Sometimes Advanced Indexing is Slow

mrdrozdov · November 8, 2018, 7:22pm

I found that sometimes index_select is much faster than advanced indexing. For a toy example it is 10x faster. Checking with the profiler reveals that many more ops are used for advanced indexing, whereas index_select is a single op.

Is this behavior intentional? Could advanced indexing easily check if it can use index_select instead? I imagine quite a few projects could benefit from using index_select.

Here’s the behavior reproduced on CPU:

gist.github.com

https://gist.github.com/mrdrozdov/37123eed34eeaa7d1c6640d7ad2c5278

profile.py



import torch
import argparse
from tqdm import tqdm

parser = argparse.ArgumentParser()
parser.add_argument('--flip', action='store_true')
options = parser.parse_args()

This file has been truncated. show original

spoiler-alert.txt

Profiler Output after 1 call
----------------------------

# Method A
-----------------  ---------------  ---------------  ---------------  ---------------  ---------------
Name                      CPU time        CUDA time            Calls        CPU total       CUDA total
-----------------  ---------------  ---------------  ---------------  ---------------  ---------------
select                    15.768us          0.000us                1         15.768us          0.000us
as_strided                10.561us          0.000us                1         10.561us          0.000us
_cast_int64_t              0.559us          0.000us                1          0.559us          0.000us

This file has been truncated. show original

Note: I somewhat recall that index_select required more memory than advanced indexing, which is what leads me to believe that this isn’t really a bug and there are cases where advanced indexing makes sense to use (maybe when you are indexing many dimensions simultaneously).