saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-10 17:37:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[CUDA] SoftmaxCrossEntropy Kernels Refactor (#12482)

Browse source 
* sce refactor

* refactor

* remove usnecessory memset
This commit is contained in:
Vincent Wang 2022-08-09 16:48:44 +08:00 committed by GitHub
parent cfa09d16d9
commit 2bed0d4abb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 311 additions and 451 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
onnxruntime/core/providers/cuda/cu_inc/elementwise_impl.cuh Normal file
View file

@ -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 <typename T, typename FuncT>
__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 <typename T, typename FuncT>
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<CUDA_LONG>(output_size);
int blocksPerGrid = CeilDiv(N, kThreadsPerBlock * kElementsPerThread);
ElementwiseKernel<T, FuncT><<<blocksPerGrid, kThreadsPerBlock, 0, stream>>>(output_data, functor, N);
}
} // namespace cuda
} // namespace onnxruntime

6
orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cc
View file

@ -122,8 +122,7 @@ Status SoftmaxCrossEntropyLoss<T, Tin>::ComputeInternal(OpKernelContext* ctx) co
IAllocatorUniquePtr<T> weight_data_nd = GetScratchBuffer<T>(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<const CudaT*>(weight_data),
N_D, C,
@ -241,8 +240,7 @@ Status SoftmaxCrossEntropyLossGrad<T, Tin>::ComputeInternal(OpKernelContext* ctx
IAllocatorUniquePtr<T> weight_data_nd = GetScratchBuffer<T>(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<const CudaT*>(weight_data),
N_D, C,

341
orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.cu
View file

@ -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 <typename T, typename Tin>
__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 <typename T, typename Tin, bool IsWeighted>
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 <typename T, typename Tin>
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<float>(count) / GridDim::maxThreadsPerBlock));
CUDA_LONG N_D = static_cast<CUDA_LONG>(count);
CUDA_LONG C = static_cast<CUDA_LONG>(label_depth);
CUDA_LONG II = static_cast<CUDA_LONG>(ignore_index);
_ComputeWeightsSoftmaxCrossEntropy<T, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
weight_data_nd,
label,
weight,
N_D,
C,
II);
}
template <typename T, typename TAcc, typename Tin>
__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<T, Tin, true> op(label, weight, static_cast<Tin>(label_depth),
static_cast<Tin>(ignore_index));
LaunchElementwiseKernel<T, decltype(op)>(stream, weight_data_nd, op, count);
} else {
CUDA_KERNEL_ASSERT(label_data[i] >= 0 && label_data[i] < C);
output_data[i] = static_cast<T>(static_cast<TAcc>(-log_prob_data[i * C + label_data[i]] * weight_data[i]) /
*normalize_factor_data);
OpSoftmaxCrossEntropyWeights<T, Tin, false> op(label, nullptr, static_cast<Tin>(label_depth),
static_cast<Tin>(ignore_index));
LaunchElementwiseKernel<T, decltype(op)>(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 <typename T, typename TAcc, typename Tin>
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<float>(count) / GridDim::maxThreadsPerBlock));
CUDA_LONG N_D = static_cast<CUDA_LONG>(count);
CUDA_LONG C = static_cast<CUDA_LONG>(label_depth);
CUDA_LONG II = static_cast<CUDA_LONG>(ignore_index);
_WeightedSoftmaxCrossEntropyLoss<T, TAcc, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
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<T>(static_cast<TAcc>(-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 <typename T, typename TAcc, typename Tin>
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<T, TAcc, Tin> op(log_prob, label, weight, normalize_factor,
static_cast<Tin>(label_depth), static_cast<Tin>(ignore_index));
LaunchElementwiseKernel<T, decltype(op)>(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 <typename T, typename TAcc, typename Tin, bool IsReductionNone>
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<int>(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 <typename T, typename TAcc, typename Tin>
__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<T>(static_cast<TAcc>((*dY) * weight[row]) *
(_Exp(static_cast<TAcc>(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<T>(static_cast<TAcc>((IsReductionNone ? dY_data_[row] : *dY_data_) * weight_data_[row]) *
(_Exp(static_cast<TAcc>(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 <typename T, typename TAcc, typename Tin>
__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<T>(static_cast<TAcc>(dY[row] * weight[row]) *
(_Exp(static_cast<TAcc>(log_prob[i])) - ONE_TAcc * (TAcc)(d == label[row])) /
(*normalize_factor));
}
}
template <typename T, typename TAcc, typename Tin>
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<CUDA_LONG>(count);
CUDA_LONG C = static_cast<CUDA_LONG>(label_depth);
int blocksPerGrid = (int)(ceil(static_cast<float>(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<T, TAcc, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
dY,
log_prob,
label,
weight,
normalize_factor,
output_data,
N_D,
C);
OpWeightedSoftmaxCrossEntropyLossGrad<T, TAcc, Tin, true> op(dY, log_prob, label, weight, normalize_factor,
static_cast<Tin>(label_depth));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count * label_depth);
} else {
_WeightedSoftmaxCrossEntropyLossGrad<T, TAcc, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
dY,
log_prob,
label,
weight,
normalize_factor,
output_data,
N_D,
C);
OpWeightedSoftmaxCrossEntropyLossGrad<T, TAcc, Tin, false> op(dY, log_prob, label, weight, normalize_factor,
static_cast<Tin>(label_depth));
LaunchElementwiseKernel<T, decltype(op)>(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
} // namespace onnxruntime

2
orttraining/orttraining/training_ops/cuda/loss/softmax_cross_entropy_loss_impl.h
View file

@ -36,7 +36,7 @@ void SoftmaxCrossEntropyLossGradImpl(
T* output_data);
template <typename T, typename Tin>
void ComputeWeightsSoftmaxCrossEntropyImpl(
void ComputeSoftmaxCrossEntropyWeightsImpl(
cudaStream_t stream,
const Tin* label,
const T* weight,

361
orttraining/orttraining/training_ops/cuda/loss/softmaxcrossentropy_impl.cu
View file

@ -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 <typename T>
__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 <typename T>
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<float>(count) / GridDim::maxThreadsPerBlock));
CUDA_LONG N = static_cast<CUDA_LONG>(count);
CUDA_LONG NORMALIZE_FACTOR = static_cast<CUDA_LONG>(normalize_factor);
_SoftMaxCrossEntropy<T><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
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 <typename T>
__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 <typename T>
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<float>(count) / GridDim::maxThreadsPerBlock));
CUDA_LONG N = static_cast<CUDA_LONG>(count);
CUDA_LONG NORMALIZE_FACTOR = static_cast<CUDA_LONG>(normalize_factor);
_SoftMaxCrossEntropyGrad<T><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
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 <typename T, typename Tin>
__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 <typename T>
void SoftMaxCrossEntropyImpl(cudaStream_t stream, const T* log_prob, const T* label, size_t normalize_factor,
T* output_data, size_t count) {
OpSoftmaxCrossEntropy<T> op(log_prob, label, static_cast<T>(normalize_factor));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count);
}
template <typename T, typename Tin>
__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 <typename T>
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 <typename T>
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<T> op(dY, log_prob, label, static_cast<T>(normalize_factor));
LaunchElementwiseKernel<T, decltype(op)>(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 <typename T, typename Tin, bool IsWeighted>
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 <typename T, typename Tin>
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<float>(count) / GridDim::maxThreadsPerBlock));
CUDA_LONG N = static_cast<CUDA_LONG>(count);
CUDA_LONG D = static_cast<CUDA_LONG>(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<T, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
log_prob,
label,
weight,
normalize_factor,
output_data,
N,
D);
OpSparseSoftmaxCrossEntropy<T, Tin, true> op(log_prob, label, weight, normalize_factor,
static_cast<Tin>(label_depth));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count);
} else {
_SparseSoftmaxCrossEntropy<T, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
log_prob,
label,
normalize_factor,
output_data,
N,
D);
OpSparseSoftmaxCrossEntropy<T, Tin, false> op(log_prob, label, nullptr, normalize_factor,
static_cast<Tin>(label_depth));
LaunchElementwiseKernel<T, decltype(op)>(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 <typename T, typename Tin, bool IsWeighted>
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 <typename T, typename Tin>
__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 <typename T, typename Tin>
__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 <typename T, typename Tin>
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<CUDA_LONG>(count);
CUDA_LONG D = static_cast<CUDA_LONG>(label_depth);
int blocksPerGrid = (int)(ceil(static_cast<float>(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<T, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
dY,
log_prob,
label,
weight,
normalize_factor,
output_data,
N,
D);
OpSparseSoftmaxCrossEntropyGrad<T, Tin, true> op(dY, log_prob, label, weight, normalize_factor,
fast_divmod(static_cast<int>(label_depth)));
LaunchElementwiseKernel<T, decltype(op)>(stream, output_data, op, count * label_depth);
} else {
_SparseSoftmaxCrossEntropyGrad<T, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(
dY,
log_prob,
label,
normalize_factor,
output_data,
N,
D);
OpSparseSoftmaxCrossEntropyGrad<T, Tin, false> op(dY, log_prob, label, nullptr, normalize_factor,
fast_divmod(static_cast<int>(label_depth)));
LaunchElementwiseKernel<T, decltype(op)>(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