[Paddle Tensorrt] add pd_op.stride_slice converter (#69038)

lizexu123 · web-flow · commit 91c16108cd3e · 2024-11-04T11:08:47.000+08:00
* pd_op.stride_slice

* fix

* fix

* 加一个单测

* fix

* 解决单测随机挂现象,对于普通张量,使用profile.set_shape,对于shape_tensor，使用profile.set_shape_input,重复使用profile.set_shape

* fix

* fix

* fix

* fix

* fix

* fix
diff --git a/paddle/fluid/framework/new_executor/collect_shape_manager.h b/paddle/fluid/framework/new_executor/collect_shape_manager.h
@@ -55,6 +55,19 @@ class CollectShapeManager {
                                               bool is_shape_tensor,
                                               ShapeMode shape_mode);
 
+  void ClearShapeInfo() {
+    shape_info_.clear();
+    shape_tensor_info_.clear();
+    min_shapes_.clear();
+    max_shapes_.clear();
+    opt_shapes_.clear();
+    min_values_.clear();
+    max_values_.clear();
+    opt_values_.clear();
+    op_value2kernel_value_.clear();
+    is_shape_range_info_ready_ = false;
+  }
+
  private:
   CollectShapeManager() {}
   std::unordered_map<pir::Value, pir::Value> op_value2kernel_value_;
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1648,6 +1648,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<StackOpPattern>(context));
     ps.Add(std::make_unique<TanhOpPattern>(context));
     ps.Add(std::make_unique<FullWithTensorPattern>(context));
+    ps.Add(std::make_unique<StridedSliceOpPattern>(context));
     return ps;
   }
 };
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -3279,6 +3279,7 @@ All parameter, weight, gradient are variables in Paddle.
            bool is_shape_tensor,
            paddle::framework::ShapeMode shape_mode) -> py::list {
           py::list res;
+
           paddle::framework::CollectShapeManager::Instance()
               .StatisticShapeRangeInfo();
           auto shape_result =
@@ -3289,6 +3290,9 @@ All parameter, weight, gradient are variables in Paddle.
           }
           return res;
         });
+  m.def("clear_shape_info", []() {
+    paddle::framework::CollectShapeManager::Instance().ClearShapeInfo();
+  });
 
 #if defined(PADDLE_WITH_PSLIB) && !defined(PADDLE_WITH_HETERPS)
   BindHeterWrapper(&m);
diff --git a/python/paddle/tensorrt/converter.py b/python/paddle/tensorrt/converter.py
@@ -26,7 +26,7 @@
 
 import paddle
 from paddle import pir
-from paddle.base.core import get_value_shape_range_info
+from paddle.base.core import clear_shape_info, get_value_shape_range_info
 from paddle.base.log_helper import get_logger
 
 from .impls.activation import *  # noqa: F403
@@ -123,6 +123,7 @@ def __is_output_value(value):
 
     def convert_subgraph_to_trt(self, program, group_op):
         _logger.info(f"start process {group_op}")
+
         operations = next(iter(group_op.blocks())).ops
         input_values, output_values = self.find_graph_inputs_outputs(group_op)
         builder = trt.Builder(trt.Logger(trt.Logger.ERROR))
@@ -270,13 +271,14 @@ def convert_subgraph_to_trt(self, program, group_op):
                         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:
+                if not trt_input.is_shape_tensor:
+                    _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
+                    )
+                else:
                     _logger.info(
                         f"set min_value of shape input: {value} as {min_value}"
                     )
@@ -438,3 +440,5 @@ def convert_program_to_trt(self):
                     orin_out_values[o_i].replace_all_uses_with(new_out[o_i])
 
                 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/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -734,54 +734,93 @@ def strided_slice_converter(network, paddle_op, inputs):
 
     trt_start_dims = [0] * nchw_input_dims
     trt_size_dims = [input_shape[i] for i in range(nchw_input_dims)]
-    trt_end_dims = [0] * nchw_input_dims
     trt_step_dims = [1] * nchw_input_dims
 
     has_neg_indices = False
+    trt_start_tensors = []
+    trt_end_tensors = []
+    trt_stride_tensors = []
+
     for i, axis in enumerate(axes):
-        trt_start_dims[axis] = starts[i]
-        trt_end_dims[axis] = ends[i]
-        trt_step_dims[axis] = strides[i]
-        trt_size_dims[axis] = max(
-            0, (ends[i] - starts[i] + strides[i] - 1) // strides[i]
-        )
-        if starts[i] < 0 or ends[i] < 0:
-            has_neg_indices = True
+        if isinstance(starts, trt.ITensor):
+            start_tensor = get_shape_tensor_element(network, starts, i)
+        else:
+            start_tensor = add_1D_constant_layer(network, [starts[i]])
 
-    shape_tensor = network.add_shape(input_tensor).get_output(0)
-    start_tensor = add_1D_constant_layer(network, trt_start_dims)
+        if isinstance(ends, trt.ITensor):
+            end_tensor = get_shape_tensor_element(network, ends, i)
+        else:
+            end_tensor = add_1D_constant_layer(network, [ends[i]])
 
-    if has_neg_indices:
-        start_tensor = fix_negative_indices(network, shape_tensor, start_tensor)
+        if isinstance(strides, trt.ITensor):
+            stride_tensor = get_shape_tensor_element(network, strides, i)
+        else:
+            stride_tensor = add_1D_constant_layer(network, [strides[i]])
 
-    end_vec_tensor = []
-    for i in range(len(trt_end_dims)):
-        end_vec_tensor.append(
-            get_shape_tensor_element(network, shape_tensor, i)
-        )
+        zero_tensor = add_1D_constant_layer(network, [0])
 
-    for i, axis in enumerate(axes):
-        if ends[i] >= 0:
-            end_vec_tensor[axis] = network.add_constant(
-                (1,), np.array([ends[i]], dtype=np.int32)
+        if isinstance(starts, trt.ITensor) or isinstance(ends, trt.ITensor):
+            is_start_neg = trt_less(network, start_tensor, zero_tensor)
+            is_end_neg = trt_less(network, end_tensor, zero_tensor)
+            temp_has_neg = network.add_elementwise(
+                is_start_neg, is_end_neg, trt.ElementWiseOperation.OR
             ).get_output(0)
+            if not has_neg_indices:
+                has_neg_indices = temp_has_neg
+            else:
+                has_neg_indices = network.add_elementwise(
+                    has_neg_indices, temp_has_neg, trt.ElementWiseOperation.OR
+                ).get_output(0)
         else:
-            adjusted_end = network.add_constant(
-                (1,), np.array([ends[i]], dtype=np.int32)
-            ).get_output(0)
-            end_vec_tensor[axis] = trt_sum(
-                network, end_vec_tensor[axis], adjusted_end
-            )
+            if starts[i] < 0 or ends[i] < 0:
+                has_neg_indices = True
+
+        trt_start_tensors.append(start_tensor)
+        trt_end_tensors.append(end_tensor)
+        trt_stride_tensors.append(stride_tensor)
+
+    # Concatenate the tensors for start, end, and strides
+    start_tensor = network.add_concatenation(trt_start_tensors).get_output(0)
+    end_tensor = network.add_concatenation(trt_end_tensors).get_output(0)
+    step_tensor = network.add_concatenation(trt_stride_tensors).get_output(0)
+
+    shape_tensor = network.add_shape(input_tensor).get_output(0)
+
+    if has_neg_indices is True:
+        start_tensor = fix_negative_indices(network, shape_tensor, start_tensor)
+    elif isinstance(has_neg_indices, trt.ITensor):
+        fixed_start_tensor = fix_negative_indices(
+            network, shape_tensor, start_tensor
+        )
+        start_tensor = network.add_select(
+            condition=has_neg_indices,
+            then_input=fixed_start_tensor,
+            else_input=start_tensor,
+        ).get_output(0)
+
+    # Process end_tensor similarly to handle negative indices
+    if has_neg_indices is True:
+        end_tensor = fix_negative_indices(network, shape_tensor, end_tensor)
+    elif isinstance(has_neg_indices, trt.ITensor):
+        fixed_end_tensor = fix_negative_indices(
+            network, shape_tensor, end_tensor
+        )
+        end_tensor = network.add_select(
+            condition=has_neg_indices,
+            then_input=fixed_end_tensor,
+            else_input=end_tensor,
+        ).get_output(0)
 
-    concat_end_tensor = network.add_concatenation(end_vec_tensor).get_output(0)
-    min_tensor = trt_min(network, concat_end_tensor, shape_tensor)
+    # Compute min_tensor
+    min_tensor = trt_min(network, end_tensor, shape_tensor)
+    # Correct size_tensor calculation
     size_tensor = trt_sub(network, start_tensor, min_tensor)
 
-    zero_t = add_1D_constant_layer(network, [0] * nchw_input_dims)
-    step_tensor = add_1D_constant_layer(network, trt_step_dims)
+    # floor_div_tensor computation
     floor_div_tensor = trt_floor_div(network, size_tensor, step_tensor)
-    size_tensor = trt_sub(network, zero_t, floor_div_tensor)
+    size_tensor = trt_sub(network, zero_tensor, floor_div_tensor)
 
+    # Create the slice layer
     layer = network.add_slice(
         input_tensor, trt_start_dims, trt_size_dims, trt_step_dims
     )
diff --git a/setup.py b/setup.py
@@ -2026,6 +2026,7 @@ def get_setup_parameters():
         'paddle.decomposition',
         'paddle._typing',
         'paddle._typing.libs',
+        'paddle.tensorrt',
     ]
 
     paddle_bins = ''
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -385,15 +385,15 @@ class TestStrideSliceCase1TRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.strided_slice
         self.api_args = {
-            "x": np.random.random([3, 4, 5, 6]).astype("float32"),
+            "x": np.random.random([3, 4, 10]).astype("float32"),
             "axes": [0, 1, 2],
             "starts": [1, 0, 2],
             "ends": [2, 3, 4],
             "strides": [1, 1, 1],
         }
         self.program_config = {"feed_list": ["x"]}
-        self.min_shape = {"x": [1, 4, 5, 6]}
-        self.max_shape = {"x": [5, 4, 5, 6]}
+        self.min_shape = {"x": [1, 4, 10]}
+        self.max_shape = {"x": [5, 4, 10]}
 
     def test_trt_result(self):
         self.check_trt_result()
@@ -403,15 +403,15 @@ class TestStrideSliceCase2TRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.strided_slice
         self.api_args = {
-            "x": np.random.random([3, 4, 5, 6]).astype("int32"),
+            "x": np.random.random([3, 4, 10]).astype("int32"),
             "axes": [0, 1, 2],
             "starts": [1, 0, 2],
             "ends": [2, 3, 4],
             "strides": [1, 1, 1],
         }
         self.program_config = {"feed_list": ["x"]}
-        self.min_shape = {"x": [1, 4, 5, 6]}
-        self.max_shape = {"x": [5, 4, 5, 6]}
+        self.min_shape = {"x": [1, 4, 10]}
+        self.max_shape = {"x": [5, 4, 10]}
 
     def test_trt_result(self):
         self.check_trt_result()
@@ -435,5 +435,41 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestStrideSliceCase4TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.strided_slice
+        self.api_args = {
+            "x": np.random.random([5, 5, 5]).astype("float32"),
+            "axes": [0, 1, 2],
+            "starts": np.array([1, 0, 0]).astype("int32"),
+            "ends": np.array([2, 1, 3]).astype("int32"),
+            "strides": np.array([1, 1, 1]).astype("int32"),
+        }
+        self.program_config = {"feed_list": ["x", "starts", "ends", "strides"]}
+        self.min_shape = {"x": [1, 5, 5]}
+        self.max_shape = {"x": [6, 5, 5]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestStrideSliceCase5TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.strided_slice
+        self.api_args = {
+            "x": np.random.random([3, 4, 10]).astype("float32"),
+            "axes": [0, 1, 2],
+            "starts": np.array([0, -1, 0]).astype("int32"),
+            "ends": np.array([2, -3, 5]).astype("int32"),
+            "strides": np.array([1, -1, 1]).astype("int32"),
+        }
+        self.program_config = {"feed_list": ["x", "starts", "ends", "strides"]}
+        self.min_shape = {"x": [1, 4, 10]}
+        self.max_shape = {"x": [5, 4, 10]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -1648,6 +1648,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {`
`1648`	`1648`	`ps.Add(std::make_unique<StackOpPattern>(context));`
`1649`	`1649`	`ps.Add(std::make_unique<TanhOpPattern>(context));`
`1650`	`1650`	`ps.Add(std::make_unique<FullWithTensorPattern>(context));`
	`1651`	`+ ps.Add(std::make_unique<StridedSliceOpPattern>(context));`
`1651`	`1652`	`return ps;`
`1652`	`1653`	`}`
`1653`	`1654`	`};`
Original file line number	Diff line number	Diff line change
`@@ -2026,6 +2026,7 @@ def get_setup_parameters():`
`2026`	`2026`	`'paddle.decomposition',`
`2027`	`2027`	`'paddle._typing',`
`2028`	`2028`	`'paddle._typing.libs',`
	`2029`	`+ 'paddle.tensorrt',`
`2029`	`2030`	`]`
`2030`	`2031`
`2031`	`2032`	`paddle_bins = ''`