Repository: pytorch
Issue: Very slow backward speed when using gather with small-range indices
Labels: module: performance,triaged

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
When the range of `index` is small in gather, the backward speed becomes very slow.


## To Reproduce

Steps to reproduce the behavior:
```python
import torch
import math
import torch.utils.benchmark as benchmark


bsz=32
num_head=12
cur_len=500
seq_len=512

data = torch.randn(bsz, num_head, cur_len, seq_len, device=torch.device("cuda"), dtype=torch.half)
data2 = torch.randn(bsz, num_head, cur_len, seq_len, device=torch.device("cuda"), dtype=torch.half, requires_grad=True)

index1 = torch.randint(low=0, high=512, size=(bsz, num_head, cur_len, cur_len ), device=torch.device("cuda"))
index2 = torch.randint(low=10, high=12, size=(bsz, num_head, cur_len, cur_len ), device=torch.device("cuda"))


def test(data, index):
    a = torch.gather(data, dim=-1, index=index)

def test_bp(data, index):
    a = torch.gather(data, dim=-1, index=index)
    s = a.sum()
    s.backward()


t0 = benchmark.Timer(
    stmt='test(data, index)',
    setup='from __main__ import test',
    globals={'data': data, 'index': index1})

t1 = benchmark.Timer(
    stmt='test(data, index)',
    setup='from __main__ import test',
    globals={'data': data, 'index': index2})

print('forward speed for large-range index,', t0.timeit(100))
print('forward speed for small-range index,', t1.timeit(100))


t2 = benchmark.Timer(
    stmt='test_bp(data, index)',
    setup='from __main__ import test_bp',
    globals={'data': data2, 'index': index1})

t3 = benchmark.Timer(
    stmt='test_bp(data, index)',
    setup='from __main__ import test_bp',
    globals={'data': data2, 'index': index2})

print('forward+backward speed for large-range index,', t2.timeit(100))
print('forward+backward speed for small-range index,', t3.timeit(100))

```

outputs:
```python
forward speed for large-range index, <torch.utils.benchmark.utils.common.Measurement object at 0x7fb92fe4bb10>
test(data, index)
  1.46 ms
  1 measurement, 100 runs , 1 thread

forward speed for small-range index, <torch.utils.benchmark.utils.common.Measurement object at 0x7fb92fe4bcd0>
test(data, index)
  1.21 ms
  1 measurement, 100 runs , 1 thread

forward+backward speed for large-range index, <torch.utils.benchmark.utils.common.Measurement object at 0x7fb92fe4be90>
test_bp(data, index)
  7.84 ms
  1 measurement, 100 runs , 1 thread

forward+backward speed for small-range index, <torch.utils.benchmark.utils.common.Measurement object at 0x7fb92fe4bb10>
test_bp(data, index)
  138.51 ms
  1 measurement, 100 runs , 1 thread

```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

PyTorch version: 1.7.1
Is debug build: False
CUDA used to build PyTorch: 11.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.19.3

Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 11.0.194
GPU models and configuration:
GPU 0: Tesla V100-PCIE-16GB
GPU 1: Tesla V100-PCIE-16GB
GPU 2: Tesla V100-PCIE-16GB
GPU 3: Tesla V100-PCIE-16GB

Nvidia driver version: 450.51.05
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn.so.8
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.0
[pip3] torch==1.7.1
[pip3] torchaudio==0.7.0a0+a853dff
[pip3] torchvision==0.8.2
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.0.221             h6bb024c_0
[conda] mkl                       2020.2                      256
[conda] mkl-service               2.3.0            py37he8ac12f_0
[conda] mkl_fft                   1.2.0            py37h23d657b_0
[conda] mkl_random                1.1.1            py37h0573a6f_0
[conda] numpy                     1.20.0                   pypi_0    pypi
[conda] pytorch                   1.7.1           py3.7_cuda11.0.221_cudnn8.0.5_0    pytorch
[conda] torchaudio                0.7.2                      py37    pytorch
[conda] torchvision               0.8.2                py37_cu110    pytorch


## Additional context

<!-- Add any other context about the problem here. -->


cc @VitalyFedyunin @ngimel