[Inference]Fix PIR-TRT bugs (#68776)

YuanRisheng · web-flow · commit cdc9d69b05c1 · 2024-11-05T20:30:09.000+08:00
* fix model bug

* re-trigger

* fix matmul bug

* fix bugs
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -611,11 +611,12 @@ class ScaleOpPattern : public pir::OpRewritePattern<paddle::dialect::ScaleOp> {
     }
     pir::Value x = op.operand_source(0);
     auto x_dtype = pir::GetDataTypeFromValue(x);
+    // TODO(YuanRisheng): The trt(<=8.5) can't support cast layer, we need
+    // support int32 and int64 after we upgrade our trt version
     if (!(x_dtype.isa<pir::Float32Type>() || x_dtype.isa<pir::Float64Type>() ||
-          x_dtype.isa<pir::Float16Type>() || x_dtype.isa<pir::Int32Type>() ||
-          x_dtype.isa<pir::Int64Type>())) {
+          x_dtype.isa<pir::Float16Type>())) {
       VLOG(3) << "At present, ScaleOp only support float32 or float16 or "
-                 "float64 or int32 or int64 into trt.";
+                 "float64 into trt.";
       return false;
     }
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
diff --git a/python/paddle/tensorrt/converter.py b/python/paddle/tensorrt/converter.py
@@ -68,6 +68,16 @@ def get_trt_version():
     return trt.__version__
 
 
+def remove_duplicate_value(value_list):
+    ret_list = []
+    ret_list_id = []
+    for value in value_list:
+        if value.id not in ret_list_id:
+            ret_list.append(value)
+            ret_list_id.append(value.id)
+    return ret_list
+
+
 class PaddleToTensorRTConverter:
     def __init__(self, paddle_program, scope):
         self.scope = scope
@@ -155,6 +165,7 @@ def convert_subgraph_to_trt(self, program, group_op):
             else:
                 origin_input_value.append(value)
 
+        origin_input_value = remove_duplicate_value(origin_input_value)
         # create TRT Weight and TRT Input
         for value in origin_input_value:
             defining_op = value.get_defining_op()
@@ -242,8 +253,16 @@ def convert_subgraph_to_trt(self, program, group_op):
                 min_value = []
                 opt_value = []
                 max_value = []
-                if value.get_defining_op().name() == "builtin.split":
-                    # TODO if the input value is generated by the other trt_engine_op, so the shape is searched by origin value
+
+                value_define_op = value.get_defining_op()
+                # if the input value is generated by the other trt_engine_op, so the shape is searched by origin value
+                if (
+                    value_define_op.name() == "builtin.split"
+                    and value_define_op.operand_source(0)
+                    .get_defining_op()
+                    .name()
+                    == "pd_op.tensorrt_engine"
+                ):
                     min_shape = self.input_info[value.id]["min_shape"]
                     opt_shape = self.input_info[value.id]["opt_shape"]
                     max_shape = self.input_info[value.id]["max_shape"]
diff --git a/python/paddle/tensorrt/converter_utils.py b/python/paddle/tensorrt/converter_utils.py
@@ -19,8 +19,6 @@
 import numpy as np
 import tensorrt as trt
 
-from .util import get_trt_version_list
-
 current_dir = os.path.dirname(os.path.abspath(__file__))
 parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
 if parent_dir not in sys.path:
@@ -33,7 +31,8 @@
     __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
 )
 
-version_list = get_trt_version_list()
+version = trt.__version__
+version_list = list(map(int, version.split('.')))
 
 
 def has_dynamic_shape(shape):
diff --git a/python/paddle/tensorrt/export.py b/python/paddle/tensorrt/export.py
@@ -35,6 +35,7 @@
 from paddle.tensorrt.converter import PaddleToTensorRTConverter
 from paddle.tensorrt.util import (
     forbid_op_lower_trt,
+    mark_buitlin_op,
     run_pir_pass,
     warmup_shape_infer,
 )
@@ -231,6 +232,9 @@ def convert_to_trt(program, trt_config, scope):
         if trt_config.disable_ops:
             forbid_op_lower_trt(program, trt_config.disable_ops)
 
+        # Adding marker labels to builtin ops facilitates convert processing, but they ultimately do not enter the TensorRT subgraph.
+        mark_buitlin_op(program)
+
         # run pir pass (including trt_sub_graph_extract_pass)
         program_with_pir = run_pir_pass(program, partition_mode=True)
 
diff --git a/python/paddle/tensorrt/impls/linalg.py b/python/paddle/tensorrt/impls/linalg.py
@@ -42,6 +42,12 @@ def matmul_converter(network, paddle_op, inputs):
         weight_tensor = network.add_constant(
             weight_shape, inputs[1]
         ).get_output(0)
+
+    if len(weight_shape) == 1:
+        layer = network.add_shuffle(weight_tensor)
+        layer.reshape_dims = (*tuple(weight_shape), 1)
+        weight_tensor = layer.get_output(0)
+
     lhs_val, rhs_val = broadcast(
         network, inputs[0], weight_tensor, inputs[0].name, weight_tensor.name
     )
