From 2bed0d4abb33178da7361213f68ba68939a39a95 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Tue, 9 Aug 2022 16:48:44 +0800 Subject: [PATCH] [CUDA] SoftmaxCrossEntropy Kernels Refactor (#12482) * sce refactor * refactor * remove usnecessory memset --- .../cuda/cu_inc/elementwise_impl.cuh | 52 +++ .../loss/softmax_cross_entropy_loss_impl.cc | 6 +- .../loss/softmax_cross_entropy_loss_impl.cu | 341 +++++++---------- .../loss/softmax_cross_entropy_loss_impl.h | 2 +- .../cuda/loss/softmaxcrossentropy_impl.cu | 361 ++++++------------ 5 files changed, 311 insertions(+), 451 deletions(-) create mode 100644 onnxruntime/core/providers/cuda/cu_inc/elementwise_impl.cuh diff --git a/onnxruntime/core/providers/cuda/cu_inc/elementwise_impl.cuh b/onnxruntime/core/providers/cuda/cu_inc/elementwise_impl.cuh new file mode 100644 index 0000000000..790a612790 --- /dev/null +++ b/onnxruntime/core/providers/cuda/cu_inc/elementwise_impl.cuh @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/cuda/cu_inc/common.cuh" + +namespace onnxruntime { +namespace cuda { + +#ifdef USE_ROCM +constexpr int kElementsPerThread = 2; +constexpr int kThreadsPerBlock = 512; +#else +constexpr int kElementsPerThread = GridDim::maxElementsPerThread; +constexpr int kThreadsPerBlock = GridDim::maxThreadsPerBlock; +#endif + +template +__global__ void ElementwiseKernel(T* output_data, const FuncT functor, CUDA_LONG N) { + CUDA_LONG start = kElementsPerThread * kThreadsPerBlock * blockIdx.x + threadIdx.x; + T value[kElementsPerThread]; + + CUDA_LONG id = start; +#pragma unroll + for (int i = 0; i < kElementsPerThread; ++i) { + if (id < N) { + value[i] = functor(id); + id += kThreadsPerBlock; + } + } + + id = start; +#pragma unroll + for (int i = 0; i < kElementsPerThread; ++i) { + if (id < N) { + output_data[id] = value[i]; + id += kThreadsPerBlock; + } + } +} + +template +void LaunchElementwiseKernel(cudaStream_t stream, T* output_data, const FuncT& functor, size_t output_size) { + if (output_size == 0) return; + CUDA_LONG N = static_cast(output_size); + int blocksPerGrid = CeilDiv(N, kThreadsPerBlock * kElementsPerThread); + ElementwiseKernel<<>>(output_data, functor, N); +} + +} // namespace cuda +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc index 03b446b1e6..3f4f09058c 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc @@ -122,8 +122,7 @@ Status SoftmaxCrossEntropyLoss::ComputeInternal(OpKernelContext* ctx) co IAllocatorUniquePtr weight_data_nd = GetScratchBuffer(N_D); T* weight_data_nd_data = weight_data_nd.get(); - CUDA_RETURN_IF_ERROR(cudaMemsetAsync(weight_data_nd_data, 0, N_D * sizeof(T), Stream())); - ComputeWeightsSoftmaxCrossEntropyImpl(Stream(), + ComputeSoftmaxCrossEntropyWeightsImpl(Stream(), label_data, reinterpret_cast(weight_data), N_D, C, @@ -241,8 +240,7 @@ Status SoftmaxCrossEntropyLossGrad::ComputeInternal(OpKernelContext* ctx IAllocatorUniquePtr weight_data_nd = GetScratchBuffer(N_D); T* weight_data_nd_data = weight_data_nd.get(); - CUDA_RETURN_IF_ERROR(cudaMemsetAsync(weight_data_nd_data, 0, N_D * sizeof(T), Stream())); - ComputeWeightsSoftmaxCrossEntropyImpl(Stream(), + ComputeSoftmaxCrossEntropyWeightsImpl(Stream(), label_data, reinterpret_cast(weight_data), N_D, C, diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu index c40a656f08..f536054e7b 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu @@ -2,242 +2,161 @@ // Licensed under the MIT License. #include "core/providers/cuda/cuda_common.h" -#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/cu_inc/elementwise_impl.cuh" namespace onnxruntime { namespace cuda { -template -__global__ void _ComputeWeightsSoftmaxCrossEntropy( - T* weight_data_nd, - const Tin* label_data, - const T* weight_data, - CUDA_LONG N_D, - CUDA_LONG C, - CUDA_LONG ignore_index) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(i, N_D); - const T ONE_T = 1; - if (label_data[i] != ignore_index) { - CUDA_KERNEL_ASSERT(label_data[i] >= 0 && label_data[i] < C); - weight_data_nd[i] = weight_data != nullptr ? weight_data[label_data[i]] : ONE_T; +template +struct OpSoftmaxCrossEntropyWeights { + OpSoftmaxCrossEntropyWeights(const Tin* label_data, const T* weight_data, Tin C, Tin ignore_index) + : label_data_(label_data), weight_data_(weight_data), C_(C), ignore_index_(ignore_index) {} + + __device__ __inline__ T operator()(CUDA_LONG idx) const { + if (label_data_[idx] != ignore_index_) { + if (IsWeighted) { + CUDA_KERNEL_ASSERT(label_data_[idx] >= 0 && label_data_[idx] < C_); + return weight_data_[label_data_[idx]]; + } + return T(1.f); + } + return T(0.f); } -} + + const Tin* label_data_; + const T* weight_data_; + Tin C_; + Tin ignore_index_; +}; template -void ComputeWeightsSoftmaxCrossEntropyImpl( - cudaStream_t stream, - const Tin* label, - const T* weight, - size_t count, - size_t label_depth, - int64_t ignore_index, - T* weight_data_nd) { - int blocksPerGrid = (int)(ceil(static_cast(count) / GridDim::maxThreadsPerBlock)); - CUDA_LONG N_D = static_cast(count); - CUDA_LONG C = static_cast(label_depth); - CUDA_LONG II = static_cast(ignore_index); - _ComputeWeightsSoftmaxCrossEntropy<<>>( - weight_data_nd, - label, - weight, - N_D, - C, - II); -} - -template -__global__ void _WeightedSoftmaxCrossEntropyLoss( - const T* log_prob_data, - const Tin* label_data, - const T* weight_data, - const TAcc* normalize_factor_data, - T* output_data, - CUDA_LONG N_D, - CUDA_LONG C, - CUDA_LONG II) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(i, N_D); - if (II == label_data[i]) { - output_data[i] = 0; +void ComputeSoftmaxCrossEntropyWeightsImpl(cudaStream_t stream, const Tin* label, const T* weight, size_t count, + size_t label_depth, int64_t ignore_index, T* weight_data_nd) { + if (weight) { + OpSoftmaxCrossEntropyWeights op(label, weight, static_cast(label_depth), + static_cast(ignore_index)); + LaunchElementwiseKernel(stream, weight_data_nd, op, count); } else { - CUDA_KERNEL_ASSERT(label_data[i] >= 0 && label_data[i] < C); - output_data[i] = static_cast(static_cast(-log_prob_data[i * C + label_data[i]] * weight_data[i]) / - *normalize_factor_data); + OpSoftmaxCrossEntropyWeights op(label, nullptr, static_cast(label_depth), + static_cast(ignore_index)); + LaunchElementwiseKernel(stream, weight_data_nd, op, count); } } +#define INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(T, Tin) \ + template void ComputeSoftmaxCrossEntropyWeightsImpl(cudaStream_t stream, const Tin* label, const T* weight, \ + size_t count, size_t label_depth, int64_t ignore_index, \ + T* weight_data_nd) + +INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(float, int32_t); +INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(float, int64_t); +INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(half, int64_t); +INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL(BFloat16, int64_t); + +#undef INSTANTIATE_COMPUTE_SCE_WEIGHTS_IMPL + template -void SoftmaxCrossEntropyLossImpl( - cudaStream_t stream, - const T* log_prob, - const Tin* label, - const T* weight, - const TAcc* normalize_factor, - size_t count, - size_t label_depth, - int64_t ignore_index, - T* output_data) { - int blocksPerGrid = (int)(ceil(static_cast(count) / GridDim::maxThreadsPerBlock)); - CUDA_LONG N_D = static_cast(count); - CUDA_LONG C = static_cast(label_depth); - CUDA_LONG II = static_cast(ignore_index); - _WeightedSoftmaxCrossEntropyLoss<<>>( - log_prob, - label, - weight, - normalize_factor, - output_data, - N_D, - C, - II); +struct OpWeightedSoftmaxCrossEntropyLoss { + OpWeightedSoftmaxCrossEntropyLoss(const T* log_prob_data, const Tin* label_data, const T* weight_data, + const TAcc* normalize_factor_data, Tin C, Tin ignore_index) + : log_prob_data_(log_prob_data), + label_data_(label_data), + weight_data_(weight_data), + normalize_factor_data_(normalize_factor_data), + C_(C), + ignore_index_(ignore_index) {} + + __device__ __inline__ T operator()(CUDA_LONG idx) const { + if (label_data_[idx] != ignore_index_) { + CUDA_KERNEL_ASSERT(label_data_[idx] >= 0 && label_data_[idx] < C_); + return static_cast(static_cast(-log_prob_data_[idx * C_ + label_data_[idx]] * weight_data_[idx]) / + (*normalize_factor_data_)); + } + return T(0.f); + } + + const T* log_prob_data_; + const Tin* label_data_; + const T* weight_data_; + const TAcc* normalize_factor_data_; + Tin C_; + Tin ignore_index_; +}; + +template +void SoftmaxCrossEntropyLossImpl(cudaStream_t stream, const T* log_prob, const Tin* label, const T* weight, + const TAcc* normalize_factor, size_t count, size_t label_depth, int64_t ignore_index, + T* output_data) { + OpWeightedSoftmaxCrossEntropyLoss op(log_prob, label, weight, normalize_factor, + static_cast(label_depth), static_cast(ignore_index)); + LaunchElementwiseKernel(stream, output_data, op, count); } -#define INSTANTIATE_IMPL_SoftMaxEntropyLossImpl(T, TAcc, Tin) \ - template void SoftmaxCrossEntropyLossImpl( \ - cudaStream_t stream, \ - const T* log_prob, \ - const Tin* label, \ - const T* weight, \ - const TAcc* normalize_factor, \ - size_t count, \ - size_t label_depth, \ - int64_t ignore_index, \ - T* output_data); +template +struct OpWeightedSoftmaxCrossEntropyLossGrad { + OpWeightedSoftmaxCrossEntropyLossGrad(const T* dY_data, const T* log_prob_data, const Tin* label_data, + const T* weight_data, const TAcc* normalize_factor_data, Tin C) + : dY_data_(dY_data), + log_prob_data_(log_prob_data), + label_data_(label_data), + weight_data_(weight_data), + normalize_factor_data_(normalize_factor_data), + C_(C) { + C_fdm_ = fast_divmod(static_cast(C)); + } -INSTANTIATE_IMPL_SoftMaxEntropyLossImpl(float, float, int32_t) -INSTANTIATE_IMPL_SoftMaxEntropyLossImpl(float, float, int64_t) -INSTANTIATE_IMPL_SoftMaxEntropyLossImpl(half, float, int64_t) -INSTANTIATE_IMPL_SoftMaxEntropyLossImpl(BFloat16, float, int64_t) - -template -__global__ void _WeightedSoftmaxCrossEntropyLossGrad( - const T* dY, - const T* log_prob, - const Tin* label, - const T* weight, - const TAcc* normalize_factor, - T* output_data, - CUDA_LONG N_D, - CUDA_LONG C) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(i, N_D * C); - - int row = i / C; - int d = i % C; - const T ZERO_T = 0; - const TAcc ZERO_TAcc = 0; - const TAcc ONE_TAcc = 1; - CUDA_KERNEL_ASSERT(weight[row] == ZERO_T || (label[row] >= 0 && label[row] < C)); - if (ZERO_TAcc == *normalize_factor) { - // normalize_factor is sum of labels' weights. Because zero - // sum implies all weights are 0, the loss function should + __device__ __inline__ T operator()(CUDA_LONG idx) const { + // normalize_factor is sum of labels' weights. Because zero sum implies all weights are 0, the loss function should // be constant 0 and its corresponding gradient should be 0 as well. - output_data[i] = ZERO_T; - } else { - output_data[i] = static_cast(static_cast((*dY) * weight[row]) * - (_Exp(static_cast(log_prob[i])) - ONE_TAcc * (TAcc)(d == label[row])) / - (*normalize_factor)); + if (*normalize_factor_data_ != TAcc(0.f)) { + int row, d; + C_fdm_.divmod(idx, row, d); + CUDA_KERNEL_ASSERT(weight_data_[row] == T(0.f) || (label_data_[row] >= 0 && label_data_[row] < C_)); + return static_cast(static_cast((IsReductionNone ? dY_data_[row] : *dY_data_) * weight_data_[row]) * + (_Exp(static_cast(log_prob_data_[idx])) - (TAcc)(d == label_data_[row])) / + (*normalize_factor_data_)); + } + return T(0.f); } -} + + const T* dY_data_; + const T* log_prob_data_; + const Tin* label_data_; + const T* weight_data_; + const TAcc* normalize_factor_data_; + Tin C_; + fast_divmod C_fdm_; +}; template -__global__ void _WeightedReductionNoneSoftmaxCrossEntropyLossGrad( - const T* dY, - const T* log_prob, - const Tin* label, - const T* weight, - const TAcc* normalize_factor, - T* output_data, - CUDA_LONG N_D, - CUDA_LONG C) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(i, N_D * C); - - int row = i / C; - int d = i % C; - const T ZERO_T = 0; - const TAcc ZERO_TAcc = 0; - const TAcc ONE_TAcc = 1; - CUDA_KERNEL_ASSERT(weight[row] == ZERO_T || (label[row] >= 0 && label[row] < C)); - if (ZERO_TAcc == *normalize_factor) { - // normalize_factor is sum of labels' weights. Because zero - // sum implies all weights are 0, the loss function should - // be constant 0 and its corresponding gradient should be 0 as well. - output_data[i] = ZERO_T; - } else { - output_data[i] = static_cast(static_cast(dY[row] * weight[row]) * - (_Exp(static_cast(log_prob[i])) - ONE_TAcc * (TAcc)(d == label[row])) / - (*normalize_factor)); - } -} - -template -void SoftmaxCrossEntropyLossGradImpl( - cudaStream_t stream, - const T* dY, - const T* log_prob, - const Tin* label, - const T* weight, - const TAcc* normalize_factor, - size_t count, - size_t label_depth, - bool reduction_none, - T* output_data) { - CUDA_LONG N_D = static_cast(count); - CUDA_LONG C = static_cast(label_depth); - int blocksPerGrid = (int)(ceil(static_cast(N_D * C) / GridDim::maxThreadsPerBlock)); +void SoftmaxCrossEntropyLossGradImpl(cudaStream_t stream, const T* dY, const T* log_prob, const Tin* label, + const T* weight, const TAcc* normalize_factor, size_t count, size_t label_depth, + bool reduction_none, T* output_data) { if (reduction_none) { - _WeightedReductionNoneSoftmaxCrossEntropyLossGrad<<>>( - dY, - log_prob, - label, - weight, - normalize_factor, - output_data, - N_D, - C); + OpWeightedSoftmaxCrossEntropyLossGrad op(dY, log_prob, label, weight, normalize_factor, + static_cast(label_depth)); + LaunchElementwiseKernel(stream, output_data, op, count * label_depth); } else { - _WeightedSoftmaxCrossEntropyLossGrad<<>>( - dY, - log_prob, - label, - weight, - normalize_factor, - output_data, - N_D, - C); + OpWeightedSoftmaxCrossEntropyLossGrad op(dY, log_prob, label, weight, normalize_factor, + static_cast(label_depth)); + LaunchElementwiseKernel(stream, output_data, op, count * label_depth); } } -#define INSTANTIATE_IMPL_SoftMaxEntropyLossGradImpl(T, TAcc, Tin) \ - template void SoftmaxCrossEntropyLossGradImpl( \ - cudaStream_t stream, \ - const T* dY, \ - const T* log_prob, \ - const Tin* label, \ - const T* weight, \ - const TAcc* normalize_factor, \ - size_t count, \ - size_t label_depth, \ - bool reducation_none, \ - T* output_data); +#define INSTANTIATE_SCE_LOSS_IMPL(T, TAcc, Tin) \ + template void SoftmaxCrossEntropyLossImpl(cudaStream_t stream, const T* log_prob, const Tin* label, const T* weight, \ + const TAcc* normalize_factor, size_t count, size_t label_depth, \ + int64_t ignore_index, T* output_data); \ + template void SoftmaxCrossEntropyLossGradImpl(cudaStream_t stream, const T* dY, const T* log_prob, const Tin* label, \ + const T* weight, const TAcc* normalize_factor, size_t count, \ + size_t label_depth, bool reducation_none, T* output_data) -INSTANTIATE_IMPL_SoftMaxEntropyLossGradImpl(float, float, int32_t) -INSTANTIATE_IMPL_SoftMaxEntropyLossGradImpl(float, float, int64_t) -INSTANTIATE_IMPL_SoftMaxEntropyLossGradImpl(half, float, int64_t) -INSTANTIATE_IMPL_SoftMaxEntropyLossGradImpl(BFloat16, float, int64_t) +INSTANTIATE_SCE_LOSS_IMPL(float, float, int32_t); +INSTANTIATE_SCE_LOSS_IMPL(float, float, int64_t); +INSTANTIATE_SCE_LOSS_IMPL(half, float, int64_t); +INSTANTIATE_SCE_LOSS_IMPL(BFloat16, float, int64_t); -#define INSTANTIATE_IMPL_ComputeWeightsSoftmaxCrossEntropyImpl(T, Tin) \ - template void ComputeWeightsSoftmaxCrossEntropyImpl( \ - cudaStream_t stream, \ - const Tin* label, \ - const T* weight, \ - size_t count, \ - size_t label_depth, \ - int64_t ignore_index, \ - T* weight_data_nd); - -INSTANTIATE_IMPL_ComputeWeightsSoftmaxCrossEntropyImpl(float, int32_t) -INSTANTIATE_IMPL_ComputeWeightsSoftmaxCrossEntropyImpl(float, int64_t) -INSTANTIATE_IMPL_ComputeWeightsSoftmaxCrossEntropyImpl(half, int64_t) -INSTANTIATE_IMPL_ComputeWeightsSoftmaxCrossEntropyImpl(BFloat16, int64_t) +#undef INSTANTIATE_SCE_LOSS_IMPL } // namespace cuda -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.h b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.h index 26604b00cc..a9ba04f077 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.h +++ b/orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.h @@ -36,7 +36,7 @@ void SoftmaxCrossEntropyLossGradImpl( T* output_data); template -void ComputeWeightsSoftmaxCrossEntropyImpl( +void ComputeSoftmaxCrossEntropyWeightsImpl( cudaStream_t stream, const Tin* label, const T* weight, diff --git a/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cu b/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cu index 9fd01720b5..9f58073256 100644 --- a/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cu @@ -2,269 +2,160 @@ // Licensed under the MIT License. #include "core/providers/cuda/cuda_common.h" -#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/cu_inc/elementwise_impl.cuh" namespace onnxruntime { namespace cuda { template -__global__ void _SoftMaxCrossEntropy( - const T* log_prob_data, - const T* label_data, - CUDA_LONG NORMALIZE_FACTOR, - T* output_data, - CUDA_LONG N) { +struct OpSoftmaxCrossEntropy { + OpSoftmaxCrossEntropy(const T* log_prob_data, const T* label_data, T normalize_factor) + : log_prob_data_(log_prob_data), label_data_(label_data), normalize_factor_(normalize_factor) {} - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); - output_data[id] = -log_prob_data[id] * label_data[id] / NORMALIZE_FACTOR; -} - -template -void SoftMaxCrossEntropyImpl( - cudaStream_t stream, - const T* log_prob, - const T* label, - size_t normalize_factor, - T* output_data, - size_t count) { - int blocksPerGrid = (int)(ceil(static_cast(count) / GridDim::maxThreadsPerBlock)); - CUDA_LONG N = static_cast(count); - CUDA_LONG NORMALIZE_FACTOR = static_cast(normalize_factor); - _SoftMaxCrossEntropy<<>>( - log_prob, - label, - NORMALIZE_FACTOR, - output_data, - N); -} - -#define SPECIALIZED_IMPL_SoftMaxEntropyImpl(T) \ - template void SoftMaxCrossEntropyImpl( \ - cudaStream_t stream, \ - const T* log_prob, \ - const T* label, \ - size_t normalize_factor, \ - T* output_data, \ - size_t count); - -SPECIALIZED_IMPL_SoftMaxEntropyImpl(float) - -template -__global__ void _SoftMaxCrossEntropyGrad( - const T* dY, - const T* log_prob, - const T* label, - CUDA_LONG NORMALIZE_FACTOR, - T* output_data, - CUDA_LONG N) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); - output_data[id] = (_Exp(log_prob[id]) - label[id]) * (*dY) / NORMALIZE_FACTOR; -} - -template -void SoftMaxCrossEntropyGradImpl( - cudaStream_t stream, - const T* dY, - const T* log_prob, - const T* label, - size_t normalize_factor, - T* output_data, - size_t count) { - int blocksPerGrid = (int)(ceil(static_cast(count) / GridDim::maxThreadsPerBlock)); - CUDA_LONG N = static_cast(count); - CUDA_LONG NORMALIZE_FACTOR = static_cast(normalize_factor); - _SoftMaxCrossEntropyGrad<<>>( - dY, - log_prob, - label, - NORMALIZE_FACTOR, - output_data, - N); -} - -#define SPECIALIZED_IMPL_SoftMaxEntropyGradImpl(T) \ - template void SoftMaxCrossEntropyGradImpl( \ - cudaStream_t stream, \ - const T* dY, \ - const T* log_prob, \ - const T* label, \ - size_t normalize_factor, \ - T* output_data, \ - size_t count); - -SPECIALIZED_IMPL_SoftMaxEntropyGradImpl(float) - -template -__global__ void _SparseSoftmaxCrossEntropy( - const T* log_prob_data, - const Tin* label_data, - const T* normalize_factor_data, - T* output_data, - CUDA_LONG N, - CUDA_LONG D) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(i, N); - CUDA_KERNEL_ASSERT(label_data[i] >= 0 && label_data[i] < D); - if (*normalize_factor_data == 0) { - output_data[i] = 0; - } else { - output_data[i] = -log_prob_data[i * D + label_data[i]] / (*normalize_factor_data); + __device__ __inline__ T operator()(CUDA_LONG idx) const { + return -log_prob_data_[idx] * label_data_[idx] / normalize_factor_; } + + const T* log_prob_data_; + const T* label_data_; + T normalize_factor_; +}; + +template +void SoftMaxCrossEntropyImpl(cudaStream_t stream, const T* log_prob, const T* label, size_t normalize_factor, + T* output_data, size_t count) { + OpSoftmaxCrossEntropy op(log_prob, label, static_cast(normalize_factor)); + LaunchElementwiseKernel(stream, output_data, op, count); } -template -__global__ void _WeightedSparseSoftmaxCrossEntropy( - const T* log_prob_data, - const Tin* label_data, - const T* weight_data, - const T* normalize_factor_data, - T* output_data, - CUDA_LONG N, - CUDA_LONG D) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(i, N); - CUDA_KERNEL_ASSERT(label_data[i] >= 0 && label_data[i] < D); - if (*normalize_factor_data == 0) { - output_data[i] = 0; - } else { - output_data[i] = -log_prob_data[i * D + label_data[i]] * weight_data[i] / (*normalize_factor_data); +template void SoftMaxCrossEntropyImpl(cudaStream_t stream, const float* log_prob, const float* label, + size_t normalize_factor, float* output_data, size_t count); + +template +struct OpSoftmaxCrossEntropyGrad { + OpSoftmaxCrossEntropyGrad(const T* dY_data, const T* log_prob_data, const T* label_data, T normalize_factor) + : dY_data_(dY_data), + log_prob_data_(log_prob_data), + label_data_(label_data), + normalize_factor_(normalize_factor) {} + + __device__ __inline__ T operator()(CUDA_LONG idx) const { + return (_Exp(log_prob_data_[idx]) - label_data_[idx]) * (*dY_data_) / normalize_factor_; } + + const T* dY_data_; + const T* log_prob_data_; + const T* label_data_; + T normalize_factor_; +}; + +template +void SoftMaxCrossEntropyGradImpl(cudaStream_t stream, const T* dY, const T* log_prob, const T* label, + size_t normalize_factor, T* output_data, size_t count) { + OpSoftmaxCrossEntropyGrad op(dY, log_prob, label, static_cast(normalize_factor)); + LaunchElementwiseKernel(stream, output_data, op, count); } +template void SoftMaxCrossEntropyGradImpl(cudaStream_t stream, const float* dY, const float* log_prob, + const float* label, size_t normalize_factor, float* output_data, + size_t count); + +template +struct OpSparseSoftmaxCrossEntropy { + OpSparseSoftmaxCrossEntropy(const T* log_prob_data, const Tin* label_data, const T* weight_data, + const T* normalize_factor_data, Tin D) + : log_prob_data_(log_prob_data), + label_data_(label_data), + weight_data_(weight_data), + normalize_factor_data_(normalize_factor_data), + D_(D) {} + + __device__ __inline__ T operator()(CUDA_LONG idx) const { + if (*normalize_factor_data_ != T(0.f)) { + CUDA_KERNEL_ASSERT(label_data_[idx] >= 0 && label_data_[idx] < D_); + return -log_prob_data_[idx * D_ + label_data_[idx]] * (IsWeighted ? weight_data_[idx] : T(1.f)) / + (*normalize_factor_data_); + } + return T(0.f); + } + + const T* log_prob_data_; + const Tin* label_data_; + const T* weight_data_; + const T* normalize_factor_data_; + Tin D_; +}; + template -void SparseSoftmaxCrossEntropyImpl( - cudaStream_t stream, - const T* log_prob, - const Tin* label, - const T* weight, - const T* normalize_factor, - T* output_data, - size_t count, - size_t label_depth) { - int blocksPerGrid = (int)(ceil(static_cast(count) / GridDim::maxThreadsPerBlock)); - CUDA_LONG N = static_cast(count); - CUDA_LONG D = static_cast(label_depth); +void SparseSoftmaxCrossEntropyImpl(cudaStream_t stream, const T* log_prob, const Tin* label, const T* weight, + const T* normalize_factor, T* output_data, size_t count, size_t label_depth) { if (weight) { - _WeightedSparseSoftmaxCrossEntropy<<>>( - log_prob, - label, - weight, - normalize_factor, - output_data, - N, - D); + OpSparseSoftmaxCrossEntropy op(log_prob, label, weight, normalize_factor, + static_cast(label_depth)); + LaunchElementwiseKernel(stream, output_data, op, count); } else { - _SparseSoftmaxCrossEntropy<<>>( - log_prob, - label, - normalize_factor, - output_data, - N, - D); + OpSparseSoftmaxCrossEntropy op(log_prob, label, nullptr, normalize_factor, + static_cast(label_depth)); + LaunchElementwiseKernel(stream, output_data, op, count); } } -#define SPECIALIZED_IMPL_SparseSoftMaxEntropyImpl(T, Tin) \ - template void SparseSoftmaxCrossEntropyImpl( \ - cudaStream_t stream, \ - const T* log_prob, \ - const Tin* label, \ - const T* weight, \ - const T* normalize_factor, \ - T* output_data, \ - size_t count, \ - size_t label_depth); +template +struct OpSparseSoftmaxCrossEntropyGrad { + OpSparseSoftmaxCrossEntropyGrad(const T* dY_data, const T* log_prob_data, const Tin* label_data, const T* weight_data, + const T* normalize_factor_data, fast_divmod D_fdm) + : dY_data_(dY_data), + log_prob_data_(log_prob_data), + label_data_(label_data), + weight_data_(weight_data), + normalize_factor_data_(normalize_factor_data), + D_fdm_(D_fdm) {} -SPECIALIZED_IMPL_SparseSoftMaxEntropyImpl(float, int32_t) -SPECIALIZED_IMPL_SparseSoftMaxEntropyImpl(float, int64_t) - -template -__global__ void _SparseSoftmaxCrossEntropyGrad( - const T* dY, - const T* log_prob, - const Tin* label, - const T* normalize_factor, - T* output_data, - CUDA_LONG N, - CUDA_LONG D) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(i, N * D); - int row = i / D; - int d = i % D; - if (*normalize_factor == 0) { - output_data[i] = 0; - } else { - output_data[i] = (*dY) * (_Exp(log_prob[i]) - 1.0 * (d == label[row])) / (*normalize_factor); + __device__ __inline__ T operator()(CUDA_LONG idx) const { + if (*normalize_factor_data_ != T(0.f)) { + int row, d; + D_fdm_.divmod(idx, row, d); + return (*dY_data_) * (IsWeighted ? weight_data_[row] : T(1.f)) * + (_Exp(log_prob_data_[idx]) - (T)(d == label_data_[row])) / (*normalize_factor_data_); + } + return T(0.f); } -} + + const T* dY_data_; + const T* log_prob_data_; + const Tin* label_data_; + const T* weight_data_; + const T* normalize_factor_data_; + fast_divmod D_fdm_; +}; template -__global__ void _WeightedSparseSoftmaxCrossEntropyGrad( - const T* dY, - const T* log_prob, - const Tin* label, - const T* weight, - const T* normalize_factor, - T* output_data, - CUDA_LONG N, - CUDA_LONG D) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(i, N * D); - int row = i / D; - int d = i % D; - if (*normalize_factor == 0) { - output_data[i] = 0; - } else { - output_data[i] = (*dY) * weight[row] * (_Exp(log_prob[i]) - 1.0 * (d == label[row])) / (*normalize_factor); - } -} - -template -void SparseSoftmaxCrossEntropyGradImpl( - cudaStream_t stream, - const T* dY, - const T* log_prob, - const Tin* label, - const T* weight, - const T* normalize_factor, - T* output_data, - size_t count, - size_t label_depth) { - CUDA_LONG N = static_cast(count); - CUDA_LONG D = static_cast(label_depth); - int blocksPerGrid = (int)(ceil(static_cast(N * D) / GridDim::maxThreadsPerBlock)); +void SparseSoftmaxCrossEntropyGradImpl(cudaStream_t stream, const T* dY, const T* log_prob, const Tin* label, + const T* weight, const T* normalize_factor, T* output_data, size_t count, + size_t label_depth) { if (weight) { - _WeightedSparseSoftmaxCrossEntropyGrad<<>>( - dY, - log_prob, - label, - weight, - normalize_factor, - output_data, - N, - D); + OpSparseSoftmaxCrossEntropyGrad op(dY, log_prob, label, weight, normalize_factor, + fast_divmod(static_cast(label_depth))); + LaunchElementwiseKernel(stream, output_data, op, count * label_depth); } else { - _SparseSoftmaxCrossEntropyGrad<<>>( - dY, - log_prob, - label, - normalize_factor, - output_data, - N, - D); + OpSparseSoftmaxCrossEntropyGrad op(dY, log_prob, label, nullptr, normalize_factor, + fast_divmod(static_cast(label_depth))); + LaunchElementwiseKernel(stream, output_data, op, count * label_depth); } } -#define SPECIALIZED_IMPL_SparseSoftMaxEntropyGradImpl(T, Tin) \ - template void SparseSoftmaxCrossEntropyGradImpl( \ - cudaStream_t stream, \ - const T* dY, \ - const T* log_prob, \ - const Tin* label, \ - const T* weight, \ - const T* normalize_factor, \ - T* output_data, \ - size_t count, \ - size_t label_depth); +#define SPECIALIZED_SPARSE_SOFTMAX_ENTROPY_IMPL(T, Tin) \ + template void SparseSoftmaxCrossEntropyImpl(cudaStream_t stream, const T* log_prob, const Tin* label, \ + const T* weight, const T* normalize_factor, T* output_data, \ + size_t count, size_t label_depth); \ + template void SparseSoftmaxCrossEntropyGradImpl(cudaStream_t stream, const T* dY, const T* log_prob, \ + const Tin* label, const T* weight, const T* normalize_factor, \ + T* output_data, size_t count, size_t label_depth) -SPECIALIZED_IMPL_SparseSoftMaxEntropyGradImpl(float, int32_t) -SPECIALIZED_IMPL_SparseSoftMaxEntropyGradImpl(float, int64_t) +SPECIALIZED_SPARSE_SOFTMAX_ENTROPY_IMPL(float, int32_t); +SPECIALIZED_SPARSE_SOFTMAX_ENTROPY_IMPL(float, int64_t); + +#undef SPECIALIZED_SPARSE_SOFTMAX_ENTROPY_IMPL } // namespace cuda } // namespace onnxruntime