194 lines
7.3 KiB
C++
194 lines
7.3 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 ResNetUnitXPUKernel(const Context &dev_ctx,
|
|
const DenseTensor &x_in,
|
|
const DenseTensor &filter_x_in,
|
|
const DenseTensor &scale_x_in,
|
|
const DenseTensor &bias_x_in,
|
|
const DenseTensor &mean_x_in,
|
|
const DenseTensor &var_x_in,
|
|
const optional<DenseTensor> &z_in,
|
|
const optional<DenseTensor> &filter_z_in,
|
|
const optional<DenseTensor> &scale_z_in,
|
|
const optional<DenseTensor> &bias_z_in,
|
|
const optional<DenseTensor> &mean_z_in,
|
|
const optional<DenseTensor> &var_z_in,
|
|
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 *out,
|
|
DenseTensor *bit_mask,
|
|
DenseTensor *conv_x,
|
|
DenseTensor *saved_mean_x,
|
|
DenseTensor *saved_invstd_x,
|
|
DenseTensor *running_mean_x,
|
|
DenseTensor *running_var_x,
|
|
DenseTensor *conv_z,
|
|
DenseTensor *saved_mean_z,
|
|
DenseTensor *saved_invstd_z,
|
|
DenseTensor *running_mean_z,
|
|
DenseTensor *running_var_z) {
|
|
using XPUType = typename XPUTypeTrait<T>::Type;
|
|
|
|
bool is_nchw = (data_format == "NCHW");
|
|
// input x
|
|
const DenseTensor *input_x = &x_in;
|
|
const DenseTensor *filter_x = &filter_x_in;
|
|
const DenseTensor *scale_x = &scale_x_in;
|
|
const DenseTensor *bias_x = &bias_x_in;
|
|
|
|
// output x
|
|
DenseTensor *conv_out_x = conv_x;
|
|
|
|
DenseTensor *output = out;
|
|
|
|
// attrs
|
|
float eps = epsilon;
|
|
float momentum = momentum_in;
|
|
bool is_train = !is_test && !use_global_stats;
|
|
|
|
std::vector<const XPUType *> x_list = {
|
|
reinterpret_cast<const XPUType *>(input_x->data<T>())};
|
|
std::vector<const XPUType *> w_list = {
|
|
reinterpret_cast<const XPUType *>(filter_x->data<T>())};
|
|
std::vector<XPUType *> conv_y_list = {
|
|
reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(conv_out_x))};
|
|
|
|
std::vector<std::vector<int64_t>> x_shape_list = {
|
|
vectorize<int64_t>(input_x->dims())};
|
|
|
|
auto filter_x_shape = vectorize<int64_t>(filter_x->dims());
|
|
std::vector<int64_t> ksize = {filter_x_shape[2], filter_x_shape[3]};
|
|
if (!is_nchw) {
|
|
ksize[0] = filter_x_shape[1];
|
|
ksize[1] = filter_x_shape[2];
|
|
}
|
|
std::vector<int64_t> strides = {stride, stride};
|
|
std::vector<std::vector<int64_t>> ksize_list = {ksize};
|
|
std::vector<std::vector<int64_t>> stride_list = {strides};
|
|
std::vector<int64_t> paddings = {padding, padding};
|
|
std::vector<int64_t> dilations = {dilation, dilation};
|
|
std::vector<const float *> scale_list = {scale_x->data<float>()};
|
|
std::vector<const float *> bias_list = {bias_x->data<float>()};
|
|
std::vector<float *> batch_mean_list = {
|
|
dev_ctx.template Alloc<float>(saved_mean_x)};
|
|
std::vector<float *> batch_invstd_list = {
|
|
dev_ctx.template Alloc<float>(saved_invstd_x)};
|
|
std::vector<float *> global_mean_list = {
|
|
dev_ctx.template Alloc<float>(running_mean_x)};
|
|
std::vector<float *> global_var_list = {
|
|
dev_ctx.template Alloc<float>(running_var_x)};
|
|
|
|
std::vector<const float *> x_maxlist = {nullptr};
|
|
std::vector<const float *> w_maxlist = {nullptr};
|
|
if (has_shortcut) {
|
|
// input z
|
|
const DenseTensor *input_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 *bias_z = bias_z_in.get_ptr();
|
|
|
|
DenseTensor *conv_out_z = conv_z;
|
|
|
|
x_list.push_back(reinterpret_cast<const XPUType *>(input_z->data<T>()));
|
|
w_list.push_back(reinterpret_cast<const XPUType *>(filter_z->data<T>()));
|
|
conv_y_list.push_back(
|
|
reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(conv_out_z)));
|
|
|
|
x_shape_list.push_back(vectorize<int64_t>(input_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});
|
|
scale_list.push_back(scale_z->data<float>());
|
|
bias_list.push_back(bias_z->data<float>());
|
|
batch_mean_list.push_back(dev_ctx.template Alloc<float>(saved_mean_z));
|
|
batch_invstd_list.push_back(dev_ctx.template Alloc<float>(saved_invstd_z));
|
|
global_mean_list.push_back(dev_ctx.template Alloc<float>(running_mean_z));
|
|
global_var_list.push_back(dev_ctx.template Alloc<float>(running_var_z));
|
|
x_maxlist.push_back(nullptr);
|
|
w_maxlist.push_back(nullptr);
|
|
} else {
|
|
if (fuse_add) {
|
|
const DenseTensor *input_z = z_in.get_ptr();
|
|
auto input_z_shape = vectorize<int64_t>(input_z->dims());
|
|
x_list.push_back(reinterpret_cast<const XPUType *>(input_z->data<T>()));
|
|
x_shape_list.push_back(input_z_shape);
|
|
x_maxlist.push_back(nullptr);
|
|
}
|
|
}
|
|
int r = xpu::resnet_unit_fusion<XPUType, XPUType, XPUType, int16_t>(
|
|
dev_ctx.x_context(),
|
|
x_list,
|
|
w_list,
|
|
conv_y_list,
|
|
reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(output)),
|
|
x_shape_list,
|
|
filter_x_shape[0],
|
|
ksize_list,
|
|
stride_list,
|
|
paddings,
|
|
dilations,
|
|
group,
|
|
eps,
|
|
momentum,
|
|
x_maxlist,
|
|
w_maxlist,
|
|
scale_list,
|
|
bias_list,
|
|
batch_mean_list,
|
|
batch_invstd_list,
|
|
global_mean_list,
|
|
global_var_list,
|
|
xpu::Activation_t::RELU,
|
|
is_nchw,
|
|
has_shortcut,
|
|
fuse_add,
|
|
is_train);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "resnet_unit_fusion");
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(resnet_unit,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::ResNetUnitXPUKernel,
|
|
phi::float16,
|
|
float) {}
|