diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh b/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh new file mode 100644 index 0000000000..93534318f4 --- /dev/null +++ b/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh @@ -0,0 +1,68 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include "core/providers/rocm/cu_inc/common.cuh" + +namespace onnxruntime { +namespace rocm { + +__forceinline__ __host__ __device__ int least_pow2_bound(int value) { + unsigned int value_ = static_cast(value); + --value_; + value_ |= value_ >> 1; + value_ |= value_ >> 2; + value_ |= value_ >> 4; + value_ |= value_ >> 8; + value_ |= value_ >> 16; + return static_cast(++value_); +} + +template +struct Cast { + __forceinline__ __device__ TAccumulated operator()(const TValue& value) { + return TAccumulated(value); + } +}; + +template +struct Square { + __forceinline__ __device__ TAccumulated operator()(const TValue& value) { + return TAccumulated(value) * TAccumulated(value); + } +}; + +template +struct Abs { + __forceinline__ __device__ TAccumulated operator()(const TValue& value) { + TAccumulated value_ = TAccumulated(value); + return value_ > TAccumulated(0) ? value_ : -value_; + } +}; + +template +struct Sqrt { + __forceinline__ __device__ T operator()(const T& value) { + return _Sqrt(value); + } +}; + +template +struct Identity { + __forceinline__ __device__ T operator()(const T& value) { + return value; + } +}; + +template +struct ToBuffer { + typedef T Type; +}; + +template <> +struct ToBuffer { + typedef float Type; +}; + +} // namespace rocm +} // namespace onnxruntime \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu b/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu new file mode 100644 index 0000000000..51b9d11ef4 --- /dev/null +++ b/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu @@ -0,0 +1,114 @@ +#include "hip/hip_runtime.h" +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "orttraining/training_ops/rocm/reduction/reduction_all.h" +#include "core/providers/rocm/cu_inc/common.cuh" +#include "core/providers/rocm/rocm_common.h" +#include "core/providers/rocm/atomic/common.cuh" +#include "core/providers/rocm/reduction/reduction_utils.cuh" + +namespace onnxruntime { +namespace rocm { + +template +__global__ void _ScalarSqrtImpl(Tin* input, Tout* output) { + *output = (Tout)_Sqrt(*input); +}; + +template +void ScalarSqrt(Tin* input, Tout* output) { + hipLaunchKernelGGL(_ScalarSqrtImpl, dim3(1), dim3(1), 0, 0, input, output); +}; + +template void ScalarSqrt(float* input, float* output); +template void ScalarSqrt(half* input, half* output); +template void ScalarSqrt(float* input, half* output); + +template +__global__ void _MultiTensorReduceImpl(ChunkGroup<1> chunk_group, TOut* output) { + const int group_index = chunk_group.block_index_to_tensor_group_index[blockIdx.x]; + const int tensor_size = chunk_group.tensor_sizes[group_index]; + const int chunk_size = chunk_group.chunk_size; + const int chunk_start = chunk_group.block_index_to_chunk_start_index[blockIdx.x]; + const TIn* w = reinterpret_cast(chunk_group.tensor_ptrs[0][group_index]) + chunk_start; + TOut* w_norm = output; + + TBuf w_sum = TBuf(0.f); + constexpr int load_count_per_thread = 4; + for (int i = threadIdx.x; i < chunk_size && i + chunk_start < tensor_size; i += blockDim.x * load_count_per_thread) { +#pragma unroll + for (int j = 0; j < load_count_per_thread; ++j) { + const int index_in_chunk = i + j * blockDim.x; + const int index_in_tensor = chunk_start + index_in_chunk; + if (index_in_chunk < chunk_size && index_in_tensor < tensor_size) { + const TBuf w_element = TBuf(w[index_in_chunk]); + w_sum += TInOp()(w_element); + } + } + } + + // Thread count in a block must be a multiple of GPU_WARP_SIZE. +#pragma unroll + for (int stride = GPU_WARP_SIZE / 2; stride > 0; stride /= 2) { + w_sum += WARP_SHFL_DOWN(w_sum, stride); + } + + const int warp_count_in_block = blockDim.x / GPU_WARP_SIZE; + const int lid = threadIdx.x % GPU_WARP_SIZE; + const int wid = threadIdx.x / GPU_WARP_SIZE; + + // Shape is 2 x warp_count_in_block. + HIP_DYNAMIC_SHARED( unsigned char, shared_memory_) + TBuf* shared_memory = reinterpret_cast(shared_memory_); + + if (lid == 0) { + shared_memory[wid] = w_sum; + } + + __syncthreads(); + +#pragma unroll + for (int stride = warp_count_in_block / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { + shared_memory[threadIdx.x] += shared_memory[threadIdx.x + stride]; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + atomic_add(w_norm, TOutOp()(shared_memory[0])); + } +}; + +template +void MultiTensorReduce(ChunkGroup<1> chunk_group, TOut* output) { + // thread count per block. + constexpr int thread_count = ChunkGroup<1>::thread_count_per_block; + // shared memory's size per block. + const int shared_memory_size = thread_count / GPU_WARP_SIZE * sizeof(TBuf); + + // Enforce assumptions used inside this reduction ROCM kernel. + ORT_ENFORCE(thread_count % GPU_WARP_SIZE == 0); + ORT_ENFORCE((thread_count & (thread_count - 1)) == 0); + + hipLaunchKernelGGL(HIP_KERNEL_NAME(_MultiTensorReduceImpl), dim3(chunk_group.chunk_count), dim3(thread_count), shared_memory_size, 0, chunk_group, output); +} + +template +void MultiTensorReduceL2::operator()(ChunkGroup<1> chunk_group, TOut* output) { + typedef typename ToBuffer::Type TBuf; + MultiTensorReduce, Cast>(chunk_group, output); +} + +#define INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(TIn, TOut) \ + template void MultiTensorReduceL2::operator()(ChunkGroup<1> chunk_group, TOut* output); + +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(double, float) +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, float) +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(half, float) +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, half) +INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(half, half) + +} // namespace rocm +} // namespace onnxruntime \ No newline at end of file diff --git a/tools/ci_build/amd_hipify.py b/tools/ci_build/amd_hipify.py index c765655a6b..935b209229 100644 --- a/tools/ci_build/amd_hipify.py +++ b/tools/ci_build/amd_hipify.py @@ -133,10 +133,12 @@ core_ops_files = [ 'object_detection/roialign.h', 'object_detection/roialign_impl.cu', 'object_detection/roialign_impl.h', + 'reduction/reduction_functions.cc', 'reduction/reduction_functions.cu', 'reduction/reduction_functions.h', 'reduction/reduction_ops.cc', 'reduction/reduction_ops.h', + 'reduction/reduction_utils.cuh', 'rnn/cudnn_rnn_base.cc', 'rnn/cudnn_rnn_base.h', 'rnn/gru.cc', @@ -258,6 +260,7 @@ training_ops_files = [ 'optimizer/adam.cu', 'optimizer/lamb.cc', 'reduction/reduction_all.cc', + 'reduction/reduction_all.cu', 'reduction/reduction_ops.cc', 'tensor/gather_elements_grad.cc', 'tensor/gather_elements_grad.h',