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

softmax perf improvement pr2 - import softmax bw (#15199)

Browse source 
when dimension to do softmax is 2048, original ort code will fallback to
cudnn, while with some optimization on ort's softmax_warp_backward, we
can be faster than cudnn implementation.

the ideas to optimize softmax_warp_backward is:
1. instead of saving intermediate result in register, we just recompute
to save resource
2. save the input data in fp16 instead of fp32 to further save resource

the perf numbers:

![image](https://user-images.githubusercontent.com/43435212/227476335-ae0b61c4-cd15-40b7-b743-a956fadaedda.png)

please be noted that when dim to do softmax is less than 2048, nothing
will be changed, so only gives perf number of 2048 case.


add more perf number for smaller batch size

![image](https://user-images.githubusercontent.com/43435212/231676120-c8944b09-a664-43f3-a1e8-dfe729c6e816.png)
This commit is contained in:
zhijiang 2023-04-13 14:57:01 +08:00 committed by GitHub
parent 67ac36101c
commit 05ec22330f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 134 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

157
orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu
View file

@ -28,20 +28,119 @@
namespace onnxruntime {
namespace cuda {
template <typename input_t, typename output_t, typename acc_t, int log2_elements, bool is_log_softmax>
__global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad, const input_t* output,
int element_count, int batch_count) {
// WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method
// warp_softmax_backward_kernel.
constexpr int next_power_of_two = 1 << log2_elements;
constexpr int WARP_SIZE = (next_power_of_two < GPU_WARP_SIZE) ? next_power_of_two : GPU_WARP_SIZE;
constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
#ifdef USE_ROCM
constexpr int WARP_BATCH = 1;
#else
constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
#endif
int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
// batch_count might not be a multiple of WARP_BATCH. Check how
// many batches have to computed within this WARP.
int local_batches = batch_count - first_batch;
if (local_batches > WARP_BATCH)
local_batches = WARP_BATCH;
// there might be multiple batches per warp. compute the index within the batch
int local_idx = threadIdx.x % WARP_SIZE;
// the first element to process by the current thread
int thread_offset = first_batch * element_count + local_idx;
grad += thread_offset;
output += thread_offset;
gradInput += thread_offset;
// The nested loops over WARP_BATCH and then WARP_ITERATIONS can be simplified to one loop,
// but I think doing so would obfuscate the logic of the algorithm, thus I chose to keep
// the nested loops.
// This should have no impact on performance because the loops are unrolled anyway.
// load data from global memory
acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS];
acc_t output_reg[WARP_BATCH][WARP_ITERATIONS];
acc_t grad_output_reg[WARP_BATCH][WARP_ITERATIONS];
for (int i = 0; i < WARP_BATCH; ++i) {
int batch_element_count = (i >= local_batches) ? 0 : element_count;
for (int it = 0; it < WARP_ITERATIONS; ++it) {
int element_index = local_idx + it * WARP_SIZE;
if (element_index < batch_element_count) {
grad_reg[i][it] = grad[i * element_count + it * WARP_SIZE];
output_reg[i][it] = output[i * element_count + it * WARP_SIZE];
grad_output_reg[i][it] = grad_reg[i][it] * output_reg[i][it];
} else {
grad_reg[i][it] = acc_t(0);
output_reg[i][it] = acc_t(0);
grad_output_reg[i][it] = acc_t(0);
}
}
}
acc_t sum[WARP_BATCH];
if (!is_log_softmax) {
#pragma unroll
for (int i = 0; i < WARP_BATCH; ++i) {
sum[i] = grad_output_reg[i][0];
#pragma unroll
for (int it = 1; it < WARP_ITERATIONS; ++it) {
sum[i] += grad_output_reg[i][it];
}
}
warp_reduce<acc_t, WARP_BATCH, WARP_SIZE, Add>(sum);
}
else {
#pragma unroll
for (int i = 0; i < WARP_BATCH; ++i) {
sum[i] = grad_reg[i][0];
#pragma unroll
for (int it = 1; it < WARP_ITERATIONS; ++it) {
sum[i] += grad_reg[i][it];
}
}
warp_reduce<acc_t, WARP_BATCH, WARP_SIZE, Add>(sum);
}
// store result
#pragma unroll
for (int i = 0; i < WARP_BATCH; ++i) {
if (i >= local_batches)
break;
#pragma unroll
for (int it = 0; it < WARP_ITERATIONS; ++it) {
int element_index = local_idx + it * WARP_SIZE;
if (element_index < element_count) {
// compute gradients
if (is_log_softmax) {
gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - std::exp(output_reg[i][it]) * sum[i]);
} else {
gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - sum[i] ) * output_reg[i][it];
}
}
}
}
}
// The function "softmax_warp_backward" saves intermediate results in float32 using registers to prevent recomputing, which can be beneficial for small shapes.
// However, for larger shapes, the usage of a large register resource can lead to low CUDA warp occupancy and poor performance.
// In contrast, "softmax_warp_backward_register_efficicent" recomputes intermediate results instead of saving them and also saves the inputs in float16 format to further reduce register usage.
// TODO: If the dimension to do softmax is greater than 2048, saving the input into shared memory can further reduce register usage.
template <typename input_t, typename output_t, typename acc_t, int log2_elements, bool is_log_softmax>
__global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad, const input_t* output,
__global__ void softmax_warp_backward_register_efficicent(output_t* gradInput, const input_t* grad, const input_t* output,
int element_count, int batch_count) {
// WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method
// warp_softmax_backward_kernel.
constexpr int next_power_of_two = 1 << log2_elements;
constexpr int WARP_SIZE = (next_power_of_two < GPU_WARP_SIZE) ? next_power_of_two : GPU_WARP_SIZE;
constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
#ifdef USE_ROCM
constexpr int WARP_BATCH = 1;
#else
constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
#endif
int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
@ -66,9 +165,8 @@ __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad,
// This should have no impact on performance because the loops are unrolled anyway.
// load data from global memory
acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS];
acc_t output_reg[WARP_BATCH][WARP_ITERATIONS];
acc_t grad_output_reg[WARP_BATCH][WARP_ITERATIONS];
input_t grad_reg[WARP_BATCH][WARP_ITERATIONS];
input_t output_reg[WARP_BATCH][WARP_ITERATIONS];
for (int i = 0; i < WARP_BATCH; ++i) {
int batch_element_count = (i >= local_batches) ? 0 : element_count;
for (int it = 0; it < WARP_ITERATIONS; ++it) {
@ -76,11 +174,9 @@ __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad,
if (element_index < batch_element_count) {
grad_reg[i][it] = grad[i * element_count + it * WARP_SIZE];
output_reg[i][it] = output[i * element_count + it * WARP_SIZE];
grad_output_reg[i][it] = grad_reg[i][it] * output_reg[i][it];
} else {
grad_reg[i][it] = acc_t(0);
output_reg[i][it] = acc_t(0);
grad_output_reg[i][it] = acc_t(0);
grad_reg[i][it] = input_t(0);
output_reg[i][it] = input_t(0);
}
}
}
@ -89,10 +185,10 @@ __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad,
if (!is_log_softmax) {
#pragma unroll
for (int i = 0; i < WARP_BATCH; ++i) {
sum[i] = grad_output_reg[i][0];
sum[i] = (acc_t)(grad_reg[i][0]) * (acc_t)(output_reg[i][0]);
#pragma unroll
for (int it = 1; it < WARP_ITERATIONS; ++it) {
sum[i] += grad_output_reg[i][it];
sum[i] += (acc_t)(grad_reg[i][it]) * (acc_t)(output_reg[i][it]);
}
}
warp_reduce<acc_t, WARP_BATCH, WARP_SIZE, Add>(sum);
@ -100,10 +196,10 @@ __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad,
else {
#pragma unroll
for (int i = 0; i < WARP_BATCH; ++i) {
sum[i] = grad_reg[i][0];
sum[i] = (acc_t)grad_reg[i][0];
#pragma unroll
for (int it = 1; it < WARP_ITERATIONS; ++it) {
sum[i] += grad_reg[i][it];
sum[i] += (acc_t)grad_reg[i][it];
}
}
warp_reduce<acc_t, WARP_BATCH, WARP_SIZE, Add>(sum);
@ -120,9 +216,9 @@ __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad,
if (element_index < element_count) {
// compute gradients
if (is_log_softmax) {
gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - std::exp(output_reg[i][it]) * sum[i]);
gradInput[i * element_count + it * WARP_SIZE] = ((acc_t)(grad_reg[i][it]) - std::exp((acc_t)output_reg[i][it]) * sum[i]);
} else {
gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - sum[i] ) * output_reg[i][it];
gradInput[i * element_count + it * WARP_SIZE] = ((acc_t)grad_reg[i][it] - sum[i] ) * (acc_t)output_reg[i][it];
}
}
}
@ -133,7 +229,11 @@ template <typename T>
Status SoftmaxGradImpl(cudaStream_t stream, cudnnHandle_t cudnn_handle, T* input_grad, const T* output_grad,
const T* softmax_output, int element_count, int batch_count, bool is_log_softmax) {
if (element_count == 0) return Status::OK();
#ifdef USE_ROCM
if (element_count <= 1024 && element_count * sizeof(T) <= 4096) {
#else
if (element_count <= 2048 && element_count * sizeof(T) <= 4096) {
#endif
typedef AccumulationType_t<T> AccT;
int log2_elements = log2_ceil(element_count);
const int next_power_of_two = 1 << log2_elements;
@ -155,17 +255,26 @@ Status SoftmaxGradImpl(cudaStream_t stream, cudnnHandle_t cudnn_handle, T* input
int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
dim3 threads(warp_size, warps_per_block, 1);
// Launch code would be more elegant if C++ supported FOR CONSTEXPR
constexpr int start_to_use_register_efficient_func = 11;
switch (log2_elements) {
#define LAUNCH_KERNEL(log2_elements_value, kernel_name) \
if (is_log_softmax) { \
kernel_name<T, T, AccT, log2_elements_value, true> \
<<<blocks, threads, 0, stream>>>(input_grad, output_grad, softmax_output, element_count, batch_count); \
} else { \
kernel_name<T, T, AccT, log2_elements_value, false> \
<<<blocks, threads, 0, stream>>>(input_grad, output_grad, softmax_output, element_count, batch_count); \
}
#define CASE_LOG2_ELEMENTS(log2_elements_value) \
case log2_elements_value: { \
if (is_log_softmax) { \
softmax_warp_backward<T, T, AccT, log2_elements_value, true> \
<<<blocks, threads, 0, stream>>>(input_grad, output_grad, softmax_output, element_count, batch_count); \
if (log2_elements_value < start_to_use_register_efficient_func) { \
LAUNCH_KERNEL(log2_elements_value, softmax_warp_backward); \
} else { \
softmax_warp_backward<T, T, AccT, log2_elements_value, false> \
<<<blocks, threads, 0, stream>>>(input_grad, output_grad, softmax_output, element_count, batch_count); \
LAUNCH_KERNEL(log2_elements_value, softmax_warp_backward_register_efficicent); \
} \
} break
CASE_LOG2_ELEMENTS(0); // 1
CASE_LOG2_ELEMENTS(1); // 2
CASE_LOG2_ELEMENTS(2); // 4
@ -177,6 +286,8 @@ Status SoftmaxGradImpl(cudaStream_t stream, cudnnHandle_t cudnn_handle, T* input
CASE_LOG2_ELEMENTS(8); // 256
CASE_LOG2_ELEMENTS(9); // 512
CASE_LOG2_ELEMENTS(10); // 1024
CASE_LOG2_ELEMENTS(11); // 2048, start to use softmax_warp_backward_register_efficicent, decided by value of start_to_use_register_efficient_func
#undef LAUNCH_KERNEL
#undef CASE_LOG2_ELEMENTS
}
return Status::OK();