[CUDA] Optimize Perf for AtomicAdd of Half Type (#14992)

onnxruntime/core/providers/cuda/atomic/common.cuh

@@ -84,5 +84,43 @@ __device__ __forceinline__ void atomic_add(BFloat16* address, BFloat16 value) {
   } while (assumed != old);
 }
 
+// CUDA's atomic_add for half type is too slow. Using half2 will be much faster. To avoid address out of bound,
+// we need to pass in the numel so that the element at the edge can be handled separately.
+// Ideally we need to deprecate above atomic_add function and use below one for better performance.
+// But since the signature is different, we can change it for specific Op kernel once we find it is slow.
+// TODO: need to add same logic for BF16.
+template <typename T>
+__device__ __forceinline__ void AtomicAdd(T *start_addr, size_t index, const size_t numel, T value) {
+  ORT_UNUSED_PARAMETER(numel);
+  atomic_add(start_addr + index, value);
+}
+
+template <>
+__device__ __forceinline__ void AtomicAdd<half>(half* start_addr, size_t index, const size_t numel, half value) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 700
+  atomic_add(start_addr + index, value);
+#else
+  // Accounts for the chance tensor falls on an odd 16 bit alignment (ie, not 32 bit aligned)
+  half* target_addr = reinterpret_cast<half*>(start_addr + index);
+  bool low_byte = (reinterpret_cast<std::uintptr_t>(target_addr) % sizeof(__half2) == 0);
+
+  if (low_byte && index < (numel - 1)) {
+    __half2 value2;
+    value2.x = value;
+    value2.y = __int2half_rz(0);
+    atomicAdd(reinterpret_cast<__half2*>(target_addr), value2);
+
+  } else if (!low_byte && index > 0) {
+    __half2 value2;
+    value2.x = __int2half_rz(0);
+    value2.y = value;
+    atomicAdd(reinterpret_cast<__half2*>(target_addr - 1), value2);
+
+  } else {
+    atomicAdd(start_addr + index, value);
+  }
+#endif
+}
+
 }  // namespace cuda
-}  // namespace onnxruntime
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu

@@ -14,7 +14,11 @@ namespace onnxruntime {
 namespace cuda {
 
 namespace {
-constexpr int kThreadsPerBlock = GridDim::maxThreadsPerBlock;
+#ifdef USE_ROCM
+constexpr int kThreadsPerBlock = 256;
+#else
+constexpr int kThreadsPerBlock = GPU_WARP_SIZE * 4;
+#endif
 constexpr int kThreadWorkSize = 4;
 
 // General case to compute the input(for Gather)/output(for Scatter) and indices data offset given the thread ID
@@ -91,13 +95,13 @@ struct OffsetCalculatorFor2D {
 
 template <class T>
 struct FuncAssignment {
-  __device__ __inline__ void operator()(T* a, const T* b) const { *a = *b; }
+  __device__ __inline__ void operator()(T* start_addr, size_t index, T value) const { start_addr[index] = value; }
 };
 
 template <typename T, typename TIndex, bool IsGather, typename OffsetCalcT, typename TFunc>
 __global__ void _GatherScatterElementsKernel(const T* src_data, const TIndex* indices_data, T* output_data,
                                              const int64_t input_dim_along_axis, const int64_t input_stride_along_axis,
-                                             const OffsetCalcT offset_calc, const TFunc& func, CUDA_LONG N) {
+                                             const OffsetCalcT offset_calc, const TFunc func, CUDA_LONG N) {
   CUDA_LONG start = kThreadsPerBlock * kThreadWorkSize * blockIdx.x + threadIdx.x;
   CUDA_LONG id;
   T value[kThreadWorkSize];
@@ -124,9 +128,9 @@ __global__ void _GatherScatterElementsKernel(const T* src_data, const TIndex* in
         CUDA_LONG input_offset = offsets[0] + static_cast<CUDA_LONG>(input_offset_along_axis * input_stride_along_axis);
 
         if (IsGather) {
-          func(value + work, src_data + input_offset);
+          func(value, static_cast<size_t>(work), src_data[input_offset]);
         } else {
-          func(output_data + input_offset, value + work);
+          func(output_data, static_cast<size_t>(input_offset), value[work]);
         }
       }
 
@@ -196,7 +200,7 @@ void GatherElementsImpl(cudaStream_t stream, const T* input_data, const TIndex*
 template <typename T, typename TIndex, typename TFunc>
 Status ScatterElementsImplInternal(cudaStream_t stream, const T* input_data, const TIndex* indices_data,
                                    const T* updates_data, T* output_data, const GatherScatterElementsArgs& args,
-                                   const TFunc& func) {
+                                   const TFunc func) {
   if (input_data != output_data) {
     CUDA_RETURN_IF_ERROR(
         cudaMemcpyAsync(output_data, input_data, args.input_size * sizeof(T), cudaMemcpyDeviceToDevice, stream));
@@ -263,7 +267,12 @@ GATHER_SCATTER_ELEMENTS_SPECIALIZED_IMPL(double)
 
 template <class T>
 struct FuncAtomicAdd {
-  __device__ __inline__ void operator()(T* a, const T* b) const { atomic_add(a, *b); }
+  FuncAtomicAdd(const size_t numel) : numel_(numel) {}
+  __device__ __inline__ void operator()(T* start_addr, size_t index, T value) const {
+    AtomicAdd(start_addr, index, numel_, value);
+  }
+
+  const size_t numel_;
 };
 
 template <typename T, typename TIndex>
@@ -271,8 +280,9 @@ Status GatherElementsGradImpl(cudaStream_t stream, const TIndex* indices_data, c
                               const GatherScatterElementsArgs& args) {
   // Give output_data as the input_data parameter by intention,
   // to skip input_data copy, which is not applicable for GatherElementsGrad.
+  // output_data's numel is same as input_data's numel.
   return ScatterElementsImplInternal(stream, output_data, indices_data, updates_data, output_data, args,
-                                     FuncAtomicAdd<T>());
+                                     FuncAtomicAdd<T>(static_cast<size_t>(args.input_size)));
 }
 
 #define GATHER_ELEMENTS_GRAD_SPECIALIZED_TINDEX_IMPL(T, TIndex)                                      \

onnxruntime/core/providers/rocm/atomic/common.cuh

@@ -52,5 +52,12 @@ __device__ __forceinline__ void atomic_add(BFloat16* address, BFloat16 value) {
   } while (assumed != old);
 }
 
+// This function is added to speed up atomic add for half/bf16 type on CUDA. For ROCM, use default implementation.
+template <typename T>
+__device__ __forceinline__ void AtomicAdd(T *start_addr, size_t index, const size_t numel, T value) {
+  ORT_UNUSED_PARAMETER(numel);
+  atomic_add(start_addr + index, value);
+}
+
 }  // namespace rocm
 }  // namespace onnxruntime