wehub-resource-sync
378 lines
12 KiB
Plaintext
378 lines
12 KiB
Plaintext
// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "paddle/phi/kernels/row_conv_grad_kernel.h"
|
|
#include "paddle/phi/backends/gpu/gpu_device_function.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/mixed_vector.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
|
|
namespace phi {
|
|
|
|
namespace {
|
|
|
|
static inline int DivUp(int x, int y) { return (x + y - 1) / y; }
|
|
|
|
// Compute input gradient (shared memory version, for small future_context)
|
|
// TODO(large-tensor): to support large tensor
|
|
template <typename T>
|
|
__global__ void RowConvGradInputSharedMemory(const T *dout,
|
|
const T *wt,
|
|
int num_sequence,
|
|
int input_dim,
|
|
int future_context,
|
|
const size_t *batch_indices,
|
|
T *din) {
|
|
int blx = blockDim.x;
|
|
int bly = blockDim.y;
|
|
int thx = threadIdx.x;
|
|
int thy = threadIdx.y;
|
|
int d = blockIdx.x * blx + thx; // index along input dim
|
|
|
|
extern __shared__ T mem[];
|
|
T *sw = mem;
|
|
if (thy < future_context) {
|
|
sw[thy * blx + thx] =
|
|
(d < input_dim) ? wt[thy * input_dim + d] : static_cast<T>(0);
|
|
}
|
|
__syncthreads();
|
|
|
|
int current_timesteps = 0;
|
|
for (int i = 0; i < num_sequence; i++) {
|
|
int start = batch_indices[i];
|
|
int end = batch_indices[i + 1];
|
|
current_timesteps = end - start;
|
|
|
|
for (int k = thy; k < current_timesteps; k += bly) {
|
|
T sum = 0;
|
|
for (int w = 0; (w < future_context) && ((k - w) >= 0); w++) {
|
|
sum += (d < input_dim)
|
|
? (sw[w * blx + thx] * dout[(k + start - w) * input_dim + d])
|
|
: static_cast<T>(0);
|
|
}
|
|
if (d < input_dim) {
|
|
din[(k + start) * input_dim + d] = sum;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Compute input gradient (Naive version)
|
|
// TODO(large-tensor): to support large tensor
|
|
template <typename T>
|
|
__global__ void RowConvGradInput(const T *dout,
|
|
const T *wt,
|
|
int num_sequence,
|
|
int input_dim,
|
|
int future_context,
|
|
const size_t *batch_indices,
|
|
T *din) {
|
|
int64_t d =
|
|
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
|
static_cast<int64_t>(threadIdx.x); // index along input_dim
|
|
int bly = blockDim.y;
|
|
int thy = threadIdx.y;
|
|
|
|
if (d >= input_dim) return;
|
|
int current_timesteps = 0;
|
|
for (int i = 0; i < num_sequence; i++) {
|
|
int start = batch_indices[i];
|
|
int end = batch_indices[i + 1];
|
|
current_timesteps = end - start;
|
|
|
|
for (int k = thy; k < current_timesteps; k += bly) {
|
|
T sum = 0;
|
|
for (int w = 0; (w < future_context) && ((k - w) >= 0); w++) {
|
|
sum += (wt[w * input_dim + d] * dout[(k + start - w) * input_dim + d]);
|
|
}
|
|
din[(k + start) * input_dim + d] = sum;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Compute W gradient (small future_context version)
|
|
// TODO(large-tensor): to support large tensor
|
|
template <typename T>
|
|
__global__ void RowConvGradFilterImproved(const T *in,
|
|
const T *dout,
|
|
int num_sequence,
|
|
int input_dim,
|
|
int future_context,
|
|
int block_x,
|
|
int block_y,
|
|
const size_t *batch_indices,
|
|
T *dfilter) {
|
|
int blx = blockDim.x;
|
|
int bly = blockDim.y;
|
|
int thx = threadIdx.x;
|
|
int thy = threadIdx.y;
|
|
int gx = blockIdx.x * blx;
|
|
int d = gx + thx; // index along input dim
|
|
|
|
extern __shared__ T mem[];
|
|
|
|
int xdim_sh_in = block_y;
|
|
int xdim_sh_dout = block_y;
|
|
int ydim_sh_in = block_x;
|
|
int ydim_sh_dout = block_x + future_context - 1;
|
|
int ydim_sh_dfilter = block_y;
|
|
|
|
T *sh_in = mem;
|
|
T *sh_dout = &mem[xdim_sh_in * ydim_sh_in];
|
|
T *sh_dfilter = &mem[xdim_sh_in * ydim_sh_in + xdim_sh_dout * ydim_sh_dout];
|
|
|
|
if (thy < future_context) {
|
|
sh_dfilter[thy * ydim_sh_dfilter + thx] = static_cast<T>(0);
|
|
}
|
|
__syncthreads();
|
|
|
|
// NOTE(zcd): temporary solution
|
|
unsigned mask = 0u;
|
|
CREATE_SHFL_MASK(mask, true);
|
|
|
|
for (int i = 0; i < num_sequence; i++) {
|
|
int start = batch_indices[i];
|
|
int end = batch_indices[i + 1];
|
|
int current_timesteps = end - start;
|
|
|
|
int scaled_cur_steps =
|
|
((current_timesteps + block_x - 1) / block_x) * block_x;
|
|
|
|
for (int k = thy; k < scaled_cur_steps; k += block_x) {
|
|
int pos = start + k;
|
|
sh_in[thx * ydim_sh_in + thy] =
|
|
(d < input_dim && pos < end) ? in[pos * input_dim + d] : T(0);
|
|
sh_dout[thx * ydim_sh_dout + thy + future_context - 1] =
|
|
(d < input_dim && pos < end) ? dout[pos * input_dim + d] : T(0);
|
|
__syncthreads();
|
|
|
|
if (thy < future_context - 1) {
|
|
int pos_offset = pos - future_context + 1;
|
|
sh_dout[thx * ydim_sh_dout + thy] =
|
|
(d < input_dim && pos_offset >= start)
|
|
? dout[pos_offset * input_dim + d]
|
|
: T(0);
|
|
}
|
|
__syncthreads();
|
|
|
|
for (int w = 0; w < future_context; w++) {
|
|
T val = sh_in[thy * ydim_sh_in + thx] *
|
|
sh_dout[thy * ydim_sh_dout + thx + future_context - 1 - w];
|
|
__syncthreads();
|
|
|
|
for (int offset = 16; offset > 0;
|
|
offset = offset / 2) { // blockDim.x is 32.
|
|
val += backends::gpu::CudaShuffleDownSync(mask, val, offset);
|
|
}
|
|
__syncthreads();
|
|
|
|
if (thx == 0) {
|
|
sh_dfilter[w * ydim_sh_dfilter + thy] += val;
|
|
}
|
|
__syncthreads();
|
|
}
|
|
}
|
|
}
|
|
for (int w = thy; (w < future_context) && (d < input_dim); w += bly) {
|
|
dfilter[w * input_dim + d] += sh_dfilter[w * ydim_sh_dfilter + thx];
|
|
}
|
|
}
|
|
|
|
// Compute weight(filter) gradient
|
|
// TODO(large-tensor): to support large tensor
|
|
template <typename T>
|
|
__global__ void RowConvGradFilter(const T *in,
|
|
const T *dout,
|
|
int num_sequence,
|
|
int input_dim,
|
|
int future_context,
|
|
int block_x,
|
|
int block_y,
|
|
const size_t *batch_indices,
|
|
T *dfilter) {
|
|
int blx = blockDim.x;
|
|
int thx = threadIdx.x;
|
|
int thy = threadIdx.y;
|
|
int gx = blockIdx.x * blx;
|
|
int d = gx + thx; // index along input dim
|
|
extern __shared__ T mem[];
|
|
T *sh_in = mem;
|
|
T *sh_dout = &mem[block_x * block_y];
|
|
|
|
// NOTE(zcd): temporary solution
|
|
unsigned mask = 0u;
|
|
CREATE_SHFL_MASK(mask, true);
|
|
for (int i = 0; i < num_sequence; i++) {
|
|
int start = batch_indices[i];
|
|
int end = batch_indices[i + 1];
|
|
int current_timesteps = end - start;
|
|
|
|
int scaled_cur_steps =
|
|
((current_timesteps + block_x - 1) / block_x) * block_x;
|
|
|
|
for (int k = thy; k < scaled_cur_steps; k += block_x) {
|
|
int pos = start + k;
|
|
sh_in[thx * block_y + thy] =
|
|
(d < input_dim && pos < end) ? in[pos * input_dim + d] : 0.0;
|
|
__syncthreads();
|
|
|
|
for (int w = 0; w < future_context; w++) {
|
|
sh_dout[thx * block_y + thy] =
|
|
(d < input_dim && (k - w) >= 0 && (k - w) < current_timesteps)
|
|
? dout[(pos - w) * input_dim + d]
|
|
: 0.0;
|
|
__syncthreads();
|
|
|
|
T val = sh_in[thy * block_y + thx] * sh_dout[thy * block_y + thx];
|
|
__syncthreads();
|
|
|
|
for (int offset = 16; offset > 0;
|
|
offset = offset / 2) { // blockDim.x is 32.
|
|
val += backends::gpu::CudaShuffleDownSync(mask, val, offset);
|
|
}
|
|
__syncthreads();
|
|
|
|
if (thx == 0 && (gx + thy) < input_dim) {
|
|
dfilter[w * input_dim + gx + thy] += val;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
template <typename T, typename Context>
|
|
void RowConvGradKernel(const Context &dev_ctx,
|
|
const DenseTensor &x_in,
|
|
const DenseTensor &filter_in,
|
|
const DenseTensor &out_grad,
|
|
DenseTensor *x_grad,
|
|
DenseTensor *filter_grad) {
|
|
auto *X = &x_in;
|
|
auto *Filter = &filter_in;
|
|
auto *dOut = &out_grad;
|
|
const T *in = X->data<T>();
|
|
const T *weights = Filter->data<T>();
|
|
const T *dout = dOut->data<T>();
|
|
|
|
DenseTensor *dX = x_grad;
|
|
DenseTensor *dFilter = filter_grad;
|
|
int batch_size = 0;
|
|
bool is_tensor = X->lod().empty();
|
|
if (is_tensor) {
|
|
batch_size = X->dims()[0];
|
|
} else {
|
|
batch_size = X->lod()[0].size() - 1;
|
|
}
|
|
|
|
int input_dim = 0;
|
|
Vector<size_t> batch_indices(batch_size + 1);
|
|
int64_t timesteps = X->dims()[1];
|
|
|
|
if (is_tensor) {
|
|
for (int i = 0; i < batch_size + 1; i++) {
|
|
batch_indices[i] = i * timesteps;
|
|
}
|
|
input_dim = X->dims()[2];
|
|
} else {
|
|
batch_indices = X->lod()[0];
|
|
input_dim = X->dims()[1];
|
|
}
|
|
// int input_dim = X->dims()[1];
|
|
int num_sequence = batch_indices.size() - 1;
|
|
int64_t future_context = Filter->dims()[0];
|
|
// TODO(large-tensor): CUDA kernel future_context not support int64
|
|
PADDLE_ENFORCE_LE_INT_MAX(future_context, "future_context");
|
|
int future_context_int = static_cast<int>(future_context);
|
|
|
|
MixVector<size_t> mixv_batch_indices(&batch_indices);
|
|
size_t *idx = mixv_batch_indices.CUDAMutableData(dev_ctx.GetPlace());
|
|
|
|
funcs::SetConstant<GPUContext, T> zero;
|
|
|
|
if (dFilter) {
|
|
T *dfilter = dev_ctx.template Alloc<T>(dFilter);
|
|
zero(dev_ctx, dFilter, static_cast<T>(0.0));
|
|
|
|
if (future_context_int <= 32) {
|
|
dim3 block_dim = dim3(32, 32);
|
|
dim3 grid_dim = dim3(DivUp(input_dim, block_dim.x), 1);
|
|
int block_x = block_dim.x;
|
|
int block_y = block_dim.y;
|
|
int mem_per_block =
|
|
(block_y * block_x + block_y * (block_x + future_context_int - 1) +
|
|
future_context_int * block_y) *
|
|
sizeof(T);
|
|
RowConvGradFilterImproved<T>
|
|
<<<grid_dim, block_dim, mem_per_block, dev_ctx.stream()>>>(
|
|
in,
|
|
dout,
|
|
num_sequence,
|
|
input_dim,
|
|
future_context_int,
|
|
block_x,
|
|
block_y,
|
|
idx,
|
|
dfilter);
|
|
} else {
|
|
dim3 block_dim = dim3(32, 32);
|
|
dim3 grid_dim = dim3(DivUp(input_dim, block_dim.x), 1);
|
|
int block_x = block_dim.x;
|
|
int block_y = block_dim.y;
|
|
int mem_per_block =
|
|
(block_x * block_y * 2) * sizeof(T); // For 2 arrays of size 32x32
|
|
RowConvGradFilter<T>
|
|
<<<grid_dim, block_dim, mem_per_block, dev_ctx.stream()>>>(
|
|
in,
|
|
dout,
|
|
num_sequence,
|
|
input_dim,
|
|
future_context_int,
|
|
block_x,
|
|
block_y,
|
|
idx,
|
|
dfilter);
|
|
}
|
|
}
|
|
|
|
if (dX) {
|
|
T *din = dev_ctx.template Alloc<T>(dX);
|
|
if (future_context_int <= 32) {
|
|
dim3 block_dim = dim3(32, 32);
|
|
dim3 grid_dim = dim3(DivUp(input_dim, block_dim.x), 1);
|
|
int mem_per_block = (future_context_int * block_dim.x) * sizeof(T);
|
|
RowConvGradInputSharedMemory<T>
|
|
<<<grid_dim, block_dim, mem_per_block, dev_ctx.stream()>>>(
|
|
dout,
|
|
weights,
|
|
num_sequence,
|
|
input_dim,
|
|
future_context_int,
|
|
idx,
|
|
din);
|
|
} else {
|
|
dim3 block_dim = dim3(32, 32);
|
|
dim3 grid_dim = dim3(DivUp(input_dim, block_dim.x), 1);
|
|
RowConvGradInput<T><<<grid_dim, block_dim, 0, dev_ctx.stream()>>>(
|
|
dout, weights, num_sequence, input_dim, future_context_int, idx, din);
|
|
}
|
|
}
|
|
mixv_batch_indices.CopyToCPU();
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(
|
|
row_conv_grad, GPU, ALL_LAYOUT, phi::RowConvGradKernel, float) {}
|