wehub-resource-sync
214 lines
7.4 KiB
C++
214 lines
7.4 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/pad3d_kernel.h"
|
|
|
|
#include "paddle/phi/backends/xpu/enforce_xpu.h"
|
|
#include "paddle/phi/backends/xpu/xpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void Pad3dKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const IntArray& paddings,
|
|
const std::string& mode,
|
|
double pad_value,
|
|
const std::string& data_format,
|
|
DenseTensor* out) {
|
|
std::vector<int64_t> pads = paddings.GetData();
|
|
|
|
auto in_dims = x.dims();
|
|
const T* in_data = x.data<T>();
|
|
|
|
bool is_ncdhw = true;
|
|
if (data_format == "NCDHW") {
|
|
out->Resize({in_dims[0],
|
|
in_dims[1],
|
|
in_dims[2] + pads[4] + pads[5],
|
|
in_dims[3] + pads[2] + pads[3],
|
|
in_dims[4] + pads[0] + pads[1]});
|
|
} else {
|
|
// NDHWC
|
|
is_ncdhw = false;
|
|
out->Resize({in_dims[0],
|
|
in_dims[1] + pads[4] + pads[5],
|
|
in_dims[2] + pads[2] + pads[3],
|
|
in_dims[3] + pads[0] + pads[1],
|
|
in_dims[4]});
|
|
}
|
|
|
|
T* out_data = dev_ctx.template Alloc<T>(out);
|
|
if (x.numel() == 0) {
|
|
Full<T, Context>(dev_ctx, out->dims(), pad_value, out);
|
|
return;
|
|
}
|
|
|
|
const int64_t num = in_dims[0]; // n
|
|
int64_t channels = in_dims[1]; // c
|
|
int64_t in_depth = in_dims[2]; // xd
|
|
int64_t in_height = in_dims[3]; // xh
|
|
int64_t in_width = in_dims[4]; // xw
|
|
if (data_format == "NDHWC") {
|
|
channels = in_dims[4]; // c
|
|
in_depth = in_dims[1]; // xd
|
|
in_height = in_dims[2]; // xh
|
|
in_width = in_dims[3]; // xw
|
|
}
|
|
|
|
if (mode == "circular") {
|
|
PADDLE_THROW(common::errors::External(
|
|
"XPU is not support circular padding mode in pad3d"));
|
|
}
|
|
|
|
if (mode == "reflect") {
|
|
PADDLE_ENFORCE_GT(
|
|
in_depth,
|
|
pads[4],
|
|
errors::InvalidArgument("The depth of Input(X)'s dimension should be "
|
|
"greater than pad_front"
|
|
" in reflect mode"
|
|
", but received depth(%d) and pad_front(%d).",
|
|
in_depth,
|
|
pads[4]));
|
|
PADDLE_ENFORCE_GT(
|
|
in_depth,
|
|
pads[5],
|
|
errors::InvalidArgument("The depth of Input(X)'s dimension should be "
|
|
"greater than pad_back"
|
|
" in reflect mode"
|
|
", but received depth(%d) and pad_back(%d).",
|
|
in_depth,
|
|
pads[5]));
|
|
|
|
PADDLE_ENFORCE_GT(
|
|
in_height,
|
|
pads[2],
|
|
errors::InvalidArgument("The height of Input(X)'s dimension should be "
|
|
"greater than pad_top"
|
|
" in reflect mode"
|
|
", but received depth(%d) and pad_top(%d).",
|
|
in_height,
|
|
pads[2]));
|
|
PADDLE_ENFORCE_GT(
|
|
in_height,
|
|
pads[3],
|
|
errors::InvalidArgument("The height of Input(X)'s dimension should be "
|
|
"greater than pad_bottom"
|
|
" in reflect mode"
|
|
", but received depth(%d) and pad_bottom(%d).",
|
|
in_height,
|
|
pads[3]));
|
|
|
|
PADDLE_ENFORCE_GT(
|
|
in_width,
|
|
pads[0],
|
|
errors::InvalidArgument("The width of Input(X)'s dimension should be "
|
|
"greater than pad_left"
|
|
" in reflect mode"
|
|
", but received depth(%d) and pad_left(%d).",
|
|
in_width,
|
|
pads[0]));
|
|
PADDLE_ENFORCE_GT(
|
|
in_width,
|
|
pads[1],
|
|
errors::InvalidArgument("The width of Input(X)'s dimension should be "
|
|
"greater than pad_right"
|
|
" in reflect mode"
|
|
", but received depth(%d) and pad_right(%d).",
|
|
in_width,
|
|
pads[1]));
|
|
} else if (mode == "replicate") {
|
|
PADDLE_ENFORCE_NE(in_depth * in_height * in_width,
|
|
0,
|
|
errors::InvalidArgument(
|
|
"The input tensor size can not be 0 for circular "
|
|
"or replicate padding mode."));
|
|
}
|
|
|
|
std::vector<int64_t> pads_xpu(6);
|
|
pads_xpu[0] = pads[4]; // pf
|
|
pads_xpu[1] = pads[5]; // pb
|
|
pads_xpu[2] = pads[2]; // pt
|
|
pads_xpu[3] = pads[3]; // pd
|
|
pads_xpu[4] = pads[0]; // pl
|
|
pads_xpu[5] = pads[1]; // pr
|
|
|
|
using XPUType = typename XPUTypeTrait<T>::Type;
|
|
using XPUTypeFP16 = typename XPUTypeTrait<phi::float16>::Type;
|
|
using XPUTypeBF16 = typename XPUTypeTrait<phi::bfloat16>::Type;
|
|
// Because the xpu api do not support pad3d with bf16 type, we use fp16
|
|
// temporarily. This would not cause problem because it is a memcpy-only
|
|
// operator.
|
|
using XPURealType = std::
|
|
conditional_t<std::is_same_v<XPUType, XPUTypeBF16>, XPUTypeFP16, XPUType>;
|
|
|
|
if (mode == "reflect") {
|
|
int r = xpu::reflection_pad3d<XPURealType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPURealType*>(in_data),
|
|
reinterpret_cast<XPURealType*>(out_data),
|
|
num,
|
|
channels,
|
|
in_depth,
|
|
in_height,
|
|
in_width,
|
|
pads_xpu,
|
|
is_ncdhw);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "reflection_pad3d");
|
|
} else if (mode == "replicate") {
|
|
int r = xpu::replication_pad3d<XPURealType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPURealType*>(in_data),
|
|
reinterpret_cast<XPURealType*>(out_data),
|
|
num,
|
|
channels,
|
|
in_depth,
|
|
in_height,
|
|
in_width,
|
|
pads_xpu,
|
|
is_ncdhw);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "replication_pad3d");
|
|
} else if (mode == "constant") {
|
|
XPUType value = static_cast<XPUType>(pad_value);
|
|
XPURealType real_value;
|
|
std::memcpy(&real_value, &value, sizeof(XPURealType));
|
|
int r = xpu::constant_pad3d<XPURealType>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPURealType*>(in_data),
|
|
reinterpret_cast<XPURealType*>(out_data),
|
|
num,
|
|
channels,
|
|
in_depth,
|
|
in_height,
|
|
in_width,
|
|
pads_xpu,
|
|
real_value,
|
|
is_ncdhw);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant_pad3d");
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(pad3d,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::Pad3dKernel,
|
|
float,
|
|
phi::float16,
|
|
phi::bfloat16) {}
|