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

Optimize GatherGrad for AMD GPU (#6381)

Browse source 
* optimize gathergrad

* address comments

Co-authored-by: Weixing Zhang <wezhan@microsoft.com>
This commit is contained in:
Weixing Zhang 2021-01-29 13:50:08 -08:00 committed by GitHub
parent 76bc0e479c
commit 7f5731741d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 71 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

84
orttraining/orttraining/training_ops/rocm/tensor/gather_grad_impl.cu
View file

@ -20,7 +20,7 @@ __global__ void _Iota(
output[idx] = input[idx];
}
template <typename T, typename Tin>
template <typename T, typename Tin, int NumElementsPerThread>
__global__ void _GatherGradImpl(
const Tin* input,
const Tin* indices,
@ -32,19 +32,18 @@ __global__ void _GatherGradImpl(
int64_t stride) {
int idx = blockIdx.x * 4 + threadIdx.y;
const int SZ = 4;
if (idx < numel && (idx == 0 || input[idx] != input[idx - 1])) {
do {
for (int itr = 0; itr < param_itrs; ++itr) {
const int start_feature = threadIdx.x + blockIdx.y * blockDim.x * SZ;
const int start_feature = threadIdx.x + blockIdx.y * blockDim.x * NumElementsPerThread;
const int weight_row = itr * input_numel + ((int)input[idx]) * stride; //the offset of the input
const int grad_row = (itr * numel + ((int)indices[idx])) * stride; //the offset of the gradient
float gradient[SZ];
float weight[SZ];
float gradient[NumElementsPerThread];
float weight[NumElementsPerThread];
#pragma unroll
for (int ii = 0; ii < SZ; ii++) {
for (int ii = 0; ii < NumElementsPerThread; ii++) {
int feature_dim = start_feature + ii * GPU_WARP_SIZE;
if (feature_dim < stride) {
gradient[ii] = static_cast<float>(grad_output[grad_row + feature_dim]);
@ -53,12 +52,12 @@ __global__ void _GatherGradImpl(
}
#pragma unroll
for (int ii = 0; ii < SZ; ii++) {
for (int ii = 0; ii < NumElementsPerThread; ii++) {
weight[ii] += gradient[ii];
}
#pragma unroll
for (int ii = 0; ii < SZ; ii++) {
for (int ii = 0; ii < NumElementsPerThread; ii++) {
int feature_dim = start_feature + ii * GPU_WARP_SIZE;
if (feature_dim < stride) {
grad_weight[weight_row + feature_dim] = static_cast<T>(weight[ii]);
@ -70,6 +69,54 @@ __global__ void _GatherGradImpl(
}
}
// Special optimization for the case which the gather is on axis=0
template <typename T, typename Tin, int NumElementsPerThread>
__global__ void _GatherAxis0GradImpl(
const Tin* input,
const Tin* indices,
const T* grad_output,
T* grad_weight,
int64_t numel,
int64_t input_numel,
int64_t stride)
{
int idx = blockIdx.x * 4 + threadIdx.y;
if (idx < numel && (idx == 0 || input[idx] != input[idx - 1])) {
const int start_feature = threadIdx.x + blockIdx.y * blockDim.x * NumElementsPerThread;
const int weight_row = ((int)input[idx]) * stride; //the offset of the input
float weight[NumElementsPerThread];
for (int ii = 0; ii < NumElementsPerThread; ii++) {
int feature_dim = start_feature + ii * GPU_WARP_SIZE/4;
if (feature_dim < stride)
weight[ii] = static_cast<float>(grad_weight[weight_row + feature_dim]);
}
do {
const int grad_row = ((int)indices[idx]) * stride; //the offset of the gradient
float gradient[NumElementsPerThread];
#pragma unroll
for (int ii = 0; ii < NumElementsPerThread; ii++) {
int feature_dim = start_feature + ii * GPU_WARP_SIZE/4;
if (feature_dim < stride) {
gradient[ii] = static_cast<float>(grad_output[grad_row + feature_dim]);
weight[ii] += gradient[ii];
}
}
idx++;
} while (idx < numel && input[idx] == input[idx - 1]);
#pragma unroll
for (int ii = 0; ii < NumElementsPerThread; ii++) {
int feature_dim = start_feature + ii * GPU_WARP_SIZE/4;
if (feature_dim < stride)
grad_weight[weight_row + feature_dim] = static_cast<T>(weight[ii]);
}
}
}
template <typename T, typename Tin>
void GatherGradImpl(
const RocmKernel& rocm_kernel,
@ -113,17 +160,28 @@ void GatherGradImpl(
num_indices));
dim3 block(GPU_WARP_SIZE, 4);
dim3 grid(CeilDiv(num_indices, 4), CeilDiv(stride, 128));
hipLaunchKernelGGL(_GatherGradImpl, dim3(grid), dim3(block), 0, 0,
dim3 grid(CeilDiv(num_indices, 4), CeilDiv(stride, GridDim::maxElementsPerThread * GPU_WARP_SIZE));
if (param_itrs == 1)
{
hipLaunchKernelGGL(HIP_KERNEL_NAME(_GatherAxis0GradImpl<T, Tin, GridDim::maxElementsPerThread>), dim3(grid), dim3(block), 0, 0,
indices_data_sorted.get(),
original_indices_sorted.get(),
grad_data,
output_data,
num_indices,
num_inputs,
param_itrs,
stride);
stride);
} else {
hipLaunchKernelGGL(HIP_KERNEL_NAME(_GatherGradImpl<T, Tin, GridDim::maxElementsPerThread>), dim3(grid), dim3(block), 0, 0,
indices_data_sorted.get(),
original_indices_sorted.get(),
grad_data,
output_data,
num_indices,
num_inputs,
param_itrs,
stride);
}
}
#define SPECIALIZED_GRAD_IMPL2(T) \