From 79501899206524e6913bf3a375125c07c8f4a247 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Sat, 11 Mar 2023 08:52:01 +0800 Subject: [PATCH] [CUDA] Optimize Perf for AtomicAdd of Half Type (#14992) --- .../core/providers/cuda/atomic/common.cuh | 40 ++++++++++++++++++- .../cuda/tensor/gather_elements_impl.cu | 26 ++++++++---- .../core/providers/rocm/atomic/common.cuh | 7 ++++ 3 files changed, 64 insertions(+), 9 deletions(-) diff --git a/onnxruntime/core/providers/cuda/atomic/common.cuh b/onnxruntime/core/providers/cuda/atomic/common.cuh index d6751f9bd0..14fa2d0706 100644 --- a/onnxruntime/core/providers/cuda/atomic/common.cuh +++ b/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 +__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* 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(start_addr + index); + bool low_byte = (reinterpret_cast(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 \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu b/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu index 51389483ad..914da00f0d 100644 --- a/onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu +++ b/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 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 __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(input_offset_along_axis * input_stride_along_axis); if (IsGather) { - func(value + work, src_data + input_offset); + func(value, static_cast(work), src_data[input_offset]); } else { - func(output_data + input_offset, value + work); + func(output_data, static_cast(input_offset), value[work]); } } @@ -196,7 +200,7 @@ void GatherElementsImpl(cudaStream_t stream, const T* input_data, const TIndex* template 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 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 @@ -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()); + FuncAtomicAdd(static_cast(args.input_size))); } #define GATHER_ELEMENTS_GRAD_SPECIALIZED_TINDEX_IMPL(T, TIndex) \ diff --git a/onnxruntime/core/providers/rocm/atomic/common.cuh b/onnxruntime/core/providers/rocm/atomic/common.cuh index 7cf4e31e9b..4e23570202 100644 --- a/onnxruntime/core/providers/rocm/atomic/common.cuh +++ b/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 +__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