From 86ac11af1a04bd5c60e7b28e725ea66a9dadcc7c Mon Sep 17 00:00:00 2001 From: Jesse Benson Date: Tue, 19 Jan 2021 12:51:25 -0800 Subject: [PATCH] Delete ROCM-specific reduction code that is identical to CUDA reduction code. --- .../rocm/reduction/reduction_functions.cu | 23 +- .../rocm/reduction/reduction_functions.h | 107 ------- .../providers/rocm/reduction/reduction_ops.h | 6 +- .../rocm/reduction/reduction_utils.cuh | 43 --- .../loss/softmax_cross_entropy_loss_impl.cc | 272 ---------------- .../rocm/loss/softmaxcrossentropy_impl.cc | 294 ------------------ .../rocm/reduction/reduction_all.cu | 116 ------- .../rocm/reduction/reduction_ops.cc | 10 +- tools/ci_build/amd_hipify.py | 5 - 9 files changed, 13 insertions(+), 863 deletions(-) delete mode 100644 onnxruntime/core/providers/rocm/reduction/reduction_functions.h delete mode 100644 onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh delete mode 100644 orttraining/orttraining/training_ops/rocm/loss/softmax_cross_entropy_loss_impl.cc delete mode 100644 orttraining/orttraining/training_ops/rocm/loss/softmaxcrossentropy_impl.cc delete mode 100644 orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_functions.cu b/onnxruntime/core/providers/rocm/reduction/reduction_functions.cu index d6e1ee4181..8089d8e5e4 100644 --- a/onnxruntime/core/providers/rocm/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/rocm/reduction/reduction_functions.cu @@ -13,13 +13,6 @@ #include "core/providers/rocm/shared_inc/rocm_utils.h" #include "core/providers/rocm/reduction/reduction_utils.cuh" -#define NUM_ELEMENTS_PER_THREAD 4 -#define NUM_WARPS_PER_BLOCK 8 -#define MAX_NUM_BLOCKS 256 - -#define ALL_ONE_MASK 0xFFFFFFFF -#define ONE_MASK 0x00000001 - namespace onnxruntime { namespace rocm { @@ -115,7 +108,7 @@ template (shared_memory_bytes); // Thread-level indices: // Linear index of thread in block. @@ -175,11 +168,7 @@ __device__ void reduce_all( } } -#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 @@ -312,9 +301,8 @@ Status call_reduce_matrix_columns( } 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); + hipLaunchKernelGGL(HIP_KERNEL_NAME(reduce_matrix_columns_kernel), dim3(grid_dim), dim3(block_dim), shared_mem_size, 0, + num_rows, num_cols, input, output, block_reductions_buffer, block_done_counts_buffer); return Status::OK(); } @@ -390,7 +378,7 @@ __global__ void reduce_matrix_rows_kernel(const TIn* input, TOut* output, int m, 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[]; + HIP_DYNAMIC_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 @@ -454,8 +442,7 @@ Status call_reduce_matrix_rows(const TIn* input, TOut* output, int m, int n, boo 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, + hipLaunchKernelGGL(HIP_KERNEL_NAME(reduce_matrix_rows_kernel), dim3(grid), dim3(block), block.y * block.x * sizeof(TBuf), 0, input, output, m, n); return Status::OK(); diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_functions.h b/onnxruntime/core/providers/rocm/reduction/reduction_functions.h deleted file mode 100644 index 2f677c3c4b..0000000000 --- a/onnxruntime/core/providers/rocm/reduction/reduction_functions.h +++ /dev/null @@ -1,107 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// 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 = 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); - -} // namespace rocm -} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_ops.h b/onnxruntime/core/providers/rocm/reduction/reduction_ops.h index 3cacb4367e..eb3c9d9e45 100644 --- a/onnxruntime/core/providers/rocm/reduction/reduction_ops.h +++ b/onnxruntime/core/providers/rocm/reduction/reduction_ops.h @@ -3,8 +3,8 @@ #pragma once #include "core/common/optional.h" -#include "core/providers/cpu/reduction/reduction_ops.h" #include "core/providers/rocm/rocm_kernel.h" +#include "core/providers/cpu/reduction/reduction_ops.h" #include "core/providers/rocm/reduction/reduction_functions.h" namespace onnxruntime { @@ -61,7 +61,7 @@ class ReduceKernel : public RocmKernel, public ReduceKernelBase Status ComputeImplEx(OpKernelContext* ctx, miopenReduceTensorOp_t miopen_reduce_op) const; - template + template Status ReduceKernelShared( const T* X, const TensorShape& input_shape, @@ -269,4 +269,4 @@ class MiopenReduceDescriptor final { }; } // namespace rocm -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh b/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh deleted file mode 100644 index 4ac5710d10..0000000000 --- a/onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh +++ /dev/null @@ -1,43 +0,0 @@ -// 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_); -} - -struct Square { - template - __forceinline__ __device__ T operator()(const T& value) { - return value * value; - } -}; - -struct Sqrt { - template - __forceinline__ __device__ T operator()(const T& value) { - return _Sqrt(value); - } -}; - -struct Identity { - template - __forceinline__ __device__ T operator()(const T& value) { - return value; - } -}; - -} // namespace rocm -} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/rocm/loss/softmax_cross_entropy_loss_impl.cc b/orttraining/orttraining/training_ops/rocm/loss/softmax_cross_entropy_loss_impl.cc deleted file mode 100644 index d4b62fc0f3..0000000000 --- a/orttraining/orttraining/training_ops/rocm/loss/softmax_cross_entropy_loss_impl.cc +++ /dev/null @@ -1,272 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "core/providers/rocm/math/softmax.h" -#include "core/providers/rocm/reduction/reduction_functions.h" -#include "core/providers/rocm/tensor/transpose.h" -#include "core/providers/cpu/controlflow/scan_utils.h" -#include "orttraining/training_ops/cpu/loss/softmax_cross_entropy_loss.h" -#include "orttraining/training_ops/rocm/loss/softmax_cross_entropy_loss_impl.h" - -namespace onnxruntime { -namespace rocm { - -#define REGISTER_KERNEL_VERSIONED_TYPED_TWO_TYPES(Class, T, Tin, domain, startver, endver) \ - ONNX_OPERATOR_VERSIONED_TWO_TYPED_KERNEL_EX( \ - Class, \ - domain, \ - startver, endver, \ - T, Tin, \ - kRocmExecutionProvider, \ - KernelDefBuilder() \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("Tin", DataTypeImpl::GetTensorType()), \ - Class); - -#define REGISTER_KERNEL_TYPED_TWO_TYPES(Class, T, Tin, domain, version) \ - ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \ - Class, \ - domain, \ - version, \ - T, Tin, \ - kRocmExecutionProvider, \ - KernelDefBuilder() \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("Tin", DataTypeImpl::GetTensorType()), \ - Class); - -template -Status SoftmaxCrossEntropyLoss::ComputeInternal(OpKernelContext* ctx) const { - const Tensor& logit = *ctx->Input(0); - const Tensor& label = *ctx->Input(1); - const TensorShape logit_shape{logit.Shape()}; - const TensorShape label_shape{label.Shape()}; - onnxruntime::contrib::VerifyLogitWeightAndLabelShape(logit_shape, label_shape, - OpKernel::Node().InputDefs().size() == 3 ? &(*(ctx->Input(2))).Shape() : nullptr); - - // N_D = N * D1 * D2...D*K - int64_t N_D; - int64_t C; - onnxruntime::contrib::GetNDCFromLogitAndLabelShape(logit_shape, label_shape, N_D, C); - const TensorShape logit_reshape({N_D, C}); - const TensorShape label_reshape({N_D}); - Tensor* total_loss = ctx->Output(0, reduction_ == ReductionType::NONE ? TensorShape(label.Shape()) : TensorShape({})); - T* total_loss_data = total_loss->template MutableData(); - T* tmp_loss_sample_buffer = nullptr; - IAllocatorUniquePtr tmp_loss_sample; - if (reduction_ == ReductionType::NONE) { - tmp_loss_sample_buffer = total_loss_data; - } else { - tmp_loss_sample = GetScratchBuffer(N_D); - tmp_loss_sample_buffer = tmp_loss_sample.get(); - } - - const T* logit_data = logit.template Data(); - const Tin* label_data = label.template Data(); - - T* log_prob_data = nullptr; - Tensor* log_prob = nullptr; - IAllocatorUniquePtr log_prob_scratch_buffer; - if (ctx->OutputCount() > 1) { - log_prob = ctx->Output(1, logit_shape); - log_prob_data = log_prob->template MutableData(); - } else { - log_prob_scratch_buffer = GetScratchBuffer(logit_shape.Size()); - log_prob_data = log_prob_scratch_buffer.get(); - } - - OrtValue transpose_output; - Tensor transpose_tensor; - std::vector new_shape; - std::vector permutations; - AllocatorPtr alloc; - const OpKernelInfo& info = OpKernel::Info(); - - // Transpose logit from [N, C, D1, D2 .. Dk] to [N, D1, D2...Dk, C] - if (logit_shape.NumDimensions() > 2) { - ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc)); - onnxruntime::contrib::GetPermutationAndShape(true, logit_shape, new_shape, permutations); - transpose_output = scan::detail::AllocateTensorInMLValue(logit.DataType(), new_shape, alloc); - ORT_RETURN_IF_ERROR(rocm::Transpose::DoTranspose(rocm::Transpose(info), permutations, logit, *transpose_output.GetMutable())); - logit_data = (*transpose_output.GetMutable()).template Data(); - } - - // calculate logsoftmax - auto status = SoftMaxComputeHelper(logit_data, - logit_reshape, - log_prob_data, - MiopenHandle(), - 1); - ORT_RETURN_IF_ERROR(status); - - const T* weight_data = nullptr; - if (OpKernel::Node().InputDefs().size() == 3) { - const Tensor& weight = *ctx->Input(2); - weight_data = weight.template Data(); - } - - IAllocatorUniquePtr weight_data_nd = GetScratchBuffer(N_D); - T* weight_data_nd_data = weight_data_nd.get(); - HIP_RETURN_IF_ERROR(hipMemsetAsync(weight_data_nd_data, 0, N_D * sizeof(T))); - ComputeWeightsSoftmaxCrossEntropyImpl(label_data, weight_data, N_D, C, ignore_index_, weight_data_nd_data); - - auto normalize_factor_data = GetScratchBuffer(1); - if (reduction_ == ReductionType::MEAN) { - // Compute buffer size in byte for reduction APIs. - const auto buffer_size = - compute_reduction_buffer_size(static_cast(N_D)); - // Allocate reduction buffer whose size is buffer_size bytes. - IAllocatorUniquePtr reduction_buffer = GetScratchBuffer( - buffer_size); - ORT_RETURN_IF_ERROR(reduce_sum( - weight_data_nd_data, - normalize_factor_data.get(), - static_cast(N_D), - reduction_buffer.get(), - buffer_size)); - } else { - const T normalize_factor = static_cast(1); - HIP_RETURN_IF_ERROR(hipMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(T), hipMemcpyHostToDevice)); - } - - SoftmaxCrossEntropyLossImpl(log_prob_data, - label_data, - weight_data_nd_data, - normalize_factor_data.get(), - N_D, - C, - ignore_index_, - tmp_loss_sample_buffer); - - // Transpose log probability from [N, D1, D2...Dk, C] to [N, C, D1, D2 .. Dk]. - if (logit_shape.NumDimensions() > 2 && log_prob != nullptr) { - TensorShape log_prob_shape = new_shape; - new_shape.clear(); - permutations.clear(); - onnxruntime::contrib::GetPermutationAndShape(false, log_prob_shape, new_shape, permutations); - auto* transposed_data = (*transpose_output.GetMutable()).template MutableData(); - transpose_output.GetMutable()->Reshape(log_prob->Shape()); - log_prob->Reshape(log_prob_shape); - ORT_RETURN_IF_ERROR(rocm::Transpose::DoTranspose(rocm::Transpose(info), permutations, *log_prob, *transpose_output.GetMutable())); - HIP_RETURN_IF_ERROR(hipMemcpyAsync(log_prob_data, transposed_data, sizeof(T) * logit_shape.Size(), hipMemcpyDeviceToDevice)); - log_prob->Reshape(new_shape); - } - - if (reduction_ != ReductionType::NONE) { - // ReduceSum on loss_per_sample - std::vector output_dims(1, 1); - ReduceKernelShared( - tmp_loss_sample_buffer, - label_reshape, - total_loss_data, - TensorShape({}), - MIOPEN_REDUCE_TENSOR_ADD, - output_dims); - } - - return Status::OK(); -} - -template -Status SoftmaxCrossEntropyLossGrad::ComputeInternal(OpKernelContext* ctx) const { - const Tensor& dY = *ctx->Input(0); - const Tensor& log_prob = *ctx->Input(1); - const Tensor& label = *ctx->Input(2); - const TensorShape probability_shape{log_prob.Shape()}; - const TensorShape label_shape{label.Shape()}; - onnxruntime::contrib::VerifyLogitWeightAndLabelShape(probability_shape, label_shape, - OpKernel::Node().InputDefs().size() == 4 ? &(*(ctx->Input(3))).Shape() : nullptr); - - // N_D = N * D1 * D2...D*K - int64_t N_D; - int64_t C; - onnxruntime::contrib::GetNDCFromLogitAndLabelShape(probability_shape, label_shape, N_D, C); - Tensor* d_logit = ctx->Output(0, probability_shape); - const T* dY_data = dY.template Data(); - const T* log_prob_data = log_prob.template Data(); - const Tin* label_data = label.template Data(); - T* d_logit_data = d_logit->template MutableData(); - const T* weight_data = nullptr; - OrtValue transpose_output; - std::vector new_shape; - std::vector permutations; - AllocatorPtr alloc; - const OpKernelInfo& info = OpKernel::Info(); - - // Transpose logit from [N, C, D1, D2 .. Dk] to [N, D1, D2...Dk, C] - if (probability_shape.NumDimensions() > 2) { - ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc)); - onnxruntime::contrib::GetPermutationAndShape(true, probability_shape, new_shape, permutations); - transpose_output = scan::detail::AllocateTensorInMLValue(log_prob.DataType(), new_shape, alloc); - ORT_RETURN_IF_ERROR(rocm::Transpose::DoTranspose(rocm::Transpose(info), permutations, log_prob, *transpose_output.GetMutable())); - log_prob_data = (*transpose_output.GetMutable()).template Data(); - } - - if (OpKernel::Node().InputDefs().size() == 4) { - const Tensor& weight = *ctx->Input(3); - weight_data = weight.template Data(); - } - - IAllocatorUniquePtr weight_data_nd = GetScratchBuffer(N_D); - T* weight_data_nd_data = weight_data_nd.get(); - HIP_RETURN_IF_ERROR(hipMemsetAsync(weight_data_nd_data, 0, N_D * sizeof(T))); - ComputeWeightsSoftmaxCrossEntropyImpl(label_data, weight_data, N_D, C, ignore_index_, weight_data_nd_data); - auto normalize_factor_data = GetScratchBuffer(1); - if (reduction_ == ReductionType::MEAN) { - // Compute buffer size in byte for reduction APIs. - const auto buffer_size = - compute_reduction_buffer_size(static_cast(N_D)); - // Allocate reduction buffer whose size is buffer_size bytes. - IAllocatorUniquePtr reduction_buffer = GetScratchBuffer( - buffer_size); - ORT_RETURN_IF_ERROR(reduce_sum( - weight_data_nd_data, - normalize_factor_data.get(), - static_cast(N_D), - reduction_buffer.get(), - buffer_size)); - } else { - const T normalize_factor = static_cast(1); - HIP_RETURN_IF_ERROR(hipMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(T), hipMemcpyHostToDevice)); - } - - SoftmaxCrossEntropyLossGradImpl(dY_data, - log_prob_data, - label_data, - weight_data_nd_data, - normalize_factor_data.get(), - N_D, - C, - ReductionType::NONE == reduction_, - d_logit_data); - - // Transpose logit from [N, D1, D2...Dk, C] to [N, C, D1, D2 .. Dk] - if (probability_shape.NumDimensions() > 2) { - TensorShape logit_shape = new_shape; - new_shape.clear(); - permutations.clear(); - onnxruntime::contrib::GetPermutationAndShape(false, logit_shape, new_shape, permutations); - transpose_output.GetMutable()->Reshape(d_logit->Shape()); - d_logit->Reshape(logit_shape); - ORT_RETURN_IF_ERROR(rocm::Transpose::DoTranspose(rocm::Transpose(info), permutations, *d_logit, *transpose_output.GetMutable())); - auto* transposed_data = (*transpose_output.GetMutable()).template Data(); - HIP_RETURN_IF_ERROR(hipMemcpyAsync(d_logit_data, transposed_data, sizeof(T) * probability_shape.Size(), hipMemcpyDeviceToDevice)); - d_logit->Reshape(new_shape); - } - - return Status::OK(); -} - -#define SPECIALIZED_VERSIONED_COMPUTE_SPARSE(Class, T, Tin, domain, startver, endvar) \ - REGISTER_KERNEL_VERSIONED_TYPED_TWO_TYPES(Class, T, Tin, domain, startver, endvar) - -#define SPECIALIZED_COMPUTE_SPARSE(Class, T, Tin, domain, version) \ - REGISTER_KERNEL_TYPED_TWO_TYPES(Class, T, Tin, domain, version) \ - template Status Class::ComputeInternal(OpKernelContext* ctx) const; - -SPECIALIZED_VERSIONED_COMPUTE_SPARSE(SoftmaxCrossEntropyLoss, float, int64_t, kOnnxDomain, 12, 12) -SPECIALIZED_COMPUTE_SPARSE(SoftmaxCrossEntropyLoss, float, int64_t, kOnnxDomain, 13) -SPECIALIZED_COMPUTE_SPARSE(SoftmaxCrossEntropyLossGrad, float, int64_t, kMSDomain, 1) - -} // namespace rocm -} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/rocm/loss/softmaxcrossentropy_impl.cc b/orttraining/orttraining/training_ops/rocm/loss/softmaxcrossentropy_impl.cc deleted file mode 100644 index 7cc017f360..0000000000 --- a/orttraining/orttraining/training_ops/rocm/loss/softmaxcrossentropy_impl.cc +++ /dev/null @@ -1,294 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#include "core/providers/rocm/reduction/reduction_functions.h" -#include "core/providers/rocm/math/softmax.h" -#include "orttraining/training_ops/rocm/loss/softmaxcrossentropy_impl.h" - -namespace onnxruntime { -namespace rocm { -#define REGISTER_KERNEL_TYPED(Class, T, domain, version) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - Class, \ - domain, \ - version, \ - T, \ - kRocmExecutionProvider, \ - KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), \ - Class); - -#define REGISTER_KERNEL_TYPED_TWO_TYPES(Class, T, Tin, domain, version) \ - ONNX_OPERATOR_TWO_TYPED_KERNEL_EX( \ - Class, \ - domain, \ - version, \ - T, Tin, \ - kRocmExecutionProvider, \ - KernelDefBuilder() \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("Tin", DataTypeImpl::GetTensorType()), \ - Class); - -template -Status SoftmaxCrossEntropy::ComputeInternal(OpKernelContext* ctx) const { - const Tensor& logit = *ctx->Input(0); - const Tensor& label = *ctx->Input(1); - - const TensorShape logit_shape{logit.Shape()}; - const TensorShape label_shape{label.Shape()}; - ORT_ENFORCE(label_shape == logit_shape, "The shape in logits and labels is not identical"); - - int64_t N = logit_shape.SizeToDimension(logit_shape.NumDimensions() - 1); - int64_t D = logit_shape[logit_shape.NumDimensions() - 1]; - const TensorShape logit_reshape({N, D}); - - Tensor* log_prob = ctx->Output(1, logit_shape); - - const T* logit_data = logit.template Data(); - const T* label_data = label.template Data(); - T* log_prob_data = log_prob->template MutableData(); - - // calculate logsoftmax - auto status = SoftMaxComputeHelper(logit_data, - logit_reshape, - log_prob_data, - MiopenHandle(), - 1 /*axis default*/); - ORT_RETURN_IF_ERROR(status); - - size_t normalize_factor = N; - if (reduction_ == ReductionType::SUM) { - normalize_factor = static_cast(1); - } - - // calculate (label * log(softmax)) for each element - IAllocatorUniquePtr temp_X = GetScratchBuffer(N * D); - SoftMaxCrossEntropyImpl( - log_prob_data, // logsoftmax result - label_data, // label - normalize_factor, // normalize_factor - temp_X.get(), // -(label * log(softmax)) - N * D); - - std::vector output_dims(2, 1); - Tensor* Y = ctx->Output(0, TensorShape({})); - // Sum((label * log(softmax)) using Reduction - return ReduceKernelShared( - temp_X.get(), - logit_reshape, - Y->template MutableData(), - TensorShape({}), - MIOPEN_REDUCE_TENSOR_ADD, - output_dims); -} - -template -Status SoftmaxCrossEntropyGrad::ComputeInternal(OpKernelContext* ctx) const { - const Tensor& dY = *ctx->Input(0); - const Tensor& log_prob = *ctx->Input(1); - const Tensor& label = *ctx->Input(2); - - const TensorShape probability_shape{log_prob.Shape()}; - const TensorShape label_shape{label.Shape()}; - ORT_ENFORCE(label_shape == probability_shape, "The shape in probability and label is not identical"); - - int64_t N = probability_shape.SizeToDimension(probability_shape.NumDimensions() - 1); - int64_t ND = probability_shape.Size(); - - Tensor* d_logits = ctx->Output(0, probability_shape); - - const T* dY_data = dY.template Data(); - const T* log_prob_data = log_prob.template Data(); - const T* label_data = label.template Data(); - - size_t normalize_factor = N; - if (reduction_ == ReductionType::SUM) { - normalize_factor = static_cast(1); - } - - T* d_logits_data = d_logits->template MutableData(); - - SoftMaxCrossEntropyGradImpl( - dY_data, // Dy - log_prob_data, // log(pi) - label_data, // Label - normalize_factor, // normalize_factor - d_logits_data, // gradient - ND); - - return Status::OK(); -} - -template -Status SparseSoftmaxCrossEntropy::ComputeInternal(OpKernelContext* ctx) const { - const Tensor& logit = *ctx->Input(0); - const Tensor& label = *ctx->Input(1); - - const TensorShape logit_shape{logit.Shape()}; - const TensorShape label_shape{label.Shape()}; - ORT_ENFORCE(logit_shape.NumDimensions() == label_shape.NumDimensions() + 1, - "logits_shape must be (1 + label_shape)"); - for (size_t i = 0; i < label_shape.NumDimensions(); i++) { - ORT_ENFORCE(label_shape[i] == logit_shape[i], "The shape in logits and labels does not match"); - } - - int64_t N = label_shape.Size(); - int64_t D = logit_shape[logit_shape.NumDimensions() - 1]; - const TensorShape logit_reshape({N, D}); - const TensorShape label_reshape({N}); - - IAllocatorUniquePtr tmp_loss_sample = GetScratchBuffer(N); - Tensor* total_loss = ctx->Output(0, TensorShape({})); - Tensor* log_prob = ctx->Output(1, logit_shape); - - const T* logit_data = logit.template Data(); - const Tin* label_data = label.template Data(); - T* total_loss_data = total_loss->template MutableData(); - T* log_prob_data = log_prob->template MutableData(); - - // calculate logsoftmax - auto status = SoftMaxComputeHelper(logit_data, - logit_reshape, - log_prob_data, - MiopenHandle(), - 1 /*axis default*/); - ORT_RETURN_IF_ERROR(status); - - // calculate (label * log(softmax)) for each sample - const T* weight_data = nullptr; - if (OpKernel::Node().InputDefs().size() == 3) { - const Tensor& weight = *ctx->Input(2); - const TensorShape weight_shape{weight.Shape()}; - ORT_ENFORCE(weight_shape == label_shape, "The shape in weights and labels is different"); - weight_data = weight.template Data(); - } - - auto normalize_factor_data = GetScratchBuffer(1); - if (reduction_ == ReductionType::SUM) { - const T normalize_factor = static_cast(1); - hipMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(T), hipMemcpyHostToDevice); - } else if (reduction_ == ReductionType::MEAN) { - if (weight_data == nullptr) { - const T normalize_factor = static_cast(N); - hipMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(T), hipMemcpyHostToDevice); - } else { - // Compute buffer size in byte for reduction APIs. - const auto buffer_size = - compute_reduction_buffer_size(static_cast(N)); - // Allocate reduction buffer whose size is buffer_size bytes. - IAllocatorUniquePtr reduction_buffer = GetScratchBuffer( - buffer_size); - ORT_RETURN_IF_ERROR(reduce_sum( - weight_data, - normalize_factor_data.get(), - static_cast(N), - reduction_buffer.get(), - buffer_size)); - } - } - - SparseSoftmaxCrossEntropyImpl(log_prob_data, - label_data, - weight_data, - normalize_factor_data.get(), - tmp_loss_sample.get(), - N, - D); - - // ReduceSum on loss_per_sample - std::vector output_dims(1, 1); - return ReduceKernelShared( - tmp_loss_sample.get(), - label_reshape, - total_loss_data, - TensorShape({}), - MIOPEN_REDUCE_TENSOR_ADD, - output_dims); -} - -template -Status SparseSoftmaxCrossEntropyGrad::ComputeInternal(OpKernelContext* ctx) const { - const Tensor& dY = *ctx->Input(0); - const Tensor& log_prob = *ctx->Input(1); - const Tensor& label = *ctx->Input(2); - - const TensorShape probability_shape{log_prob.Shape()}; - const TensorShape label_shape{label.Shape()}; - ORT_ENFORCE(probability_shape.NumDimensions() == label_shape.NumDimensions() + 1, - "probability_shape must be (1 + label_shape)"); - for (size_t i = 0; i < label_shape.NumDimensions(); i++) { - ORT_ENFORCE(label_shape[i] == probability_shape[i], "The shape in probability and labels does not match"); - } - - int64_t N = label_shape.Size(); - int64_t D = probability_shape[probability_shape.NumDimensions() - 1]; - - Tensor* d_logit = ctx->Output(0, probability_shape); - - const T* dY_data = dY.template Data(); - const T* log_prob_data = log_prob.template Data(); - const Tin* label_data = label.template Data(); - T* d_logit_data = d_logit->template MutableData(); - - const T* weight_data = nullptr; - if (OpKernel::Node().InputDefs().size() == 4) { - const Tensor& weight = *ctx->Input(3); - const TensorShape weight_shape{weight.Shape()}; - ORT_ENFORCE(weight_shape == label_shape, "The shape in weights and labels is different"); - weight_data = weight.template Data(); - } - - auto normalize_factor_data = GetScratchBuffer(1); - if (reduction_ == ReductionType::SUM) { - const T normalize_factor = static_cast(1); - hipMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(T), hipMemcpyHostToDevice); - } else if (reduction_ == ReductionType::MEAN) { - if (weight_data == nullptr) { - const T normalize_factor = static_cast(N); - hipMemcpyAsync(normalize_factor_data.get(), &normalize_factor, sizeof(T), hipMemcpyHostToDevice); - } else { - // Compute buffer size in byte for reduction APIs. - const auto buffer_size = - compute_reduction_buffer_size(static_cast(N)); - // Allocate reduction buffer whose size is buffer_size bytes. - IAllocatorUniquePtr reduction_buffer = GetScratchBuffer( - buffer_size); - ORT_RETURN_IF_ERROR(reduce_sum( - weight_data, - normalize_factor_data.get(), - static_cast(N), - reduction_buffer.get(), - buffer_size)); - } - } - - SparseSoftmaxCrossEntropyGradImpl(dY_data, - log_prob_data, - label_data, - weight_data, - normalize_factor_data.get(), - d_logit_data, - N, - D); - - return Status::OK(); -} - -#define SPECIALIZED_COMPUTE(Class, T, domain, version) \ - REGISTER_KERNEL_TYPED(Class, T, domain, version) \ - template Status Class::ComputeInternal(OpKernelContext* ctx) const; - -SPECIALIZED_COMPUTE(SoftmaxCrossEntropy, float, kMSDomain, 1) -SPECIALIZED_COMPUTE(SoftmaxCrossEntropyGrad, float, kMSDomain, 1) - -#define SPECIALIZED_COMPUTE_SPARSE(Class, T, Tin, domain, version) \ - REGISTER_KERNEL_TYPED_TWO_TYPES(Class, T, Tin, domain, version) \ - template Status Class::ComputeInternal(OpKernelContext* ctx) const; - -// SPECIALIZED_COMPUTE_SPARSE(SparseSoftmaxCrossEntropy, float, int32_t, kOnnxDomain, 9) -SPECIALIZED_COMPUTE_SPARSE(SparseSoftmaxCrossEntropy, float, int64_t, kOnnxDomain, 9) -// SPECIALIZED_COMPUTE_SPARSE(SparseSoftmaxCrossEntropyGrad, float, int32_t, kOnnxDomain, 9) -SPECIALIZED_COMPUTE_SPARSE(SparseSoftmaxCrossEntropyGrad, float, int64_t, kOnnxDomain, 9) - -} // namespace rocm -} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu b/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu deleted file mode 100644 index b6bd48bd1d..0000000000 --- a/orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu +++ /dev/null @@ -1,116 +0,0 @@ -// 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 ScalarSqrtKernel(Tin* input, Tout* output) { - *output = (Tout)_Sqrt(*input); -} - -template -void ScalarSqrt(Tin* input, Tout* 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 -__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; - 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. - extern __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()(TOut(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(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) { - using TBuf = AccumulationType_t; - MultiTensorReduce(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/orttraining/orttraining/training_ops/rocm/reduction/reduction_ops.cc b/orttraining/orttraining/training_ops/rocm/reduction/reduction_ops.cc index d7a5e4477f..2bafe92209 100644 --- a/orttraining/orttraining/training_ops/rocm/reduction/reduction_ops.cc +++ b/orttraining/orttraining/training_ops/rocm/reduction/reduction_ops.cc @@ -131,17 +131,17 @@ Status ReduceKernel::ComputeImplEx indices_miopen = GetScratchBuffer(indices_bytes); - IAllocatorUniquePtr workspace_miopen = GetScratchBuffer(workspace_bytes); + IAllocatorUniquePtr indices_rocm = GetScratchBuffer(indices_bytes); + IAllocatorUniquePtr workspace_rocm = GetScratchBuffer(workspace_bytes); const auto one = Consts::One; const auto zero = Consts::Zero; auto temp_Y = GetScratchBuffer(output_count); MIOPEN_RETURN_IF_ERROR(miopenReduceTensor(MiopenHandle(), reduce_desc, - indices_miopen.get(), + indices_rocm.get(), indices_bytes, - workspace_miopen.get(), + workspace_rocm.get(), workspace_bytes, &one, input_tensor, @@ -156,4 +156,4 @@ Status ReduceKernel::ComputeImplEx