【Paddle TensorRT】fix PaddleX pir-trt 3 (#71287)

lizexu123 · YuanRisheng · web-flow · commit 97a59db7509b · 2025-03-07T11:20:46.000+08:00
* fix trt model bugs

* fix bugs

* fix inplace bugs

* fix

* fix

* fix

* fix

* fix slice

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* code-style

* fix

* fix

---------

Co-authored-by: YuanRisheng &lt;yuanrisheng@baidu.com&gt;
diff --git a/paddle/fluid/inference/tensorrt/convert/slice_op.cc b/paddle/fluid/inference/tensorrt/convert/slice_op.cc
@@ -99,8 +99,13 @@ class SliceOpConverter : public OpConverter {
     if (slice_inputs.find("EndsTensor") != slice_inputs.end() &&
         !op_desc.Input("EndsTensor").empty()) {  // has EndsTensor input
       for (size_t i = 0; i < axes.size(); ++i) {
-        ends_tensor[axes[i]] = GetEleTensorOfShape(
-            engine_->GetITensor(op_desc.Input("EndsTensor")[0]), i);
+        auto axis = axes[i];
+        auto input_dim = GetEleTensorOfShape(shape_tensor, axis);
+        ends_tensor[axes[i]] =
+            Min(Max(GetEleTensorOfShape(
+                        engine_->GetITensor(op_desc.Input("EndsTensor")[0]), i),
+                    Add1DConstantLayer(0)),
+                input_dim);
       }
     } else if (slice_inputs.find("EndsTensorList") != slice_inputs.end() &&
                !op_desc.Input("EndsTensorList").empty()) {
diff --git a/paddle/fluid/inference/tensorrt/convert/yolo_box_op.cc b/paddle/fluid/inference/tensorrt/convert/yolo_box_op.cc
@@ -50,7 +50,8 @@ class YoloBoxOpConverter : public OpConverter {
             : 0.5;
 
     int type_id = static_cast<int>(engine_->WithFp16());
-    auto* yolo_box_plugin = new plugin::YoloBoxPluginDynamic(
+    auto input_dim = X_tensor->getDimensions();
+    auto* yolo_box_plugin = new plugin::YoloBoxPlugin(
         type_id ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT,
         anchors,
         class_num,
@@ -59,7 +60,9 @@ class YoloBoxOpConverter : public OpConverter {
         clip_bbox,
         scale_x_y,
         iou_aware,
-        iou_aware_factor);
+        iou_aware_factor,
+        input_dim.d[1],
+        input_dim.d[2]);
 
     std::vector<nvinfer1::ITensor*> yolo_box_inputs;
     yolo_box_inputs.push_back(X_tensor);
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -955,6 +955,31 @@ struct SimpleOpTypeSetTeller : public Teller {
       if (resize_inputs.find("SizeTensor") != resize_inputs.end()) {
 #if IS_TRT_VERSION_GE(8200)
         if (desc.Input("SizeTensor").size() == 2) {
+          // TODO(lizexu123): When SizeTensor exists, at least one of the input
+          // variable names must contain 'shape' in order for TRT conversion to
+          // proceed; otherwise, TRT conversion will be disallowed."
+          auto* block = desc.Block();
+          if (block == nullptr) {
+            VLOG(3)
+                << "The block desc is nullptr,we can't continue to analyze.";
+            return false;
+          }
+          bool valid_source = false;
+          //
+          std::vector<std::string> size_tensor_names = desc.Input("SizeTensor");
+          for (const auto& tensor_name : size_tensor_names) {
+            auto* var_desc = block->FindVarRecursive(tensor_name);
+            if (!var_desc) continue;
+            if (tensor_name.find("shape") != std::string::npos) {
+              valid_source = true;
+              break;
+            }
+          }
+          if (!valid_source) {
+            VLOG(3) << "The SizeTensor for bilinear_interp_v2 doesn't come "
+                       "from a valid source.";
+            return false;
+          }
           return true;
         }
 #else
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -421,11 +421,11 @@ class Pool2dOpPattern
         op->attribute<pir::StrAttribute>("padding_algorithm").AsString();
 
     auto adaptive = op->attribute<pir::BoolAttribute>("adaptive").data();
-    // 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.
+    // TODO(lizexu123): 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 input_type = input.type().dyn_cast<paddle::dialect::DenseTensorType>();
     auto input_dims = input_type.dims();
@@ -1173,6 +1173,13 @@ class CastOpPattern : public pir::OpRewritePattern<paddle::dialect::CastOp> {
       return false;
 #endif
     }
+    if (dtype != phi::DataType::BOOL && dtype != phi::DataType::FLOAT32 &&
+        dtype != phi::DataType::FLOAT64 && dtype != phi::DataType::FLOAT16 &&
+        dtype != phi::DataType::INT32 && dtype != phi::DataType::INT64) {
+      VLOG(3) << "the cast op does not support type: "
+              << phi::DataTypeToString(dtype);
+      return false;
+    }
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
     return true;
   }
@@ -1691,11 +1698,45 @@ class BilinearInterpV2Pattern
       return false;
     }
 #if IS_TRT_VERSION_GE(8200)
+    // TODO(lizexu123): Starting from the size_tensor, traverse up three levels.
+    // If a pd_op.shape64 operator is found within those three levels, then
+    // allow it to enter TRT; otherwise, prohibit TRT conversion to avoid
+    // potential bugs.
     auto size_tensor = op.operand_source(2);
     if (size_tensor.impl()) {
+      auto *first_def_op = size_tensor.defining_op();
+      std::vector<std::string> upstream_op_names;
+      upstream_op_names.push_back(first_def_op->name());
+      if (first_def_op->num_operands() > 0 &&
+          first_def_op->operand_source(0).impl()) {
+        auto second_input = first_def_op->operand_source(0);
+        auto *second_def_op = second_input.defining_op();
+        upstream_op_names.push_back(second_def_op->name());
+        if (second_def_op->num_operands() > 0 &&
+            second_def_op->operand_source(0).impl()) {
+          auto third_input = second_def_op->operand_source(0);
+          auto *third_def_op = third_input.defining_op();
+          upstream_op_names.push_back(third_def_op->name());
+        }
+      }
+      bool found_shape = false;
+      for (const auto &name : upstream_op_names) {
+        if (name.find("shape64") != std::string::npos) {
+          found_shape = true;
+        }
+      }
+      if (!found_shape) {
+        VLOG(3) << "BilinearInterpV2: Upstream ops do not contain 'shape':";
+        for (const auto &name : upstream_op_names) {
+          VLOG(3) << "\t" << name;
+        }
+        return false;
+      }
+
+      // 同时检查 size_tensor 类型为 VectorType 且大小为2
       auto size_tensor_type = size_tensor.type();
       if (size_tensor_type.isa<pir::VectorType>()) {
-        auto vector_type = size_tensor.type().dyn_cast<pir::VectorType>();
+        auto vector_type = size_tensor_type.dyn_cast<pir::VectorType>();
         if (vector_type.size() == 2) {
           op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
           return true;
@@ -1711,14 +1752,12 @@ class BilinearInterpV2Pattern
     }
 #endif
     pir::Value scale_tensor = op.operand_source(3);
-
     bool has_scale_input = false;
     if (scale_tensor) {
       has_scale_input = true;
     }
-
     if (has_scale_input) {
-      VLOG(3) << "BilinearInterpV2 has scale input can not into trt,support "
+      VLOG(3) << "BilinearInterpV2 has scale input can not into trt, support "
                  "scale attribute into trt";
       return false;
     }
@@ -1760,7 +1799,6 @@ class BilinearInterpV2Pattern
         }
       }
     }
-
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
     return true;
   }
@@ -2907,14 +2945,14 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Dropout)
     ADD_PATTERN(Bmm)
     ADD_PATTERN(Concat)
+    ADD_PATTERN(Nonzero)
     ADD_PATTERN(Full)
     ADD_PATTERN(Fused_gemm_epilogue)
     ADD_PATTERN(Add)
     ADD_PATTERN(Silu)
     ADD_PATTERN(Conv2d)
     ADD_PATTERN(FusedConv2dAddAct)
     ADD_PATTERN(DepthwiseConv2d)
-    ADD_PATTERN(Nonzero)
     ADD_PATTERN(Gelu)
     ADD_PATTERN(Relu6)
     ADD_PATTERN(Shape)
diff --git a/python/paddle/tensorrt/converter.py b/python/paddle/tensorrt/converter.py
@@ -257,7 +257,7 @@ def convert_subgraph_to_trt(self, program, group_op):
                             f'{source_id} not found in value_to_trt_tensor'
                         )
 
-            if precision_mode == PrecisionMode.INT8:
+            if precision_mode.value == PrecisionMode.INT8.value:
                 set_dynamic_range(op, operands)
             trt_outs = self.convert(network, op, operands)
 
@@ -439,15 +439,15 @@ def convert_subgraph_to_trt(self, program, group_op):
             trt.MemoryPoolType.WORKSPACE, self.trt_config.workspace_size
         )
 
-        if precision_mode == PrecisionMode.FP16:
+        if precision_mode.value == PrecisionMode.FP16.value:
             if builder.platform_has_fast_fp16:
                 config.set_flag(trt.BuilderFlag.FP16)
                 _logger.info("Run Paddle-TRT FP16 mode")
             else:
                 _logger.warning(
                     "Hardware does not support FP16. Continuing in FP32 mode."
                 )
-        elif precision_mode == PrecisionMode.BF16:
+        elif precision_mode.value == PrecisionMode.BF16.value:
             if version_list[0] >= 9:
                 if builder.platform_has_fast_bfp16 and hasattr(
                     builder, 'plateform_has_fast_bf16'
@@ -468,7 +468,7 @@ def convert_subgraph_to_trt(self, program, group_op):
                     _logger.warning(
                         "Hardware does not support FP16. Continuing in FP32 mode."
                     )
-        elif precision_mode == PrecisionMode.INT8:
+        elif precision_mode.value == PrecisionMode.INT8.value:
             config.set_flag(trt.BuilderFlag.INT8)
             _logger.info("Run Paddle-TRT INT8 mode")
         elif self.trt_config is not None:
diff --git a/python/paddle/tensorrt/impls/creation.py b/python/paddle/tensorrt/impls/creation.py
@@ -121,6 +121,8 @@ def arange_converter(network, paddle_op, inputs):
         network, delta, step, name=[paddle_op.name(), 'f_quotient_tensor']
     )
 
+    dtype = paddle_op.attrs().get("dtype")
+
     if start.dtype == trt.DataType.FLOAT:
         quotient_tensor = trt_cast(
             network,
@@ -150,7 +152,8 @@ def arange_converter(network, paddle_op, inputs):
     set_layer_name(fill_layer, paddle_op)
 
     output_tensor = fill_layer.get_output(0)
-    if start.dtype == trt.DataType.FLOAT:
+
+    if dtype == paddle.int64 or dtype == paddle.int32:
         output_tensor = trt_cast(
             network,
             output_tensor,
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -517,6 +517,7 @@ def slice_converter(network, paddle_op, inputs):
                 idx,
                 name=[paddle_op.name(), f'starts_tensor_{idx}'],
             )
+    start_tensor = trt_concat(network, starts_tensor)
 
     ends = get_input_constant_value(paddle_op, inputs, 2)
     if ends is not None:
@@ -568,21 +569,36 @@ def slice_converter(network, paddle_op, inputs):
     else:
         ends = inputs[2]
         for idx in range(len(axes)):
-            ends_tensor[axes[idx]] = get_shape_tensor_element(
+            axis = axes[idx]
+            input_dim = get_shape_tensor_element(
+                network,
+                input_shape_tensor,
+                axis,
+                name=[paddle_op.name(), f'input_dim_{idx}'],
+            )
+            end_element = get_shape_tensor_element(
                 network,
                 ends,
                 idx,
-                name=[paddle_op.name(), f'ends_tensor_{idx}'],
+                name=[paddle_op.name(), f'end_element_{idx}'],
             )
 
-    start_tensor_layer = network.add_concatenation(starts_tensor)
-    start_tensor_layer.axis = 0
-    set_layer_name(start_tensor_layer, paddle_op)
-    start_tensor = start_tensor_layer.get_output(0)
-    end_tensor_layer = network.add_concatenation(ends_tensor)
-    end_tensor_layer.axis = 0
-    set_layer_name(end_tensor_layer, paddle_op)
-    end_tensor = end_tensor_layer.get_output(0)
+            ends_tensor[axes[idx]] = trt_min(
+                network,
+                trt_max(
+                    network,
+                    end_element,
+                    add_1D_constant_layer(
+                        network, 0, name=[paddle_op.name(), 'zero_tensor_{idx}']
+                    ),
+                    name=[paddle_op.name(), 'trt_max_{idx}'],
+                ),
+                input_dim,
+                name=[paddle_op.name(), 'trt_min_{idx}'],
+            )
+    end_tensor = trt_concat(
+        network, ends_tensor, name=[paddle_op.name(), 'end_tensor']
+    )
     size_tensor = trt_sub(
         network,
         end_tensor,
@@ -613,7 +629,12 @@ def slice_converter(network, paddle_op, inputs):
             shuffle_layer = network.add_shuffle(output_tensor)
             shuffle_layer.reshape_dims = ()
         else:
-            real_size_tensor = trt_gather(network, size_tensor, gather_indices)
+            real_size_tensor = trt_gather(
+                network,
+                size_tensor,
+                gather_indices,
+                name=[paddle_op.name(), 'real_size_tensor'],
+            )
             shuffle_layer = network.add_shuffle(output_tensor)
             shuffle_layer.set_input(1, real_size_tensor)
 
diff --git a/test/ir/inference/CMakeLists.txt b/test/ir/inference/CMakeLists.txt
@@ -171,6 +171,7 @@ if(WITH_GPU AND TENSORRT_FOUND)
   set_tests_properties(test_trt_ops_fp32_mix_precision PROPERTIES TIMEOUT 300)
   set_tests_properties(test_trt_convert_unary PROPERTIES TIMEOUT 600)
   set_tests_properties(test_trt_convert_pool2d PROPERTIES TIMEOUT 600)
+  set_tests_properties(test_trt_convert_slice PROPERTIES TIMEOUT 600)
 
   if(NOT WIN32)
 
diff --git a/test/ir/inference/test_trt_convert_slice.py b/test/ir/inference/test_trt_convert_slice.py
@@ -106,13 +106,15 @@ def generate_input1(attrs: list[dict[str, Any]]):
 
                             yield program_config
 
+    def generate_dynamic_shape(self):
+        self.dynamic_shape.min_input_shape = {"input_data": [1, 3, 32, 32]}
+        self.dynamic_shape.max_input_shape = {"input_data": [8, 8, 64, 64]}
+        self.dynamic_shape.opt_input_shape = {"input_data": [6, 6, 64, 64]}
+        return self.dynamic_shape
+
     def sample_predictor_configs(
-        self, program_config
+        self, program_config, run_pir=False
     ) -> tuple[paddle_infer.Config, list[int], float]:
-        def generate_dynamic_shape(attrs):
-            self.dynamic_shape.min_input_shape = {"input_data": [1, 3, 32, 32]}
-            self.dynamic_shape.max_input_shape = {"input_data": [8, 8, 64, 64]}
-            self.dynamic_shape.opt_input_shape = {"input_data": [6, 6, 64, 64]}
 
         def clear_dynamic_shape():
             self.dynamic_shape.min_input_shape = {}
@@ -144,7 +146,7 @@ def generate_trt_nodes_num(attrs, dynamic_shape):
         ), 1e-3
 
         # for dynamic_shape
-        generate_dynamic_shape(attrs)
+        self.generate_dynamic_shape()
         self.trt_param.precision = paddle_infer.PrecisionType.Float32
         program_config.set_input_type(np.float32)
         yield self.create_inference_config(), generate_trt_nodes_num(
@@ -156,12 +158,15 @@ def generate_trt_nodes_num(attrs, dynamic_shape):
             attrs, True
         ), 1e-3
 
-    def test(self):
+    def test_old_ir(self):
         # TODO(inference): fix.
         # trt6 and trt7.1 has bug.
         # trt7.2 deserialize has bug.
         self.run_test()
 
+    def test_pir(self):
+        self.run_test(run_pir=True)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/tensorrt/test_converter_common.py b/test/tensorrt/test_converter_common.py
