wehub-resource-sync
326 lines
11 KiB
C++
326 lines
11 KiB
C++
// 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_add_grad_kernel.h"
|
|
|
|
#include <memory>
|
|
#include <string>
|
|
|
|
#include "paddle/phi/backends/xpu/enforce_xpu.h"
|
|
#include "paddle/phi/backends/xpu/xpu_context.h"
|
|
#include "paddle/phi/backends/xpu/xpu_header.h"
|
|
#include "paddle/phi/backends/xpu/xpu_info.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/tensor_utils.h"
|
|
#include "paddle/phi/kernels/complex_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
|
|
|
namespace phi {
|
|
template <typename YType, typename Context>
|
|
void MixedPrecisionAddGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& y,
|
|
const DenseTensor& dout,
|
|
int axis,
|
|
DenseTensor* dx,
|
|
DenseTensor* dy) {
|
|
using T = float;
|
|
using XPUType = typename XPUTypeTrait<T>::Type;
|
|
using XPUYType = typename XPUTypeTrait<YType>::Type;
|
|
|
|
if (dout.numel() == 0) {
|
|
if (dx) {
|
|
dev_ctx.template Alloc<T>(dx);
|
|
if (dx->numel() > 0) {
|
|
int ret =
|
|
xpu::constant<XPUType>(dev_ctx.x_context(),
|
|
reinterpret_cast<XPUType*>(dx->data<T>()),
|
|
dx->numel(),
|
|
static_cast<XPUType>(0));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "constant");
|
|
}
|
|
}
|
|
if (dy) {
|
|
dev_ctx.template Alloc<YType>(dy);
|
|
if (dy->numel() > 0) {
|
|
int ret = xpu::constant<XPUYType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<XPUYType*>(dy->data<YType>()),
|
|
dy->numel(),
|
|
static_cast<XPUYType>(0));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "constant");
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
funcs::ElementwiseGradPreProcess(dout, dx);
|
|
auto* dz = &dout;
|
|
const DDim& dz_dims = dz->dims();
|
|
const T* dz_data = dz->data<T>();
|
|
|
|
if (dx != nullptr) {
|
|
T* dx_data = dev_ctx.template Alloc<T>(dx);
|
|
if (dx->dims() == dz_dims) {
|
|
if (dx_data != dz_data) {
|
|
int ret = xpu::copy(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(dz_data),
|
|
reinterpret_cast<XPUType*>(dx_data),
|
|
dx->numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "copy");
|
|
}
|
|
} else {
|
|
// For inplace strategy, dx will be stored in addr of dz, which makes
|
|
// the result of dy wrong.
|
|
if (dx->IsSharedBufferWith(*dz)) {
|
|
dx->clear();
|
|
dx->Resize(x.dims());
|
|
dx_data = dev_ctx.template Alloc<T>(dx);
|
|
}
|
|
std::vector<int> reduce_dims =
|
|
funcs::GetReduceDim(dx->dims(), dz_dims, axis);
|
|
std::vector<int64_t> dz_vector = vectorize<int64_t>(dz_dims);
|
|
|
|
int ret = xpu::reduce_sum<XPUType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(dz_data),
|
|
reinterpret_cast<XPUType*>(dx_data),
|
|
dz_vector,
|
|
std::vector<int64_t>(reduce_dims.begin(), reduce_dims.end()));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "reduce_sum");
|
|
}
|
|
}
|
|
|
|
if (dy != nullptr) {
|
|
YType* dy_data = dev_ctx.template Alloc<YType>(dy);
|
|
if (dy->dims() == dz_dims) {
|
|
int ret = xpu::cast<XPUType, XPUYType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(dz_data),
|
|
reinterpret_cast<XPUYType*>(dy_data),
|
|
dout.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "cast");
|
|
} else {
|
|
std::vector<int> reduce_dims =
|
|
funcs::GetReduceDim(dy->dims(), dz_dims, axis);
|
|
std::vector<int64_t> dz_vector = vectorize<int64_t>(dz_dims);
|
|
|
|
DenseTensor casted_dz;
|
|
casted_dz.Resize(dz_dims);
|
|
YType* casted_dz_data = dev_ctx.template Alloc<YType>(&casted_dz);
|
|
|
|
int ret_cast = xpu::cast<XPUType, XPUYType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(dz_data),
|
|
reinterpret_cast<XPUYType*>(casted_dz_data),
|
|
dout.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret_cast, "cast");
|
|
|
|
int ret_reduce = xpu::reduce_sum<XPUYType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUYType*>(casted_dz_data),
|
|
reinterpret_cast<XPUYType*>(dy_data),
|
|
dz_vector,
|
|
std::vector<int64_t>(reduce_dims.begin(), reduce_dims.end()));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret_reduce, "reduce_sum");
|
|
}
|
|
}
|
|
}
|
|
|
|
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) {
|
|
// special case for "float32 + bfloat16", or "float32 + float16"
|
|
if (x.dtype() == DataType::FLOAT32) {
|
|
if (y.dtype() == DataType::FLOAT16) {
|
|
MixedPrecisionAddGradKernel<phi::float16>(
|
|
dev_ctx, x, y, dout, axis, dx, dy);
|
|
return;
|
|
}
|
|
if (y.dtype() == DataType::BFLOAT16) {
|
|
MixedPrecisionAddGradKernel<phi::bfloat16>(
|
|
dev_ctx, x, y, dout, axis, dx, dy);
|
|
return;
|
|
}
|
|
}
|
|
|
|
using XPUType = typename XPUTypeTrait<T>::Type;
|
|
if (dout.numel() == 0) {
|
|
if (dx) {
|
|
dev_ctx.template Alloc<T>(dx);
|
|
if (dx->numel() > 0) {
|
|
int ret =
|
|
xpu::constant<XPUType>(dev_ctx.x_context(),
|
|
reinterpret_cast<XPUType*>(dx->data<T>()),
|
|
dx->numel(),
|
|
static_cast<XPUType>(0));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "constant");
|
|
}
|
|
}
|
|
if (dy) {
|
|
dev_ctx.template Alloc<T>(dy);
|
|
if (dy->numel() > 0) {
|
|
int ret =
|
|
xpu::constant<XPUType>(dev_ctx.x_context(),
|
|
reinterpret_cast<XPUType*>(dy->data<T>()),
|
|
dy->numel(),
|
|
static_cast<XPUType>(0));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "constant");
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
funcs::ElementwiseGradPreProcess(dout, dx);
|
|
auto* dz = &dout;
|
|
const DDim& dz_dims = dz->dims();
|
|
const T* dz_data = dz->data<T>();
|
|
|
|
if (dx != nullptr) {
|
|
T* dx_data = dev_ctx.template Alloc<T>(dx);
|
|
if (dx->dims() == dz_dims) {
|
|
if (dx_data != dz_data) {
|
|
int ret = xpu::copy(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(dz_data),
|
|
reinterpret_cast<XPUType*>(dx_data),
|
|
dx->numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "copy");
|
|
}
|
|
} else {
|
|
// For inplace strategy, dx will be stored in addr of dz, which makes
|
|
// the result of dy wrong.
|
|
if (dx->IsSharedBufferWith(*dz)) {
|
|
dx->clear();
|
|
dx->Resize(x.dims());
|
|
dx_data = dev_ctx.template Alloc<T>(dx);
|
|
}
|
|
std::vector<int> reduce_dims =
|
|
funcs::GetReduceDim(dx->dims(), dz_dims, axis);
|
|
std::vector<int64_t> dz_vector = vectorize<int64_t>(dz_dims);
|
|
|
|
int ret = xpu::reduce_sum<XPUType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(dz_data),
|
|
reinterpret_cast<XPUType*>(dx_data),
|
|
dz_vector,
|
|
std::vector<int64_t>(reduce_dims.begin(), reduce_dims.end()));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "reduce_sum");
|
|
}
|
|
}
|
|
|
|
if (dy != nullptr) {
|
|
T* dy_data = dev_ctx.template Alloc<T>(dy);
|
|
if (dy->dims() == dz_dims) {
|
|
if (dy_data != dz_data) {
|
|
int ret = xpu::copy(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(dz_data),
|
|
reinterpret_cast<XPUType*>(dy_data),
|
|
dy->numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "copy");
|
|
}
|
|
} else {
|
|
std::vector<int> reduce_dims =
|
|
funcs::GetReduceDim(dy->dims(), dz_dims, axis);
|
|
std::vector<int64_t> dz_vector = vectorize<int64_t>(dz_dims);
|
|
int ret = xpu::reduce_sum<XPUType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(dz_data),
|
|
reinterpret_cast<XPUType*>(dy_data),
|
|
dz_vector,
|
|
std::vector<int64_t>(reduce_dims.begin(), reduce_dims.end()));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(ret, "reduce_sum");
|
|
}
|
|
}
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_XPU_FFT
|
|
template <>
|
|
void AddGradKernel<phi::complex64, XPUContext>(const XPUContext& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& y,
|
|
const DenseTensor& dout,
|
|
int axis,
|
|
DenseTensor* dx,
|
|
DenseTensor* dy) {
|
|
using T = phi::complex64;
|
|
const bool compute_dx = (dx != nullptr);
|
|
const bool compute_dy = (dy != nullptr);
|
|
|
|
// The current complex number implementation uses separate real/imaginary
|
|
// parts,resulting in redundant operations and performance
|
|
// penalties.Optimization should address this in future iterations.
|
|
DenseTensor dout_real = Real<T, XPUContext>(dev_ctx, dout);
|
|
DenseTensor dout_imag = Imag<T, XPUContext>(dev_ctx, dout);
|
|
|
|
if (compute_dx || compute_dy) {
|
|
DenseTensor dx_real, dx_imag, dy_real, dy_imag;
|
|
DenseTensor tmp_real, tmp_imag;
|
|
|
|
if (compute_dx) {
|
|
dx_real.Resize(dx->dims());
|
|
dx_imag.Resize(dx->dims());
|
|
}
|
|
if (compute_dy) {
|
|
dy_real.Resize(dy->dims());
|
|
dy_imag.Resize(dy->dims());
|
|
}
|
|
|
|
AddGradKernel<float, XPUContext>(dev_ctx,
|
|
tmp_real, // unused
|
|
tmp_imag, // unused
|
|
dout_real,
|
|
axis,
|
|
compute_dx ? &dx_real : nullptr,
|
|
compute_dy ? &dy_real : nullptr);
|
|
|
|
AddGradKernel<float, XPUContext>(dev_ctx,
|
|
tmp_real, // unused
|
|
tmp_imag, // unused
|
|
dout_imag,
|
|
axis,
|
|
compute_dx ? &dx_imag : nullptr,
|
|
compute_dy ? &dy_imag : nullptr);
|
|
|
|
if (compute_dx) {
|
|
dev_ctx.template Alloc<T>(dx);
|
|
phi::ComplexKernel<float>(dev_ctx, dx_real, dx_imag, dx);
|
|
}
|
|
if (compute_dy) {
|
|
dev_ctx.template Alloc<T>(dy);
|
|
phi::ComplexKernel<float>(dev_ctx, dy_real, dy_imag, dy);
|
|
}
|
|
}
|
|
}
|
|
#endif
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(add_grad,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::AddGradKernel,
|
|
phi::float16,
|
|
phi::bfloat16,
|
|
#ifdef PADDLE_WITH_XPU_FFT
|
|
phi::complex64,
|
|
#endif
|
|
float,
|
|
int,
|
|
int64_t) {
|
|
}
|