Repository: pytorch
Issue: pytorch index_select is too slow
Labels: module: performance,module: cuda,triaged

### 🐛 Describe the bug

pytorch's `index_select` function could be faster.

I tested multiple versions of pytorch's `index_select` function against two possibly implementation using a CUDA extension that I wrote.  The extension outperforms all pytorch implementations.  I'm not yet sure why.

index_select.py:
```python
#!/usr/bin/env python3

import torch
import time
from torch.utils.cpp_extension import load

index_select = load(
  name='index_select',
  sources=["index_select.cu"],
  build_directory='.',
  verbose=True,
  extra_cuda_cflags=["--expt-extended-lambda"],
)

def main():
  randomize_trials = 3;
  index_select_trials = 3;
  key_type = torch.int32
  indices_type = torch.int32
  array_size = 1000;
  indices_count = 10000000;

  for _ in range(randomize_trials):
    keys = torch.randint(0, 100000, (array_size,), device="cuda", dtype=key_type)
    indices = torch.randint(0, array_size, (indices_count,), device="cuda", dtype=indices_type)
    for _ in range(index_select_trials):

      out = torch.empty((indices.size(0),), device="cuda", dtype=keys.dtype)
      torch.cuda.nvtx.range_push("torch.index_select_out")
      torch.index_select(keys, 0, indices, out=out)
      torch.cuda.nvtx.range_pop()

      torch.cuda.nvtx.range_push("cuda 1024x1024")
      out2 = index_select.index_select(keys, indices, 1024, 1024)
      torch.cuda.nvtx.range_pop()

      torch.cuda.nvtx.range_push("cuda 864x128")
      out3 = index_select.index_select(keys, indices, 864, 128)
      torch.cuda.nvtx.range_pop()

      torch.cuda.nvtx.range_push("torch.index_select")
      out4 = torch.index_select(keys, 0, indices)
      torch.cuda.nvtx.range_pop()

      torch.cuda.nvtx.range_push("keys[indcies]")
      out5 = keys[indices]
      torch.cuda.nvtx.range_pop()

      assert torch.sum(out == out2, 0) == indices_count
      assert torch.sum(out == out3, 0) == indices_count
      assert torch.sum(out == out4, 0) == indices_count
      assert torch.sum(out == out5, 0) == indices_count

if __name__ == "__main__":
  main()
```

index_select.cu:
```cuda
#include <pybind11/pybind11.h>
#include <torch/extension.h>
#include <chrono>

template <typename hash_index_type, typename key_type>
__global__ void index_select_kernel(hash_index_type* __restrict__ out,
                                    const hash_index_type* __restrict__ input,
                                    const key_type* __restrict__ indices,
                                    int64_t indices_count) {
  // _D is the size of data, 30 million
  // _N is the size of mask, 30 million

  // data[_N] will contain the id for each element when we're done. HASH
  // ids[_N] are the elements that we want to hash. INPUT
  // hash[_M] is the size of the hash, like 180 HASH
  // hash element _M will be used for atomics.

  auto block_count = gridDim.x;
  auto threads_per_block = blockDim.x;
  auto block_index = blockIdx.x;
  auto thread_in_block_index = threadIdx.x;
  auto thread_in_grid_index = threads_per_block * block_index + thread_in_block_index;
  // Finally, for each element, we look up the position of its
  // representative in the hash and write it into the output.
  for (uint64_t i = thread_in_grid_index; i < indices_count; i += threads_per_block*block_count) {
    auto current_index = indices[i];
    assert(current_index < 1000000);
    out[i] = input[current_index];
  }
}

torch::Tensor index_select(torch::Tensor keys, torch::Tensor indices, uint blocks, uint threads_per_block) {
  auto key_tensor_options = torch::TensorOptions()
                            .dtype(keys.dtype())
                            .device(torch::kCUDA);
  auto out = torch::empty(indices.size(0), key_tensor_options);
  index_select_kernel<<<blocks, threads_per_block>>>(
      out.template data_ptr<int32_t>(), keys.data_ptr<int32_t>(), indices.data_ptr<int32_t>(),
      indices.size(0));
  return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("index_select", &index_select, "blah blah");
}
```

### Versions

```
Collecting environment information...
PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.27.5
Libc version: glibc-2.31

Python version: 3.8.10 (default, May 26 2023, 14:05:08)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-1049-azure-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100 80GB PCIe
Nvidia driver version: 535.113.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
Address sizes:                      48 bits physical, 48 bits virtual
CPU(s):                             24
On-line CPU(s) list:                0-23
Thread(s) per core:                 1
Core(s) per socket:                 24
Socket(s):                          1
NUMA node(s):                       1
Vendor ID:                          AuthenticAMD
CPU family:                         25
Model:                              1
Model name:                         AMD EPYC 7V13 64-Core Processor
Stepping:                           1
CPU MHz:                            2445.437
BogoMIPS:                           4890.87
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          768 KiB
L1i cache:                          768 KiB
L2 cache:                           12 MiB
L3 cache:                           96 MiB
NUMA node0 CPU(s):                  0-23
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Mitigation; safe RET, no microcode
Vulnerability Spec store bypass:    Vulnerable
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl tsc_reliable nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core invpcid_single vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves clzero xsaveerptr rdpru arat umip vaes vpclmulqdq rdpid fsrm

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] torch==2.0.0
[pip3] triton==2.0.0
[conda] Could not collect
```

cc @ptrblck