Files
paddlepaddle--paddle/paddle/fluid/inference/tensorrt/dynamic_shape_infermeta.cc
T
2026-07-13 12:40:42 +08:00

957 lines
34 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/fluid/inference/tensorrt/dynamic_shape_infermeta_factory.h"
#include "paddle/fluid/inference/tensorrt/helper.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/phi/kernels/funcs/unfold_functor.h"
namespace paddle::inference::tensorrt {
class ExprWrapper {
public:
ExprWrapper() = default;
ExprWrapper(const nvinfer1::IDimensionExpr* expr,
nvinfer1::IExprBuilder* expr_builder) {
this->expr = expr;
this->expr_builder = expr_builder;
}
ExprWrapper(int value, nvinfer1::IExprBuilder* expr_builder) {
this->expr = expr_builder->constant(value);
this->expr_builder = expr_builder;
}
const nvinfer1::IDimensionExpr* extract_expr() const { return expr; }
public:
friend ExprWrapper BinaryOp(const ExprWrapper& a,
const ExprWrapper& b,
nvinfer1::DimensionOperation op) {
ExprWrapper result = {};
assert(a.expr);
assert(b.expr);
if (a.expr_builder) {
result.expr_builder = a.expr_builder;
}
if (b.expr_builder) {
result.expr_builder = b.expr_builder;
}
assert(result.expr_builder);
assert(result.expr);
result.expr = result.expr_builder->operation(op, *a.expr, *b.expr);
return result;
}
friend ExprWrapper BinaryOp(const ExprWrapper& a,
int b_value,
nvinfer1::DimensionOperation op) {
assert(a.expr_builder);
ExprWrapper b = {};
b.expr_builder = a.expr_builder;
b.expr = b.expr_builder->constant(b_value);
return BinaryOp(a, b, op);
}
friend ExprWrapper operator+(const ExprWrapper& a, const ExprWrapper& b) {
return BinaryOp(a, b, nvinfer1::DimensionOperation::kSUM);
}
friend ExprWrapper operator+(const ExprWrapper& a, int b_value) {
return BinaryOp(a, b_value, nvinfer1::DimensionOperation::kSUM);
}
friend ExprWrapper operator+(int a_value, const ExprWrapper& b) {
return b + a_value;
}
friend ExprWrapper operator-(const ExprWrapper& a, const ExprWrapper& b) {
return BinaryOp(a, b, nvinfer1::DimensionOperation::kSUB);
}
friend ExprWrapper operator-(const ExprWrapper& a, int b_value) {
return BinaryOp(a, b_value, nvinfer1::DimensionOperation::kSUB);
}
friend ExprWrapper operator*(const ExprWrapper& a, const ExprWrapper& b) {
return BinaryOp(a, b, nvinfer1::DimensionOperation::kPROD);
}
friend ExprWrapper operator*(const ExprWrapper& a, int b_value) {
return BinaryOp(a, b_value, nvinfer1::DimensionOperation::kPROD);
}
friend ExprWrapper operator*(int a_value, const ExprWrapper& b) {
return b * a_value;
}
friend ExprWrapper operator/(const ExprWrapper& a, const ExprWrapper& b) {
return BinaryOp(a, b, nvinfer1::DimensionOperation::kFLOOR_DIV);
}
friend ExprWrapper operator/(const ExprWrapper& a, int b_value) {
return BinaryOp(a, b_value, nvinfer1::DimensionOperation::kFLOOR_DIV);
}
friend ExprWrapper max(const ExprWrapper& a, const ExprWrapper& b) {
return BinaryOp(a, b, nvinfer1::DimensionOperation::kMAX);
}
friend ExprWrapper max(const ExprWrapper& a, int b_value) {
return BinaryOp(a, b_value, nvinfer1::DimensionOperation::kMAX);
}
public:
const nvinfer1::IDimensionExpr* expr;
nvinfer1::IExprBuilder* expr_builder;
};
static std::vector<ExprWrapper> DimsExprs2VecExprWrapper(
const nvinfer1::DimsExprs& x_dims,
nvinfer1::IExprBuilder& expr_builder // NOLINT
) {
std::vector<ExprWrapper> x_dims_wrap;
x_dims_wrap.reserve(x_dims.nbDims);
for (int i = 0; i < x_dims.nbDims; i++) {
x_dims_wrap.emplace_back(x_dims.d[i], &expr_builder);
}
return x_dims_wrap;
}
static nvinfer1::DimsExprs VecExprWrapper2DimsExprs(
const std::vector<ExprWrapper>& output_dims_wrapper) {
nvinfer1::DimsExprs output_dims = {};
output_dims.nbDims = output_dims_wrapper.size();
for (int i = 0; i < output_dims.nbDims; i++) {
output_dims.d[i] = output_dims_wrapper[i].extract_expr();
}
return output_dims;
}
nvinfer1::DimsExprs GatherNdInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
const nvinfer1::DimsExprs x_dims = inputs[0];
const int x_dims_size = inputs[0].nbDims;
const nvinfer1::DimsExprs index_dims = inputs[1];
const int index_dims_size = inputs[1].nbDims;
std::vector<const nvinfer1::IDimensionExpr*> result_dims;
// The result dims is
// Index.shape[:-1] + X.shape[Index.shape[-1]:]
result_dims.reserve(index_dims_size - 1);
for (int i = 0; i < index_dims_size - 1; ++i) {
result_dims.emplace_back(index_dims.d[i]);
}
if (index_dims.d[index_dims_size - 1]->isConstant()) {
for (int i = index_dims.d[index_dims_size - 1]->getConstantValue();
i < x_dims_size;
++i) {
result_dims.emplace_back(x_dims.d[i]);
}
}
nvinfer1::DimsExprs output = {};
output.nbDims = result_dims.size();
for (int i = 0; i < output.nbDims; i++) {
output.d[i] = result_dims[i];
}
return output;
}
nvinfer1::DimsExprs YoloBoxInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
PADDLE_ENFORCE_EQ(nb_inputs,
2,
common::errors::InvalidArgument(
"inputs of yolo_box should be equal to 2, "
"But received (%s)",
nb_inputs));
const nvinfer1::DimsExprs dim_x = inputs[0];
auto anchors = PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("anchors"));
int anchor_num = anchors.size() / 2;
// box_num = dim_x[2] * dim_x[3] * anchor_num;
const nvinfer1::IDimensionExpr* box_num = expr_builder.operation(
nvinfer1::DimensionOperation::kPROD,
*expr_builder.operation(
nvinfer1::DimensionOperation::kPROD, *dim_x.d[2], *dim_x.d[3]),
*expr_builder.constant(anchor_num));
nvinfer1::DimsExprs output = {};
output.nbDims = 3;
if (output_index == 0) {
output.d[0] = dim_x.d[0];
output.d[1] = box_num;
output.d[2] = expr_builder.constant(4);
} else {
auto class_num = PADDLE_GET_CONST(int, op_desc.GetAttr("class_num"));
output.d[0] = dim_x.d[0];
output.d[1] = box_num;
output.d[2] = expr_builder.constant(class_num);
}
return output;
}
nvinfer1::DimsExprs InstanceNormInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
nvinfer1::DimsExprs x_dims = inputs[0];
return x_dims;
}
inline const nvinfer1::IDimensionExpr* CalcOutputSize(
const nvinfer1::IDimensionExpr* input_size,
const nvinfer1::IDimensionExpr* filter_size,
const nvinfer1::IDimensionExpr* dilation,
const nvinfer1::IDimensionExpr* padding1,
const nvinfer1::IDimensionExpr* padding2,
const nvinfer1::IDimensionExpr* stride,
nvinfer1::IExprBuilder& expr_builder // NOLINT
) {
// dkernel = dilation * (filter_size - 1) + 1;
const nvinfer1::IDimensionExpr* dkernel = expr_builder.operation(
nvinfer1::DimensionOperation::kSUM,
*expr_builder.operation(
nvinfer1::DimensionOperation::kPROD,
*dilation,
*expr_builder.operation(nvinfer1::DimensionOperation::kSUB,
*filter_size,
*expr_builder.constant(1))),
*expr_builder.constant(1));
// output_size = (input_size + padding1 + padding2 - dkernel) / stride + 1;
const nvinfer1::IDimensionExpr* tmp = expr_builder.operation(
nvinfer1::DimensionOperation::kSUB,
*expr_builder.operation(
nvinfer1::DimensionOperation::kSUM,
*expr_builder.operation(
nvinfer1::DimensionOperation::kSUM, *input_size, *padding1),
*padding2),
*dkernel);
const nvinfer1::IDimensionExpr* output_size = expr_builder.operation(
nvinfer1::DimensionOperation::kSUM,
*expr_builder.operation(
nvinfer1::DimensionOperation::kFLOOR_DIV, *tmp, *stride),
*expr_builder.constant(1));
return output_size;
}
nvinfer1::DimsExprs UnfoldInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
PADDLE_ENFORCE_EQ(
nb_inputs,
1,
common::errors::InvalidArgument("inputs of unfold should be equal to 1, "
"But received (%s)",
nb_inputs));
const nvinfer1::DimsExprs in_dims = inputs[0];
std::vector<const nvinfer1::IDimensionExpr*> out_dims;
out_dims.push_back(in_dims.d[0]);
auto kernel_sizes =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("kernel_sizes"));
auto dilations =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("dilations"));
auto paddings =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("paddings"));
auto strides = PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("strides"));
// output_channels = in_dims[1] * kernel_sizes[0] * kernel_sizes[1];
const nvinfer1::IDimensionExpr* output_channels = expr_builder.operation(
nvinfer1::DimensionOperation::kPROD,
*in_dims.d[1],
*expr_builder.operation(nvinfer1::DimensionOperation::kPROD,
*expr_builder.constant(kernel_sizes[0]),
*expr_builder.constant(kernel_sizes[1])));
out_dims.push_back(output_channels);
const nvinfer1::IDimensionExpr* output_height =
CalcOutputSize(in_dims.d[2],
expr_builder.constant(kernel_sizes[0]),
expr_builder.constant(dilations[0]),
expr_builder.constant(paddings[0]),
expr_builder.constant(paddings[2]),
expr_builder.constant(strides[0]),
expr_builder);
const nvinfer1::IDimensionExpr* output_width =
CalcOutputSize(in_dims.d[3],
expr_builder.constant(kernel_sizes[1]),
expr_builder.constant(dilations[1]),
expr_builder.constant(paddings[1]),
expr_builder.constant(paddings[3]),
expr_builder.constant(strides[1]),
expr_builder);
const nvinfer1::IDimensionExpr* output_col_length = expr_builder.operation(
nvinfer1::DimensionOperation::kPROD, *output_height, *output_width);
out_dims.push_back(output_col_length);
nvinfer1::DimsExprs output = {};
output.nbDims = out_dims.size();
for (size_t i = 0; i < out_dims.size(); i++) output.d[i] = out_dims[i];
return output;
}
nvinfer1::DimsExprs ScatterNdAddInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
PADDLE_ENFORCE_EQ(nb_inputs,
3,
common::errors::InvalidArgument(
"inputs of scatter_nd_add should be equal to 3, "
"But received (%s)",
nb_inputs));
const nvinfer1::DimsExprs ref_dims = inputs[0];
return ref_dims;
}
nvinfer1::DimsExprs UnchangedInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
PADDLE_ENFORCE_EQ(nb_inputs,
1,
common::errors::InvalidArgument(
"inputs of UnchangedInferMeta should be equal to 1, "
"But received (%s)",
nb_inputs));
return inputs[0];
}
nvinfer1::DimsExprs Pad3dInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
const nvinfer1::DimsExprs x_dim = inputs[0];
nvinfer1::DimsExprs out_dims = {};
out_dims.nbDims = x_dim.nbDims;
out_dims.d[0] = x_dim.d[0];
auto paddings =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("paddings"));
auto data_format =
PADDLE_GET_CONST(std::string, op_desc.GetAttr("data_format"));
if (data_format == "NCDHW") {
out_dims.d[1] = x_dim.d[1];
} else {
out_dims.d[4] = x_dim.d[4];
}
if (data_format == "NCDHW") {
// depth
out_dims.d[2] = expr_builder.operation(
nvinfer1::DimensionOperation::kSUM,
*expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
*x_dim.d[2],
*expr_builder.constant(paddings[4])),
*expr_builder.constant(paddings[5]));
// height
out_dims.d[3] = expr_builder.operation(
nvinfer1::DimensionOperation::kSUM,
*expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
*x_dim.d[3],
*expr_builder.constant(paddings[2])),
*expr_builder.constant(paddings[3]));
// width
out_dims.d[4] = expr_builder.operation(
nvinfer1::DimensionOperation::kSUM,
*expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
*x_dim.d[4],
*expr_builder.constant(paddings[0])),
*expr_builder.constant(paddings[1]));
} else { // NDHWC
// depth
out_dims.d[1] = expr_builder.operation(
nvinfer1::DimensionOperation::kSUM,
*expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
*x_dim.d[1],
*expr_builder.constant(paddings[4])),
*expr_builder.constant(paddings[5]));
// height
out_dims.d[2] = expr_builder.operation(
nvinfer1::DimensionOperation::kSUM,
*expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
*x_dim.d[2],
*expr_builder.constant(paddings[2])),
*expr_builder.constant(paddings[3]));
// width
out_dims.d[3] = expr_builder.operation(
nvinfer1::DimensionOperation::kSUM,
*expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
*x_dim.d[3],
*expr_builder.constant(paddings[0])),
*expr_builder.constant(paddings[1]));
}
return out_dims;
}
nvinfer1::DimsExprs PNormInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
bool asvector = PADDLE_GET_CONST(bool, op_desc.GetAttr("asvector"));
bool keepdim = PADDLE_GET_CONST(bool, op_desc.GetAttr("keepdim"));
int axis = PADDLE_GET_CONST(int, op_desc.GetAttr("axis"));
auto x_dim = inputs[0];
auto x_rank = x_dim.nbDims;
PADDLE_ENFORCE_GE(axis,
-x_rank,
common::errors::InvalidArgument(
"Attr(axis) value should be in range [-R, R-1], R is "
"the rank of Input(X). But received axis: %d, R: %d. "
"Current Input(X)'s shape is=[%s].",
axis,
x_rank,
x_dim.d));
PADDLE_ENFORCE_LT(axis,
x_rank,
common::errors::InvalidArgument(
"Attr(axis) value should be in range [-R, R-1], R is "
"the rank of Input(X). But received axis: %d, R: %d. "
"Current Input(X)'s shape is=[%s].",
axis,
x_rank,
x_dim.d));
// TODO(liuyuanle): support asvector = True
PADDLE_ENFORCE_EQ(
asvector,
false,
common::errors::InvalidArgument(
"p_norm only support asvector=false, but received asvector: %d.",
asvector));
std::vector<const nvinfer1::IDimensionExpr*> reduce_dims;
if (asvector) {
reduce_dims.emplace_back(expr_builder.constant(1));
if (keepdim) {
for (int i = 1; i < x_dim.nbDims; ++i) {
reduce_dims.emplace_back(expr_builder.constant(1));
}
x_dim.nbDims = reduce_dims.size();
for (size_t i = 0; i < reduce_dims.size(); i++) {
x_dim.d[i] = reduce_dims[i];
}
}
} else {
if (axis < 0) axis = x_dim.nbDims + axis;
for (int i = 0; i < x_dim.nbDims; ++i) {
if (i != axis) reduce_dims.emplace_back(x_dim.d[i]);
}
if (reduce_dims.empty()) {
reduce_dims.emplace_back(expr_builder.constant(1));
}
}
x_dim.d[axis] = expr_builder.constant(1);
nvinfer1::DimsExprs output = {};
if (keepdim) {
output = x_dim;
} else {
output.nbDims = reduce_dims.size();
for (int i = 0; i < output.nbDims; i++) output.d[i] = reduce_dims[i];
}
return output;
}
nvinfer1::DimsExprs GridSamplerInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
const nvinfer1::DimsExprs x_dims = inputs[0];
const nvinfer1::DimsExprs grid_dims = inputs[1];
nvinfer1::DimsExprs output = {};
if (grid_dims.nbDims == 4) {
output.nbDims = 4;
output.d[0] = x_dims.d[0];
output.d[1] = x_dims.d[1];
output.d[2] = grid_dims.d[1];
output.d[3] = grid_dims.d[2];
} else {
output.nbDims = 5;
output.d[0] = x_dims.d[0];
output.d[1] = x_dims.d[1];
output.d[2] = grid_dims.d[1];
output.d[3] = grid_dims.d[2];
output.d[4] = grid_dims.d[3];
}
return output;
}
inline const void UpdatePaddingAndDilation(
std::vector<ExprWrapper>* paddings_wrap,
std::vector<int>* dilation,
const std::string padding_algorithm,
const std::vector<ExprWrapper>& hw_dims,
const std::vector<int>& strides,
const std::vector<ExprWrapper>& k_dims,
nvinfer1::IExprBuilder& expr_builder // NOLINT
) {
if (paddings_wrap->size() == hw_dims.size()) {
for (size_t i = 0; i < hw_dims.size(); ++i) {
auto copy_pad = *(paddings_wrap->begin() + 2 * i);
paddings_wrap->insert(paddings_wrap->begin() + 2 * i + 1, copy_pad);
}
} else {
PADDLE_ENFORCE_EQ(
hw_dims.size(),
paddings_wrap->size(),
common::errors::InvalidArgument(
"Required hw_dims.size() should be equal to paddings_wrap->size(), "
"But received hw_dims.size() = %d, paddings_wrap->size() = %d",
hw_dims.size(),
paddings_wrap->size()));
}
// when padding_algorithm is "VALID" or "SAME"
if (padding_algorithm == "SAME") {
for (size_t i = 0; i < hw_dims.size(); ++i) {
auto out_size = (hw_dims[i] + strides[i] - 1) / strides[i];
auto pad_sum =
max((out_size - 1) * strides[i] + k_dims[i] - hw_dims[i], 0);
auto pad_0 = pad_sum / 2;
auto pad_1 = pad_sum - pad_0;
*(paddings_wrap->begin() + i * 2) = pad_0;
*(paddings_wrap->begin() + i * 2 + 1) = pad_1;
// dilation
*(dilation->begin() + i) = 1;
}
} else if (padding_algorithm == "VALID") {
for (auto& val : *paddings_wrap) {
val = ExprWrapper(0, &expr_builder);
}
}
}
// Here are all examples of using h(height), ok for weight too.
inline ExprWrapper ConvOutputSize(ExprWrapper ih,
ExprWrapper kh,
int dilation_h,
ExprWrapper pad_h0,
ExprWrapper pad_h1,
int stride_h) {
ExprWrapper oh =
(ih + pad_h0 + pad_h1 - dilation_h * (kh - 1) - 1) / stride_h + 1;
return oh;
}
nvinfer1::DimsExprs FusedConv2dAddActInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
// we may update dilations.
std::vector<int> dilations =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("dilations"));
const std::vector<int> strides =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("strides"));
std::vector<int> paddings =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("paddings"));
std::string padding_algorithm = "EXPLICIT";
if (op_desc.HasAttr("padding_algorithm"))
padding_algorithm =
PADDLE_GET_CONST(std::string, op_desc.GetAttr("padding_algorithm"));
if (padding_algorithm == "VALID") {
for (auto& padding : paddings) {
padding = 0;
}
}
// TODO(zhangjun): nhwc support
bool channel_last = false;
// conv_fusion: input, filter, bias
const nvinfer1::DimsExprs input_dims = inputs[0];
const nvinfer1::DimsExprs filter_dims = inputs[1];
auto input_dims_wrap = DimsExprs2VecExprWrapper(input_dims, expr_builder);
auto filter_dims_wrap = DimsExprs2VecExprWrapper(filter_dims, expr_builder);
std::vector<ExprWrapper> hw_dims_wrap; // d, h, w
if (channel_last) {
for (int i = 1; i < input_dims.nbDims - 1; ++i) {
hw_dims_wrap.emplace_back(input_dims_wrap[i]);
}
} else {
for (int i = 2; i < input_dims.nbDims; ++i) {
hw_dims_wrap.emplace_back(input_dims_wrap[i]);
}
}
std::vector<ExprWrapper> filter_hw_dims_wrap; // filter_h, filter_w
if (channel_last) {
for (int i = 1; i < filter_dims.nbDims - 1; ++i) {
filter_hw_dims_wrap.emplace_back(filter_dims_wrap[i]);
}
} else {
for (int i = 2; i < filter_dims.nbDims; ++i) {
filter_hw_dims_wrap.emplace_back(filter_dims_wrap[i]);
}
}
std::vector<ExprWrapper> paddings_wrap;
for (const auto& padding : paddings) {
paddings_wrap.emplace_back(padding, &expr_builder);
}
UpdatePaddingAndDilation(&paddings_wrap,
&dilations,
padding_algorithm,
hw_dims_wrap,
strides,
filter_hw_dims_wrap,
expr_builder);
std::vector<ExprWrapper> output_dims_wrap(input_dims.nbDims);
int out_idx = 0;
output_dims_wrap[out_idx++] = input_dims_wrap[0];
if (!channel_last) {
output_dims_wrap[out_idx++] = filter_dims_wrap[0];
}
for (size_t i = 0; i < hw_dims_wrap.size(); ++i) {
output_dims_wrap[out_idx++] = ConvOutputSize(hw_dims_wrap[i],
filter_hw_dims_wrap[i],
dilations[i],
paddings_wrap[2 * i],
paddings_wrap[2 * i + 1],
strides[i]);
}
if (channel_last) {
output_dims_wrap[out_idx++] = filter_dims_wrap[0];
}
return VecExprWrapper2DimsExprs(output_dims_wrap);
}
nvinfer1::DimsExprs LookupTableV2InferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
const auto x_dims = inputs[0];
const auto weight_dims = inputs[1];
nvinfer1::DimsExprs output = {};
output.nbDims = x_dims.nbDims + 1;
for (int i = 0; i < x_dims.nbDims; ++i) {
output.d[i] = x_dims.d[i];
}
output.d[x_dims.nbDims] = weight_dims.d[1];
return output;
}
nvinfer1::DimsExprs MemoryEfficientAttentionInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
PADDLE_ENFORCE_LE(output_index,
2,
common::errors::InvalidArgument(
"memory_efficient_attention only has three "
"output, but received asvector: %d.",
output_index));
PADDLE_ENFORCE_EQ(
nb_inputs,
8,
common::errors::InvalidArgument("memory_efficient_attention has three "
"input, but received asvector: %d.",
nb_inputs));
if (output_index == 0) {
return inputs[0];
} else if (output_index == 1) {
nvinfer1::DimsExprs output = {};
output.nbDims = 2;
output.d[0] = inputs[0].d[0];
output.d[1] = inputs[0].d[2];
return output;
} else {
nvinfer1::DimsExprs output = {};
output.nbDims = 1;
output.d[0] = expr_builder.constant(2);
return output;
}
}
nvinfer1::DimsExprs Conv2dTransposeInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
auto x_dims = inputs[0];
auto filter_dims = inputs[1];
std::vector<ExprWrapper> x_dims_wrap =
DimsExprs2VecExprWrapper(x_dims, expr_builder);
std::vector<ExprWrapper> filter_dims_wrap =
DimsExprs2VecExprWrapper(filter_dims, expr_builder);
const std::vector<int> dilations =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("dilations"));
const std::vector<int> strides =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("strides"));
std::vector<int> paddings =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("paddings"));
std::vector<int> output_size =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("output_size"));
std::vector<int> output_padding =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("output_padding"));
auto data_format =
PADDLE_GET_CONST(std::string, op_desc.GetAttr("data_format"));
int groups = PADDLE_GET_CONST(int, op_desc.GetAttr("groups"));
std::string padding_algorithm = "EXPLICIT";
if (op_desc.HasAttr("padding_algorithm")) {
padding_algorithm =
PADDLE_GET_CONST(std::string, op_desc.GetAttr("padding_algorithm"));
}
PADDLE_ENFORCE_EQ(padding_algorithm,
"EXPLICIT",
common::errors::InvalidArgument(
"Required padding_algorithm should be 'EXPLICIT', "
"but received padding_algorithm: %s.",
padding_algorithm));
PADDLE_ENFORCE_EQ(
data_format,
"NCHW",
common::errors::InvalidArgument("Required data_format should be 'NCHW', "
"but received data_format: %s.",
data_format));
PADDLE_ENFORCE_EQ(output_size.empty(),
true,
common::errors::InvalidArgument(
"output_size is not empty! Please Check!"));
PADDLE_ENFORCE_EQ(paddings.size(),
2,
common::errors::InvalidArgument(
"Required paddings.size() should be equal to 2, "
"but received paddings.size() = %d.",
paddings.size()));
PADDLE_ENFORCE_EQ(x_dims.nbDims,
4,
common::errors::InvalidArgument(
"Required x_dims.nbDims should be equal to 4, "
"but received x_dims.nbDims = %d.",
x_dims.nbDims));
PADDLE_ENFORCE_EQ(
x_dims.nbDims,
filter_dims.nbDims,
common::errors::InvalidArgument(
"Required x_dims.nbDims should be equal to filter_dims.nbDims, "
"but received x_dims.nbDims = %d, filter_dims.nbDims = %d",
x_dims.nbDims,
filter_dims.nbDims));
PADDLE_ENFORCE_EQ(output_padding.empty(),
true,
common::errors::InvalidArgument(
"output_padding is not empty! Please Check!"));
int stride_size = strides.size();
for (int i = 0; i < stride_size; ++i) {
PADDLE_ENFORCE_EQ(strides[i] > 0,
true,
common::errors::InvalidArgument(
"Required strides[i] should be greater than 0, "
"but received strides[i] = %d",
strides[i]));
}
int in_sub_stride_size = x_dims.nbDims - stride_size;
PADDLE_ENFORCE_EQ(in_sub_stride_size,
2,
common::errors::InvalidArgument(
"Required in_sub_stride_size should be equal to 2, "
"but received in_sub_stride_size = %d",
in_sub_stride_size));
if (!output_size.empty()) {
PADDLE_ENFORCE_EQ(
output_size.size(),
strides.size(),
common::errors::InvalidArgument(
"Required output_size.size() should be equal to strides.size(), "
"but received output_size.size() = %d, strides.size() = %d",
output_size.size(),
strides.size()));
}
if (!output_padding.empty()) {
PADDLE_ENFORCE_EQ(
strides.size(),
output_padding.size(),
common::errors::InvalidArgument(
"Required strides.size should be equal to output_padding.size, "
"but received strides.size() = %d, output_padding.size() = %d",
strides.size(),
output_padding.size()));
}
std::vector<ExprWrapper> output_dims_wrap(x_dims.nbDims);
output_dims_wrap[0] = x_dims_wrap[0];
output_dims_wrap[1] = filter_dims_wrap[1] * groups;
auto ih = x_dims_wrap[2];
auto iw = x_dims_wrap[3];
auto kh = filter_dims_wrap[2];
auto kw = filter_dims_wrap[3];
int pad_h0 = paddings[0];
int pad_h1 = paddings[0];
int pad_w0 = paddings[1];
int pad_w1 = paddings[1];
output_dims_wrap[2] =
(ih - 1) * strides[0] - pad_h0 - pad_h1 + (kh - 1) * dilations[0] + 1;
output_dims_wrap[3] =
(iw - 1) * strides[1] - pad_w0 - pad_w1 + (kw - 1) * dilations[1] + 1;
return VecExprWrapper2DimsExprs(output_dims_wrap);
}
nvinfer1::DimsExprs PadInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
const auto x_dims = inputs[0];
auto paddings =
PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("paddings"));
nvinfer1::DimsExprs output = {};
output.nbDims = x_dims.nbDims;
for (int i = 0; i < x_dims.nbDims; ++i) {
output.d[i] = expr_builder.operation(
nvinfer1::DimensionOperation::kSUM,
*expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
*x_dims.d[i],
*expr_builder.constant(paddings[2 * i])),
*expr_builder.constant(paddings[2 * i + 1]));
}
return output;
}
nvinfer1::DimsExprs ScatterInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
PADDLE_ENFORCE_EQ(
nb_inputs,
3,
common::errors::InvalidArgument("inputs of scatter should be equal to 3, "
"But received (%s)",
nb_inputs));
const nvinfer1::DimsExprs ref_dims = inputs[0];
return ref_dims;
}
nvinfer1::DimsExprs ArgsortInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
const nvinfer1::DimsExprs input_dims = inputs[0];
nvinfer1::DimsExprs output = {};
output.nbDims = input_dims.nbDims;
for (int i = 0; i < input_dims.nbDims; ++i) {
output.d[i] = input_dims.d[i];
}
return output;
}
nvinfer1::DimsExprs SolveInferMeta(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder, // NOLINT
const framework::OpDesc& op_desc) {
PADDLE_ENFORCE_EQ(
nb_inputs,
2,
common::errors::InvalidArgument("inputs of solve should be equal to 2, "
"But received (%s)",
nb_inputs));
const nvinfer1::DimsExprs ref_dims = inputs[1];
return ref_dims;
}
PD_REGISTER_DYNAMIC_INFER_META_FN(gather_nd, GatherNdInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(yolo_box, YoloBoxInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(instance_norm, InstanceNormInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(unfold, UnfoldInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(scatter_nd_add, ScatterNdAddInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(inverse, UnchangedInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(pad3d, Pad3dInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(grid_sampler, GridSamplerInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(fused_conv2d_add_act,
FusedConv2dAddActInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(conv2d, FusedConv2dAddActInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(conv2d_transpose, Conv2dTransposeInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(p_norm, PNormInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(memory_efficient_attention,
MemoryEfficientAttentionInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(pad, PadInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(argsort, ArgsortInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(scatter, ScatterInferMeta);
PD_REGISTER_DYNAMIC_INFER_META_FN(solve, SolveInferMeta);
} // namespace paddle::inference::tensorrt