saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-20 19:12:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Delete ROCM-specific reduction code that is identical to CUDA reduction code.

Browse source
This commit is contained in:
Jesse Benson 2021-01-19 12:51:25 -08:00 committed by Jesse Benson
parent 5d8792705b
commit 86ac11af1a
9 changed files with 13 additions and 863 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
onnxruntime/core/providers/rocm/reduction/reduction_functions.cu
View file

@ -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 <typename TIn, typename TOut, typename TBuf, typename TOp, typename TFi
__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[];
HIP_DYNAMIC_SHARED( unsigned char, shared_memory_bytes)
TBuf* shared_memory = reinterpret_cast<TBuf*>(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<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);
hipLaunchKernelGGL(HIP_KERNEL_NAME(reduce_matrix_columns_kernel<TIn, TOut, TBuf, TOp, TFinalOp, DivideResultBySize>), 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<TBuf*>(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<TIn, TOut, TBuf>),
grid, block, block.y * block.x * sizeof(TBuf), 0,
hipLaunchKernelGGL(HIP_KERNEL_NAME(reduce_matrix_rows_kernel<TIn, TOut, TBuf>), dim3(grid), dim3(block), block.y * block.x * sizeof(TBuf), 0,
input, output, m, n);
return Status::OK();

107
onnxruntime/core/providers/rocm/reduction/reduction_functions.h
View file

@ -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 <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 = 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);
} // namespace rocm
} // namespace onnxruntime

6
onnxruntime/core/providers/rocm/reduction/reduction_ops.h
View file

@ -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<allow_multi_axes
template <typename T, miopenReduceTensorIndices_t ReduceTensorIndices = MIOPEN_REDUCE_TENSOR_NO_INDICES>
Status ComputeImplEx(OpKernelContext* ctx, miopenReduceTensorOp_t miopen_reduce_op) const;
template <typename T, typename OutT, miopenReduceTensorIndices_t ReduceTensorIndices = MIOPEN_REDUCE_TENSOR_NO_INDICES>
template <typename T, typename OutT, miopenReduceTensorIndices_t ReduceTensorIndices>
Status ReduceKernelShared(
const T* X,
const TensorShape& input_shape,
@ -269,4 +269,4 @@ class MiopenReduceDescriptor final {
};
} // namespace rocm
} // namespace onnxruntime
} // namespace onnxruntime

43
onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh
View file

@ -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<unsigned int>(value);
--value_;
value_ |= value_ >> 1;
value_ |= value_ >> 2;
value_ |= value_ >> 4;
value_ |= value_ >> 8;
value_ |= value_ >> 16;
return static_cast<int>(++value_);
}
struct Square {
template <typename T>
__forceinline__ __device__ T operator()(const T& value) {
return value * value;
}
};
struct Sqrt {
template <typename T>
__forceinline__ __device__ T operator()(const T& value) {
return _Sqrt(value);
}
};
struct Identity {
template <typename T>
__forceinline__ __device__ T operator()(const T& value) {
return value;
}
};
} // namespace rocm
} // namespace onnxruntime

272
orttraining/orttraining/training_ops/rocm/loss/softmax_cross_entropy_loss_impl.cc
View file

@ -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<T>()) \
.TypeConstraint("Tin", DataTypeImpl::GetTensorType<Tin>()), \
Class<T, Tin>);
#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<T>()) \
.TypeConstraint("Tin", DataTypeImpl::GetTensorType<Tin>()), \
Class<T, Tin>);
template <typename T, typename Tin>
Status SoftmaxCrossEntropyLoss<T, Tin>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor& logit = *ctx->Input<Tensor>(0);
const Tensor& label = *ctx->Input<Tensor>(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<Tensor>(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>();
T* tmp_loss_sample_buffer = nullptr;
IAllocatorUniquePtr<T> tmp_loss_sample;
if (reduction_ == ReductionType::NONE) {
tmp_loss_sample_buffer = total_loss_data;
} else {
tmp_loss_sample = GetScratchBuffer<T>(N_D);
tmp_loss_sample_buffer = tmp_loss_sample.get();
}
const T* logit_data = logit.template Data<T>();
const Tin* label_data = label.template Data<Tin>();
T* log_prob_data = nullptr;
Tensor* log_prob = nullptr;
IAllocatorUniquePtr<T> log_prob_scratch_buffer;
if (ctx->OutputCount() > 1) {
log_prob = ctx->Output(1, logit_shape);
log_prob_data = log_prob->template MutableData<T>();
} else {
log_prob_scratch_buffer = GetScratchBuffer<T>(logit_shape.Size());
log_prob_data = log_prob_scratch_buffer.get();
}
OrtValue transpose_output;
Tensor transpose_tensor;
std::vector<int64_t> new_shape;
std::vector<size_t> 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<Tensor>()));
logit_data = (*transpose_output.GetMutable<Tensor>()).template Data<T>();
}
// calculate logsoftmax
auto status = SoftMaxComputeHelper<T, true>(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<Tensor>(2);
weight_data = weight.template Data<T>();
}
IAllocatorUniquePtr<T> weight_data_nd = GetScratchBuffer<T>(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<T>(1);
if (reduction_ == ReductionType::MEAN) {
// Compute buffer size in byte for reduction APIs.
const auto buffer_size =
compute_reduction_buffer_size<T>(static_cast<int>(N_D));
// Allocate reduction buffer whose size is buffer_size bytes.
IAllocatorUniquePtr<void> reduction_buffer = GetScratchBuffer<void>(
buffer_size);
ORT_RETURN_IF_ERROR(reduce_sum(
weight_data_nd_data,
normalize_factor_data.get(),
static_cast<int>(N_D),
reduction_buffer.get(),
buffer_size));
} else {
const T normalize_factor = static_cast<T>(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<Tensor>()).template MutableData<T>();
transpose_output.GetMutable<Tensor>()->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<Tensor>()));
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<int64_t> output_dims(1, 1);
ReduceKernelShared<T, T, MIOPEN_REDUCE_TENSOR_NO_INDICES>(
tmp_loss_sample_buffer,
label_reshape,
total_loss_data,
TensorShape({}),
MIOPEN_REDUCE_TENSOR_ADD,
output_dims);
}
return Status::OK();
}
template <typename T, typename Tin>
Status SoftmaxCrossEntropyLossGrad<T, Tin>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor& dY = *ctx->Input<Tensor>(0);
const Tensor& log_prob = *ctx->Input<Tensor>(1);
const Tensor& label = *ctx->Input<Tensor>(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<Tensor>(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<T>();
const T* log_prob_data = log_prob.template Data<T>();
const Tin* label_data = label.template Data<Tin>();
T* d_logit_data = d_logit->template MutableData<T>();
const T* weight_data = nullptr;
OrtValue transpose_output;
std::vector<int64_t> new_shape;
std::vector<size_t> 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<Tensor>()));
log_prob_data = (*transpose_output.GetMutable<Tensor>()).template Data<T>();
}
if (OpKernel::Node().InputDefs().size() == 4) {
const Tensor& weight = *ctx->Input<Tensor>(3);
weight_data = weight.template Data<T>();
}
IAllocatorUniquePtr<T> weight_data_nd = GetScratchBuffer<T>(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<T>(1);
if (reduction_ == ReductionType::MEAN) {
// Compute buffer size in byte for reduction APIs.
const auto buffer_size =
compute_reduction_buffer_size<T>(static_cast<int>(N_D));
// Allocate reduction buffer whose size is buffer_size bytes.
IAllocatorUniquePtr<void> reduction_buffer = GetScratchBuffer<void>(
buffer_size);
ORT_RETURN_IF_ERROR(reduce_sum(
weight_data_nd_data,
normalize_factor_data.get(),
static_cast<int>(N_D),
reduction_buffer.get(),
buffer_size));
} else {
const T normalize_factor = static_cast<T>(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<Tensor>()->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<Tensor>()));
auto* transposed_data = (*transpose_output.GetMutable<Tensor>()).template Data<T>();
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<T, Tin>::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

294
orttraining/orttraining/training_ops/rocm/loss/softmaxcrossentropy_impl.cc
View file

@ -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<T>()), \
Class<T>);
#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<T>()) \
.TypeConstraint("Tin", DataTypeImpl::GetTensorType<Tin>()), \
Class<T, Tin>);
template <typename T>
Status SoftmaxCrossEntropy<T>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor& logit = *ctx->Input<Tensor>(0);
const Tensor& label = *ctx->Input<Tensor>(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<T>();
const T* label_data = label.template Data<T>();
T* log_prob_data = log_prob->template MutableData<T>();
// calculate logsoftmax
auto status = SoftMaxComputeHelper<T, true>(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<size_t>(1);
}
// calculate (label * log(softmax)) for each element
IAllocatorUniquePtr<T> temp_X = GetScratchBuffer<T>(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<int64_t> output_dims(2, 1);
Tensor* Y = ctx->Output(0, TensorShape({}));
// Sum((label * log(softmax)) using Reduction
return ReduceKernelShared<T, T, MIOPEN_REDUCE_TENSOR_NO_INDICES>(
temp_X.get(),
logit_reshape,
Y->template MutableData<T>(),
TensorShape({}),
MIOPEN_REDUCE_TENSOR_ADD,
output_dims);
}
template <typename T>
Status SoftmaxCrossEntropyGrad<T>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor& dY = *ctx->Input<Tensor>(0);
const Tensor& log_prob = *ctx->Input<Tensor>(1);
const Tensor& label = *ctx->Input<Tensor>(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<T>();
const T* log_prob_data = log_prob.template Data<T>();
const T* label_data = label.template Data<T>();
size_t normalize_factor = N;
if (reduction_ == ReductionType::SUM) {
normalize_factor = static_cast<size_t>(1);
}
T* d_logits_data = d_logits->template MutableData<T>();
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 <typename T, typename Tin>
Status SparseSoftmaxCrossEntropy<T, Tin>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor& logit = *ctx->Input<Tensor>(0);
const Tensor& label = *ctx->Input<Tensor>(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<T> tmp_loss_sample = GetScratchBuffer<T>(N);
Tensor* total_loss = ctx->Output(0, TensorShape({}));
Tensor* log_prob = ctx->Output(1, logit_shape);
const T* logit_data = logit.template Data<T>();
const Tin* label_data = label.template Data<Tin>();
T* total_loss_data = total_loss->template MutableData<T>();
T* log_prob_data = log_prob->template MutableData<T>();
// calculate logsoftmax
auto status = SoftMaxComputeHelper<T, true>(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<Tensor>(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<T>();
}
auto normalize_factor_data = GetScratchBuffer<T>(1);
if (reduction_ == ReductionType::SUM) {
const T normalize_factor = static_cast<T>(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<T>(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<T>(static_cast<int>(N));
// Allocate reduction buffer whose size is buffer_size bytes.
IAllocatorUniquePtr<void> reduction_buffer = GetScratchBuffer<void>(
buffer_size);
ORT_RETURN_IF_ERROR(reduce_sum(
weight_data,
normalize_factor_data.get(),
static_cast<int>(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<int64_t> output_dims(1, 1);
return ReduceKernelShared<T, T, MIOPEN_REDUCE_TENSOR_NO_INDICES>(
tmp_loss_sample.get(),
label_reshape,
total_loss_data,
TensorShape({}),
MIOPEN_REDUCE_TENSOR_ADD,
output_dims);
}
template <typename T, typename Tin>
Status SparseSoftmaxCrossEntropyGrad<T, Tin>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor& dY = *ctx->Input<Tensor>(0);
const Tensor& log_prob = *ctx->Input<Tensor>(1);
const Tensor& label = *ctx->Input<Tensor>(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<T>();
const T* log_prob_data = log_prob.template Data<T>();
const Tin* label_data = label.template Data<Tin>();
T* d_logit_data = d_logit->template MutableData<T>();
const T* weight_data = nullptr;
if (OpKernel::Node().InputDefs().size() == 4) {
const Tensor& weight = *ctx->Input<Tensor>(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<T>();
}
auto normalize_factor_data = GetScratchBuffer<T>(1);
if (reduction_ == ReductionType::SUM) {
const T normalize_factor = static_cast<T>(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<T>(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<T>(static_cast<int>(N));
// Allocate reduction buffer whose size is buffer_size bytes.
IAllocatorUniquePtr<void> reduction_buffer = GetScratchBuffer<void>(
buffer_size);
ORT_RETURN_IF_ERROR(reduce_sum(
weight_data,
normalize_factor_data.get(),
static_cast<int>(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<T>::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<T, Tin>::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

116
orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu
View file

@ -1,116 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <hip/hip_runtime.h>
#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 <typename Tin, typename Tout>
__global__ void ScalarSqrtKernel(Tin* input, Tout* output) {
*output = (Tout)_Sqrt(*input);
}
template <typename Tin, typename Tout>
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 <typename TIn, typename TOut, typename TBuf, typename TInOp, typename TOutOp>
__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<const TIn*>(chunk_group.tensor_ptrs[0][group_index]) + chunk_start;
TOut* w_norm = output;
TBuf w_sum = TBuf(0.f);
constexpr int load_count_per_thread = 4;
for (int i = threadIdx.x; i < chunk_size && i + chunk_start < tensor_size; i += blockDim.x * load_count_per_thread) {
#pragma unroll
for (int j = 0; j < load_count_per_thread; ++j) {
const int index_in_chunk = i + j * blockDim.x;
const int index_in_tensor = chunk_start + index_in_chunk;
if (index_in_chunk < chunk_size && index_in_tensor < tensor_size) {
const TBuf w_element = TBuf(w[index_in_chunk]);
w_sum += TInOp()(w_element);
}
}
}
// Thread count in a block must be a multiple of GPU_WARP_SIZE.
#pragma unroll
for (int stride = GPU_WARP_SIZE / 2; stride > 0; stride /= 2) {
w_sum += WARP_SHFL_DOWN(w_sum, stride);
}
const int warp_count_in_block = blockDim.x / GPU_WARP_SIZE;
const int lid = threadIdx.x % GPU_WARP_SIZE;
const int wid = threadIdx.x / GPU_WARP_SIZE;
// Shape is 2 x warp_count_in_block.
extern __shared__ unsigned char shared_memory_[];
TBuf* shared_memory = reinterpret_cast<TBuf*>(shared_memory_);
if (lid == 0) {
shared_memory[wid] = w_sum;
}
__syncthreads();
#pragma unroll
for (int stride = warp_count_in_block / 2; stride > 0; stride /= 2) {
if (threadIdx.x < stride) {
shared_memory[threadIdx.x] += shared_memory[threadIdx.x + stride];
}
__syncthreads();
}
if (threadIdx.x == 0) {
atomic_add(w_norm, TOutOp()(TOut(shared_memory[0])));
}
}
template <typename TIn, typename TOut, typename TBuf, typename TInOp, typename TOutOp>
void MultiTensorReduce(ChunkGroup<1> chunk_group, TOut* output) {
// thread count per block.
constexpr int thread_count = ChunkGroup<1>::thread_count_per_block;
// shared memory's size per block.
const int shared_memory_size = thread_count / GPU_WARP_SIZE * sizeof(TBuf);
// Enforce assumptions used inside this reduction ROCM kernel.
ORT_ENFORCE(thread_count % GPU_WARP_SIZE == 0);
ORT_ENFORCE((thread_count & (thread_count - 1)) == 0);
hipLaunchKernelGGL(HIP_KERNEL_NAME(MultiTensorReduceKernel<TIn, TOut, TBuf, TInOp, TOutOp>), dim3(chunk_group.chunk_count), dim3(thread_count), shared_memory_size, 0, chunk_group, output);
}
template <typename TIn, typename TOut>
void MultiTensorReduceL2<TIn, TOut>::operator()(ChunkGroup<1> chunk_group, TOut* output) {
using TBuf = AccumulationType_t<TIn>;
MultiTensorReduce<TIn, TOut, TBuf, Square, Identity>(chunk_group, output);
}
#define INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(TIn, TOut) \
template void MultiTensorReduceL2<TIn, TOut>::operator()(ChunkGroup<1> chunk_group, TOut* output);
INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(double, float)
INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, float)
INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(half, float)
INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(float, half)
INSTANTIATE_MULTI_TENSOR_REDUCTION_L2_FUNCTOR(half, half)
} // namespace rocm
} // namespace onnxruntime

10
orttraining/orttraining/training_ops/rocm/reduction/reduction_ops.cc
View file

@ -131,17 +131,17 @@ Status ReduceKernel<true>::ComputeImplEx<int32_t, MIOPEN_REDUCE_TENSOR_NO_INDICE
ORT_RETURN_IF_ERROR(output_tensor.Set(output_dims_miopen, miopen_type_X));
MIOPEN_RETURN_IF_ERROR(miopenGetReductionIndicesSize(MiopenHandle(), reduce_desc, input_tensor, output_tensor, &indices_bytes));
MIOPEN_RETURN_IF_ERROR(miopenGetReductionWorkspaceSize(MiopenHandle(), reduce_desc, input_tensor, output_tensor, &workspace_bytes));
IAllocatorUniquePtr<uint32_t> indices_miopen = GetScratchBuffer<uint32_t>(indices_bytes);
IAllocatorUniquePtr<HipT> workspace_miopen = GetScratchBuffer<HipT>(workspace_bytes);
IAllocatorUniquePtr<uint32_t> indices_rocm = GetScratchBuffer<uint32_t>(indices_bytes);
IAllocatorUniquePtr<HipT> workspace_rocm = GetScratchBuffer<HipT>(workspace_bytes);
const auto one = Consts<float>::One;
const auto zero = Consts<float>::Zero;
auto temp_Y = GetScratchBuffer<float>(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<true>::ComputeImplEx<int32_t, MIOPEN_REDUCE_TENSOR_NO_INDICE
}
} // namespace rocm
} // namespace onnxruntime
} // namespace onnxruntime

5
tools/ci_build/amd_hipify.py
View file

@ -132,10 +132,8 @@ provider_excluded_files = [
'object_detection/roialign_impl.cu',
'object_detection/roialign_impl.h',
'reduction/reduction_functions.cu',
'reduction/reduction_functions.h',
'reduction/reduction_ops.cc',
'reduction/reduction_ops.h',
'reduction/reduction_utils.cuh',
'rnn/cudnn_rnn_base.cc',
'rnn/cudnn_rnn_base.h',
'rnn/gru.cc',
@ -237,8 +235,6 @@ training_ops_excluded_files = [
'controlflow/record.h',
'controlflow/wait.cc',
'controlflow/wait.h',
'loss/softmax_cross_entropy_loss_impl.cc',
'loss/softmaxcrossentropy_impl.cc',
'math/div_grad.cc',
'math/div_grad.h',
'math/div_grad_impl.cu',
@ -254,7 +250,6 @@ training_ops_excluded_files = [
'optimizer/adam.cu',
'optimizer/lamb.cc',
'reduction/reduction_all.cc',
'reduction/reduction_all.cu',
'reduction/reduction_ops.cc',
'tensor/gather_elements_grad.cc',
'tensor/gather_elements_grad.h',