PaddlePaddle
diff --git a/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+33-61Lines changed: 33 additions & 61 deletions b/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+33-61Lines changed: 33 additions & 61 deletions
diff --git a/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
+3-11Lines changed: 3 additions & 11 deletions b/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
+3-11Lines changed: 3 additions & 11 deletions
diff --git a/‎python/paddle/tensorrt/export.py
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions b/‎python/paddle/tensorrt/export.py
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions
diff --git a/‎python/paddle/tensorrt/impls/manipulation.py
Copy file name to clipboard
+47-11Lines changed: 47 additions & 11 deletions b/‎python/paddle/tensorrt/impls/manipulation.py
Copy file name to clipboard
+47-11Lines changed: 47 additions & 11 deletions
@@ -87,6 +87,7 @@ DEFINE_GENERAL_PATTERN(Roll, paddle::dialect::RollOp)
 DEFINE_GENERAL_PATTERN(Softplus, paddle::dialect::SoftplusOp)
 DEFINE_GENERAL_PATTERN(ThresholdedRelu, paddle::dialect::ThresholdedReluOp)
 DEFINE_GENERAL_PATTERN(Flip, paddle::dialect::FlipOp)
+DEFINE_GENERAL_PATTERN(Mish, paddle::dialect::MishOp)
 
 #undef DEFINE_GENERAL_PATTERN
 
@@ -259,7 +260,6 @@ class ActOpPattern : public pir::OpRewritePattern<OpType> {
 };
 using TanhOpPattern = ActOpPattern<paddle::dialect::TanhOp>;
 using CeluOpPattern = ActOpPattern<paddle::dialect::CeluOp>;
-using MishOpPattern = ActOpPattern<paddle::dialect::MishOp>;
 
 class Pool2dOpPattern
     : public pir::OpRewritePattern<paddle::dialect::Pool2dOp> {
@@ -278,6 +278,13 @@ class Pool2dOpPattern
       VLOG(3) << "Cannot find FullIntArrayOp";
       return false;
     }
+    auto attr_value =
+        full_int_array_op->attribute<pir::ArrayAttribute>("value");
+    std::vector<int64_t> kernel_size;
+    for (const auto &attr : attr_value.AsVector()) {
+      kernel_size.push_back(attr.dyn_cast<pir::Int64Attribute>().data());
+    }
+
     auto padding_attr = op->attribute<pir::ArrayAttribute>("paddings");
     std::vector<int32_t> paddings;
     for (const auto &attr : padding_attr.AsVector()) {
@@ -298,33 +305,25 @@ class Pool2dOpPattern
     if (!op->HasAttribute("pooling_type")) {
       VLOG(3) << "The pooling_type attribute does not exist";
       return false;
-    } else {
-      std::string pool_type =
-          op->attribute<pir::StrAttribute>("pooling_type").AsString();
-      if (pool_type != "max" && pool_type != "avg") {
-        VLOG(3) << "Wrong pool op type, the trt do not support the "
-                << pool_type << " pool type.";
-        return false;
-      }
-      if (pool_type == "avg") {
-        if (op->HasAttribute("global_pooling")) {
-          if (!op->attribute<pir::BoolAttribute>("global_pooling").data()) {
-            if (op->HasAttribute("exclusive")) {
-              if (op->attribute<pir::BoolAttribute>("exclusive").data()) {
-                auto attr_value =
-                    full_int_array_op->attribute<pir::ArrayAttribute>("value");
-                std::vector<int64_t> kernel_size;
-                for (const auto &attr : attr_value.AsVector()) {
-                  kernel_size.push_back(
-                      attr.dyn_cast<pir::Int64Attribute>().data());
-                }
-                for (size_t i = 0; i < kernel_size.size(); ++i) {
-                  if (kernel_size[i] <= paddings[i]) {
-                    VLOG(3) << "the padding size should be less than the "
-                               "filter size "
-                               "for exclusive-counting pooling.";
-                    return false;
-                  }
+    }
+    std::string pool_type =
+        op->attribute<pir::StrAttribute>("pooling_type").AsString();
+    if (pool_type != "max" && pool_type != "avg") {
+      VLOG(3) << "Wrong pool op type, the trt do not support the " << pool_type
+              << " pool type.";
+      return false;
+    }
+    if (pool_type == "avg") {
+      if (op->HasAttribute("global_pooling")) {
+        if (!op->attribute<pir::BoolAttribute>("global_pooling").data()) {
+          if (op->HasAttribute("exclusive")) {
+            if (op->attribute<pir::BoolAttribute>("exclusive").data()) {
+              for (size_t i = 0; i < kernel_size.size(); ++i) {
+                if (kernel_size[i] <= paddings[i]) {
+                  VLOG(3) << "the padding size should be less than the "
+                             "filter size "
+                             "for exclusive-counting pooling.";
+                  return false;
                 }
               }
             }
@@ -338,27 +337,14 @@ class Pool2dOpPattern
         op->attribute<pir::BoolAttribute>("global_pooling").data();
     std::string padding_algorithm =
         op->attribute<pir::StrAttribute>("padding_algorithm").AsString();
-    // TODO(Lizexu): The general plugin approach for entering TensorRT has not
-    // been supported yet.
+
     auto adaptive = op->attribute<pir::BoolAttribute>("adaptive").data();
-    if (adaptive) {
-      VLOG(3)
-          << "The adaptive is true pd_op.pool2d is not supported by trt now";
-      return false;
-    }
     // TODO(Lizexu): This piece of code exists in the old IR-TRT implementation
     // but is not covered by unit tests, raising suspicions about its
     // correctness. In the PIR-TRT implementation, following the same approach
     // causes precision issues. For now, we will exclude it from entering
     // TensorRT.
     pir::Value input = op.operand_source(0);
-    auto kernel_size_attr =
-        full_int_array_op->attribute<pir::ArrayAttribute>("value");
-    std::vector<int64_t> kernel_size;
-    for (const auto &attr : kernel_size_attr.AsVector()) {
-      kernel_size.push_back(attr.dyn_cast<pir::Int64Attribute>().data());
-    }
-
     auto input_type = input.type().dyn_cast<paddle::dialect::DenseTensorType>();
     auto input_dims = input_type.dims();
     int g_post_pad_h = 0;
@@ -960,35 +946,21 @@ class FlattenOpPattern
       return false;
     }
     int start_axis = op->attribute<pir::Int32Attribute>("start_axis").data();
-    int stop_axis = op->attribute<pir::Int32Attribute>("stop_axis").data();
 
     pir::Value x = op.operand_source(0);
     auto x_type = x.type().dyn_cast<paddle::dialect::DenseTensorType>();
     auto x_shape = x_type.dims();
     int dims = x_shape.size();
-    if (dims == 0) {
-      VLOG(3) << "Flatten op does not support input's dim is 0 in tensorrt "
-                 "static shape mode.";
-    }
+
     if (start_axis < 0) {
       start_axis += dims;
     }
-
     if (start_axis == 0) {
-      VLOG(3) << "TRT flatten_contiguous_range not support the "
-                 "batch-dimension being changed";
+      VLOG(3)
+          << "TRT pd_op.flatten not support the batch-dimension being changed";
       return false;
     }
-    if (stop_axis < 0) {
-      stop_axis += dims;
-    }
-    for (int i = start_axis; i <= stop_axis; ++i) {
-      if (x_shape[i] < 0) {
-        VLOG(3) << "On TRT static shape,flatten_contiguous_range input dim "
-                   "should be > 0";
-        return false;
-      }
-    }
+
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
     return true;
   }
@@ -2197,6 +2169,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Softplus)
     ADD_PATTERN(ThresholdedRelu)
     ADD_PATTERN(Flip)
+    ADD_PATTERN(Mish)
 #if IS_TRT_VERSION_GE(8600)
     ADD_PATTERN(Layer_norm)
 #endif
@@ -2261,7 +2234,6 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<NotEqualOpPattern>(context));
     ps.Add(std::make_unique<TanhOpPattern>(context));
     ps.Add(std::make_unique<CeluOpPattern>(context));
-    ps.Add(std::make_unique<MishOpPattern>(context));
     ps.Add(std::make_unique<OneHotOpPattern>(context));
     ps.Add(std::make_unique<AssignValueOpPattern>(context));
     ps.Add(std::make_unique<AssignValue_OpPattern>(context));
 
@@ -81,13 +81,8 @@ def __init__(self, paddle_program, scope, trt_config=None):
         self.param_dict = param_dict
 
         trt_manager = TensorRTConfigManager()
-        if (
-            self.trt_config is not None
-            and self.trt_config.tensorrt_ops_run_float
-        ):
-            trt_manager.set_force_fp32_ops(
-                self.trt_config.tensorrt_ops_run_float
-            )
+        if self.trt_config is not None and self.trt_config.ops_run_float:
+            trt_manager.set_force_fp32_ops(self.trt_config.ops_run_float)
             _logger.info(f"force_fp32_ops: {trt_manager.get_force_fp32_ops()}")
 
         self.input_info = {}
@@ -441,10 +436,7 @@ def convert_subgraph_to_trt(self, program, group_op):
             and version_list[1] >= 2
             and version_list[2] >= 1
         ):
-            if (
-                self.trt_config is not None
-                and self.trt_config.tensorrt_ops_run_float
-            ):
+            if self.trt_config is not None and self.trt_config.ops_run_float:
                 config.set_flag(trt.BuilderFlag.PREFER_PRECISION_CONSTRAINTS)
 
         trt_engine = builder.build_serialized_network(network, config)
 
@@ -170,7 +170,7 @@ def __init__(
         save_model_dir: str | None = None,
         disable_ops: str | list | None = None,
         precision_mode: PrecisionMode = PrecisionMode.FP32,
-        tensorrt_ops_run_float: str | list | None = None,
+        ops_run_float: str | list | None = None,
     ) -> None:
         """
         A class for configuring TensorRT optimizations.
@@ -190,7 +190,7 @@ def __init__(
                 - PrecisionMode.FP16: 16-bit floating point precision.
                 - PrecisionMode.INT8: 8-bit integer precision.
                 - PrecisionMode.BFP16: 16-bit Brain Floating Point precision. Only supported in TensorRT versions greater than 9.0.
-            tensorrt_ops_run_float (str|list, optional):
+            ops_run_float (str|list, optional):
                 A set of operation names that should be executed using FP32 precision regardless of the `tensorrt_precision_mode` setting.
                  The directory where the optimized model will be saved (default is None).
         Returns:
@@ -215,13 +215,13 @@ def __init__(
             >>> trt_config = TensorRTConfig(inputs=[input])
             >>> trt_config.disable_ops = "pd_op.dropout"
             >>> trt_config.precision_mode = PrecisionMode.FP16
-            >>> trt_config.tensorrt_ops_run_float = "pd_op.conv2d"
+            >>> trt_config.ops_run_float = "pd_op.conv2d"
         """
         self.inputs = inputs
         self.min_subgraph_size = min_subgraph_size
         self.save_model_dir = save_model_dir
         self.precision_mode = precision_mode
-        self.tensorrt_ops_run_float = tensorrt_ops_run_float
+        self.ops_run_float = ops_run_float
         self.disable_ops = disable_ops
         paddle.framework.set_flags(
             {'FLAGS_trt_min_group_size': min_subgraph_size}
 
@@ -22,7 +22,6 @@
     cast_tensor,
     fix_negative_indices,
     get_axes_for_reduce_op,
-    get_positive_dim,
     get_shape_tensor_element,
     has_dynamic_shape,
     resize_to_1d,
@@ -246,21 +245,58 @@ def squeeze_converter(network, paddle_op, inputs):
     input_shape = input_val.shape
     input_shape_size = len(input_shape)
 
-    if type(input_val) == trt.Weights:
+    # If input is weights, convert to TensorRT tensor
+    if isinstance(input_val, trt.Weights):
         input_val = network.add_constant(input_shape, input_val).get_output(0)
 
-    axis = paddle_op.operands()[1].source().get_defining_op().attrs()["value"]
-    axis = axis[0]
+    # Get axis
+    axis = (
+        paddle_op.operands()[1]
+        .source()
+        .get_defining_op()
+        .attrs()
+        .get("value", [])
+    )
 
-    axis = get_positive_dim(axis, input_shape_size + 1)
-    output_shape = []
-    for i, s in enumerate(input_shape):
-        if i == axis and s == 1:
-            continue
-        output_shape.append(s)
+    if not axis:
+        for i in range(input_shape_size):
+            if input_shape[i] == -1:
+                raise RuntimeError(
+                    "The necessary attributes of the squeeze operator axis is missing"
+                )
+            elif input_shape[i] == 1:
+                axis.append(i)
+    else:
+        # Verify that each axis to squeeze has size 1
+        for a in axis:
+            if a < 0:
+                a += input_shape_size
+            if input_shape[a] != 1:
+                raise RuntimeError(
+                    f"Cannot squeeze dimension {a} with size {input_shape[a]}. Only dimensions with size 1 can be squeezed."
+                )
 
+    axes_size = len(axis)
+    if axes_size == 0:
+        raise RuntimeError(
+            f"axis.size should be >0 in pd_op.squeeze op in TensorRT, but received {axes_size}"
+        )
+    # Mark which dimensions to squeeze
+    should_squeeze = [False] * input_shape_size
+    for a in axis:
+        should_squeeze[a] = True
+
+    # Get dimensions to keep
+    gather_indices = [
+        i for i, squeeze in enumerate(should_squeeze) if not squeeze
+    ]
+
+    # Add Shuffle layer
     layer = network.add_shuffle(input_val)
-    layer.reshape_dims = tuple(output_shape)
+    shape_tensor = trt_shape(network, input_val)
+    real_shape_tensor = trt_gather(network, shape_tensor, gather_indices)
+    layer.set_input(1, real_shape_tensor)
+
     return layer.get_output(0)