wehub-resource-sync
399 lines
15 KiB
Plaintext
399 lines
15 KiB
Plaintext
// Copyright (c) 2021 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/lerp_grad_kernel.h"
|
|
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/common/flags.h"
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
#include "paddle/phi/common/amp_type_traits.h"
|
|
#include "paddle/phi/kernels/broadcast_tensors_kernel.h"
|
|
#include "paddle/phi/kernels/empty_kernel.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/common_shape.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
|
#include "paddle/phi/kernels/funcs/reduce_function.h"
|
|
#include "paddle/phi/kernels/gpu/reduce.h"
|
|
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
|
|
|
COMMON_DECLARE_bool(use_accuracy_compatible_kernel);
|
|
|
|
namespace phi {
|
|
|
|
template <typename T>
|
|
__global__ void LerpGradKernelImpl(const T* weight,
|
|
const T* dout,
|
|
T* dx,
|
|
T* dy,
|
|
const int64_t out_size,
|
|
const int64_t x_size,
|
|
const int64_t y_size) {
|
|
using MT = typename MPTypeTrait<T>::Type;
|
|
CUDA_KERNEL_LOOP_TYPE(idx, out_size, int64_t) {
|
|
MT temp_dx = static_cast<MT>(weight[idx]) * static_cast<MT>(dout[idx]);
|
|
if (dx) {
|
|
if (idx < x_size) {
|
|
dx[idx] = static_cast<T>(static_cast<MT>(dout[idx]) - temp_dx);
|
|
}
|
|
}
|
|
if (dy) {
|
|
if (idx < y_size) {
|
|
dy[idx] = static_cast<T>(temp_dx);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void LerpGradKernelCompatibleImpl(const T* weight,
|
|
const T* dout,
|
|
T* dx,
|
|
T* dy,
|
|
const int64_t out_size,
|
|
const int64_t x_size,
|
|
const int64_t y_size) {
|
|
CUDA_KERNEL_LOOP_TYPE(idx, out_size, int64_t) {
|
|
T weight_value = weight[idx];
|
|
T remaining_weight_value = static_cast<T>(1) - weight[idx];
|
|
if (dx) {
|
|
if (idx < x_size) {
|
|
dx[idx] = remaining_weight_value * dout[idx];
|
|
}
|
|
}
|
|
if (dy) {
|
|
if (idx < y_size) {
|
|
dy[idx] = weight_value * dout[idx];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename WeightT = T>
|
|
__global__ void LerpGradScalarKernelImpl(const WeightT* weight,
|
|
const T* dout,
|
|
T* dx,
|
|
T* dy,
|
|
const int64_t out_size,
|
|
const int64_t x_size,
|
|
const int64_t y_size) {
|
|
double weight_scalar = static_cast<double>(weight[0]);
|
|
CUDA_KERNEL_LOOP_TYPE(idx, out_size, int64_t) {
|
|
double temp_dx = weight_scalar * static_cast<double>(dout[idx]);
|
|
if (dx) {
|
|
if (idx < x_size) {
|
|
dx[idx] = static_cast<T>(static_cast<double>(dout[idx]) - temp_dx);
|
|
}
|
|
}
|
|
if (dy) {
|
|
if (idx < y_size) {
|
|
dy[idx] = static_cast<T>(temp_dx);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename WeightT = T>
|
|
__global__ void LerpGradScalarKernelCompatibleImpl(const WeightT* weight,
|
|
const T* dout,
|
|
T* dx,
|
|
T* dy,
|
|
const int64_t out_size,
|
|
const int64_t x_size,
|
|
const int64_t y_size) {
|
|
T weight_scalar = static_cast<T>(weight[0]);
|
|
T remaining_weight_scalar =
|
|
static_cast<T>(1 - static_cast<double>(weight[0]));
|
|
CUDA_KERNEL_LOOP_TYPE(idx, out_size, int64_t) {
|
|
if (dx) {
|
|
if (idx < x_size) {
|
|
dx[idx] = remaining_weight_scalar * dout[idx];
|
|
}
|
|
}
|
|
if (dy) {
|
|
if (idx < y_size) {
|
|
dy[idx] = weight_scalar * dout[idx];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
bool XYNeedReduce(const DenseTensor& x,
|
|
const DenseTensor& y,
|
|
const DenseTensor& out) {
|
|
auto x_dims =
|
|
x.dims().size() ? x.dims() : make_ddim(std::vector<int64_t>(1, 1));
|
|
auto y_dims =
|
|
y.dims().size() ? y.dims() : make_ddim(std::vector<int64_t>(1, 1));
|
|
|
|
auto out_dims = out.dims();
|
|
if (out_dims.size() == 0) {
|
|
return false;
|
|
}
|
|
int x_rank = x_dims.size();
|
|
int y_rank = y_dims.size();
|
|
int out_rank = out_dims.size();
|
|
int smaller_rank = std::min(x_rank, y_rank);
|
|
if (std::max(x_rank, y_rank) < out_rank) {
|
|
return true;
|
|
}
|
|
for (int i = 1; i <= smaller_rank; ++i) {
|
|
int x_idx = x_rank - i;
|
|
int y_idx = y_rank - i;
|
|
int out_idx = out_rank - i;
|
|
if (x_dims[x_idx] != y_dims[y_idx]) {
|
|
return true;
|
|
}
|
|
if (x_dims[x_idx] == 1 && y_dims[y_idx] == 1 && out_dims[out_idx] != 1) {
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void SwitchKernel(const Context& dev_ctx,
|
|
const DenseTensor& weight,
|
|
const DenseTensor& out_grad,
|
|
const int64_t x_grad_size,
|
|
const int64_t y_grad_size,
|
|
T* x_grad_data,
|
|
T* y_grad_data) {
|
|
if (weight.numel() == 1) {
|
|
// condition when weight is a scalar
|
|
const T* out_grad_data = out_grad.data<T>();
|
|
const int64_t out_size = out_grad.numel();
|
|
const int64_t weight_size = weight.numel();
|
|
|
|
auto gpu_config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_size);
|
|
const size_t grid_size = gpu_config.GetGridSize();
|
|
const size_t block_size = gpu_config.GetBlockSize();
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_size, "grid");
|
|
const uint32_t grid = static_cast<uint32_t>(grid_size);
|
|
const uint32_t block = static_cast<uint32_t>(block_size);
|
|
|
|
if (weight.dtype() == DataType::FLOAT64) {
|
|
const double* weight_data = weight.data<double>();
|
|
if (FLAGS_use_accuracy_compatible_kernel) {
|
|
LerpGradScalarKernelCompatibleImpl<T, double>
|
|
<<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
|
|
out_grad_data,
|
|
x_grad_data,
|
|
y_grad_data,
|
|
out_size,
|
|
x_grad_size,
|
|
y_grad_size);
|
|
} else {
|
|
LerpGradScalarKernelImpl<T, double>
|
|
<<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
|
|
out_grad_data,
|
|
x_grad_data,
|
|
y_grad_data,
|
|
out_size,
|
|
x_grad_size,
|
|
y_grad_size);
|
|
}
|
|
} else {
|
|
const T* weight_data = weight.data<T>();
|
|
if (FLAGS_use_accuracy_compatible_kernel) {
|
|
LerpGradScalarKernelCompatibleImpl<T>
|
|
<<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
|
|
out_grad_data,
|
|
x_grad_data,
|
|
y_grad_data,
|
|
out_size,
|
|
x_grad_size,
|
|
y_grad_size);
|
|
} else {
|
|
LerpGradScalarKernelImpl<T>
|
|
<<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
|
|
out_grad_data,
|
|
x_grad_data,
|
|
y_grad_data,
|
|
out_size,
|
|
x_grad_size,
|
|
y_grad_size);
|
|
}
|
|
}
|
|
} else {
|
|
// broadcast weight with out_grad's dimensions
|
|
const std::vector<const DenseTensor*> in_tensors = {&weight, &out_grad};
|
|
DenseTensor b_weight = EmptyLike<T>(dev_ctx, out_grad);
|
|
DenseTensor b_out = EmptyLike<T>(dev_ctx, out_grad);
|
|
std::vector<DenseTensor*> out_tensors = {&b_weight, &b_out};
|
|
|
|
BroadcastTensorsKernel<T, Context>(dev_ctx, in_tensors, out_tensors);
|
|
|
|
const T* weight_data = b_weight.data<T>();
|
|
const T* out_grad_data = b_out.data<T>();
|
|
const int64_t out_size = out_grad.numel();
|
|
const int64_t weight_size = weight.numel();
|
|
auto gpu_config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, out_size);
|
|
const int64_t grid_size = gpu_config.GetGridSize();
|
|
const int64_t block_size = gpu_config.GetBlockSize();
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_size, "grid");
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(block_size, "block");
|
|
const uint32_t grid = static_cast<uint32_t>(grid_size);
|
|
const uint32_t block = static_cast<uint32_t>(block_size);
|
|
if (FLAGS_use_accuracy_compatible_kernel) {
|
|
LerpGradKernelCompatibleImpl<T>
|
|
<<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
|
|
out_grad_data,
|
|
x_grad_data,
|
|
y_grad_data,
|
|
out_size,
|
|
x_grad_size,
|
|
y_grad_size);
|
|
} else {
|
|
LerpGradKernelImpl<T><<<grid, block, 0, dev_ctx.stream()>>>(weight_data,
|
|
out_grad_data,
|
|
x_grad_data,
|
|
y_grad_data,
|
|
out_size,
|
|
x_grad_size,
|
|
y_grad_size);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void LerpGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& y,
|
|
const DenseTensor& weight,
|
|
const DenseTensor& out,
|
|
const DenseTensor& out_grad,
|
|
DenseTensor* x_grad,
|
|
DenseTensor* y_grad) {
|
|
if (out_grad.numel() == 0) {
|
|
if (x_grad) {
|
|
Full<T, Context>(dev_ctx, x_grad->dims(), 0, x_grad);
|
|
}
|
|
if (y_grad) {
|
|
Full<T, Context>(dev_ctx, y_grad->dims(), 0, y_grad);
|
|
}
|
|
return;
|
|
}
|
|
|
|
const int rank = out.dims().size();
|
|
PADDLE_ENFORCE_GE(
|
|
rank,
|
|
0,
|
|
common::errors::InvalidArgument(
|
|
"The number of dimensions for LerpGradOp must be "
|
|
"greater than or equal to 0, but the value received is %d.",
|
|
rank));
|
|
PADDLE_ENFORCE_LE(
|
|
rank,
|
|
6,
|
|
common::errors::InvalidArgument(
|
|
"The number of dimensions for LerpGradOp must be "
|
|
"less than or equal to 6, but the value received is %d.",
|
|
rank));
|
|
|
|
// check if x_grad and y_grad need to be reduced
|
|
// if x has a different dimension with y or weight in the middle axis, then
|
|
// they need to be broadcast and then reduced.
|
|
bool reduce_flag = XYNeedReduce(x, y, out);
|
|
if (!reduce_flag) {
|
|
int64_t x_grad_size = 0, y_grad_size = 0;
|
|
T* x_grad_data = NULL;
|
|
T* y_grad_data = NULL;
|
|
|
|
if (x_grad) {
|
|
x_grad_data = dev_ctx.template Alloc<T>(x_grad);
|
|
x_grad_size = x.numel();
|
|
}
|
|
|
|
if (y_grad) {
|
|
y_grad_data = dev_ctx.template Alloc<T>(y_grad);
|
|
y_grad_size = y.numel();
|
|
}
|
|
|
|
SwitchKernel<T, Context>(dev_ctx,
|
|
weight,
|
|
out_grad,
|
|
x_grad_size,
|
|
y_grad_size,
|
|
x_grad_data,
|
|
y_grad_data);
|
|
|
|
} else {
|
|
int64_t x_grad_size = 0, y_grad_size = 0;
|
|
DenseTensor b_xgrad = EmptyLike<T, Context>(dev_ctx, out_grad);
|
|
DenseTensor b_ygrad = EmptyLike<T, Context>(dev_ctx, out_grad);
|
|
T* x_grad_data = NULL;
|
|
T* y_grad_data = NULL;
|
|
|
|
if (x_grad) {
|
|
x_grad_data = dev_ctx.template Alloc<T>(&b_xgrad);
|
|
x_grad_size = out.numel();
|
|
}
|
|
|
|
if (y_grad) {
|
|
y_grad_data = dev_ctx.template Alloc<T>(&b_ygrad);
|
|
y_grad_size = out.numel();
|
|
}
|
|
|
|
SwitchKernel<T, Context>(dev_ctx,
|
|
weight,
|
|
out_grad,
|
|
x_grad_size,
|
|
y_grad_size,
|
|
x_grad_data,
|
|
y_grad_data);
|
|
|
|
auto zero_dim = make_ddim(std::vector<int64_t>(1, 1));
|
|
if (x_grad) {
|
|
std::vector<int> reduce_axis_x =
|
|
funcs::GetReduceDim(x_grad->dims().size() ? x_grad->dims() : zero_dim,
|
|
b_xgrad.dims(),
|
|
-1);
|
|
if (!reduce_axis_x.empty()) {
|
|
SumKernel<T, Context>(
|
|
dev_ctx, b_xgrad, reduce_axis_x, b_xgrad.dtype(), false, x_grad);
|
|
} else {
|
|
x_grad->ShareDataWith(b_xgrad);
|
|
}
|
|
}
|
|
|
|
if (y_grad) {
|
|
std::vector<int> reduce_axis_y =
|
|
funcs::GetReduceDim(y_grad->dims().size() ? y_grad->dims() : zero_dim,
|
|
b_ygrad.dims(),
|
|
-1);
|
|
if (!reduce_axis_y.empty()) {
|
|
SumKernel<T, Context>(
|
|
dev_ctx, b_ygrad, reduce_axis_y, b_ygrad.dtype(), false, y_grad);
|
|
} else {
|
|
y_grad->ShareDataWith(b_ygrad);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(lerp_grad,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::LerpGradKernel,
|
|
phi::float16,
|
|
phi::bfloat16,
|
|
float,
|
|
double) {}
|