PaddlePaddle
diff --git a/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+35-8Lines changed: 35 additions & 8 deletions b/‎paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
Copy file name to clipboard
+35-8Lines changed: 35 additions & 8 deletions
diff --git a/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
+109-35Lines changed: 109 additions & 35 deletions b/‎python/paddle/tensorrt/converter.py
Copy file name to clipboard
+109-35Lines changed: 109 additions & 35 deletions
diff --git a/‎python/paddle/tensorrt/converter_utils.py
Copy file name to clipboard
+38Lines changed: 38 additions & 0 deletions b/‎python/paddle/tensorrt/converter_utils.py
Copy file name to clipboard
+38Lines changed: 38 additions & 0 deletions
@@ -637,21 +637,48 @@ class SliceOpPattern : public pir::OpRewritePattern<paddle::dialect::SliceOp> {
     }
 
     auto axes_attr = op->attribute<pir::ArrayAttribute>("axes");
-
     std::vector<int64_t> axes;
     for (const auto &attr : axes_attr.AsVector()) {
       axes.push_back(attr.dyn_cast<pir::Int64Attribute>().data());
     }
-    pir::Value input = op.operand_source(0);
 
-    auto inputs = input.type().dyn_cast<paddle::dialect::DenseTensorType>();
-    auto inputs_shape = inputs.dims();
-    if (axes.size() !=
-        static_cast<std::vector<int64_t>::size_type>(inputs_shape.size())) {
-      VLOG(3) << "The shape of attributes of the slice operator axes "
-                 "and starts are not equal.";
+    size_t starts_size = axes.size();
+    size_t ends_size = axes.size();
+    if (pir::GetDefiningOpForInput(op, 1)
+            ->isa<paddle::dialect::FullIntArrayOp>()) {
+      paddle::dialect::FullIntArrayOp full_int_array_op_start =
+          pir::GetDefiningOpForInput(op, 1)
+              ->dyn_cast<paddle::dialect::FullIntArrayOp>();
+      auto starts_attr =
+          full_int_array_op_start->attribute<pir::ArrayAttribute>("value");
+      std::vector<int64_t> starts;
+      for (const auto &attr : starts_attr.AsVector()) {
+        starts.push_back(attr.dyn_cast<pir::Int64Attribute>().data());
+      }
+      starts_size = starts.size();
+    }
+
+    if (pir::GetDefiningOpForInput(op, 2)
+            ->isa<paddle::dialect::FullIntArrayOp>()) {
+      paddle::dialect::FullIntArrayOp full_int_array_op_end =
+          pir::GetDefiningOpForInput(op, 2)
+              ->dyn_cast<paddle::dialect::FullIntArrayOp>();
+      auto ends_attr =
+          full_int_array_op_end->attribute<pir::ArrayAttribute>("value");
+      std::vector<int64_t> ends;
+      for (const auto &attr : ends_attr.AsVector()) {
+        ends.push_back(attr.dyn_cast<pir::Int64Attribute>().data());
+      }
+      ends_size = ends.size();
+    }
+    if (starts_size != axes.size() || ends_size != axes.size()) {
+      VLOG(3) << "The size of axes and starts are not equal. "
+                 "Axes size: "
+              << axes.size() << ", Starts size: " << starts_size
+              << ", Ends size: " << ends_size;
       return false;
     }
+
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
     return true;
   }
 
@@ -80,7 +80,7 @@ def __init__(self, paddle_program, scope):
             # weights = trt.Weights(weight_array)
             param_dict.update({name: weight_array})
         self.param_dict = param_dict
-        self.shape_map = {}
+        self.input_info = {}
         self.trt_output_value_map = {}
 
     def find_graph_inputs_outputs(self, group_op):
@@ -135,6 +135,9 @@ def convert_subgraph_to_trt(self, program, group_op):
         min_shape_map = {}
         opt_shape_map = {}
         max_shape_map = {}
+        min_value_map = {}
+        opt_value_map = {}
+        max_value_map = {}
         input_names = []
 
         # Because one of the inputs to pd_op.concat is builtin.combine,
@@ -150,6 +153,7 @@ def convert_subgraph_to_trt(self, program, group_op):
             else:
                 origin_input_value.append(value)
 
+        # create TRT Weight and TRT Input
         for value in origin_input_value:
             defining_op = value.get_defining_op()
             if defining_op.name() == "builtin.parameter":
@@ -160,39 +164,10 @@ def convert_subgraph_to_trt(self, program, group_op):
             else:
                 shape = value.shape
                 dtype = map_dtype(value.dtype.name)
-                _logger.info(
-                    f"set shape of {value}, op is: {value.get_defining_op()}"
-                )
-                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
-                    min_shape = self.shape_map[value.id]["min_shape"]
-                    opt_shape = self.shape_map[value.id]["opt_shape"]
-                    max_shape = self.shape_map[value.id]["max_shape"]
-                else:
-                    min_shape = get_value_shape_range_info(
-                        value, False, paddle.base.core.ShapeMode.kMIN
-                    )
-                    opt_shape = get_value_shape_range_info(
-                        value, False, paddle.base.core.ShapeMode.kOPT
-                    )
-                    max_shape = get_value_shape_range_info(
-                        value, False, paddle.base.core.ShapeMode.kMAX
-                    )
-
                 input_name = f"input_{value.id}"
-
                 input_tensor = network.add_input(
                     name=input_name, dtype=dtype, shape=shape
                 )
-                _logger.info(f"set min_shape of {value} as {min_shape}")
-                _logger.info(f"set opt_shape of {value} as {opt_shape}")
-                _logger.info(f"set max_shape of {value} as {max_shape}")
-                profile.set_shape(
-                    input_name, min=min_shape, opt=opt_shape, max=max_shape
-                )
-                min_shape_map[input_name] = min_shape
-                opt_shape_map[input_name] = opt_shape
-                max_shape_map[input_name] = max_shape
                 input_names.append(input_name)
                 value_to_trt_tensor[value.id] = input_tensor
 
@@ -232,6 +207,79 @@ def convert_subgraph_to_trt(self, program, group_op):
                     value_to_trt_tensor[result.id] = trt_outs[idx]
                 else:
                     value_to_trt_tensor[result.id] = None
+
+        # Set TRT min/opt/max input shape and the value of shape tensor
+        for value in origin_input_value:
+            trt_input = value_to_trt_tensor[value.id]
+            if isinstance(trt_input, trt.Weights):
+                continue
+            input_name = trt_input.name
+            if input_name != "":
+                _logger.info(
+                    f"set shape of {value}, op is: {value.get_defining_op()}"
+                )
+                min_shape = []
+                opt_shape = []
+                max_shape = []
+                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
+                    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"]
+                    if trt_input.is_shape_tensor:
+                        min_value = self.input_info[value.id]["min_value"]
+                        opt_value = self.input_info[value.id]["opt_value"]
+                        max_value = self.input_info[value.id]["max_value"]
+                else:
+                    min_shape = get_value_shape_range_info(
+                        value, False, paddle.base.core.ShapeMode.kMIN
+                    )
+                    opt_shape = get_value_shape_range_info(
+                        value, False, paddle.base.core.ShapeMode.kOPT
+                    )
+                    max_shape = get_value_shape_range_info(
+                        value, False, paddle.base.core.ShapeMode.kMAX
+                    )
+                    if trt_input.is_shape_tensor:
+                        min_value = get_value_shape_range_info(
+                            value, True, paddle.base.core.ShapeMode.kMIN
+                        )
+                        opt_value = get_value_shape_range_info(
+                            value, True, paddle.base.core.ShapeMode.kOPT
+                        )
+                        max_value = get_value_shape_range_info(
+                            value, True, paddle.base.core.ShapeMode.kMAX
+                        )
+                _logger.info(f"set min_shape of {value} as {min_shape}")
+                _logger.info(f"set opt_shape of {value} as {opt_shape}")
+                _logger.info(f"set max_shape of {value} as {max_shape}")
+                profile.set_shape(
+                    input_name, min=min_shape, opt=opt_shape, max=max_shape
+                )
+                if trt_input.is_shape_tensor:
+                    _logger.info(
+                        f"set min_value of shape input: {value} as {min_value}"
+                    )
+                    _logger.info(
+                        f"set max_value of shape input: {value} as {opt_value}"
+                    )
+                    _logger.info(
+                        f"set opt_value of shape input: {value} as {max_value}"
+                    )
+                    profile.set_shape_input(
+                        input_name, min=min_value, opt=opt_value, max=max_value
+                    )
+
+                min_shape_map[input_name] = min_shape
+                opt_shape_map[input_name] = opt_shape
+                max_shape_map[input_name] = max_shape
+                min_value_map[input_name] = min_value
+                opt_value_map[input_name] = opt_value
+                max_value_map[input_name] = max_value
+
         out_shapes = []
         out_names = []
         out_types = []
@@ -259,10 +307,27 @@ def convert_subgraph_to_trt(self, program, group_op):
             max_shape = get_value_shape_range_info(
                 result_value, False, paddle.base.core.ShapeMode.kMAX
             )
-            self.shape_map[result_value.id] = {
+            min_value = []
+            opt_value = []
+            max_value = []
+            if output_tensor.is_shape_tensor:
+                min_value = get_value_shape_range_info(
+                    result_value, True, paddle.base.core.ShapeMode.kMIN
+                )
+                opt_value = get_value_shape_range_info(
+                    result_value, True, paddle.base.core.ShapeMode.kOPT
+                )
+                max_value = get_value_shape_range_info(
+                    result_value, True, paddle.base.core.ShapeMode.kMAX
+                )
+
+            self.input_info[result_value.id] = {
                 "min_shape": min_shape,
                 "opt_shape": opt_shape,
                 "max_shape": max_shape,
+                "min_value": min_value,
+                "opt_value": opt_value,
+                "max_value": max_value,
             }
 
         config = builder.create_builder_config()
@@ -277,6 +342,9 @@ def convert_subgraph_to_trt(self, program, group_op):
         trt_params.min_input_shape = min_shape_map
         trt_params.max_input_shape = max_shape_map
         trt_params.optim_input_shape = opt_shape_map
+        trt_params.min_shape_tensor = min_value_map
+        trt_params.max_shape_tensor = max_value_map
+        trt_params.optim_shape_tensor = opt_value_map
         group_str = str(group_op)
         engine_name = (
             int(hashlib.sha256(group_str.encode('utf-8')).hexdigest(), 16)
@@ -309,13 +377,19 @@ def convert_subgraph_to_trt(self, program, group_op):
                     continue
                 ori_value = output_values[out_index]
                 current_value = out[out_index]
-                orin_min_shape = self.shape_map[ori_value.id]["min_shape"]
-                orin_opt_shape = self.shape_map[ori_value.id]["opt_shape"]
-                orin_max_shape = self.shape_map[ori_value.id]["max_shape"]
-                self.shape_map[current_value.id] = {
+                orin_min_shape = self.input_info[ori_value.id]["min_shape"]
+                orin_opt_shape = self.input_info[ori_value.id]["opt_shape"]
+                orin_max_shape = self.input_info[ori_value.id]["max_shape"]
+                orin_min_value = self.input_info[ori_value.id]["min_value"]
+                orin_opt_value = self.input_info[ori_value.id]["opt_value"]
+                orin_max_value = self.input_info[ori_value.id]["max_value"]
+                self.input_info[current_value.id] = {
                     "min_shape": orin_min_shape,
                     "opt_shape": orin_opt_shape,
                     "max_shape": orin_max_shape,
+                    "min_value": orin_min_value,
+                    "opt_value": orin_opt_value,
+                    "max_value": orin_max_value,
                 }
 
         return out
 
@@ -155,3 +155,41 @@ def add_elementwise_layer(network, paddle_op, inputs, op_type):
     )
     layer = network.add_elementwise(lhs_val, rhs_val, op_type)
     return layer.get_output(0)
+
+
+# Create and add 1D constant layer
+def add_1D_constant_layer(network, data, dtype=np.int32):
+    constant_data = np.array([data], dtype=dtype)
+    constant_layer = network.add_constant(constant_data.shape, constant_data)
+    return constant_layer.get_output(0)
+
+
+# Get element tensor of 1D shape tensor
+def get_shape_tensor_element(network, x, index):
+    assert index >= 0, (
+        "The index should be greater or equal than 0, but got %d" % index
+    )
+    gather_layer = network.add_gather(
+        input=x, indices=add_1D_constant_layer(network, index), axis=0
+    )
+    return gather_layer.get_output(0)
+
+
+def trt_sum(network, a, b):
+    layer = network.add_elementwise(a, b, trt.ElementWiseOperation.SUM)
+    return layer.get_output(0)
+
+
+def trt_max(network, a, b):
+    layer = network.add_elementwise(a, b, trt.ElementWiseOperation.MAX)
+    return layer.get_output(0)
+
+
+def trt_sub(network, a, b):
+    layer = network.add_elementwise(a, b, trt.ElementWiseOperation.SUB)
+    return layer.get_output(0)
+
+
+def trt_min(network, a, b):
+    layer = network.add_elementwise(a, b, trt.ElementWiseOperation.MIN)
+    return layer.get_output(0)