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.

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<unsigned int>(value);
+  --value_;
+  value_ |= value_ >> 1;
+  value_ |= value_ >> 2;
+  value_ |= value_ >> 4;
+  value_ |= value_ >> 8;
+  value_ |= value_ >> 16;
+  return static_cast<int>(++value_);
+}
+
+template <typename TAccumulated, typename TValue>
+struct Cast {
+  __forceinline__ __device__ TAccumulated operator()(const TValue& value) {
+    return TAccumulated(value);
+  }
+};
+
+template <typename TAccumulated, typename TValue>
+struct Square {
+  __forceinline__ __device__ TAccumulated operator()(const TValue& value) {
+    return TAccumulated(value) * TAccumulated(value);
+  }
+};
+
+template <typename TAccumulated, typename TValue>
+struct Abs {
+  __forceinline__ __device__ TAccumulated operator()(const TValue& value) {
+    TAccumulated value_ = TAccumulated(value);
+    return value_ > TAccumulated(0) ? value_ : -value_;
+  }
+};
+
+template <typename T>
+struct Sqrt {
+  __forceinline__ __device__ T operator()(const T& value) {
+    return _Sqrt(value);
+  }
+};
+
+template <typename T>
+struct Identity {
+  __forceinline__ __device__ T operator()(const T& value) {
+    return value;
+  }
+};
+
+template <typename T>
+struct ToBuffer {
+  typedef T Type;
+};
+
+template <>
+struct ToBuffer<half> {
+  typedef float Type;
+};
+
+}  // namespace rocm
+}  // namespace onnxruntime

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<typename Tin, typename Tout>
+__global__ void _ScalarSqrtImpl(Tin* input, Tout* output) {
+  *output = (Tout)_Sqrt(*input);
+};
+
+template<typename Tin, typename Tout>
+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 <typename TIn, typename TOut, typename TBuf, typename TInOp, typename TOutOp>
+__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<const TIn*>(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<TBuf*>(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 <typename TIn, typename TOut, typename TBuf, typename TInOp, typename TOutOp>
+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<TIn, TOut, TBuf, TInOp, TOutOp>), dim3(chunk_group.chunk_count), dim3(thread_count), shared_memory_size, 0, chunk_group, output);
+}
+
+template <typename TIn, typename TOut>
+void MultiTensorReduceL2<TIn, TOut>::operator()(ChunkGroup<1> chunk_group, TOut* output) {
+  typedef typename ToBuffer<TIn>::Type TBuf;
+  MultiTensorReduce<TIn, TOut, TBuf, Square<TBuf, TIn>, Cast<TOut, TBuf>>(chunk_group, output);
+}
+
+#define INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(TIn, TOut) \
+  template void MultiTensorReduceL2<TIn, TOut>::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

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',