torch.gather can be slow on AMD with duplicated index

[Yuzhen11]

🐛 Describe the bug

We found that torch.gather can be very slow running on AMD GPU compared with NVIDIA GPU with duplicated index.

To repro:

	
import torch
import triton

device = "cuda"
input = torch.randn([100], dtype=torch.float, device=device, requires_grad=True)
indices = torch.randint(low=0, high=1, size=(74544,), dtype=torch.int64, device=device)

def gather():
    output = torch.gather(input, dim=0, index=indices)
    output.sum().backward()

triton.testing.do_bench(lambda: gather())

The slowness comes from the backward kernel: _scatter_gather_elementwise_kernel

On AMD MI300X:

I0612 223553.302 test_gather.py:41] Link to trace: https://www.internalfb.com/intern/perfdoctor/trace_view?filepath=tree/traces/test/testttt_20240612_223534_3731825843.json.gz&bucket=pyper_traces (Meta internal only)
I0612 223608.688 test_gather.py:42] Time (ms): 1887.4288330078125

Trace:

On H100:

I0612 224254.306 test_gather.py:41] Link to trace: https://www.internalfb.com/intern/perfdoctor/trace_view?filepath=tree/traces/test/testttt_20240612_224254_3319354165.json.gz&bucket=pyper_traces (Meta internal only)
I0612 224254.408 test_gather.py:42] Time (ms): 0.26988503336906433

Versions

Latest PyTorch trunk.

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang

[jerrymannil]

The issue seems to be fixed by using atomicAdd instead of atomicAddNoRet here

I was getting a deprecated warning for atomicAddNoRet; hence tried with atomicAdd.
They are implemented using different compiler intrinsics.

atomicAddNoRet: https://github.com/ROCm/clr/blob/4ec3a977b21da0025609fad0564edaa98436e077/hipamd/include/hip/amd_detail/amd_hip_atomic.h#L283

atomicAdd: https://github.com/ROCm/clr/blob/4ec3a977b21da0025609fad0564edaa98436e077/hipamd/include/hip/amd_detail/amd_hip_atomic.h#L1425