PaddlePaddle
diff --git a/‎paddle/fluid/inference/tensorrt/plugin/pool_op_plugin.h
Copy file name to clipboard
+74-4Lines changed: 74 additions & 4 deletions b/‎paddle/fluid/inference/tensorrt/plugin/pool_op_plugin.h
Copy file name to clipboard
+74-4Lines changed: 74 additions & 4 deletions
diff --git a/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+10-2Lines changed: 10 additions & 2 deletions b/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+10-2Lines changed: 10 additions & 2 deletions
diff --git a/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
+15-9Lines changed: 15 additions & 9 deletions b/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
+15-9Lines changed: 15 additions & 9 deletions
diff --git a/‎python/paddle/tensorrt/converter_utils.py
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎python/paddle/tensorrt/converter_utils.py
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎python/paddle/tensorrt/export.py
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions b/‎python/paddle/tensorrt/export.py
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions
diff --git a/‎python/paddle/tensorrt/impls/creation.py
Copy file name to clipboard
+65-55Lines changed: 65 additions & 55 deletions b/‎python/paddle/tensorrt/impls/creation.py
Copy file name to clipboard
+65-55Lines changed: 65 additions & 55 deletions
@@ -206,10 +206,11 @@ class PoolPluginDynamic : public DynamicPluginTensorRT {
   size_t getSerializationSize() const TRT_NOEXCEPT override;
   void serialize(void* buffer) const TRT_NOEXCEPT override;
 
-  nvinfer1::DimsExprs getOutputDimensions(int output_index,
-                                          const nvinfer1::DimsExprs* inputs,
-                                          int nb_inputs,
-                                          nvinfer1::IExprBuilder& expr_builder)
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
       TRT_NOEXCEPT override;
 
   bool supportsFormatCombination(int pos,
@@ -268,7 +269,76 @@ class PoolPluginDynamicCreator : public TensorRTPluginCreator {
     return new PoolPluginDynamic(serial_data, serial_length);
   }
 };
+
+class PIRPoolPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "pir_pool_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new PoolPluginDynamic(serial_data, serial_length);
+  }
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    bool ceil_mode = false;
+    std::string pool_type;
+    bool adaptive = false;
+    bool exclusive = false;
+    std::vector<int> ksize;
+    std::vector<int> strides;
+    std::vector<int> paddings;
+    bool global_pooling = false;
+
+    for (int i = 0; i < fc->nbFields; ++i) {
+      const nvinfer1::PluginField& field = fc->fields[i];
+      const std::string field_name(field.name);
+      if (field_name.compare("ceil_mode") == 0) {
+        ceil_mode = *static_cast<const bool*>(field.data);
+      } else if (field_name.compare("pool_type") == 0) {
+        pool_type = std::string(static_cast<const char*>(fc->fields[i].data),
+                                fc->fields[i].length);
+      } else if (field_name.compare("adaptive") == 0) {
+        adaptive = *static_cast<const bool*>(field.data);
+      } else if (field_name.compare("exclusive") == 0) {
+        exclusive = *static_cast<const bool*>(field.data);
+      } else if (field_name.compare("ksize") == 0) {
+        const int length = fc->fields[i].length;
+        const int* data = static_cast<const int*>(fc->fields[i].data);
+        ksize.insert(ksize.end(), data, data + length);
+      } else if (field_name.compare("strides") == 0) {
+        const int length = fc->fields[i].length;
+        const int* data = static_cast<const int*>(fc->fields[i].data);
+        strides.insert(strides.end(), data, data + length);
+      } else if (field_name.compare("paddings") == 0) {
+        const int length = fc->fields[i].length;
+        const int* data = static_cast<const int*>(fc->fields[i].data);
+        paddings.insert(paddings.end(), data, data + length);
+      } else if (field_name.compare("global_pooling") == 0) {
+        global_pooling = *static_cast<const bool*>(field.data);
+      } else {
+        assert(false && "unknown plugin field name.");
+      }
+    }
+    return new PoolPluginDynamic(ceil_mode,
+                                 pool_type,
+                                 adaptive,
+                                 exclusive,
+                                 ksize,
+                                 strides,
+                                 paddings,
+                                 global_pooling);
+  }
+};
+
 REGISTER_TRT_PLUGIN_V2(PoolPluginDynamicCreator);
+REGISTER_TRT_PLUGIN_V2(PIRPoolPluginDynamicCreator);
 #endif
 
 }  // namespace plugin
 
@@ -108,6 +108,7 @@ DEFINE_GENERAL_PATTERN(Cosh, paddle::dialect::CoshOp)
 DEFINE_GENERAL_PATTERN(Asinh, paddle::dialect::AsinhOp)
 DEFINE_GENERAL_PATTERN(Acosh, paddle::dialect::AcoshOp)
 DEFINE_GENERAL_PATTERN(Atanh, paddle::dialect::AtanhOp)
+DEFINE_GENERAL_PATTERN(Tanh, paddle::dialect::TanhOp)
 DEFINE_GENERAL_PATTERN(Ceil, paddle::dialect::CeilOp)
 DEFINE_GENERAL_PATTERN(Rsqrt, paddle::dialect::RsqrtOp)
 DEFINE_GENERAL_PATTERN(Reciprocal, paddle::dialect::ReciprocalOp)
@@ -117,6 +118,10 @@ DEFINE_GENERAL_PATTERN(Sign, paddle::dialect::SignOp)
 DEFINE_GENERAL_PATTERN(Round, paddle::dialect::RoundOp)
 DEFINE_GENERAL_PATTERN(Numel, paddle::dialect::NumelOp)
 DEFINE_GENERAL_PATTERN(Pool3d, paddle::dialect::Pool3dOp)
+DEFINE_GENERAL_PATTERN(Tan, paddle::dialect::TanOp)
+DEFINE_GENERAL_PATTERN(Asin, paddle::dialect::AsinOp)
+DEFINE_GENERAL_PATTERN(Acos, paddle::dialect::AcosOp)
+DEFINE_GENERAL_PATTERN(Atan, paddle::dialect::AtanOp)
 DEFINE_GENERAL_PATTERN(ShuffleChannel, paddle::dialect::ShuffleChannelOp)
 
 #undef DEFINE_GENERAL_PATTERN
@@ -287,7 +292,6 @@ class ActOpPattern : public pir::OpRewritePattern<OpType> {
     return true;
   }
 };
-using TanhOpPattern = ActOpPattern<paddle::dialect::TanhOp>;
 using CeluOpPattern = ActOpPattern<paddle::dialect::CeluOp>;
 using TanhShrinkOpPattern = ActOpPattern<paddle::dialect::TanhShrinkOp>;
 
@@ -2575,6 +2579,11 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Round)
     ADD_PATTERN(Numel)
     ADD_PATTERN(Pool3d)
+    ADD_PATTERN(Tanh)
+    ADD_PATTERN(Tan)
+    ADD_PATTERN(Asin)
+    ADD_PATTERN(Acos)
+    ADD_PATTERN(Atan)
     ADD_PATTERN(ShuffleChannel)
 #if IS_TRT_VERSION_GE(8600)
     ADD_PATTERN(Layer_norm)
@@ -2654,7 +2663,6 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<EqualOpPattern>(context));
     ps.Add(std::make_unique<NotEqualOpPattern>(context));
     ps.Add(std::make_unique<LogicalXorOpPattern>(context));
-    ps.Add(std::make_unique<TanhOpPattern>(context));
     ps.Add(std::make_unique<CeluOpPattern>(context));
     ps.Add(std::make_unique<OneHotOpPattern>(context));
     ps.Add(std::make_unique<TemporalShiftOpPattern>(context));
 
@@ -380,15 +380,21 @@ def convert_subgraph_to_trt(self, program, group_op):
             if group_op.result(out_index).use_empty():
                 # if result value is not used, it doesn't need get shape, continue
                 continue
-            min_shape = get_value_shape_range_info(
-                result_value, False, paddle.base.core.ShapeMode.kMIN
-            )
-            opt_shape = get_value_shape_range_info(
-                result_value, False, paddle.base.core.ShapeMode.kOPT
-            )
-            max_shape = get_value_shape_range_info(
-                result_value, False, paddle.base.core.ShapeMode.kMAX
-            )
+            if not is_shape_tensor(result_value):
+                if len(result_value.shape) == 0:
+                    min_shape = []
+                    opt_shape = []
+                    max_shape = []
+                else:
+                    min_shape = get_value_shape_range_info(
+                        result_value, False, paddle.base.core.ShapeMode.kMIN
+                    )
+                    opt_shape = get_value_shape_range_info(
+                        result_value, False, paddle.base.core.ShapeMode.kOPT
+                    )
+                    max_shape = get_value_shape_range_info(
+                        result_value, False, paddle.base.core.ShapeMode.kMAX
+                    )
             min_value = []
             opt_value = []
             max_value = []
 
@@ -725,6 +725,10 @@ def unary_op_converter(network, paddle_op, inputs):
         "pd_op.round": [trt.UnaryOperation.ROUND],
         "pd_op.logical_not": [trt.UnaryOperation.NOT],
         "pd_op.rsqrt": [trt.UnaryOperation.SQRT, trt.UnaryOperation.RECIP],
+        "pd_op.tan": [trt.UnaryOperation.TAN],
+        "pd_op.asin": [trt.UnaryOperation.ASIN],
+        "pd_op.acos": [trt.UnaryOperation.ACOS],
+        "pd_op.atan": [trt.UnaryOperation.ATAN],
     }
 
     input_tensor = inputs[0]
 
@@ -36,6 +36,7 @@
 from paddle.nn import Layer
 from paddle.tensorrt.converter import PaddleToTensorRTConverter
 from paddle.tensorrt.util import (
+    enforce_op_lower_trt,
     forbid_op_lower_trt,
     mark_builtin_op,
     run_pir_pass,
@@ -175,6 +176,7 @@ def __init__(
         optimization_level: int | None = 3,
         disable_passes: list = [],
         workspace_size: int | None = 1 << 30,
+        allow_only_specified_trt_ops: str | list | None = None,
     ) -> None:
         """
         A class for configuring TensorRT optimizations.
@@ -203,6 +205,8 @@ def __init__(
                 A list of string representing the names of pass that should not be used for origin program (default is []).
             workspace_size (int, optional):
                 Specifies the maximum GPU memory (in bytes) that TensorRT can use for the optimization process (default is 1 << 30).
+            allow_only_specified_trt_ops(str|list,optional):
+                Allows a specified operation or a list of operations to be passed into TRT, while excluding all others.(default is None).
         Returns:
             None
 
@@ -237,6 +241,7 @@ def __init__(
         self.disable_passes = disable_passes
         self.optimization_level = optimization_level
         self.workspace_size = workspace_size
+        self.allow_only_specified_trt_ops = allow_only_specified_trt_ops
         paddle.framework.set_flags(
             {'FLAGS_trt_min_group_size': min_subgraph_size}
         )
@@ -282,6 +287,14 @@ def convert_to_trt(program, trt_config, scope):
             max_shape_feed=max_shape_feed,
             scope=scope,
         )
+        if trt_config.allow_only_specified_trt_ops and trt_config.disable_ops:
+            raise ValueError(
+                "The parameters 'allow_only_specified_trt_ops' and 'disable_ops' are mutually exclusive. Please specify only one."
+            )
+        if trt_config.allow_only_specified_trt_ops:
+            enforce_op_lower_trt(
+                program, trt_config.allow_only_specified_trt_ops
+            )
 
         # specify certain operators to be excluded from entering TensorRT
         if trt_config.disable_ops:
 
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 import numpy as np
 import tensorrt as trt
 
@@ -25,8 +26,10 @@
     trt_cast,
     trt_floor_div,
     trt_max,
+    trt_min,
     trt_reduce_to_scalar,
     trt_reshape,
+    trt_shape,
     trt_sub,
 )
 from paddle.tensorrt.register import converter_registry
@@ -129,65 +132,72 @@ def arange_converter(network, paddle_op, inputs):
 
 @converter_registry.register("pd_op.full_like", trt_version="8.x")
 def full_like_converter(network, paddle_op, inputs):
-    shape = inputs[0].shape
+    input_tensor = inputs[0]
+    shape = input_tensor.shape
     ndims = len(shape)
 
-    out_dtype = int(paddle_op.attrs().get("dtype", None))
-    # Reference paddle/phi/common/data_type.h enum DataType
-    if out_dtype == 1:  # paddle.bool
-        out_dtype = trt.int32
-    elif out_dtype == 7:  # paddle.int32
-        out_dtype = trt.int32
-    elif out_dtype == 9:  # paddle.int64
-        out_dtype = trt.int32
-    elif out_dtype == 10:  # paddle.float32
-        out_dtype = trt.float32
-    elif out_dtype == 11:  # paddle.float64
-        out_dtype = trt.float32
-    else:
-        raise RuntimeError(
-            f"cast converter currently doesn't support dtype: {out_dtype}"
-        )
+    dtype = int(paddle_op.attrs().get("dtype", -1))
+
+    dtype_map = {
+        0: None,  # Undefined
+        1: trt.bool,  # bool
+        2: trt.int32,  # int32
+        3: trt.int32,  # int64 -> int32
+        4: trt.int32,  # int16 -> int32
+        5: trt.float32,  # float16 -> float32
+        6: trt.float32,  # float64 -> float32
+        7: trt.float32,  # float32
+        8: trt.int32,  # uint8 -> int32
+        11: trt.float32,  # float32
+    }
+
+    target_dtype = dtype_map.get(dtype, None)
+    if target_dtype is None:
+        target_dtype = input_tensor.dtype
 
-    fill_value = get_input_constant_value(paddle_op, inputs, 1)
-    if fill_value is not None:
-        fill_value = fill_value[0]
-        value = network.add_constant(
-            (1,),
-            np.array(
-                [
-                    fill_value,
-                ],
-                dtype=np.float32,
-            ),
-        ).get_output(0)
-        value = trt_cast(network, value, out_dtype)
+    value = get_input_constant_value(paddle_op, inputs, 1)
+    if value is not None:
+        if isinstance(value, (list, tuple)):
+            value = value[0] if value else 0
+
+        if target_dtype == trt.int32:
+            value_tensor = add_1D_constant_layer(network, int(value), np.int32)
+        else:
+            value_tensor = add_1D_constant_layer(
+                network, float(value), np.float32
+            )
     else:
-        value = inputs[1]
-
-    shuffle_layer = network.add_shuffle(value)
-    shuffle_layer.reshape_dims = (1,) * ndims
-
-    start_vec = np.zeros((ndims,), dtype=np.int32)
-    start_tensor = network.add_constant((ndims,), start_vec).get_output(0)
-    shape_tensor = network.add_shape(inputs[0]).get_output(0)
-    stride_tensor = network.add_constant(
-        (ndims,), np.ones((ndims,), dtype=np.int32)
-    ).get_output(0)
-
-    slice_layer = network.add_slice(
-        shuffle_layer.get_output(0),
-        start_vec,
-        [1] * ndims,
-        np.ones((ndims,), dtype=np.int32),
-    )
-    slice_layer.mode = trt.SliceMode.FILL
-    slice_layer.set_input(1, start_tensor)
-    slice_layer.set_input(2, shape_tensor)
-    slice_layer.set_input(3, stride_tensor)
-    value = trt_cast(network, value, out_dtype)
-    slice_layer.set_input(4, value)
-    return slice_layer.get_output(0)
+        value_tensor = inputs[1]
+        if value_tensor.dtype != target_dtype:
+            value_tensor = trt_cast(network, value_tensor, target_dtype)
+
+    shape_tensor = trt_shape(network, input_tensor)
+    one_rank_tensor = add_1D_constant_layer(network, [1] * ndims)
+    input_shape_tensor = one_rank_tensor
+
+    shuffle_layer = network.add_shuffle(value_tensor)
+    shuffle_layer.set_input(1, input_shape_tensor)
+    start = trt.Dims([0] * ndims)
+    size = trt.Dims([1] * ndims)
+    stride = trt.Dims([1] * ndims)
+
+    starts_tensor = add_1D_constant_layer(network, [0] * ndims)
+    one_tensor = add_1D_constant_layer(network, 1)
+    sizes_tensor = trt_max(network, input_shape_tensor, shape_tensor)
+    input_sub_tensor = trt_sub(network, input_shape_tensor, one_tensor)
+    strides_tensor = trt_min(network, one_tensor, input_sub_tensor)
+
+    layer = network.add_slice(shuffle_layer.get_output(0), start, size, stride)
+    layer.set_input(1, starts_tensor)
+    layer.set_input(2, sizes_tensor)
+    layer.set_input(3, strides_tensor)
+
+    output = layer.get_output(0)
+
+    if output.dtype != target_dtype:
+        output = trt_cast(network, output, target_dtype)
+
+    return output
 
 
 @converter_registry.register("pd_op.full_with_tensor", trt_version="8.x")