wehub-resource-sync
202 lines
8.0 KiB
C++
202 lines
8.0 KiB
C++
// 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/backends/xpu/enforce_xpu.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/utils/optional.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void ResNetUnitGradXPUKernel(const Context &dev_ctx,
|
|
const DenseTensor &x_in,
|
|
const DenseTensor &filter_x_in,
|
|
const DenseTensor &conv_x_in,
|
|
const DenseTensor &scale_x_in,
|
|
const DenseTensor &bias_x_in,
|
|
const DenseTensor &saved_mean_x_in,
|
|
const DenseTensor &saved_invstd_x_in,
|
|
const optional<DenseTensor> &z_in,
|
|
const optional<DenseTensor> &filter_z_in,
|
|
const optional<DenseTensor> &conv_z_in,
|
|
const optional<DenseTensor> &scale_z_in,
|
|
const optional<DenseTensor> &bias_z_in,
|
|
const optional<DenseTensor> &saved_mean_z_in,
|
|
const optional<DenseTensor> &saved_invstd_z_in,
|
|
const DenseTensor &out,
|
|
const DenseTensor &bit_mask,
|
|
const DenseTensor &out_grad,
|
|
int stride,
|
|
int stride_z,
|
|
int padding,
|
|
int dilation,
|
|
int group,
|
|
float momentum_in,
|
|
float epsilon,
|
|
const std::string &data_format,
|
|
bool fuse_add,
|
|
bool has_shortcut,
|
|
bool use_global_stats,
|
|
bool is_test,
|
|
bool use_addto,
|
|
const std::string &act_type,
|
|
DenseTensor *x_grad,
|
|
DenseTensor *filter_x_grad,
|
|
DenseTensor *scale_x_grad,
|
|
DenseTensor *bias_x_grad,
|
|
DenseTensor *z_grad,
|
|
DenseTensor *filter_z_grad,
|
|
DenseTensor *scale_z_grad,
|
|
DenseTensor *bias_z_grad) {
|
|
using XPUType = typename XPUTypeTrait<T>::Type;
|
|
|
|
bool is_nchw = (data_format == "NCHW");
|
|
const DenseTensor *y_grad = &out_grad;
|
|
const DenseTensor *x = &x_in;
|
|
const DenseTensor *filter_x = &filter_x_in;
|
|
const DenseTensor *scale_x = &scale_x_in;
|
|
const DenseTensor *saved_mean_x = &saved_mean_x_in;
|
|
const DenseTensor *saved_invstd_x = &saved_invstd_x_in;
|
|
const DenseTensor *conv_out_x = &conv_x_in;
|
|
const DenseTensor *output = &out;
|
|
|
|
float eps = epsilon;
|
|
|
|
std::vector<const XPUType *> x_list = {
|
|
reinterpret_cast<const XPUType *>(x->data<T>())};
|
|
std::vector<const XPUType *> w_list = {
|
|
reinterpret_cast<const XPUType *>(filter_x->data<T>())};
|
|
std::vector<const XPUType *> conv_y_list = {
|
|
reinterpret_cast<const XPUType *>(conv_out_x->data<T>())};
|
|
std::vector<XPUType *> dx_list = {
|
|
reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(x_grad))};
|
|
std::vector<XPUType *> dw_list = {
|
|
reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(filter_x_grad))};
|
|
|
|
std::vector<std::vector<int64_t>> x_shape_list = {
|
|
vectorize<int64_t>(x->dims())};
|
|
|
|
auto filter_x_shape = vectorize<int64_t>(filter_x->dims());
|
|
std::vector<int64_t> x_ksize = {filter_x_shape[2], filter_x_shape[3]};
|
|
if (!is_nchw) {
|
|
x_ksize[0] = filter_x_shape[1];
|
|
x_ksize[1] = filter_x_shape[2];
|
|
}
|
|
std::vector<std::vector<int64_t>> ksize_list = {x_ksize};
|
|
std::vector<std::vector<int64_t>> stride_list = {{stride, stride}};
|
|
std::vector<int64_t> paddings = {padding, padding};
|
|
std::vector<int64_t> dilations = {dilation, dilation};
|
|
|
|
std::vector<const float *> x_maxlist = {nullptr};
|
|
std::vector<const float *> w_maxlist = {nullptr};
|
|
|
|
std::vector<const float *> scale_list = {scale_x->data<float>()};
|
|
std::vector<const float *> batch_mean_list = {saved_mean_x->data<float>()};
|
|
std::vector<const float *> batch_invstd_list = {
|
|
saved_invstd_x->data<float>()};
|
|
std::vector<float *> dscale_list = {
|
|
dev_ctx.template Alloc<float>(scale_x_grad)};
|
|
std::vector<float *> dbias_list = {
|
|
dev_ctx.template Alloc<float>(bias_x_grad)};
|
|
|
|
if (has_shortcut) {
|
|
// X Z
|
|
// | |
|
|
// NormConv NormConv
|
|
// | |
|
|
// BNStatsFinalize BNStatsFinalize
|
|
// \ /
|
|
// ScaleBiasAddRelu
|
|
// |
|
|
// Y
|
|
const DenseTensor *z = z_in.get_ptr();
|
|
const DenseTensor *filter_z = filter_z_in.get_ptr();
|
|
const DenseTensor *scale_z = scale_z_in.get_ptr();
|
|
const DenseTensor *saved_mean_z = saved_mean_z_in.get_ptr();
|
|
const DenseTensor *saved_invstd_z = saved_invstd_z_in.get_ptr();
|
|
const DenseTensor *conv_out_z = conv_z_in.get_ptr();
|
|
|
|
x_list.push_back(reinterpret_cast<const XPUType *>(z->data<T>()));
|
|
w_list.push_back(reinterpret_cast<const XPUType *>(filter_z->data<T>()));
|
|
conv_y_list.push_back(
|
|
reinterpret_cast<const XPUType *>(conv_out_z->data<T>()));
|
|
dx_list.push_back(
|
|
reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(z_grad)));
|
|
dw_list.push_back(
|
|
reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(filter_z_grad)));
|
|
x_shape_list.push_back(vectorize<int64_t>(z->dims()));
|
|
|
|
auto filter_z_shape = vectorize<int64_t>(filter_z->dims());
|
|
std::vector<int64_t> ksize_z = {filter_z_shape[2], filter_z_shape[3]};
|
|
if (!is_nchw) {
|
|
ksize_z[0] = filter_z_shape[1];
|
|
ksize_z[1] = filter_z_shape[2];
|
|
}
|
|
ksize_list.push_back(ksize_z);
|
|
stride_list.push_back({stride_z, stride_z});
|
|
x_maxlist.push_back(nullptr);
|
|
w_maxlist.push_back(nullptr);
|
|
|
|
scale_list.push_back(scale_z->data<float>());
|
|
batch_mean_list.push_back(saved_mean_z->data<float>());
|
|
batch_invstd_list.push_back(saved_invstd_z->data<float>());
|
|
dscale_list.push_back(dev_ctx.template Alloc<float>(scale_z_grad));
|
|
dbias_list.push_back(dev_ctx.template Alloc<float>(bias_z_grad));
|
|
} else {
|
|
if (fuse_add) {
|
|
auto z_grad_tmp = z_grad;
|
|
dx_list.push_back(
|
|
reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(z_grad_tmp)));
|
|
}
|
|
}
|
|
|
|
int r = xpu::resnet_unit_grad_fusion<XPUType, XPUType, XPUType, int16_t>(
|
|
dev_ctx.x_context(),
|
|
x_list,
|
|
w_list,
|
|
reinterpret_cast<const XPUType *>(y_grad->data<T>()),
|
|
reinterpret_cast<const XPUType *>(output->data<T>()),
|
|
conv_y_list,
|
|
dx_list,
|
|
dw_list,
|
|
x_shape_list,
|
|
filter_x_shape[0],
|
|
ksize_list,
|
|
stride_list,
|
|
paddings,
|
|
dilations,
|
|
group,
|
|
x_maxlist,
|
|
w_maxlist,
|
|
scale_list,
|
|
batch_mean_list,
|
|
batch_invstd_list,
|
|
dscale_list,
|
|
dbias_list,
|
|
xpu::Activation_t::RELU,
|
|
eps,
|
|
is_nchw,
|
|
has_shortcut,
|
|
fuse_add);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "resnet_unit_grad_fusion");
|
|
}
|
|
|
|
} // namespace phi
|
|
PD_REGISTER_KERNEL(resnet_unit_grad,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::ResNetUnitGradXPUKernel,
|
|
phi::float16,
|
|
float) {}
|