Repository: pytorch
Issue: Advanced indexing gradient is extremely slow when there are many duplicate indices
Labels: module: performance,triaged,has workaround

## 🐛 Bug

We observed a ~2x slowdown in training performance of a detection network when we did broadcast on a tensor.
This broadcasting operation launches the following kernel in the backward pass which completely eats away performance:
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/Indexing.cu#L25
The index of the broadcast has a lot of duplicates and the broadcast is done on the last dimension.

There is no performance loss when the index doesn't contain duplicates. 
Also, it doesn't matter if the index is sorted or unsorted. The code path which calls `indexing_backward_kernel` function sorts the index before computation.

## To Reproduce

Steps to reproduce the behavior:

1. Here is a python script to see the difference in performance:
```
import torch
from time import time


input_tensor = torch.randn(3, 224, 224, requires_grad=True, device='cuda')
bad_index = torch.randint(19999, 20000, (1, 224, 512), dtype=torch.long, device='cuda')
index = torch.randint(0, 20000, (1, 224, 512), dtype=torch.long, device='cuda')
x = torch.randint(0, 20, [3, 40000], dtype=torch.float32, device='cuda', requires_grad=True)

##### Good Indexing ######
#warming up cache
for i in range(10):
    y = x[:, index]
torch.cuda.synchronize()

t1 = time()
#broadcasting
y = x[:, index]
z = y.sum()
z.backward()
torch.cuda.synchronize()
t1 = time() - t1


##### Bad Indexing ######
#warming up cache
for i in range(10):
    y = x[:, bad_index]
torch.cuda.synchronize()

t2 = time()
#broadcasting
y = x[:, bad_index]
z = y.sum()
z.backward()
torch.cuda.synchronize()
t2 = time() - t2


print("Bad Indexing time: {:.3f}ms".format(t2*1000))
print("Good Indexing time: {:.3f}ms".format(t1*1000))
print("Ratio of Bad Indexing vs Good Indexing: {:.4f}x".format(t2/t1))

```
**Output:**
```
> python indexing_perf_bug.py
Bad Indexing time: 41.887ms
Good Indexing time: 2.191ms
Ratio of Bad Indexing vs Good Indexing: 19.1172x
```
Clearly with duplicate indexing @performance is ~20x worse than the non-duplicate indexing.

## 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:
```
Collecting environment information...
PyTorch version: 1.5.0
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 8.4.0-1ubuntu1~18.04) 8.4.0
CMake version: version 3.17.3

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: 10.1.243
GPU models and configuration:
GPU 0: TITAN RTX
GPU 1: TITAN RTX

Nvidia driver version: 440.100
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.16.6
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               10.2.89              hfd86e86_1
[conda] mkl                       2020.0                      166
[conda] mkl-service               2.3.0            py37he904b0f_0
[conda] mkl_fft                   1.0.15           py37ha843d7b_0
[conda] mkl_random                1.1.0            py37hd6b4f25_0
[conda] numpy                     1.18.1           py37h4f9e942_0
[conda] numpy-base                1.18.1           py37hde5b4d6_1
[conda] numpydoc                  0.9.2                      py_0
[conda] pytorch                   1.5.0           py3.7_cuda10.2.89_cudnn7.6.5_0    pytorch
[conda] torchvision               0.6.0                py37_cu102    pytorch
[conda] torchviz                  0.0.1                    pypi_0    pypi
```
@ngimel @mcarilli @csarofeen @ptrblck 

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @VitalyFedyunin @ngimel @heitorschueroff