[Inference]Improve pir-trt performance Part-2 (#71712)

YuanRisheng · web-flow · commit 38f6a5a0d6ae · 2025-03-31T14:38:45.000+08:00
* fix performance

* add meshgrid

* fix ci

* fix ci

* add coverage

* del code

* fix coverage
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -126,6 +126,7 @@ 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)
+DEFINE_GENERAL_PATTERN(Meshgrid, paddle::dialect::MeshgridOp)
 
 #undef DEFINE_GENERAL_PATTERN
 
@@ -927,10 +928,16 @@ class UnsqueezeOpPattern
         dynamic_dims.push_back(i);
       }
     }
-    if (dynamic_dims.size() > 1) {
-      VLOG(3) << "Currently we don't support unsqueeze with more than one "
-                 "dynamic dims";
-      return false;
+    if (dynamic_dims.size() == 0) {
+      std::vector<int64_t> axes;
+      for (auto &axis_ele : axis.AsVector()) {
+        axes.push_back(axis_ele.dyn_cast<pir::Int64Attribute>().data());
+      }
+      if (std::find(axes.begin(), axes.end(), 0) != axes.end()) {
+        VLOG(3) << "Invalid squeeze axes. Axes having batch axis is not "
+                   "supported in static shape";
+        return false;
+      }
     }
 
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
@@ -967,10 +974,16 @@ class Unsqueeze_OpPattern
         dynamic_dims.push_back(i);
       }
     }
-    if (dynamic_dims.size() > 1) {
-      VLOG(3) << "Currently we don't support unsqueeze with more than one "
-                 "dynamic dims";
-      return false;
+    if (dynamic_dims.size() == 0) {
+      std::vector<int64_t> axes;
+      for (auto &axis_ele : axis.AsVector()) {
+        axes.push_back(axis_ele.dyn_cast<pir::Int64Attribute>().data());
+      }
+      if (std::find(axes.begin(), axes.end(), 0) != axes.end()) {
+        VLOG(3) << "Invalid squeeze axes. Axes having batch axis is not "
+                   "supported in static shape";
+        return false;
+      }
     }
 
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
@@ -1953,8 +1966,12 @@ class StackOpPattern : public pir::OpRewritePattern<paddle::dialect::StackOp> {
     pir::Value x = op.operand_source(0);
     int rank = 1;
     auto x_type = x.type();
-    if (x_type.isa<pir::VectorType>()) {
-      rank = x_type.dyn_cast<pir::VectorType>().size();
+    if (x_type.isa<pir::VectorType>() &&
+        x_type.dyn_cast<pir::VectorType>().size() > 0) {
+      auto vec_type = x_type.dyn_cast<pir::VectorType>();
+      auto tensor_element =
+          vec_type.data()[0].dyn_cast<paddle::dialect::DenseTensorType>();
+      rank = tensor_element.dims().size();
     } else {
       auto x_shape = pir::GetShapeFromValue(x);
       rank = x_shape.size();
@@ -3004,6 +3021,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Acos)
     ADD_PATTERN(Atan)
     ADD_PATTERN(ShuffleChannel)
+    ADD_PATTERN(Meshgrid)
 #if IS_TRT_VERSION_GE(8600)
     ADD_PATTERN(Layer_norm)
 #endif
diff --git a/python/paddle/tensorrt/converter.py b/python/paddle/tensorrt/converter.py
@@ -586,8 +586,42 @@ def convert_program_to_trt(self):
                 if op.results()[0].use_empty():
                     self.program.global_block().remove_op(op)
             if op.name() == "builtin.constant":
+                # builtin.constant can't be saved/loaded, we need del it
                 if op.results()[0].use_empty():
                     self.program.global_block().remove_op(op)
+                else:
+                    constant_result = op.results()[0]
+                    constant_value_name = op.attrs()["value"]
+                    out_dtype = np.dtype(
+                        paddle.pir.core._PADDLE_PIR_DTYPE_2_NUMPY_DTYPE[
+                            constant_result.dtype
+                        ]
+                    )
+                    tensor_data = self.scope.var(
+                        constant_value_name
+                    ).get_tensor()
+                    constant_array = np.array(
+                        tensor_data, dtype=out_dtype
+                    ).tolist()
+
+                    # convert builtin.constant to pd_op.full_int_array/full and then delete it
+                    with paddle.pir.core.program_guard(self.program):
+                        paddle.base.libpaddle.pir.reset_insertion_point_to_start()
+                        if len(constant_array) == 1:
+                            full_value = paddle._C_ops.full(
+                                [1],
+                                constant_array[0],
+                                constant_result.dtype,
+                                paddle.CUDAPlace(0),
+                            )
+                        else:
+                            full_value = paddle._C_ops.full_int_array(
+                                constant_array,
+                                constant_result.dtype,
+                                paddle.CUDAPlace(0),
+                            )
+                    op.replace_all_uses_with([full_value])
+                    self.program.global_block().remove_op(op)
 
         # Call clear_shape_info to clear the previous shape information
         clear_shape_info()
diff --git a/python/paddle/tensorrt/impls/creation.py b/python/paddle/tensorrt/impls/creation.py
@@ -371,3 +371,79 @@ def full_with_tensor_converter(network, paddle_op, inputs):
     set_layer_name(fill_layer, paddle_op)
     output_tensor = fill_layer.get_output(0)
     return output_tensor
+
+
+@converter_registry.register("pd_op.meshgrid", trt_version="8.x")
+def meshgrid_converter(network, paddle_op, vec_inputs):
+    inputs = vec_inputs[0]
+    n = len(inputs)
+    outputs = []
+
+    # get all input dims (all input is 1-dim)
+    input_dims = [network.add_shape(inp).get_output(0) for inp in inputs]
+
+    for k in range(n):
+        # --------------------------------
+        # step1:reshape k input as [1,..,Dk,..,1]
+        # --------------------------------
+        x = inputs[k]
+        reshape_dims = []  # init dims as 1
+        for i in range(n):
+            one = add_1D_constant_layer(
+                network,
+                1,
+                dtype=np.int32,
+                is_scalar=False,
+                name=[paddle_op.name(), f'one_{k}'],
+            )
+            reshape_dims.append(one)
+        # replace k-th input dim as Dk
+        reshape_dims[k] = input_dims[k]
+
+        dim_concat = network.add_concatenation(reshape_dims)
+        set_layer_name(dim_concat, paddle_op)
+        x_reshaped = network.add_shuffle(x)
+        x_reshaped.set_input(1, dim_concat.get_output(0))
+
+        # --------------------------------
+        # step2: create tensor([D1, D2, ..., 1, ..., Dn]) that filled with 1
+        # --------------------------------
+        ones_shape = []
+        for i in range(n):
+            ones_shape.append(input_dims[i])
+        ones_shape[k] = add_1D_constant_layer(
+            network,
+            1,
+            dtype=np.int32,
+            is_scalar=False,
+            name=[paddle_op.name(), f'ones_shape_{k}'],
+        )
+        dim_concat = network.add_concatenation(ones_shape)
+        set_layer_name(dim_concat, paddle_op)
+
+        # Fill constant 1
+        fill_layer = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
+        fill_layer.set_input(0, dim_concat.get_output(0))
+        value_input = add_1D_constant_layer(
+            network,
+            1,
+            dtype=np.float32,
+            is_scalar=True,
+            name=[paddle_op.name(), 'one_for_fill'],
+        )
+        fill_layer.set_input(1, value_input)
+        beta_vec = [0] * n
+        fill_layer.set_input(
+            2, add_1D_constant_layer(network, beta_vec, np.float32)
+        )
+
+        # --------------------------------
+        # step3: element wise multiplication
+        # --------------------------------
+        grid = network.add_elementwise(
+            x_reshaped.get_output(0),
+            fill_layer.get_output(0),
+            trt.ElementWiseOperation.PROD,
+        ).get_output(0)
+        outputs.append(grid)
+    return outputs
diff --git a/python/paddle/tensorrt/util.py b/python/paddle/tensorrt/util.py
@@ -52,19 +52,10 @@ def map_dtype(pd_dtype):
         raise TypeError(f"Unsupported dtype: {pd_dtype}")
 
 
-def all_ops_into_trt(program):
+def support_constant_folding_pass(program):
     for op in program.global_block().ops:
-        if (
-            op.name() == "pd_op.fetch"
-            or op.name() == "pd_op.data"
-            or op.name().split('.')[0] == "builtin"
-        ):
-            continue
-        if op.has_attr("__l_trt__") is False:
+        if op.name() == "pd_op.while" or op.name() == "pd_op.if":
             return False
-        if op.attrs()["__l_trt__"] is False:
-            return False
-    _logger.info("All ops convert to trt.")
     return True
 
 
@@ -107,7 +98,7 @@ def _add_pass_(pm, passes, disable_passes):
     # run other passes
     pm.clear()
     passes = []
-    if all_ops_into_trt(program):
+    if support_constant_folding_pass(program):
         # only run constant_folding_pass when all ops into trt
         passes.append(
             {
@@ -117,18 +108,19 @@ def _add_pass_(pm, passes, disable_passes):
                 }
             }
         )
-
+        passes.append(
+            {
+                'dead_code_elimination_pass': {
+                    "__place__": place,
+                    "__param_scope__": scope,
+                }
+            }
+        )
         passes.append({'conv2d_add_fuse_pass': {}})
     passes.append({'trt_op_marker_pass': {}})  # for op that created by pass
     _add_pass_(pm, passes, disable_passes)
     pm.run(program)
 
-    # delete unused op
-    for op in program.global_block().ops:
-        if op.name() == "builtin.constant" or op.name() == "builtin.parameter":
-            if op.results()[0].use_empty():
-                program.global_block().remove_op(op)
-
     return program
 
 
@@ -282,6 +274,7 @@ def weight_to_tensor(network, paddle_value, trt_tensor, use_op_name):
     # the following op needn't cast trt.Weight to ITensor, because the layer need weight as input
     forbid_cast_op = [
         "pd_op.depthwise_conv2d",
+        "pd_op.conv2d",
         "pd_op.conv2d_transpose",
         "pd_op.conv3d",
         "pd_op.conv3d_transpose",
diff --git a/test/tensorrt/tensorrt_test_base.py b/test/tensorrt/tensorrt_test_base.py
@@ -46,6 +46,7 @@ def __init__(self, methodName='runTest'):
         self.dynamic_shape_data = {}
         self.disable_passes = [
             "constant_folding_pass",
+            "dead_code_elimination_pass",
         ]
 
     def create_fake_program(self):
@@ -267,6 +268,14 @@ def check_trt_result(self, rtol=1e-5, atol=1e-5, precision_mode="fp32"):
                 main_program,
                 disable_passes=self.disable_passes,
             )
+            # delete unused op
+            for op in main_program.global_block().ops:
+                if (
+                    op.name() == "builtin.constant"
+                    or op.name() == "builtin.parameter"
+                ):
+                    if op.results()[0].use_empty():
+                        main_program.global_block().remove_op(op)
 
             scope = paddle.static.global_scope()
             main_program = warmup_shape_infer(
diff --git a/test/tensorrt/test_converter_conv.py b/test/tensorrt/test_converter_conv.py
@@ -41,7 +41,11 @@ def setUp(self):
         self.min_shape = {"x": [1, 3, 8, 8]}
         self.opt_shape = {"x": [2, 3, 8, 8]}
         self.max_shape = {"x": [10, 3, 8, 8]}
-        self.disable_passes = ['constant_folding_pass', 'conv2d_add_fuse_pass']
+        self.disable_passes = [
+            'constant_folding_pass',
+            'conv2d_add_fuse_pass',
+            'dead_code_elimination_pass',
+        ]
 
     def test_trt_result_fp16(self):
         self.check_trt_result(precision_mode="fp16")
@@ -62,7 +66,11 @@ def setUp(self):
         self.min_shape = {"x": [1, 3, 8, 8]}
         self.opt_shape = {"x": [2, 3, 8, 8]}
         self.max_shape = {"x": [10, 3, 8, 8]}
-        self.disable_passes = ['constant_folding_pass', 'conv2d_add_fuse_pass']
+        self.disable_passes = [
+            'constant_folding_pass',
+            'conv2d_add_fuse_pass',
+            'dead_code_elimination_pass',
+        ]
 
     def test_trt_result(self):
         self.check_trt_result()
@@ -81,7 +89,11 @@ def setUp(self):
         self.min_shape = {"x": [1, 3, 8, 8]}
         self.opt_shape = {"x": [2, 3, 8, 8]}
         self.max_shape = {"x": [10, 3, 8, 8]}
-        self.disable_passes = ['constant_folding_pass', 'conv2d_add_fuse_pass']
+        self.disable_passes = [
+            'constant_folding_pass',
+            'conv2d_add_fuse_pass',
+            'dead_code_elimination_pass',
+        ]
 
     def test_trt_result(self):
         self.check_trt_result()
@@ -489,7 +501,7 @@ def setUp(self):
         self.min_shape = {"x": [1, 3, 8, 8]}
         self.opt_shape = {"x": [2, 3, 8, 8]}
         self.max_shape = {"x": [10, 3, 8, 8]}
-        self.disable_passes = []
+        self.disable_passes = ['dead_code_elimination_pass']
 
     def test_trt_result_fp16(self):
         self.check_trt_result(precision_mode="fp16")
diff --git a/test/tensorrt/test_converter_creation.py b/test/tensorrt/test_converter_creation.py
@@ -251,5 +251,23 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestMeshgridTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.meshgrid
+        self.api_args = {
+            "x": [
+                np.random.random([20]).astype("float32"),
+                np.random.random([30]).astype("float32"),
+            ],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [[10], [20]]}
+        self.opt_shape = {"x": [[20], [30]]}
+        self.max_shape = {"x": [[30], [40]]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -1101,5 +1101,41 @@ def test_fp16_result(self):
         self.check_trt_result(precision_mode="fp16")
 
 
+class TestUnsqueezeTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.unsqueeze
+        self.api_args = {
+            "x": np.random.random([5, 10]).astype("float32"),
+            "axis": 0,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {}
+        self.opt_shape = {}
+        self.max_shape = {}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
+def unsqueeze_inplace_wrapper(x, axis):
+    return _C_ops.unsqueeze_(x, axis)
+
+
+class TestUnsqueeze_TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = unsqueeze_inplace_wrapper
+        self.api_args = {
+            "x": np.random.random([5, 10]).astype("float32"),
+            "axis": 0,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {}
+        self.opt_shape = {}
+        self.max_shape = {}
+
+    def test_trt_result(self):
+        self.check_marker(expected_result=False)
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/tensorrt/test_converter_model_bert.py b/test/tensorrt/test_converter_model_bert.py
diff --git a/test/tensorrt/test_converter_model_dummy.py b/test/tensorrt/test_converter_model_dummy.py
diff --git a/test/tensorrt/test_converter_model_resnet50.py b/test/tensorrt/test_converter_model_resnet50.py
