wehub-resource-sync
355 lines
14 KiB
C++
355 lines
14 KiB
C++
/* Copyright (c) 2018 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/convert/op_converter.h"
|
||
#include "paddle/fluid/inference/tensorrt/plugin/elementwise_op_plugin.h"
|
||
|
||
namespace paddle::inference::tensorrt {
|
||
|
||
class ElementwiseTensorOpConverter : public OpConverter {
|
||
public:
|
||
ElementwiseTensorOpConverter() = default;
|
||
void operator()(const framework::proto::OpDesc& op,
|
||
const framework::Scope& scope,
|
||
bool test_mode) override {
|
||
VLOG(3) << "Convert a elementwise op to TensorRT IElementWiseLayer";
|
||
framework::OpDesc op_desc(op, nullptr);
|
||
auto* X = engine_->GetITensor(op_desc.Input("X").front());
|
||
nvinfer1::ITensor* Y = nullptr;
|
||
Y = engine_->GetITensor(op_desc.Input("Y").front());
|
||
bool swap_xy = false;
|
||
// Swap X and Y
|
||
if (X->getDimensions().nbDims < Y->getDimensions().nbDims) {
|
||
auto* tmp = X;
|
||
X = Y;
|
||
Y = tmp;
|
||
swap_xy = true;
|
||
}
|
||
nvinfer1::Dims dims_x = X->getDimensions();
|
||
nvinfer1::Dims dims_y = Y->getDimensions();
|
||
auto output_name = op_desc.Output("Out")[0];
|
||
|
||
int axis = -1;
|
||
// axis here is relative to explicit batch
|
||
if (op_type_ != "logical_or" && op_type_ != "logical_xor" &&
|
||
op_type_ != "logical_and") {
|
||
axis = PADDLE_GET_CONST(int, op_desc.GetAttr("axis"));
|
||
}
|
||
int real_x_rank = dims_x.nbDims;
|
||
int real_y_rank = dims_y.nbDims;
|
||
if (axis == -1) {
|
||
axis = real_x_rank - real_y_rank;
|
||
}
|
||
|
||
// X: - - - - - - -
|
||
// axis
|
||
// Y: - - -
|
||
// we need expand Y's rank = X's rank
|
||
int left_one_num = axis;
|
||
int right_one_num = dims_x.nbDims - axis - dims_y.nbDims;
|
||
nvinfer1::IShuffleLayer* reshape_layer;
|
||
nvinfer1::ITensor* reshape_y_tensor;
|
||
if (dims_x.nbDims != dims_y.nbDims &&
|
||
(left_one_num > 0 || right_one_num > 0)) {
|
||
auto* y_shape_tensor = Shape(Y);
|
||
auto* new_y_shape_tensor = y_shape_tensor;
|
||
if (axis > 0) {
|
||
std::vector<int32_t> left_one(left_one_num, 1);
|
||
auto* left_one_tensor = Add1DConstantLayer(left_one);
|
||
new_y_shape_tensor = Concat(std::vector<nvinfer1::ITensor*>{
|
||
left_one_tensor, new_y_shape_tensor});
|
||
}
|
||
if (right_one_num > 0) {
|
||
std::vector<int32_t> right_one(right_one_num, 1);
|
||
auto* right_one_tensor = Add1DConstantLayer(right_one);
|
||
new_y_shape_tensor = Concat(std::vector<nvinfer1::ITensor*>{
|
||
new_y_shape_tensor, right_one_tensor});
|
||
}
|
||
reshape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *Y);
|
||
reshape_layer->setInput(1, *new_y_shape_tensor);
|
||
reshape_y_tensor = reshape_layer->getOutput(0);
|
||
} else {
|
||
// In fact, we can remove this `else`, but -> rt_resnet50_test CI in trt
|
||
// 6015 failing, how ridiculous！
|
||
reshape_y_tensor = Y;
|
||
}
|
||
|
||
// We should swap X and Y back, because some operators do not have symmetry
|
||
if (swap_xy) {
|
||
auto* tmp = reshape_y_tensor;
|
||
reshape_y_tensor = X;
|
||
X = tmp;
|
||
}
|
||
|
||
if (op_type_ == "less_equal") {
|
||
auto* less_layer =
|
||
TRT_ENGINE_ADD_LAYER(engine_,
|
||
ElementWise,
|
||
*X,
|
||
*reshape_y_tensor,
|
||
nvinfer1::ElementWiseOperation::kLESS);
|
||
auto* equal_layer =
|
||
TRT_ENGINE_ADD_LAYER(engine_,
|
||
ElementWise,
|
||
*X,
|
||
*reshape_y_tensor,
|
||
nvinfer1::ElementWiseOperation::kEQUAL);
|
||
auto* layer = TRT_ENGINE_ADD_LAYER(engine_,
|
||
ElementWise,
|
||
*(less_layer->getOutput(0)),
|
||
*(equal_layer->getOutput(0)),
|
||
nvinfer1::ElementWiseOperation::kOR);
|
||
ReplenishLayerAndOutput(layer, "elementwise", {output_name}, test_mode);
|
||
} else if (op_type_ == "greater_equal") {
|
||
auto* greater_layer =
|
||
TRT_ENGINE_ADD_LAYER(engine_,
|
||
ElementWise,
|
||
*X,
|
||
*reshape_y_tensor,
|
||
nvinfer1::ElementWiseOperation::kGREATER);
|
||
auto* equal_layer =
|
||
TRT_ENGINE_ADD_LAYER(engine_,
|
||
ElementWise,
|
||
*X,
|
||
*reshape_y_tensor,
|
||
nvinfer1::ElementWiseOperation::kEQUAL);
|
||
auto* layer = TRT_ENGINE_ADD_LAYER(engine_,
|
||
ElementWise,
|
||
*(greater_layer->getOutput(0)),
|
||
*(equal_layer->getOutput(0)),
|
||
nvinfer1::ElementWiseOperation::kOR);
|
||
ReplenishLayerAndOutput(layer, "elementwise", {output_name}, test_mode);
|
||
} else if (op_type_ == "mod") {
|
||
auto* div_layer =
|
||
TRT_ENGINE_ADD_LAYER(engine_,
|
||
ElementWise,
|
||
*X,
|
||
*reshape_y_tensor,
|
||
nvinfer1::ElementWiseOperation::kFLOOR_DIV);
|
||
SupportFP32MixPrecision(output_name, op_desc.Type(), div_layer);
|
||
auto* mul_layer =
|
||
TRT_ENGINE_ADD_LAYER(engine_,
|
||
ElementWise,
|
||
*(div_layer->getOutput(0)),
|
||
*reshape_y_tensor,
|
||
nvinfer1::ElementWiseOperation::kPROD);
|
||
SupportFP32MixPrecision(output_name, op_desc.Type(), mul_layer);
|
||
auto* layer = TRT_ENGINE_ADD_LAYER(engine_,
|
||
ElementWise,
|
||
*X,
|
||
*(mul_layer->getOutput(0)),
|
||
nvinfer1::ElementWiseOperation::kSUB);
|
||
SupportFP32MixPrecision(output_name, op_desc.Type(), layer);
|
||
ReplenishLayerAndOutput(layer, "elementwise", {output_name}, test_mode);
|
||
} else {
|
||
auto op_pair = ops.find(op_type_);
|
||
PADDLE_ENFORCE_NE(
|
||
op_pair,
|
||
ops.end(),
|
||
common::errors::InvalidArgument(
|
||
"Elementwise op's type(%s) is not supported. Please "
|
||
"check if the op_type is correct.",
|
||
op_type_));
|
||
|
||
auto* layer = TRT_ENGINE_ADD_LAYER(
|
||
engine_, ElementWise, *X, *reshape_y_tensor, op_pair->second);
|
||
SupportFP32MixPrecision(output_name, op_desc.Type(), layer);
|
||
ReplenishLayerAndOutput(layer, "elementwise", {output_name}, test_mode);
|
||
}
|
||
}
|
||
|
||
protected:
|
||
static const std::unordered_map<std::string, nvinfer1::ElementWiseOperation>
|
||
ops;
|
||
std::string op_type_;
|
||
};
|
||
|
||
const std::unordered_map<std::string, nvinfer1::ElementWiseOperation>
|
||
ElementwiseTensorOpConverter::ops = {
|
||
{"add", nvinfer1::ElementWiseOperation::kSUM},
|
||
{"mul", nvinfer1::ElementWiseOperation::kPROD},
|
||
{"sub", nvinfer1::ElementWiseOperation::kSUB},
|
||
{"div", nvinfer1::ElementWiseOperation::kDIV},
|
||
{"min", nvinfer1::ElementWiseOperation::kMIN},
|
||
{"pow", nvinfer1::ElementWiseOperation::kPOW},
|
||
{"max", nvinfer1::ElementWiseOperation::kMAX},
|
||
{"floordiv", nvinfer1::ElementWiseOperation::kFLOOR_DIV},
|
||
{"less_than", nvinfer1::ElementWiseOperation::kLESS},
|
||
{"greater_than", nvinfer1::ElementWiseOperation::kGREATER},
|
||
{"logical_or", nvinfer1::ElementWiseOperation::kOR},
|
||
{"logical_xor", nvinfer1::ElementWiseOperation::kXOR},
|
||
{"logical_and", nvinfer1::ElementWiseOperation::kAND},
|
||
};
|
||
|
||
class ElementwiseTensorAddOpConverter : public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorAddOpConverter() { op_type_ = "add"; }
|
||
};
|
||
|
||
class ElementwiseTensorMulOpConverter : public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorMulOpConverter() { op_type_ = "mul"; }
|
||
};
|
||
|
||
class ElementwiseTensorSubOpConverter : public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorSubOpConverter() { op_type_ = "sub"; }
|
||
};
|
||
|
||
class ElementwiseTensorDivOpConverter : public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorDivOpConverter() { op_type_ = "div"; }
|
||
};
|
||
|
||
class ElementwiseTensorMinOpConverter : public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorMinOpConverter() { op_type_ = "min"; }
|
||
};
|
||
|
||
class ElementwiseTensorMaxOpConverter : public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorMaxOpConverter() { op_type_ = "max"; }
|
||
};
|
||
|
||
class ElementwiseTensorPowOpConverter : public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorPowOpConverter() { op_type_ = "pow"; }
|
||
};
|
||
class ElementwiseTensorFloorDivOpConverter
|
||
: public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorFloorDivOpConverter() { op_type_ = "floordiv"; }
|
||
};
|
||
class ElementwiseTensorLessThanOpConverter
|
||
: public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorLessThanOpConverter() { op_type_ = "less_than"; }
|
||
};
|
||
class ElementwiseTensorGreaterThanOpConverter
|
||
: public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorGreaterThanOpConverter() { op_type_ = "greater_than"; }
|
||
};
|
||
class ElementwiseTensorLogicalOrOpConverter
|
||
: public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorLogicalOrOpConverter() { op_type_ = "logical_or"; }
|
||
};
|
||
class ElementwiseTensorLogicalXorOpConverter
|
||
: public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorLogicalXorOpConverter() { op_type_ = "logical_xor"; }
|
||
};
|
||
class ElementwiseTensorLogicalAndOpConverter
|
||
: public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorLogicalAndOpConverter() { op_type_ = "logical_and"; }
|
||
};
|
||
class ElementwiseTensorLessEqualOpConverter
|
||
: public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorLessEqualOpConverter() { op_type_ = "less_equal"; }
|
||
};
|
||
class ElementwiseTensorGreaterEqualOpConverter
|
||
: public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorGreaterEqualOpConverter() { op_type_ = "greater_equal"; }
|
||
};
|
||
class ElementwiseTensorModOpConverter : public ElementwiseTensorOpConverter {
|
||
public:
|
||
ElementwiseTensorModOpConverter() { op_type_ = "mod"; }
|
||
};
|
||
|
||
// The diff between `pow` and `elementwise_pow` is in:
|
||
// https://github.com/PaddlePaddle/Paddle/blob/release/2.4/python/paddle/tensor/math.py#L420
|
||
class PowOpConverter : public OpConverter {
|
||
public:
|
||
PowOpConverter() = default;
|
||
void operator()(const framework::proto::OpDesc& op,
|
||
const framework::Scope& scope,
|
||
bool test_mode) override {
|
||
VLOG(3) << "Convert a pow op to TensorRT IElementWiseLayer";
|
||
framework::OpDesc op_desc(op, nullptr);
|
||
auto* X = engine_->GetITensor(op_desc.Input("X").front());
|
||
float factor = PADDLE_GET_CONST(float, op_desc.GetAttr("factor"));
|
||
nvinfer1::Dims dims_x = X->getDimensions();
|
||
auto output_name = op_desc.Output("Out")[0];
|
||
|
||
nvinfer1::Dims trt_dims_y;
|
||
trt_dims_y.nbDims = dims_x.nbDims;
|
||
for (int i = 0; i < trt_dims_y.nbDims; i++) {
|
||
trt_dims_y.d[i] = 1;
|
||
}
|
||
|
||
std::vector<float> w_data{factor};
|
||
auto* Y = AddConstantLayer(w_data.data(), trt_dims_y);
|
||
|
||
auto* layer = TRT_ENGINE_ADD_LAYER(
|
||
engine_, ElementWise, *X, *Y, nvinfer1::ElementWiseOperation::kPOW);
|
||
SupportFP32MixPrecision(output_name, op_desc.Type(), layer);
|
||
ReplenishLayerAndOutput(layer, "elementwise", {output_name}, test_mode);
|
||
}
|
||
};
|
||
|
||
} // namespace paddle::inference::tensorrt
|
||
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_add_weight,
|
||
ElementwiseTensorAddOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_mul_weight,
|
||
ElementwiseTensorMulOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_sub_weight,
|
||
ElementwiseTensorSubOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_div_weight,
|
||
ElementwiseTensorDivOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_max_weight,
|
||
ElementwiseTensorMaxOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_min_weight,
|
||
ElementwiseTensorMinOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_pow_weight,
|
||
ElementwiseTensorPowOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_floordiv_weight,
|
||
ElementwiseTensorFloorDivOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_mod_weight,
|
||
ElementwiseTensorModOpConverter);
|
||
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_add_tensor,
|
||
ElementwiseTensorAddOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_sub_tensor,
|
||
ElementwiseTensorSubOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_div_tensor,
|
||
ElementwiseTensorDivOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_mul_tensor,
|
||
ElementwiseTensorMulOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_max_tensor,
|
||
ElementwiseTensorMaxOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_min_tensor,
|
||
ElementwiseTensorMinOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_pow_tensor,
|
||
ElementwiseTensorPowOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_floordiv_tensor,
|
||
ElementwiseTensorFloorDivOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(elementwise_mod_tensor,
|
||
ElementwiseTensorModOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(less_than, ElementwiseTensorLessThanOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(greater_than,
|
||
ElementwiseTensorGreaterThanOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(logical_or, ElementwiseTensorLogicalOrOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(logical_xor, ElementwiseTensorLogicalXorOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(logical_and, ElementwiseTensorLogicalAndOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(less_equal, ElementwiseTensorLessEqualOpConverter);
|
||
REGISTER_TRT_OP_CONVERTER(greater_equal,
|
||
ElementwiseTensorGreaterEqualOpConverter);
|
||
|
||
REGISTER_TRT_OP_CONVERTER(pow, PowOpConverter);
|