diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index 8d5da0381e..69988862aa 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/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(). diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_functions.cu b/onnxruntime/core/providers/rocm/reduction/reduction_functions.cu index e6abc5e8da..cd55592330 100644 --- a/onnxruntime/core/providers/rocm/reduction/reduction_functions.cu +++ b/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 + #include #include -#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 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( + MAX_NUM_BLOCKS_IN_GRID_ROW, + std::max(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(num_rows) * grid_dim.x * element_size; + + buffer_size = round_up_to_aligned(buffer_size, alignof(int)); + buffer_size += static_cast(num_rows) * sizeof(int); + + // add padding to give us room to align + buffer_size += alignof(max_align_t) - 1; + + return buffer_size; +} + +template +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(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(num_rows) * grid_dim.x * sizeof(TBuf), alignof(int)); + const uintptr_t end_addr = + block_done_counts_buffer_addr + static_cast(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(block_reductions_addr); + block_done_counts_buffer = reinterpret_cast(block_done_counts_buffer_addr); + + return Status::OK(); +} + +template +__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(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 +__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( + num_cols, row_data, row_output, + row_block_reductions_buffer, row_block_done_counts_buffer); + } +} + +template +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; + + 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), + 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 +Status reduce_sum( + const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size) { + return detail::call_reduce_matrix_columns( + input, output, 1, size, buffer, buffer_size); +} + +template +Status reduce_square_sum( + const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size) { + return detail::call_reduce_matrix_columns( + input, output, 1, size, buffer, buffer_size); +} + +template +Status reduce_l2_norm( + const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size) { + return detail::call_reduce_matrix_columns( + input, output, 1, size, buffer, buffer_size); +} + +template +Status reduce_mean( + const TIn* input, TOut* output, int size, void* buffer, size_t buffer_size) { + return detail::call_reduce_matrix_columns( + input, output, 1, size, buffer, buffer_size); +} + +#define INSTANTIATE_REDUCE_SUM(TIn, TOut) \ + template Status reduce_sum(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(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(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(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 +__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(shared_memory_bytes); + + // to prevent int overflow in index calculation for input size m*n + const int64_t n_int64 = static_cast(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 +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), + grid, block, block.y * block.x * sizeof(TBuf), 0, + input, output, m, n); + + return Status::OK(); +} +} // namespace detail + +template +Status reduce_matrix_rows(const TIn* input, TOut* output, int m, int n, bool reset_initial_output) { + using TBuf = AccumulationType_t; + return detail::call_reduce_matrix_rows(input, output, m, n, reset_initial_output); +} + +#define INSTANTIATE_REDUCE_MATRIX_ROWS(T) \ + template Status reduce_matrix_rows(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 +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( + input, output, m, n, buffer, buffer_size); +} + +#define INSTANTIATE_REDUCE_MATRIX_COLUMNS(T) \ + template Status reduce_matrix_columns(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 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 __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(shared_memory_); // Thread-level indexes: // Linear index of thread in block. diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_functions.h b/onnxruntime/core/providers/rocm/reduction/reduction_functions.h index 5bc2df04c0..3b8c796b42 100644 --- a/onnxruntime/core/providers/rocm/reduction/reduction_functions.h +++ b/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 +size_t compute_reduce_matrix_columns_buffer_size(int m, int n) { + using TBuf = AccumulationType_t; + 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 +size_t compute_reduction_buffer_size(int size) { + using TBuf = AccumulationType_t; + return detail::compute_reduce_matrix_columns_intermediate_buffer_size( + sizeof(TBuf), 1, size); +} + +/** Computes the sum of the given elements. */ +template +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 +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 +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 +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& dims, const std::vector& 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 +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 +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 diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_ops.cc b/onnxruntime/core/providers/rocm/reduction/reduction_ops.cc index cb66a50764..b7f742eca7 100644 --- a/onnxruntime/core/providers/rocm/reduction/reduction_ops.cc +++ b/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()), \ name); -// 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()), \ name); +// 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()), \ + name); \ + ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ + name, \ + kOnnxDomain, \ + 11, 12, \ + T, \ + kRocmExecutionProvider, \ + KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + name); \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + name, \ + kOnnxDomain, \ + 13, \ + T, \ + kRocmExecutionProvider, \ + KernelDefBuilder() \ + .InputMemoryType(1) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + name); + static bool is_matrix_row_reduction( const miopenReduceTensorOp_t miopen_reduce_op, const int m, diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_ops.h b/onnxruntime/core/providers/rocm/reduction/reduction_ops.h index a2a953e25f..fcd97d3fbe 100644 --- a/onnxruntime/core/providers/rocm/reduction/reduction_ops.h +++ b/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 +Tensor ReduceCompute(ROCMExecutionProvider& rocm_ep, miopenReduceTensorOp_t miopen_reduce_op, AllocatorPtr allocator, + const Tensor& input, const std::vector& 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 squeezed_output_dims; std::vector input_dims_miopen; std::vector 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& axes, - PrepareReduceMetadata& prepare_reduce_metadata, - const TensorShape* input_shape_override = nullptr); - -template -Status ReduceComputeCore(const Tensor& input, PrepareReduceMetadata& prepare_reduce_metadata, - /*out*/ Tensor& output, miopenReduceTensorOp_t miopen_reduce_op, - const std::vector& axes, - bool calculate_log, bool calculate_sqt, bool log_sum_exp, bool fast_reduction, - const TensorShape* input_shape_override = nullptr); - template class ReduceKernel : public RocmKernel, public ReduceKernelBase { protected: @@ -217,5 +226,26 @@ class ReduceLogSumExp final : public ReduceKernel { } }; + +Status PrepareForReduce(const Tensor* X, + bool keepdims, + const std::vector& axes, + PrepareReduceMetadata& prepare_reduce_metadata, + const TensorShape* input_shape_override = nullptr); + +template +Status ReduceComputeCore(ROCMExecutionProvider& rocm_ep, const Tensor& input, PrepareReduceMetadata& prepare_reduce_metadata, + /*out*/ Tensor& output, miopenReduceTensorOp_t miopen_reduce_op, + const std::vector& axes, + bool calculate_log, bool calculate_sqt, bool log_sum_exp, bool fast_reduction, + const TensorShape* input_shape_override = nullptr); + +template +Status ReduceComputeCore(const Tensor& input, PrepareReduceMetadata& prepare_reduce_metadata, + /*out*/ Tensor& output, miopenReduceTensorOp_t miopen_reduce_op, + const std::vector& axes, + bool calculate_log, bool calculate_sqt, bool log_sum_exp, bool fast_reduction, + const TensorShape* input_shape_override = nullptr); + } // namespace rocm } // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh b/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh index 808756b295..9d9d210f3f 100644 --- a/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh +++ b/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh @@ -18,6 +18,34 @@ __forceinline__ __host__ __device__ int least_pow2_bound(int value) { return static_cast(++value_); } +struct Square2 { + template + __forceinline__ __device__ T operator()(const T& value) { + return value * value; + } +}; + +struct Sqrt2 { + template + __forceinline__ __device__ T operator()(const T& value) { + return _Sqrt(value); + } +}; + +struct Identity2 { + template + __forceinline__ __device__ T operator()(const T& value) { + return value; + } +}; + + + +// +// TODO: DELETE EVERYTHING BELOW +// TODO: RENAME STRUCTS ABOVE (no '2') +// + template struct Cast { __forceinline__ __device__ TAccumulated operator()(const TValue& value) {