Files
2026-07-13 12:40:42 +08:00

632 lines
21 KiB
Plaintext

// Copyright (c) 2022 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/elementwise_grad_kernel.h"
#include "paddle/phi/kernels/elementwise_add_grad_kernel.h"
#include "paddle/phi/kernels/elementwise_divide_grad_kernel.h"
#include "paddle/phi/kernels/elementwise_multiply_grad_kernel.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/tensor_utils.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
#include "paddle/phi/kernels/gpu/elementwise_grad.h"
#include "paddle/phi/kernels/impl/elementwise_grad_kernel_impl.h"
namespace phi {
template <typename T, typename Context>
void SubtractGradKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& dout,
int axis,
DenseTensor* dx,
DenseTensor* dy) {
// skip out
auto* out = &dout;
if (dout.numel() == 0) {
if (dx) {
dev_ctx.template Alloc<T>(dx);
if (dx->numel() != 0) {
Full<T, Context>(dev_ctx, dx->dims(), 0, dx);
}
}
if (dy) {
dev_ctx.template Alloc<T>(dy);
if (dy->numel() != 0) {
Full<T, Context>(dev_ctx, dy->dims(), 0, dy);
}
}
return;
}
if (dx != nullptr && dy != nullptr && (dx->dims() == dy->dims())) {
elementwise_sub_grad<T>(dev_ctx, x, y, *out, dout, dx, dy);
} else {
default_elementwise_sub_grad<T>(dev_ctx, x, y, *out, dout, dx, dy, axis);
}
}
template <typename T, typename Context>
void SubtractDoubleGradKernel(const Context& dev_ctx,
const DenseTensor& y,
const DenseTensor& dout,
const optional<DenseTensor>& ddx,
const optional<DenseTensor>& ddy,
int axis,
DenseTensor* ddout) {
SubtractDoubleGradImpl<T>(dev_ctx, y, ddx, ddy, dout, axis, ddout);
}
template <typename T, typename Context>
void MultiplyGradKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& dout,
int axis,
DenseTensor* dx,
DenseTensor* dy) {
funcs::ElementwiseGradPreProcess(dout, dx);
ElementwiseMulGrad<T>(dev_ctx, x, y, dout, dx, dy, axis);
}
template <typename T, typename Context>
void DivideGradKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& out,
const DenseTensor& dout,
int axis,
DenseTensor* dx,
DenseTensor* dy) {
const auto place = dev_ctx.GetPlace();
if (dx != nullptr && dy != nullptr) {
std::vector<const DenseTensor*> ins = {&dout, &x, &y};
GetGradXAndYOut<T>(dev_ctx,
place,
axis,
ins,
dout,
dx,
dy,
funcs::DivGradXYFunctor<T, T>());
} else if (dx != nullptr && dy == nullptr) {
std::vector<const DenseTensor*> ins = {&dout, &y};
GetGradXOrYOut<T>(
dev_ctx, place, axis, ins, dout, dx, funcs::DivGradXFunctor<T>());
} else if (dy != nullptr && dx == nullptr) {
std::vector<const DenseTensor*> ins = {&dout, &x, &y};
GetGradXOrYOut<T>(
dev_ctx, place, axis, ins, dout, dy, funcs::DivGradYFunctor<T>());
}
}
template <typename T>
void MixedPrecisionAddGradFunc(const GPUContext& dev_ctx,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& out,
const DenseTensor& dout,
DenseTensor* dx,
DenseTensor* dy,
int axis = -1) {
const auto& x_dtype = x.dtype();
const auto& y_dtype = y.dtype();
bool no_broadcast =
(dx && dy && dx->dims() == dy->dims() && dx->dims() == dout.dims());
if (no_broadcast) {
// Dispatch to non-broadcast (elementwise) kernels
if (x_dtype == DataType::FLOAT32 && y_dtype == DataType::FLOAT16) {
ElementwiseMixedPrecisionAddGrad<phi::float16>(dev_ctx, dout, dx, dy);
} else if (x_dtype == DataType::FLOAT32 && y_dtype == DataType::BFLOAT16) {
ElementwiseMixedPrecisionAddGrad<phi::bfloat16>(dev_ctx, dout, dx, dy);
} else {
PADDLE_THROW(common::errors::Unimplemented(
"Unsupported mixed precision combination for AddGrad non-broadcast "
"path: x_dtype=%s, y_dtype=%s",
DataTypeToString(x_dtype),
DataTypeToString(y_dtype)));
}
} else {
// Dispatch to broadcast-aware kernels
if (x_dtype == DataType::FLOAT32 && y_dtype == DataType::FLOAT16) {
DefaultMixedPrecisionAddGrad<phi::float16>(
dev_ctx, x, y, dout, dx, dy, axis);
} else if (x_dtype == DataType::FLOAT32 && y_dtype == DataType::BFLOAT16) {
DefaultMixedPrecisionAddGrad<phi::bfloat16>(
dev_ctx, x, y, dout, dx, dy, axis);
} else {
PADDLE_THROW(common::errors::Unimplemented(
"Unsupported mixed precision combination for AddGrad broadcast path: "
"x_dtype=%s, y_dtype=%s",
DataTypeToString(x_dtype),
DataTypeToString(y_dtype)));
}
}
}
template <typename T>
void AddGradFunc(const GPUContext& dev_ctx,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& out,
const DenseTensor& dout,
DenseTensor* dx,
DenseTensor* dy,
int axis = -1) {
if (dx != nullptr && dy != nullptr && (dx->dims() == dy->dims())) {
ElementwiseAddGrad<T>(dev_ctx, x, y, out, dout, dx, dy);
} else {
DefaultElementwiseAddGrad<T>(dev_ctx, x, y, out, dout, dx, dy, axis);
}
}
template <typename T, typename Context>
void AddGradKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& dout,
int axis,
DenseTensor* dx,
DenseTensor* dy) {
#ifdef PADDLE_WITH_CUDA
if (x.dtype() == DataType::FLOAT32 &&
(y.dtype() == DataType::FLOAT16 || y.dtype() == DataType::BFLOAT16)) {
MixedPrecisionAddGradImpl<float>(
dev_ctx, x, y, dout, axis, dx, dy, MixedPrecisionAddGradFunc<float>);
return;
}
#endif
AddGradImpl<T>(dev_ctx, x, y, dout, axis, dx, dy, AddGradFunc<T>);
}
template <typename T, typename Context>
void AddDoubleGradKernel(const Context& dev_ctx,
const DenseTensor& y,
const DenseTensor& dout,
const optional<DenseTensor>& ddx,
const optional<DenseTensor>& ddy,
int axis,
DenseTensor* ddout) {
AddDoubleGradImpl<T>(dev_ctx, y, ddx, ddy, dout, axis, ddout);
}
template <typename T, typename Context>
void AddTripleGradKernel(const Context& dev_ctx,
const DenseTensor& ddx,
const DenseTensor& ddy,
const DenseTensor& d_ddout,
int axis,
DenseTensor* d_ddx,
DenseTensor* d_ddy) {
AddGradImpl<T>(
dev_ctx, ddx, ddy, d_ddout, axis, d_ddx, d_ddy, AddGradFunc<T>);
}
template <typename T, typename Context>
void MaximumGradKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& dout,
DenseTensor* dx,
DenseTensor* dy) {
if (dout.numel() == 0) {
if (dx) {
if (dx->numel() == 0) {
dev_ctx.template Alloc<T>(dx);
} else {
Full<T, Context>(dev_ctx, dx->dims(), 0, dx);
}
}
if (dy) {
if (dy->numel() == 0) {
dev_ctx.template Alloc<T>(dy);
} else {
Full<T, Context>(dev_ctx, dy->dims(), 0, dy);
}
}
return;
}
const auto place = dev_ctx.GetPlace();
int axis = -1;
if (dx != nullptr && dy != nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &dout};
GetGradXAndYOut<T>(dev_ctx,
place,
axis,
ins,
dout,
dx,
dy,
funcs::MaxGradXYFunctor<T, T>());
} else if (dx != nullptr && dy == nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &dout};
GetGradXOrYOut<T>(
dev_ctx, place, axis, ins, dout, dx, funcs::MaxGradXFunctor<T>());
} else if (dy != nullptr && dx == nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &dout};
GetGradXOrYOut<T>(
dev_ctx, place, axis, ins, dout, dy, funcs::MaxGradYFunctor<T>());
}
}
template <typename T, typename Context>
void MinimumGradKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& dout,
DenseTensor* dx,
DenseTensor* dy) {
if (dout.numel() == 0) {
if (dx) {
if (dx->numel() == 0) {
dev_ctx.template Alloc<T>(dx);
} else {
Full<T, Context>(dev_ctx, dx->dims(), 0, dx);
}
}
if (dy) {
if (dy->numel() == 0) {
dev_ctx.template Alloc<T>(dy);
} else {
Full<T, Context>(dev_ctx, dy->dims(), 0, dy);
}
}
return;
}
const auto place = dev_ctx.GetPlace();
int axis = -1;
if (dx != nullptr && dy != nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &dout};
GetGradXAndYOut<T>(dev_ctx,
place,
axis,
ins,
dout,
dx,
dy,
funcs::MinGradXYFunctor<T, T>());
} else if (dx != nullptr && dy == nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &dout};
GetGradXOrYOut<T>(
dev_ctx, place, axis, ins, dout, dx, funcs::MinGradXFunctor<T>());
} else if (dy != nullptr && dx == nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &dout};
GetGradXOrYOut<T>(
dev_ctx, place, axis, ins, dout, dy, funcs::MinGradYFunctor<T>());
}
}
template <typename T, typename Context>
void RemainderGradKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& dout,
DenseTensor* dx,
DenseTensor* dy) {
if (dout.numel() == 0) {
if (dx) {
if (dx->numel() == 0) {
dev_ctx.template Alloc<T>(dx);
} else {
Full<T, Context>(dev_ctx, dx->dims(), 0, dx);
}
}
if (dy) {
if (dy->numel() == 0) {
dev_ctx.template Alloc<T>(dy);
} else {
Full<T, Context>(dev_ctx, dy->dims(), 0, dy);
}
}
return;
}
const auto place = dev_ctx.GetPlace();
int axis = -1;
if (dx != nullptr && dy != nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &dout};
GetGradXAndYOut<T>(dev_ctx,
place,
axis,
ins,
dout,
dx,
dy,
funcs::RemainderGradXYFunctor<T, T>());
} else if (dx != nullptr && dy == nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &dout};
GetGradXOrYOut<T>(
dev_ctx, place, axis, ins, dout, dx, funcs::RemainderGradXFunctor<T>());
} else if (dy != nullptr && dx == nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &dout};
GetGradXOrYOut<T>(
dev_ctx, place, axis, ins, dout, dy, funcs::RemainderGradYFunctor<T>());
}
}
template <typename T, typename Context>
void CopySignGradKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& y,
const DenseTensor& out_grad,
DenseTensor* x_grad,
DenseTensor* y_grad) {
const auto place = dev_ctx.GetPlace();
int axis = -1;
if (x_grad != nullptr && y_grad != nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &out_grad};
GetGradXAndYOut<T>(dev_ctx,
place,
axis,
ins,
out_grad,
x_grad,
y_grad,
funcs::CopySignGradXYFunctor<T, T>());
} else if (x_grad != nullptr && y_grad == nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &out_grad};
GetGradXOrYOut<T>(dev_ctx,
place,
axis,
ins,
out_grad,
x_grad,
funcs::CopySignGradXFunctor<T>());
} else if (y_grad != nullptr && x_grad == nullptr) {
std::vector<const DenseTensor*> ins = {&x, &y, &out_grad};
GetGradXOrYOut<T>(dev_ctx,
place,
axis,
ins,
out_grad,
y_grad,
funcs::CopySignGradYFunctor<T>());
}
}
} // namespace phi
PD_REGISTER_KERNEL(fmax_grad,
GPU,
ALL_LAYOUT,
phi::ElementwiseFMaxGradKernel,
float,
double,
int,
phi::float16,
phi::bfloat16,
int64_t) {}
PD_REGISTER_KERNEL(fmin_grad,
GPU,
ALL_LAYOUT,
phi::ElementwiseFMinGradKernel,
float,
double,
int,
phi::float16,
phi::bfloat16,
int64_t) {}
PD_REGISTER_KERNEL(maximum_grad,
GPU,
ALL_LAYOUT,
phi::MaximumGradKernel,
float,
double,
int,
int64_t,
phi::float16,
phi::bfloat16) {}
PD_REGISTER_KERNEL(minimum_grad,
GPU,
ALL_LAYOUT,
phi::MinimumGradKernel,
float,
double,
int,
int64_t,
phi::float16,
phi::bfloat16) {}
PD_REGISTER_KERNEL(remainder_grad,
GPU,
ALL_LAYOUT,
phi::RemainderGradKernel,
float,
double,
int,
int64_t,
phi::float16,
phi::bfloat16) {}
PD_REGISTER_KERNEL(heaviside_grad,
GPU,
ALL_LAYOUT,
phi::HeavisideGradKernel,
float,
double,
int,
phi::float16,
phi::bfloat16,
int64_t) {}
PD_REGISTER_KERNEL(elementwise_pow_grad,
GPU,
ALL_LAYOUT,
phi::ElementwisePowGradKernel,
float,
double,
int,
phi::float16,
phi::bfloat16,
int64_t,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(add_grad,
GPU,
ALL_LAYOUT,
phi::AddGradKernel,
float,
double,
int,
int64_t,
phi::float16,
phi::bfloat16,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(add_double_grad,
GPU,
ALL_LAYOUT,
phi::AddDoubleGradKernel,
float,
double,
int,
int64_t,
phi::float16,
phi::bfloat16,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(add_triple_grad,
GPU,
ALL_LAYOUT,
phi::AddTripleGradKernel,
float,
double,
int,
int64_t,
phi::float16,
phi::bfloat16,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(divide_grad,
GPU,
ALL_LAYOUT,
phi::DivideGradKernel,
float,
phi::float16,
phi::bfloat16,
double,
int8_t,
uint8_t,
int16_t,
int,
int64_t,
bool,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(divide_double_grad,
GPU,
ALL_LAYOUT,
phi::DivideDoubleGradKernel,
float,
phi::float16,
phi::bfloat16,
double,
int,
int64_t,
bool,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(multiply_grad,
GPU,
ALL_LAYOUT,
phi::MultiplyGradKernel,
float,
phi::float16,
double,
int,
int64_t,
bool,
phi::bfloat16,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(multiply_double_grad,
GPU,
ALL_LAYOUT,
phi::MultiplyDoubleGradKernel,
float,
phi::float16,
double,
int,
int64_t,
bool,
phi::bfloat16,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(multiply_triple_grad,
GPU,
ALL_LAYOUT,
phi::MultiplyTripleGradKernel,
float,
phi::float16,
double,
int,
int64_t,
bool,
phi::bfloat16,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(subtract_grad,
GPU,
ALL_LAYOUT,
phi::SubtractGradKernel,
float,
double,
int,
int64_t,
phi::float16,
phi::bfloat16,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(subtract_double_grad,
GPU,
ALL_LAYOUT,
phi::SubtractDoubleGradKernel,
float,
double,
int,
int64_t,
phi::float16,
phi::bfloat16,
phi::complex64,
phi::complex128) {}
PD_REGISTER_KERNEL(copysign_grad,
GPU,
ALL_LAYOUT,
phi::CopySignGradKernel,
bool,
uint8_t,
int8_t,
int16_t,
int,
int64_t,
float,
double,
phi::float16,
phi::bfloat16) {}