Files
paddlepaddle--paddle/paddle/phi/kernels/onednn/interpolate_kernel.cc
T
2026-07-13 12:40:42 +08:00

375 lines
13 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/interpolate_kernel.h"
#include "paddle/phi/backends/onednn/onednn_reuse.h"
#include "paddle/phi/core/compat/get_kerneltype_forvar_utils.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/interpolate_function.h"
namespace phi {
KernelKey InterpolateGetKernelTypeForVar(
const GetKernelTypeForVarContext* dev_ctx) {
const std::string& var_name = dev_ctx->GetVarName();
const DenseTensor& tensor = dev_ctx->GetTensor();
const KernelKey& expected_kernel_type = dev_ctx->GetKernelKey();
const AttributeMap& attrs = dev_ctx->GetAttrs();
// Only input require reshaping, weights and
// bias are having shape in NCHW order
if ((expected_kernel_type.layout() == DataLayout::ONEDNN) &&
(tensor.layout() != DataLayout::ONEDNN)) {
auto it = attrs.find("data_layout");
const std::string data_layout = PADDLE_GET_CONST(std::string, it->second);
auto dl = StringToDataLayout(data_layout);
// Some models may have intentionally set "AnyLayout" for pool
// op. Treat this as NCHW (default data_format value)
if (dl != DataLayout::ANY) {
return KernelKey(tensor.place(), dl, expected_kernel_type.dtype());
}
}
if (var_name == "OutSize" || var_name == "SizeTensor" ||
var_name == "Scale") {
return KernelKey(Backend::ALL_BACKEND,
expected_kernel_type.layout(),
expected_kernel_type.dtype());
}
return KernelKey(
tensor.place(), tensor.layout(), expected_kernel_type.dtype());
}
namespace funcs {
template <typename T = float>
class InterpolateOneDNNHandler
: public OneDNNHandlerNoCachingT<T, dnnl::resampling_forward> {
public:
InterpolateOneDNNHandler(const dnnl::algorithm algo,
const dnnl::engine engine,
Place cpu_place,
const DenseTensor* x,
DenseTensor* out)
: OneDNNHandlerNoCachingT<T, dnnl::resampling_forward>(engine,
cpu_place) {
const auto dst_tz = vectorize(out->dims());
const auto dst_md = dnnl::memory::desc(
dst_tz, OneDNNGetDataType<T>(), OneDNNMemoryFormat::any);
this->AcquireForwardPrimitiveDescriptor(
dnnl::prop_kind::forward_inference, algo, x->mem_desc(), dst_md);
}
};
} // namespace funcs
std::vector<int> ComputeOutputShape(
const DenseTensor* x,
const optional<DenseTensor>& out_size,
const optional<std::vector<const DenseTensor*>>& size_tensor,
const optional<DenseTensor>& scale_tensor,
const std::string& data_layout,
int out_d,
int out_h,
int out_w,
const std::vector<double>& scale_attr) {
const auto& in_dims = x->dims();
const DDim in_dhw_dims = slice_ddim(in_dims, 2, in_dims.size());
std::vector<int> out_dims;
out_dims.reserve(5);
if (in_dhw_dims.size() == 1) {
out_dims.push_back(out_w);
} else if (in_dhw_dims.size() == 2) {
out_dims.push_back(out_h);
out_dims.push_back(out_w);
} else if (in_dhw_dims.size() == 3) {
out_dims.push_back(out_d);
out_dims.push_back(out_h);
out_dims.push_back(out_w);
}
if (size_tensor && !size_tensor.get().empty()) {
auto new_size = funcs::get_new_shape(size_tensor.get());
if (new_size.size() == out_dims.size()) {
out_dims = new_size;
}
} else if (out_size) {
auto out_size_data =
funcs::get_new_data_from_tensor<int>(out_size.get_ptr());
if (out_size_data.size() == out_dims.size()) {
out_dims = out_size_data;
}
} else {
std::vector<float> scale;
scale.reserve(3);
if (scale_tensor) {
auto scale_data =
funcs::get_new_data_from_tensor<float>(scale_tensor.get_ptr());
scale.resize(3, scale_data[0]);
std::copy(scale_data.begin(), scale_data.end(), scale.begin());
} else {
if (!scale_attr.empty()) {
scale.resize(3, scale_attr[0]);
std::copy(scale_attr.begin(), scale_attr.end(), scale.begin());
}
}
if (scale.size() == 3 && scale[0] > 0.0f && scale[1] > 0.0f &&
scale[2] > 0.0f) {
int j = 0;
std::vector<int64_t> in_dhw_vec = vectorize(in_dhw_dims);
std::transform(
in_dhw_vec.begin(),
in_dhw_vec.end(),
out_dims.begin(),
[&](int64_t i) -> int { return static_cast<int>(i * scale[j++]); });
}
}
PADDLE_ENFORCE_GT(
std::all_of(
out_dims.begin(), out_dims.end(), [](int i) { return i > 0; }),
0,
errors::InvalidArgument("out_d, out_h, out_w of Op(interpolate) "
"should be greater than 0."));
const std::vector<int64_t> nc_dims = {in_dims[0], in_dims[1]};
out_dims.insert(out_dims.begin(), nc_dims.begin(), nc_dims.end());
return out_dims;
}
template <typename T, typename Context>
void InterpolateKernel(
const Context& dev_ctx,
const DenseTensor& x,
const optional<DenseTensor>& out_size,
const optional<std::vector<const DenseTensor*>>& size_tensor,
const optional<DenseTensor>& scale_tensor,
const std::string& data_layout,
int out_d,
int out_h,
int out_w,
const std::vector<double>& scale,
const std::string& interp_method,
DenseTensor* out) {
const auto& onednn_engine = dev_ctx.GetEngine();
const dnnl::algorithm algo = (interp_method == "nearest")
? dnnl::algorithm::resampling_nearest
: dnnl::algorithm::resampling_linear;
const auto out_dims_vec = ComputeOutputShape(&x,
out_size,
size_tensor,
scale_tensor,
data_layout,
out_d,
out_h,
out_w,
scale);
DDim dim_out = make_ddim(out_dims_vec);
out->Resize(dim_out);
funcs::InterpolateOneDNNHandler<T> handler(
algo, onednn_engine, dev_ctx.GetPlace(), &x, out);
auto src_memory_p = handler.AcquireSrcMemory(&x);
auto dst_memory_p = handler.AcquireDstMemory(out);
auto resampling_prim = handler.AcquireForwardPrimitive();
const std::unordered_map<int, dnnl::memory> args = {
{DNNL_ARG_SRC, *src_memory_p}, {DNNL_ARG_DST, *dst_memory_p}};
auto& astream = OneDNNContext::tls().get_stream();
resampling_prim->execute(astream, args);
astream.wait();
out->set_mem_desc(dst_memory_p->get_desc());
}
template <typename T, typename Context>
void BilinearInterpKernel(
const Context& dev_ctx,
const DenseTensor& x,
const optional<DenseTensor>& out_size,
const optional<std::vector<const DenseTensor*>>& size_tensor,
const optional<DenseTensor>& scale_tensor,
const std::string& data_layout,
int out_d,
int out_h,
int out_w,
const std::vector<double>& scale,
const std::string& interp_method,
bool align_corners UNUSED,
int align_mode UNUSED,
DenseTensor* output) {
InterpolateKernel<T, Context>(dev_ctx,
x,
out_size,
size_tensor,
scale_tensor,
data_layout,
out_d,
out_h,
out_w,
scale,
interp_method,
output);
}
template <typename T, typename Context>
void LegacyBilinearInterpKernel(
const Context& dev_ctx,
const DenseTensor& x,
const optional<DenseTensor>& out_size,
const optional<std::vector<const DenseTensor*>>& size_tensor,
const optional<DenseTensor>& scale_tensor,
const std::string& data_layout,
int out_d,
int out_h,
int out_w,
float scale,
const std::string& interp_method,
bool align_corners UNUSED,
int align_mode UNUSED,
DenseTensor* output) {
const auto& dim_x = x.dims();
std::vector<double> scale_vec;
if (scale > 0) {
for (int i = 0; i < dim_x.size() - 2; i++) {
scale_vec.push_back(scale);
}
}
InterpolateKernel<T, Context>(dev_ctx,
x,
out_size,
size_tensor,
scale_tensor,
data_layout,
out_d,
out_h,
out_w,
scale_vec,
interp_method,
output);
}
template <typename T, typename Context>
void NearestInterpKernel(
const Context& dev_ctx,
const DenseTensor& x,
const optional<DenseTensor>& out_size,
const optional<std::vector<const DenseTensor*>>& size_tensor,
const optional<DenseTensor>& scale_tensor,
const std::string& data_layout,
int out_d,
int out_h,
int out_w,
const std::vector<double>& scale,
const std::string& interp_method,
bool align_corners UNUSED,
int align_mode UNUSED,
DenseTensor* output) {
InterpolateKernel<T, Context>(dev_ctx,
x,
out_size,
size_tensor,
scale_tensor,
data_layout,
out_d,
out_h,
out_w,
scale,
interp_method,
output);
}
template <typename T, typename Context>
void LegacyNearestInterpKernel(
const Context& dev_ctx,
const DenseTensor& x,
const optional<DenseTensor>& out_size,
const optional<std::vector<const DenseTensor*>>& size_tensor,
const optional<DenseTensor>& scale_tensor,
const std::string& data_layout,
int out_d,
int out_h,
int out_w,
float scale,
const std::string& interp_method,
bool align_corners UNUSED,
int align_mode UNUSED,
DenseTensor* output) {
const auto& dim_x = x.dims();
std::vector<double> scale_vec;
if (scale > 0) {
for (int i = 0; i < dim_x.size() - 2; i++) {
scale_vec.push_back(scale);
}
}
InterpolateKernel<T, Context>(dev_ctx,
x,
out_size,
size_tensor,
scale_tensor,
data_layout,
out_d,
out_h,
out_w,
scale_vec,
interp_method,
output);
}
} // namespace phi
PD_REGISTER_KERNEL(bilinear_interp,
OneDNN,
ONEDNN,
phi::BilinearInterpKernel,
float,
phi::bfloat16,
phi::float16) {
kernel->get_kerneltype_forvar_fn_ = phi::InterpolateGetKernelTypeForVar;
}
PD_REGISTER_KERNEL(nearest_interp,
OneDNN,
ONEDNN,
phi::NearestInterpKernel,
float,
phi::bfloat16,
phi::float16,
int8_t,
uint8_t) {
kernel->get_kerneltype_forvar_fn_ = phi::InterpolateGetKernelTypeForVar;
}
PD_REGISTER_KERNEL(legacy_bilinear_interp,
OneDNN,
ONEDNN,
phi::LegacyBilinearInterpKernel,
float,
phi::bfloat16,
phi::float16) {
kernel->get_kerneltype_forvar_fn_ = phi::InterpolateGetKernelTypeForVar;
}
PD_REGISTER_KERNEL(legacy_nearest_interp,
OneDNN,
ONEDNN,
phi::LegacyNearestInterpKernel,
float,
phi::bfloat16,
phi::float16,
int8_t,
uint8_t) {
kernel->get_kerneltype_forvar_fn_ = phi::InterpolateGetKernelTypeForVar;
}