diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh b/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh index 93534318f4..808756b295 100644 --- a/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh +++ b/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh @@ -54,15 +54,5 @@ struct Identity { } }; -template -struct ToBuffer { - typedef T Type; -}; - -template <> -struct ToBuffer { - typedef float Type; -}; - } // namespace rocm -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu b/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu index 185018ed97..298cfabc90 100644 --- a/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu +++ b/orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu @@ -14,18 +14,19 @@ namespace cuda { template __global__ void ScalarSqrtKernel(Tin* input, Tout* output) { *output = (Tout)_Sqrt(*input); -}; +} template void ScalarSqrt(Tin* input, Tout* output) { ScalarSqrtKernel<<<1, 1, 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 +__launch_bounds__(ChunkGroup<1>::thread_count_per_block) __global__ void MultiTensorReduceKernel(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]; @@ -79,7 +80,7 @@ __global__ void MultiTensorReduceKernel(ChunkGroup<1> chunk_group, TOut* output) if (threadIdx.x == 0) { atomic_add(w_norm, TOutOp()(TOut(shared_memory[0]))); } -}; +} template void MultiTensorReduce(ChunkGroup<1> chunk_group, TOut* output) { diff --git a/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cc b/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cc index 18cc38bcbb..0e74a682f2 100644 --- a/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cc +++ b/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cc @@ -2,23 +2,14 @@ // Licensed under the MIT License. #include "orttraining/training_ops/rocm/reduction/reduction_all.h" -#include "core/providers/rocm/reduction/reduction_functions.h" + #include "core/framework/op_kernel_context_internal.h" +#include "core/providers/rocm/reduction/reduction_functions.h" +#include "core/providers/rocm/shared_inc/accumulation_type.h" namespace onnxruntime { namespace rocm { -template -struct AccumulateType {}; -template <> -struct AccumulateType { using type = float; }; -template <> -struct AccumulateType { using type = float; }; -template <> -struct AccumulateType { using type = double; }; -template -using AccType = typename AccumulateType::type; - #define REGISTER_REDUCE_ALL_KERNEL_TYPED(Name, TIn, TOut) \ ONNX_OPERATOR_TYPED_KERNEL_EX( \ Name, \ @@ -58,7 +49,6 @@ Status ReduceAllL2::ComputeInternal(OpKernelContext* ctx) const { // bool deterministic = ctx->GetUseDeterministicCompute(); bool deterministic = true; if (!deterministic) { - typedef MultiTensorReduceL2 TFunctor; TFunctor functor; @@ -70,11 +60,9 @@ Status ReduceAllL2::ComputeInternal(OpKernelContext* ctx) const { // *p_output is the squared sum of all elements. // Let's take a sqrt to get the actual L2-norm. ScalarSqrt(p_output, p_output); - } - else { - + } else { // alternate path only for deterministic compute .. - typedef AccType HipTAcc; + typedef AccumulationType_t HipTAcc; // find scratch buffer size needed by 'reduce_square_sum' for each tensor int scratch_size = 0; @@ -86,16 +74,16 @@ Status ReduceAllL2::ComputeInternal(OpKernelContext* ctx) const { scratch_size = std::max(scratch_size, compute_reduction_buffer_size(sizeof(HipTAcc), total_tensor_count)); // add head room for final output and square norms of each tensor - scratch_size += (1 + total_tensor_count)*sizeof(HipTAcc); + scratch_size += (1 + total_tensor_count) * sizeof(HipTAcc); // create GPU scratch space and zero target for each tensor square norm auto scratch_buffer = GetScratchBuffer(scratch_size); - HIP_RETURN_IF_ERROR(hipMemsetAsync(scratch_buffer.get(), 0, sizeof(HipTAcc)*(1 + total_tensor_count))); + HIP_RETURN_IF_ERROR(hipMemsetAsync(scratch_buffer.get(), 0, sizeof(HipTAcc) * (1 + total_tensor_count))); HipTAcc* p_global_sqnorm = reinterpret_cast(scratch_buffer.get()); HipTAcc* p_tensor_sqnorm = p_global_sqnorm + 1; HipTAcc* p_reduce_buffer = p_tensor_sqnorm + total_tensor_count; - + // perform reduction l2norm = sqrt[sum(tensor[i][j]**2)] for i,j over all tensor elements for (int i = 0; i < total_tensor_count; ++i) { HipTIn* p_tensor_i = reinterpret_cast(grouped_tensor_pointers[i][0]); diff --git a/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu b/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu index 51b9d11ef4..346e093bd4 100644 --- a/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu +++ b/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu @@ -1,32 +1,34 @@ -#include "hip/hip_runtime.h" // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #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" +#include "core/providers/rocm/shared_inc/accumulation_type.h" namespace onnxruntime { namespace rocm { template -__global__ void _ScalarSqrtImpl(Tin* input, Tout* output) { +__global__ void ScalarSqrtKernel(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); -}; + hipLaunchKernelGGL(ScalarSqrtKernel, 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) { +__launch_bounds__(ChunkGroup<1>::thread_count_per_block) +__global__ void MultiTensorReduceKernel(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; @@ -79,7 +81,7 @@ __global__ void _MultiTensorReduceImpl(ChunkGroup<1> chunk_group, TOut* output) if (threadIdx.x == 0) { atomic_add(w_norm, TOutOp()(shared_memory[0])); } -}; +} template void MultiTensorReduce(ChunkGroup<1> chunk_group, TOut* output) { @@ -92,12 +94,12 @@ void MultiTensorReduce(ChunkGroup<1> chunk_group, TOut* output) { 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); + hipLaunchKernelGGL(HIP_KERNEL_NAME(MultiTensorReduceKernel), 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; + using TBuf = AccumulationType_t; MultiTensorReduce, Cast>(chunk_group, output); }