wehub-resource-sync
206 lines
7.1 KiB
C++
206 lines
7.1 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/p_norm_grad_kernel.h"
|
|
#include "paddle/phi/backends/xpu/enforce_xpu.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
namespace phi {
|
|
|
|
inline void GetDims(
|
|
const DDim& dim, int axis, int* m, int* t, int* n, bool asvector) {
|
|
*m = 1;
|
|
*n = 1;
|
|
*t = dim[axis];
|
|
if (asvector) {
|
|
*t = product(dim);
|
|
} else {
|
|
for (int i = 0; i < axis; ++i) {
|
|
(*m) *= dim[i];
|
|
}
|
|
for (int i = axis + 1; i < dim.size(); ++i) {
|
|
(*n) *= dim[i];
|
|
}
|
|
}
|
|
}
|
|
template <typename T, typename Context>
|
|
void PNormGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& out,
|
|
const DenseTensor& out_grad,
|
|
double porder,
|
|
int axis,
|
|
float epsilon,
|
|
bool keepdim,
|
|
bool asvector,
|
|
DenseTensor* x_grad) {
|
|
using XPUType = typename XPUTypeTrait<T>::Type;
|
|
dev_ctx.template Alloc<T>(x_grad);
|
|
if (x.numel() == 0) return;
|
|
auto xdim = x.dims();
|
|
axis = axis < 0 ? xdim.size() + axis : axis;
|
|
int m, t, n;
|
|
GetDims(xdim, axis, &m, &t, &n, asvector);
|
|
|
|
std::vector<int64_t> r_dim;
|
|
std::vector<int64_t> x_dim;
|
|
std::vector<int64_t> y_dim;
|
|
|
|
x_dim.push_back(m);
|
|
x_dim.push_back(t);
|
|
x_dim.push_back(n);
|
|
|
|
y_dim.push_back(m);
|
|
y_dim.push_back(1);
|
|
y_dim.push_back(n);
|
|
|
|
int r = 0;
|
|
if (porder == 0) {
|
|
r = xpu::constant(dev_ctx.x_context(),
|
|
reinterpret_cast<XPUType*>(x_grad->data<T>()),
|
|
m * t * n,
|
|
static_cast<T>(0));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
|
|
} else if (porder == INFINITY || porder == -INFINITY) {
|
|
xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
|
|
XPUType* x_abs = RAII_GUARD.alloc_l3_or_gm<XPUType>(m * t * n);
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(x_abs);
|
|
r = xpu::abs(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(x.data<T>()),
|
|
x_abs,
|
|
m * t * n);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "abs");
|
|
|
|
bool* dx_t = RAII_GUARD.alloc_l3_or_gm<bool>(m * t * n);
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(dx_t);
|
|
|
|
XPUType* dx_mid = RAII_GUARD.alloc_l3_or_gm<XPUType>(m * t * n);
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(dx_mid);
|
|
|
|
r = xpu::broadcast_equal<XPUType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(x_abs),
|
|
reinterpret_cast<const XPUType*>(out.data<T>()),
|
|
dx_t,
|
|
x_dim,
|
|
y_dim);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast_equal");
|
|
|
|
r = xpu::cast<bool, XPUType>(dev_ctx.x_context(), dx_t, dx_mid, m * t * n);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
|
|
|
|
XPUType* x_sign = RAII_GUARD.alloc_l3_or_gm<XPUType>(m * t * n);
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(x_sign);
|
|
r = xpu::sign(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(x.data<T>()),
|
|
x_sign,
|
|
m * t * n);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "sign");
|
|
|
|
XPUType* dx_pre_dy = x_abs;
|
|
r = xpu::mul(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(dx_mid),
|
|
reinterpret_cast<const XPUType*>(x_sign),
|
|
dx_pre_dy,
|
|
m * t * n);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "mul");
|
|
|
|
r = xpu::broadcast_mul(dev_ctx.x_context(),
|
|
dx_pre_dy,
|
|
reinterpret_cast<const XPUType*>(out_grad.data<T>()),
|
|
reinterpret_cast<XPUType*>(x_grad->data<T>()),
|
|
x_dim,
|
|
y_dim);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast_mul");
|
|
|
|
} else {
|
|
xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
|
|
XPUType* x_abs = RAII_GUARD.alloc_l3_or_gm<XPUType>(m * t * n);
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(x_abs);
|
|
r = xpu::abs(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(x.data<T>()),
|
|
x_abs,
|
|
m * t * n);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "abs");
|
|
|
|
DenseTensor porder_tensor;
|
|
DDim pdim = make_ddim({1});
|
|
porder_tensor.Resize(pdim);
|
|
dev_ctx.template Alloc<float>(&porder_tensor);
|
|
r = xpu::constant(dev_ctx.x_context(),
|
|
porder_tensor.data<float>(),
|
|
1,
|
|
static_cast<float>(porder - 1.0));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
|
|
std::vector<int64_t> p_dim(1, 1);
|
|
|
|
XPUType* x_pow = RAII_GUARD.alloc_l3_or_gm<XPUType>(m * t * n);
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(x_pow);
|
|
r = xpu::broadcast_pow(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(x_abs),
|
|
reinterpret_cast<const XPUType*>(porder_tensor.data<float>()),
|
|
x_pow,
|
|
x_dim,
|
|
p_dim);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast_pow");
|
|
|
|
XPUType* y_pow = RAII_GUARD.alloc_l3_or_gm<XPUType>(m * n);
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(y_pow);
|
|
r = xpu::broadcast_pow(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(out.data<T>()),
|
|
reinterpret_cast<const XPUType*>(porder_tensor.data<float>()),
|
|
y_pow,
|
|
y_dim,
|
|
p_dim);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast_pow");
|
|
dev_ctx.Wait();
|
|
|
|
XPUType* dx_t = x_abs;
|
|
|
|
r = xpu::broadcast_div(
|
|
dev_ctx.x_context(), x_pow, y_pow, dx_t, x_dim, y_dim);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast_div");
|
|
|
|
XPUType* x_sign = x_pow;
|
|
r = xpu::sign(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(x.data<T>()),
|
|
x_sign,
|
|
m * t * n);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "sign");
|
|
|
|
XPUType* dx_mid = RAII_GUARD.alloc_l3_or_gm<XPUType>(m * t * n);
|
|
PADDLE_ENFORCE_XDNN_NOT_NULL(dx_mid);
|
|
|
|
r = xpu::broadcast_mul(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(x_sign),
|
|
reinterpret_cast<const XPUType*>(out_grad.data<T>()),
|
|
dx_mid,
|
|
x_dim,
|
|
y_dim);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast_mul");
|
|
|
|
r = xpu::broadcast_mul(dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(dx_t),
|
|
reinterpret_cast<const XPUType*>(dx_mid),
|
|
reinterpret_cast<XPUType*>(x_grad->data<T>()),
|
|
x_dim,
|
|
x_dim);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast_mul");
|
|
}
|
|
}
|
|
} // namespace phi
|
|
PD_REGISTER_KERNEL(p_norm_grad, XPU, ALL_LAYOUT, phi::PNormGradKernel, float) {}
|