Partial updating of ROCM reduction code.

onnxruntime/core/providers/cuda/reduction/reduction_functions.h

@@ -12,7 +12,7 @@ namespace cuda {
 namespace detail {
 size_t compute_reduce_matrix_columns_intermediate_buffer_size(
     int element_size, int num_rows, int num_cols);
-}
+}  // namespace detail
 
 /**
  * Computes the size in bytes of the intermediate buffer needed by reduce_matrix_columns().

onnxruntime/core/providers/rocm/reduction/reduction_functions.cu

@@ -1,12 +1,16 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "core/providers/rocm/reduction/reduction_functions.h"
+
 #include <algorithm>
+
 #include <hip/hip_runtime.h>
 #include <hip/hip_fp16.h>
-#include "core/providers/rocm/cu_inc/common.cuh"
+#include "core/common/common.h"
 #include "core/providers/rocm/atomic/common.cuh"
-#include "core/providers/rocm/reduction/reduction_functions.h"
+#include "core/providers/rocm/cu_inc/common.cuh"
+#include "core/providers/rocm/shared_inc/rocm_utils.h"
 #include "core/providers/rocm/reduction/reduction_utils.cuh"
 
 #define NUM_ELEMENTS_PER_THREAD 4
@@ -19,6 +23,481 @@
 namespace onnxruntime {
 namespace rocm {
 
+namespace detail {
+constexpr auto MAX_NUM_ELEMENTS_PER_THREAD = 4;
+constexpr auto MAX_NUM_WARPS_PER_BLOCK = 8;
+constexpr auto MAX_NUM_BLOCKS_IN_GRID_ROW = 256;
+constexpr auto MAX_NUM_GRID_ROWS = 32768;
+
+dim3 compute_block_dim(int num_cols) {
+  const int x = GPU_WARP_SIZE;
+  const int y = std::min(MAX_NUM_WARPS_PER_BLOCK, std::max(1, num_cols / (MAX_NUM_ELEMENTS_PER_THREAD * x)));
+  return dim3(x, y);
+}
+
+std::pair<dim3, dim3> compute_grid_and_block_dims(int num_rows, int num_cols) {
+  const auto block_dim = compute_block_dim(num_cols);
+  const auto grid_x =
+      std::min<int>(
+          MAX_NUM_BLOCKS_IN_GRID_ROW,
+          std::max<int>(1, num_cols / (MAX_NUM_ELEMENTS_PER_THREAD * block_dim.x * block_dim.y)));
+  const auto grid_y = std::min(MAX_NUM_GRID_ROWS, num_rows);
+  const dim3 grid_dim(grid_x, grid_y);
+  return {grid_dim, block_dim};
+}
+
+uintptr_t round_up_to_aligned(uintptr_t original, size_t alignment) {
+  assert((alignment & (alignment - 1)) == 0);
+  const size_t alignment_mask = ~(alignment - 1);
+  return (original + alignment - 1) & alignment_mask;
+}
+
+/**
+ * call_reduce_matrix_columns() intermediate buffer layout
+ *
+ * Given buffer element type TBuf, the intermediate buffer layout looks like this:
+ *
+ * -----
+ * m * num_blocks_per_row * sizeof(TBuf) bytes for block reductions per row
+ * alignment padding bytes as needed
+ * m * sizeof(int) bytes for block done counts per row
+ * -----
+ */
+
+size_t compute_reduce_matrix_columns_intermediate_buffer_size(
+    int element_size, int num_rows, int num_cols) {
+  ORT_ENFORCE(element_size >= 0 && num_rows >= 0 && num_cols >= 0);
+
+  const auto grid_dim = compute_grid_and_block_dims(num_rows, num_cols).first;
+
+  size_t buffer_size{};
+
+  // at the beginning, for sizing purposes, assume we are aligned
+  buffer_size += static_cast<size_t>(num_rows) * grid_dim.x * element_size;
+
+  buffer_size = round_up_to_aligned(buffer_size, alignof(int));
+  buffer_size += static_cast<size_t>(num_rows) * sizeof(int);
+
+  // add padding to give us room to align
+  buffer_size += alignof(max_align_t) - 1;
+
+  return buffer_size;
+}
+
+template <typename TBuf>
+Status get_reduction_buffers(
+    int num_rows, int num_cols, void* buffer, size_t buffer_size,
+    TBuf*& block_reductions_buffer, int*& block_done_counts_buffer) {
+  const auto grid_dim = compute_grid_and_block_dims(num_rows, num_cols).first;
+
+  const uintptr_t begin_addr = reinterpret_cast<uintptr_t>(buffer);
+  const uintptr_t block_reductions_addr =
+      round_up_to_aligned(begin_addr, alignof(TBuf));
+  const uintptr_t block_done_counts_buffer_addr =
+      round_up_to_aligned(
+          block_reductions_addr + static_cast<size_t>(num_rows) * grid_dim.x * sizeof(TBuf), alignof(int));
+  const uintptr_t end_addr =
+      block_done_counts_buffer_addr + static_cast<size_t>(num_rows) * sizeof(int);
+  const size_t required_size = end_addr - begin_addr;
+
+  ORT_RETURN_IF_NOT(
+      required_size <= buffer_size,
+      "Buffer size is too small (", buffer_size, " bytes). ",
+      "At least ", required_size, " bytes are needed from the given base address (", buffer, ").");
+
+  block_reductions_buffer = reinterpret_cast<TBuf*>(block_reductions_addr);
+  block_done_counts_buffer = reinterpret_cast<int*>(block_done_counts_buffer_addr);
+
+  return Status::OK();
+}
+
+template <typename TIn, typename TOut, typename TBuf, typename TOp, typename TFinalOp, bool DivideResultBySize>
+__device__ void reduce_all(
+    const int num_elements, const TIn* const input, TOut* const output,
+    TBuf* const block_reductions_buffer, int* const block_done_count_buffer) {
+  extern __shared__ unsigned char shared_memory_bytes[];
+  TBuf* shared_memory = reinterpret_cast<TBuf*>(shared_memory_bytes);
+  // Thread-level indices:
+  // Linear index of thread in block.
+  const int tid_in_block = threadIdx.y * blockDim.x + threadIdx.x;
+  // Total number of threads in a 2-D block.
+  const int num_threads_in_block = blockDim.x * blockDim.y;
+
+  // Warp-level indices:
+  // Warp index of thread.
+  const int wid_in_block = tid_in_block / GPU_WARP_SIZE;
+  // Lane index of thread.
+  const int lid_in_block = tid_in_block % GPU_WARP_SIZE;
+  // Warp count per block.
+  const int num_warps_in_block = num_threads_in_block / GPU_WARP_SIZE;
+
+  // Grid-level indices:
+  // Linear index of block in grid row.
+  const int bid_in_grid_row = blockIdx.x;
+  // Linear index of thread in grid row.
+  const int tid_in_grid_row = bid_in_grid_row * (blockDim.x * blockDim.y) + tid_in_block;
+  // Total number of blocks in a grid row.
+  const int num_blocks_in_grid_row = gridDim.x;
+  // Total number of threads in a grid row with 2-D blocks.
+  const int num_threads_in_grid_row = num_blocks_in_grid_row * num_threads_in_block;
+
+  const auto write_result = [&output, &num_elements](const TOut result) {
+    // Compilation time if-else branch controlled by template argument can be
+    // optimized out, so there will be no branch in real computation phase.
+    if (DivideResultBySize) {
+      output[0] = TFinalOp()(result / TOut(num_elements));
+    } else {
+      output[0] = TFinalOp()(result);
+    }
+  };
+
+  // Thread-level reduction (storage change: global memory -> register).
+  // One thread reduces MAX_NUM_ELEMENTS_PER_THREAD elements to a thread register
+  // in one iteration.
+  TBuf value = 0;
+  for (int id = tid_in_grid_row; id < num_elements; id += MAX_NUM_ELEMENTS_PER_THREAD * num_threads_in_grid_row) {
+    TIn v[MAX_NUM_ELEMENTS_PER_THREAD];
+
+#pragma unroll
+    for (int i = 0; i < MAX_NUM_ELEMENTS_PER_THREAD; i++) {
+      const int offset = id + i * num_threads_in_grid_row;
+      if (offset < num_elements) {
+        v[i] = input[offset];
+      }
+    }
+
+#pragma unroll
+    for (int i = 0; i < MAX_NUM_ELEMENTS_PER_THREAD; i++) {
+      const int offset = id + i * num_threads_in_grid_row;
+      if (offset < num_elements) {
+        value += TOp()(TBuf(v[i]));
+      }
+    }
+  }
+
+#if __ROCM_ARCH__ >= 700
+  __syncwarp();
+#else
+  __syncthreads();
+#endif
+
+  // Warp-level reduction (storage change: register -> register).
+  // The values in a warp will be summed up to a scalar. After warp-level
+  // reduction, each block holds num_warps_in_block values in the shared memory.
+#pragma unroll
+  for (int stride = GPU_WARP_SIZE / 2; stride > 0; stride /= 2) {
+    value += WARP_SHFL_DOWN(value, stride);
+  }
+
+  // Return early if only one warp is used for reduction.
+  // Given a fixed amount of threads, we prefer threads over warps over blocks so that we never have cases such as
+  // 1. two blocks and each of them has only 1 warp (32 threads).
+  // 2. two warps and each of them has only 2 threads.
+  if (num_warps_in_block == 1) {
+    if (tid_in_grid_row == 0) {
+      write_result(value);
+    }
+    return;
+  }
+
+  if (lid_in_block == 0) {
+    shared_memory[wid_in_block] = value;
+  }
+
+  __syncthreads();
+
+  // Block-level reduction (storage change: shared memory -> global memory).
+  // The values in a block will be summed up to a scalar.
+  // Note that the values are stored in the shared memory.
+  // Here we assume that the size of shared_memory is smaller
+  // than num_warps_in_block, so we just keep halving the number
+  // of threads in each iteration. Our assumption is always true because
+  // the size of shared_memory equals to the number of warps.
+#pragma unroll
+  for (int stride = MAX_NUM_WARPS_PER_BLOCK / 2; stride > 0; stride /= 2) {
+    if (tid_in_block + stride < num_warps_in_block) {
+      shared_memory[tid_in_block] += shared_memory[tid_in_block + stride];
+    }
+    __syncthreads();
+  }
+
+  // Return early if only one block is used for reduction.
+  if (num_blocks_in_grid_row == 1) {
+    if (tid_in_grid_row == 0) {
+      write_result(shared_memory[0]);
+    }
+    return;
+  }
+
+  if (tid_in_block == 0) {
+    block_reductions_buffer[bid_in_grid_row] = shared_memory[0];
+  }
+
+  __threadfence();
+  __syncthreads();
+
+  // Grid-level reduction. We use the last block to sum up values
+  // stored in the global block_reductions_buffer.
+  __shared__ bool is_last_block_done;
+
+  if (tid_in_block == 0) {
+    const int count = atomicAdd(block_done_count_buffer, 1);
+    is_last_block_done = (count == (num_blocks_in_grid_row - 1));
+  }
+
+  // All threads in each block see if they belong the last active block
+  // (i.e., the value of is_last_block_done).
+  __syncthreads();
+
+  // Only the block which saw that count equals to num_blocks_in_grid_row - 1 can
+  // enter the following block.
+  if (is_last_block_done) {
+    const int pow2_bound = least_pow2_bound(num_blocks_in_grid_row);
+    for (int stride = pow2_bound / 2; stride > 0; stride /= 2) {
+      if (tid_in_block < stride && tid_in_block + stride < num_blocks_in_grid_row) {
+        block_reductions_buffer[tid_in_block] += block_reductions_buffer[tid_in_block + stride];
+      }
+      __syncthreads();
+    }
+
+    // The first thread in the last block assigns the final output.
+    if (tid_in_block == 0) {
+      write_result(block_reductions_buffer[0]);
+    }
+  }
+}
+
+template <typename TIn, typename TOut, typename TBuf, typename TOp, typename TFinalOp, bool DivideResultBySize>
+__global__ void reduce_matrix_columns_kernel(
+    const int num_rows, const int num_cols, const TIn* const input, TOut* const output,
+    TBuf* const block_reductions_buffer, int* const block_done_counts_buffer) {
+  const int num_blocks_in_grid_row = gridDim.x;
+  const int row_id_in_grid = blockIdx.y;
+  const int num_grid_rows = gridDim.y;
+
+  // one row per iteration
+  // row_id is int64_t to avoid int overflow in offset calculations
+  for (int64_t row_id = row_id_in_grid; row_id < num_rows; row_id += num_grid_rows) {
+    const TIn* const row_data = input + row_id * num_cols;
+    TOut* const row_output = output + row_id;
+    TBuf* const row_block_reductions_buffer = block_reductions_buffer + row_id * num_blocks_in_grid_row;
+    int* const row_block_done_counts_buffer = block_done_counts_buffer + row_id;
+
+    reduce_all<TIn, TOut, TBuf, TOp, TFinalOp, DivideResultBySize>(
+        num_cols, row_data, row_output,
+        row_block_reductions_buffer, row_block_done_counts_buffer);
+  }
+}
+
+template <typename TIn, typename TOut, typename TOp, typename TFinalOp, bool DivideResultBySize>
+Status call_reduce_matrix_columns(
+    const TIn* input, TOut* output, const int num_rows, const int num_cols, void* buffer, size_t buffer_size) {
+  ORT_ENFORCE(num_rows >= 0 && num_cols >= 0);
+
+  using TBuf = AccumulationType_t<TIn>;
+
+  const auto grid_and_block_dims = compute_grid_and_block_dims(num_rows, num_cols);
+  const dim3& grid_dim = grid_and_block_dims.first;
+  const dim3& block_dim = grid_and_block_dims.second;
+
+  TBuf* block_reductions_buffer;
+  int* block_done_counts_buffer;
+  ORT_RETURN_IF_ERROR(get_reduction_buffers(
+      num_rows, num_cols, buffer, buffer_size,
+      block_reductions_buffer, block_done_counts_buffer));
+
+  // If more than one block is used per grid row, then inter-block reduction is needed.
+  if (grid_dim.x > 1) {
+    HIP_RETURN_IF_ERROR(hipMemsetAsync(block_done_counts_buffer, 0, num_rows * sizeof(int)));
+  }
+
+  const int shared_mem_size = sizeof(TBuf) * block_dim.x * block_dim.y / GPU_WARP_SIZE;
+  hipLaunchKernelGGL(HIP_KERNEL_NAME(reduce_matrix_columns_kernel<TIn, TOut, TBuf, TOp, TFinalOp, DivideResultBySize>),
+      grid_dim, block_dim, shared_mem_size, 0,
+      num_rows, num_cols, input, output, block_reductions_buffer, block_done_counts_buffer);
+
+  return Status::OK();
+}
+}  // namespace detail
+
+template <typename TIn, typename TOut>
+Status reduce_sum(
+    const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size) {
+  return detail::call_reduce_matrix_columns<TIn, TOut, Identity2, Identity2, false>(
+      input, output, 1, size, buffer, buffer_size);
+}
+
+template <typename TIn, typename TOut>
+Status reduce_square_sum(
+    const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size) {
+  return detail::call_reduce_matrix_columns<TIn, TOut, Square2, Identity2, false>(
+      input, output, 1, size, buffer, buffer_size);
+}
+
+template <typename TIn, typename TOut>
+Status reduce_l2_norm(
+    const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size) {
+  return detail::call_reduce_matrix_columns<TIn, TOut, Square2, Sqrt2, false>(
+      input, output, 1, size, buffer, buffer_size);
+}
+
+template <typename TIn, typename TOut>
+Status reduce_mean(
+    const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size) {
+  return detail::call_reduce_matrix_columns<TIn, TOut, Identity2, Identity2, true>(
+      input, output, 1, size, buffer, buffer_size);
+}
+
+#define INSTANTIATE_REDUCE_SUM(TIn, TOut) \
+  template Status reduce_sum<TIn, TOut>(const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size)
+INSTANTIATE_REDUCE_SUM(half, float);
+INSTANTIATE_REDUCE_SUM(float, float);
+INSTANTIATE_REDUCE_SUM(double, double);
+#undef INSTANTIATE_REDUCE_SUM
+
+#define INSTANTIATE_REDUCE_SQUARE_SUM(TIn, TOut) \
+  template Status reduce_square_sum<TIn, TOut>(const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size)
+INSTANTIATE_REDUCE_SQUARE_SUM(half, float);
+INSTANTIATE_REDUCE_SQUARE_SUM(float, float);
+INSTANTIATE_REDUCE_SQUARE_SUM(double, double);
+#undef INSTANTIATE_REDUCE_SQUARE_SUM
+
+#define INSTANTIATE_REDUCE_L2_NORM(TIn, TOut) \
+  template Status reduce_l2_norm<TIn, TOut>(const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size)
+INSTANTIATE_REDUCE_L2_NORM(half, float);
+INSTANTIATE_REDUCE_L2_NORM(float, float);
+INSTANTIATE_REDUCE_L2_NORM(double, double);
+#undef INSTANTIATE_REDUCE_L2_NORM
+
+#define INSTANTIATE_REDUCE_MEAN(TIn, TOut) \
+  template Status reduce_mean<TIn, TOut>(const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size)
+INSTANTIATE_REDUCE_MEAN(half, float);
+INSTANTIATE_REDUCE_MEAN(float, float);
+INSTANTIATE_REDUCE_MEAN(double, double);
+#undef INSTANTIATE_REDUCE_MEAN
+
+namespace detail {
+template <typename TIn, typename TOut, typename TBuf>
+__global__ void reduce_matrix_rows_kernel(const TIn* input, TOut* output, int m, int n) {
+  constexpr int x_load_count_per_thread = 1;
+  constexpr int y_load_count_per_thread = 4;
+  const int t_count_x_in_grid = blockDim.x * gridDim.x;
+  const int t_count_y_in_grid = blockDim.y * gridDim.y;
+  const int x_grid_stride = t_count_x_in_grid * x_load_count_per_thread;
+  const int y_grid_stride = t_count_y_in_grid * y_load_count_per_thread;
+  const int tid_x_in_grid = threadIdx.x + blockDim.x * blockIdx.x;
+  const int tid_y_in_grid = threadIdx.y + blockDim.y * blockIdx.y;
+  const int tid_in_block = threadIdx.x + blockDim.x * threadIdx.y;
+
+  // Shape is blockDim.y-by-blockDim.x and element type is TBuf.
+  extern __shared__ unsigned char shared_memory_bytes[];
+  TBuf* shared_memory = reinterpret_cast<TBuf*>(shared_memory_bytes);
+
+  // to prevent int overflow in index calculation for input size m*n
+  const int64_t n_int64 = static_cast<int64_t>(n);
+
+  for (int col = tid_x_in_grid; col < n; col += x_grid_stride) {
+    shared_memory[tid_in_block] = TBuf(0.0f);
+    TBuf sum = TBuf(0.0f);
+    // This loops load multiple blockDim.y-by-blockDim.x sub-tensors from the input.
+    for (int row = tid_y_in_grid; row < m; row += y_grid_stride) {
+      // Thread-level reduction. Each thread loads y_load_count_per_thread values
+      // and aggregrate them.
+#pragma unroll(y_load_count_per_thread)
+      for (int row_inner = 0; row_inner < y_load_count_per_thread; ++row_inner) {
+        int row_final = row + row_inner * t_count_y_in_grid;
+        int col_final = col;
+        if (row_final < m && col_final < n) {
+          sum += TBuf(input[row_final * n_int64 + col_final]);
+        }
+      }
+    }
+    // Write thread-level reduction result into shared memory.
+    shared_memory[tid_in_block] = sum;
+
+    // Wait all threads to finish their thread-level reductions.
+    __syncthreads();
+
+// This loop conducts reduction on elements stored in shared memory.
+// Each block reduces blockDim.y-by-blockDim.x tensor to 1-by-blockDim.x tensor.
+#pragma unroll(4)
+    for (int stride = blockDim.y / 2; stride > 0; stride /= 2) {
+      if (threadIdx.y < stride) {
+        shared_memory[tid_in_block] += shared_memory[tid_in_block + stride * blockDim.x];
+      }
+      __syncthreads();
+    }
+
+    if (threadIdx.y == 0) {
+      atomic_add(output + col, TOut(shared_memory[threadIdx.x]));
+    }
+  }
+}
+
+template <typename TIn, typename TOut, typename TBuf>
+Status call_reduce_matrix_rows(const TIn* input, TOut* output, int m, int n, bool reset_initial_output) {
+  ORT_ENFORCE(m >= 0 && n >= 0);
+
+  if (reset_initial_output) {
+    HIP_RETURN_IF_ERROR(hipMemsetAsync(output, 0, n * sizeof(TOut)));
+  }
+
+  constexpr int max_num_threads_in_block = 512;
+  constexpr int max_num_blocks_in_grid = 512;
+  constexpr int load_count_per_thread = 4;
+
+  const int block_x_dim = least_pow2_bound(std::max(1, std::min(n, GPU_WARP_SIZE)));
+  const int block_y_dim = least_pow2_bound(std::max(1, std::min(max_num_threads_in_block / block_x_dim, m / load_count_per_thread)));
+  const int grid_x_dim = std::max(1, std::min(n / block_x_dim, max_num_blocks_in_grid));
+  const int grid_y_dim = std::max(1, std::min(max_num_blocks_in_grid / grid_x_dim, m / block_y_dim / 4));
+
+  const dim3 grid(grid_x_dim, grid_y_dim, 1);
+  const dim3 block(block_x_dim, block_y_dim, 1);
+
+  hipLaunchKernelGGL(HIP_KERNEL_NAME(reduce_matrix_rows_kernel<TIn, TOut, TBuf>),
+      grid, block, block.y * block.x * sizeof(TBuf), 0,
+      input, output, m, n);
+
+  return Status::OK();
+}
+}  // namespace detail
+
+template <typename TIn, typename TOut>
+Status reduce_matrix_rows(const TIn* input, TOut* output, int m, int n, bool reset_initial_output) {
+  using TBuf = AccumulationType_t<TIn>;
+  return detail::call_reduce_matrix_rows<TIn, TOut, TBuf>(input, output, m, n, reset_initial_output);
+}
+
+#define INSTANTIATE_REDUCE_MATRIX_ROWS(T) \
+  template Status reduce_matrix_rows<T, T>(const T* input, T* output, int m, int n, bool reset_initial_output)
+INSTANTIATE_REDUCE_MATRIX_ROWS(half);
+INSTANTIATE_REDUCE_MATRIX_ROWS(float);
+INSTANTIATE_REDUCE_MATRIX_ROWS(double);
+#undef INSTANTIATE_REDUCE_MATRIX_ROWS
+
+template <typename TIn, typename TOut>
+Status reduce_matrix_columns(const TIn* input, TOut* output, int m, int n, void* buffer, size_t buffer_size) {
+  return detail::call_reduce_matrix_columns<TIn, TOut, Identity2, Identity2, false>(
+      input, output, m, n, buffer, buffer_size);
+}
+
+#define INSTANTIATE_REDUCE_MATRIX_COLUMNS(T) \
+  template Status reduce_matrix_columns<T, T>(const T* input, T* output, int m, int n, void* buffer, size_t buffer_size)
+INSTANTIATE_REDUCE_MATRIX_COLUMNS(half);
+INSTANTIATE_REDUCE_MATRIX_COLUMNS(float);
+INSTANTIATE_REDUCE_MATRIX_COLUMNS(double);
+#undef INSTANTIATE_REDUCE_MATRIX_COLUMNS
+
+
+
+
+
+
+
+//
+// TODO: DELETE EVERYTHING BELOW
+//
+
 std::pair<int, int> compute_block_size(int size) {
   int x = GPU_WARP_SIZE;
   int y = std::min(NUM_WARPS_PER_BLOCK, std::max(1, size / (NUM_ELEMENTS_PER_THREAD * GPU_WARP_SIZE)));
@@ -37,7 +516,7 @@ int compute_reduction_buffer_size(int element_size, int size) {
 
 template<typename TIn, typename TOut, typename TOp, typename TFinalOp, bool DivideResultBySize>
 __global__ void reduce_all_kernel(const int size, const TIn * data, TOut* output, TOut* buffer) {
-  HIP_DYNAMIC_SHARED( unsigned char, shared_memory_)
+  extern __shared__ unsigned char shared_memory_[];
   TOut* shared_memory = reinterpret_cast<TOut*>(shared_memory_);
   // Thread-level indexes:
   // Linear index of thread in block.

onnxruntime/core/providers/rocm/reduction/reduction_functions.h

@@ -2,11 +2,115 @@
 // Licensed under the MIT License.
 
 #pragma once
+
 #include "core/providers/rocm/rocm_common.h"
+#include "core/providers/rocm/shared_inc/accumulation_type.h"
 
 namespace onnxruntime {
 namespace rocm {
 
+namespace detail {
+size_t compute_reduce_matrix_columns_intermediate_buffer_size(
+    int element_size, int num_rows, int num_cols);
+}  // namespace detail
+
+/**
+ * Computes the size in bytes of the intermediate buffer needed by reduce_matrix_columns().
+ * @tparam TIn The input data type.
+ * @param m The number of matrix rows.
+ * @param n The number of matrix columns.
+ * @return The size of the intermediate buffer.
+ */
+template <typename TIn>
+size_t compute_reduce_matrix_columns_buffer_size(int m, int n) {
+  using TBuf = AccumulationType_t<TIn>;
+  return detail::compute_reduce_matrix_columns_intermediate_buffer_size(
+      sizeof(TBuf), m, n);
+}
+
+/**
+ * Computes the size in bytes of the intermediate buffer needed by the reduce_x() functions.
+ * @tparam TIn The input data type.
+ * @param size The number of elements.
+ * @return The size of the intermediate buffer.
+ */
+template <typename TIn>
+size_t compute_reduction_buffer_size(int size) {
+  using TBuf = AccumulationType_t<TIn>;
+  return detail::compute_reduce_matrix_columns_intermediate_buffer_size(
+      sizeof(TBuf), 1, size);
+}
+
+/** Computes the sum of the given elements. */
+template <typename TIn, typename TOut>
+Status reduce_sum(const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size);
+
+/** Computes the sum of the squares of the given elements. */
+template <typename TIn, typename TOut>
+Status reduce_square_sum(const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size);
+
+/** Computes the L2 norm of the given elements. */
+template <typename TIn, typename TOut>
+Status reduce_l2_norm(const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size);
+
+/** Computes the mean of the given elements. */
+template <typename TIn, typename TOut>
+Status reduce_mean(const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size);
+
+enum class ApplicableMatrixReduction {
+  // can use reduce_matrix_rows()
+  Rows,
+  // can use reduce_matrix_columns()
+  Columns,
+  // no optimized matrix reduction function applies
+  None,
+};
+
+/**
+ * Determines whether a cuDNN reduction can be computed by an optimized matrix reduction function.
+ * @param miopen_reduce_op The MIOpen reduction op type.
+ * @param dims The input dimensions.
+ * @param axes The reduction axes.
+ * @param[out] m If matrix reduction is possible, the number of matrix rows to use.
+ * @param[out] n If matrix reduction is possible, the number of matrix columns to use.
+ * @return The type of matrix reduction that can be done.
+ */
+ApplicableMatrixReduction get_applicable_matrix_reduction(
+    const miopenReduceTensorOp_t miopen_reduce_op,
+    const std::vector<int64_t>& dims, const std::vector<int64_t>& axes,
+    int& m, int& n);
+
+/**
+ * Reduces the rows in a row-major matrix to a single row containing the sum of each column.
+ * @param input The input data.
+ * @param output The output data.
+ * @param m The number of matrix rows.
+ * @param n The number of matrix columns.
+ * @param reset_initial_output Whether to reset (i.e., zero) the output values first.
+ */
+template <typename TIn, typename TOut>
+Status reduce_matrix_rows(const TIn* input, TOut* output, int m, int n, bool reset_initial_output /* TODO: = true*/);
+
+/**
+ * Reduces the columns in a row-major matrix to a single column containing the sum of each row.
+ * @param input The input data.
+ * @param output The output data.
+ * @param m The number of matrix rows.
+ * @param n The number of matrix columns.
+ * @param buffer The intermediate buffer.
+ * @param buffer_size The size of the intermediate buffer in bytes.
+ */
+template <typename TIn, typename TOut>
+Status reduce_matrix_columns(const TIn* input, TOut* output, int m, int n, void* buffer, size_t buffer_size);
+
+
+
+
+
+//
+// TODO: DELETE EVERYTHING BELOW
+//
+
 int compute_reduction_buffer_size(int element_size, int size);
 
 template <typename TIn, typename TOut>

onnxruntime/core/providers/rocm/reduction/reduction_ops.cc

@@ -1,14 +1,16 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#include "reduction_ops.h"
+#include "core/providers/rocm/reduction/reduction_ops.h"
+
+#include "core/framework/data_types_internal.h"
+#include "core/framework/op_kernel_context_internal.h"
 #include "core/providers/common.h"
+#include "core/providers/cpu/tensor/utils.h"
 #include "core/providers/rocm/miopen_common.h"
-#include "core/providers/rocm/math/unary_elementwise_ops_impl.h"
 #include "core/providers/rocm/math/binary_elementwise_ops_impl.h"
 #include "core/providers/rocm/math/binary_elementwise_ops.h"
-#include "core/providers/cpu/tensor/utils.h"
-#include "core/framework/op_kernel_context_internal.h"
+#include "core/providers/rocm/math/unary_elementwise_ops_impl.h"
 
 using namespace onnxruntime::common;
 namespace onnxruntime {
@@ -76,7 +78,7 @@ namespace rocm {
       KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
       name<T>);
 
-// CUDA ArgMax/ArgMin doesn't have OpSet12 implementation (with select_last_index attr), keep it in OpSet11 for now.
+// ROCM ArgMax/ArgMin doesn't have OpSet12 implementation (with select_last_index attr), keep it in OpSet11 for now.
 #define REGISTER_KERNEL_TYPED_11(name, T)                                       \
   ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX(                                      \
       name,                                                                     \
@@ -95,6 +97,35 @@ namespace rocm {
       KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
       name<T>);
 
+// Register with the latest version 13
+#define REGISTER_KERNEL_TYPED_13(name, T)                                       \
+  ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX(                                      \
+      name,                                                                     \
+      kOnnxDomain,                                                              \
+      1, 10,                                                                    \
+      T,                                                                        \
+      kRocmExecutionProvider,                                                   \
+      KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
+      name<T>);                                                                 \
+  ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX(                                      \
+      name,                                                                     \
+      kOnnxDomain,                                                              \
+      11, 12,                                                                   \
+      T,                                                                        \
+      kRocmExecutionProvider,                                                   \
+      KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
+      name<T>);                                                                 \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                                \
+      name,                                                                     \
+      kOnnxDomain,                                                              \
+      13,                                                                       \
+      T,                                                                        \
+      kRocmExecutionProvider,                                                   \
+      KernelDefBuilder()                                                        \
+          .InputMemoryType<OrtMemTypeCPUInput>(1)                               \
+          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()),               \
+      name<T>);
+
 static bool is_matrix_row_reduction(
     const miopenReduceTensorOp_t miopen_reduce_op,
     const int m,

onnxruntime/core/providers/rocm/reduction/reduction_ops.h

@@ -11,15 +11,33 @@ namespace onnxruntime {
 namespace rocm {
 
 enum miopenReduceTensorOp_t {
-  MIOPEN_REDUCE_TENSOR_MAX,
+  MIOPEN_REDUCE_TENSOR_ADD,
+  MIOPEN_REDUCE_TENSOR_MUL,
   MIOPEN_REDUCE_TENSOR_MIN,
+  MIOPEN_REDUCE_TENSOR_MAX,
+  MIOPEN_REDUCE_TENSOR_AVG,
   MIOPEN_REDUCE_TENSOR_NORM1,
   MIOPEN_REDUCE_TENSOR_NORM2,
-  MIOPEN_REDUCE_TENSOR_AVG,
-  MIOPEN_REDUCE_TENSOR_MUL,
-  MIOPEN_REDUCE_TENSOR_ADD
 };
 
+enum miopenReduceTensorIndices_t {
+  MIOPEN_REDUCE_TENSOR_NO_INDICES,
+  MIOPEN_REDUCE_TENSOR_FLATTENED_INDICES,
+};
+
+namespace ReductionOps {
+
+// Implementation that holds the core logic of reduction op processing
+// `input_shape_override` is the input shape for compute purposes (if provided)
+
+template <typename T, miopenReduceTensorIndices_t ReduceTensorIndices = MIOPEN_REDUCE_TENSOR_NO_INDICES>
+Tensor ReduceCompute(ROCMExecutionProvider& rocm_ep, miopenReduceTensorOp_t miopen_reduce_op, AllocatorPtr allocator,
+                     const Tensor& input, const std::vector<int64_t>& axes,
+                     bool keep_dims, bool calculate_log, bool calculate_sqt, bool log_sum_exp,
+                     bool fast_reduction, const TensorShape* input_shape_override = nullptr);
+
+}  // namespace ReductionOps
+
 // Holds some metadata that will be used during actual reduction op compute time
 struct PrepareReduceMetadata {
   int64_t input_count;
@@ -30,24 +48,15 @@ struct PrepareReduceMetadata {
   std::vector<int64_t> squeezed_output_dims;
   std::vector<int64_t> input_dims_miopen;
   std::vector<int64_t> output_dims_miopen;
+
+  //
+  // TODO: delete these fields
+  //
   int64_t rank;
   int64_t stride;
   bool contiguous_axes;
 };
 
-Status PrepareForReduce(const Tensor* X,
-                        bool keepdims,
-                        const std::vector<int64_t>& axes,
-                        PrepareReduceMetadata& prepare_reduce_metadata,
-                        const TensorShape* input_shape_override = nullptr);
-
-template <typename T>
-Status ReduceComputeCore(const Tensor& input, PrepareReduceMetadata& prepare_reduce_metadata,
-                         /*out*/ Tensor& output, miopenReduceTensorOp_t miopen_reduce_op,
-                         const std::vector<int64_t>& axes,
-                         bool calculate_log, bool calculate_sqt, bool log_sum_exp, bool fast_reduction,
-                         const TensorShape* input_shape_override = nullptr);
-
 template <bool allow_multi_axes>
 class ReduceKernel : public RocmKernel, public ReduceKernelBase<allow_multi_axes> {
  protected:
@@ -217,5 +226,26 @@ class ReduceLogSumExp final : public ReduceKernel<true> {
   }
 };
 
+
+Status PrepareForReduce(const Tensor* X,
+                        bool keepdims,
+                        const std::vector<int64_t>& axes,
+                        PrepareReduceMetadata& prepare_reduce_metadata,
+                        const TensorShape* input_shape_override = nullptr);
+
+template <typename T, miopenReduceTensorIndices_t ReduceTensorIndices>
+Status ReduceComputeCore(ROCMExecutionProvider& rocm_ep, const Tensor& input, PrepareReduceMetadata& prepare_reduce_metadata,
+                         /*out*/ Tensor& output, miopenReduceTensorOp_t miopen_reduce_op,
+                         const std::vector<int64_t>& axes,
+                         bool calculate_log, bool calculate_sqt, bool log_sum_exp, bool fast_reduction,
+                         const TensorShape* input_shape_override = nullptr);
+
+template <typename T>
+Status ReduceComputeCore(const Tensor& input, PrepareReduceMetadata& prepare_reduce_metadata,
+                         /*out*/ Tensor& output, miopenReduceTensorOp_t miopen_reduce_op,
+                         const std::vector<int64_t>& axes,
+                         bool calculate_log, bool calculate_sqt, bool log_sum_exp, bool fast_reduction,
+                         const TensorShape* input_shape_override = nullptr);
+
 }  // namespace rocm
 }  // namespace onnxruntime

onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh

@@ -18,6 +18,34 @@ __forceinline__ __host__ __device__ int least_pow2_bound(int value) {
   return static_cast<int>(++value_);
 }
 
+struct Square2 {
+  template <typename T>
+  __forceinline__ __device__ T operator()(const T& value) {
+    return value * value;
+  }
+};
+
+struct Sqrt2 {
+  template <typename T>
+  __forceinline__ __device__ T operator()(const T& value) {
+    return _Sqrt(value);
+  }
+};
+
+struct Identity2 {
+  template <typename T>
+  __forceinline__ __device__ T operator()(const T& value) {
+    return value;
+  }
+};
+
+
+
+//
+// TODO: DELETE EVERYTHING BELOW
+// TODO: RENAME STRUCTS ABOVE (no '2')
+//
+
 template <typename TAccumulated, typename TValue>
 struct Cast {
   __forceinline__ __device__ TAccumulated operator()(const TValue& value) {