[Paddle Inference] Add where trt converter

paddle/fluid/inference/api/analysis_predictor.cc

@@ -2258,6 +2258,7 @@ USE_TRT_CONVERTER(prelu);
 USE_TRT_CONVERTER(conv2d_transpose);
 USE_TRT_CONVERTER(leaky_relu);
 USE_TRT_CONVERTER(shuffle_channel);
+USE_TRT_CONVERTER(where);
 USE_TRT_CONVERTER(swish);
 USE_TRT_CONVERTER(silu);
 USE_TRT_CONVERTER(group_norm);

paddle/fluid/inference/tensorrt/convert/CMakeLists.txt

@@ -24,6 +24,7 @@ list(
  layer_norm_op.cc
  multihead_matmul_op.cc
  shuffle_channel_op.cc
+ where_op.cc
  swish_op.cc
  silu_op.cc
  instance_norm_op.cc

paddle/fluid/inference/tensorrt/convert/where_op.cc

@@ -0,0 +1,62 @@
+/* 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/convert/op_converter.h"
+
+namespace paddle {
+namespace framework {
+class Scope;
+
+namespace proto {
+class OpDesc;
+} // namespace proto
+} // namespace framework
+} // namespace paddle
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+/*
+ * Where Op
+ */
+class WhereOpConverter : public OpConverter {
+ public:
+ void operator()(const framework::proto::OpDesc& op,
+ const framework::Scope& scope,
+ bool test_mode) override {
+ VLOG(3) << "convert a fluid where op to tensorrt where layer";
+
+ framework::OpDesc op_desc(op, nullptr);
+ std::string input_x_name = op_desc.Input("X").front();
+ std::string condition_name = op_desc.Input("Condition").front();
+ std::string input_y_name = op_desc.Input("Y").front();
+ std::string output_name = op_desc.Output("Out").front();
+
+ const auto input_x_tensor = engine_->GetITensor(input_x_name);
+ const auto condition_tensor = engine_->GetITensor(condition_name);
+ const auto input_y_tensor = engine_->GetITensor(input_y_name);
+
+ auto layer = TRT_ENGINE_ADD_LAYER(
+ engine_, Select, *condition_tensor, *input_x_tensor, *input_y_tensor);
+
+ RreplenishLayerAndOutput(layer, "where", {output_name}, test_mode);
+ }
+};
+
+} // namespace tensorrt
+} // namespace inference
+} // namespace paddle
+
+REGISTER_TRT_OP_CONVERTER(where, WhereOpConverter);

paddle/fluid/inference/tensorrt/engine.h

@@ -63,6 +63,10 @@ TRT_DT FluidDataType2TRT(FluidDT type) {
  return TRT_DT::kINT32;
  case FluidDT::VarType_Type_FP16:
  return TRT_DT::kHALF;
+#if IS_TRT_VERSION_GE(8400)
+ case FluidDT::VarType_Type_BOOL:
+ return TRT_DT::kBOOL;
+#endif
  default:
  return TRT_DT::kINT32;
  }

paddle/fluid/inference/tensorrt/op_teller.cc

@@ -1654,6 +1654,17 @@ struct SimpleOpTypeSetTeller : public Teller {
 #endif
  }
 
+ if (op_type == "where") {
+#if !IS_TRT_VERSION_GE(8400)
+ VLOG(3) << "where is not supported when TensorRT < 8.4";
+ return false;
+#endif
+ if (!with_dynamic_shape) {
+ VLOG(3) << "the where op does not support static shape yet";
+ return false;
+ }
+ }
+
  if (op_type == "skip_layernorm") {
  if (!with_dynamic_shape) {
  VLOG(3) << "the skip_layernorm does not support static shape yet";
@@ -2225,6 +2236,7 @@ struct SimpleOpTypeSetTeller : public Teller {
  "leaky_relu",
  "fc",
  "shuffle_channel",
+ "where",
  "swish",
  "silu",
  "celu",
@@ -2346,6 +2358,7 @@ struct SimpleOpTypeSetTeller : public Teller {
  "leaky_relu",
  "fc",
  "shuffle_channel",
+ "where",
  "swish",
  "silu",
  "celu",

paddle/fluid/operators/tensorrt/tensorrt_engine_op.h

@@ -601,10 +601,14 @@ class TensorRTEngineOp : public framework::OperatorBase {
  buffers[bind_index] = static_cast<void *>(t.data<int32_t>());
  } else if (type == framework::proto::VarType::FP16) {
  buffers[bind_index] = static_cast<void *>(t.data<float16>());
+#if IS_TRT_VERSION_GE(8400)
+ } else if (type == framework::proto::VarType::BOOL) {
+ buffers[bind_index] = static_cast<void *>(t.data<bool>());
+#endif
  } else {
- PADDLE_THROW(
- platform::errors::Fatal("The TRT Engine OP only support "
-  "float/int32_t/int64_t/float16 input."));
+ PADDLE_THROW(platform::errors::Fatal(
+ "The TRT Engine OP only support "
+ "float/int32_t/int64_t/float16/bool input."));
  }
  }
 

python/paddle/fluid/tests/unittests/ir/inference/auto_scan_test.py

@@ -181,14 +181,25 @@ def generate_op_config(
  ops = []
  for i in range(len(ops_config)):
  op_config = ops_config[i]
- ops.append(
- OpConfig(
- type=op_config['op_type'],
- inputs=op_config['op_inputs'],
- outputs=op_config['op_outputs'],
- attrs=op_config['op_attrs'],
+ if 'outputs_dtype' in op_config:
+ ops.append(
+ OpConfig(
+ type=op_config['op_type'],
+ inputs=op_config['op_inputs'],
+ outputs=op_config['op_outputs'],
+ attrs=op_config['op_attrs'],
+ outputs_dtype=op_config['outputs_dtype'],
+ )
+ )
+ else:
+ ops.append(
+ OpConfig(
+ type=op_config['op_type'],
+ inputs=op_config['op_inputs'],
+ outputs=op_config['op_outputs'],
+ attrs=op_config['op_attrs'],
+ )
  )
- )
  return ops
 
  @abc.abstractmethod