From bb1af718b524f9aa819f0c5f5fa200e0fe744220 Mon Sep 17 00:00:00 2001 From: Weixing Zhang Date: Mon, 9 Nov 2020 08:41:30 -0800 Subject: [PATCH] fix build failures due to recent change(858040fa) in CUDA EP (#5736) Some part of code for reduction kernels has been changed in 858040fa, which cause failures in rocm build since ROCm EP shares some code with CUDA EP. This PR is to quick fix this failure by not sharing two files for now to unblock CI enabling on ROCm EP. Another PR for leveraging 858040fa for ROCm EP will be done later. --- .../rocm/reduction/reduction_utils.cuh | 68 +++++++++++ .../rocm/reduction/reduction_all.cu | 114 ++++++++++++++++++ tools/ci_build/amd_hipify.py | 3 + 3 files changed, 185 insertions(+) create mode 100644 onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh create mode 100644 orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu 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',