fix (#69671)

lizexu123 · web-flow · commit 6fd300b75a71 · 2024-11-26T12:52:52.000+08:00
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1800,6 +1800,16 @@ class StridedSliceOpPattern
       VLOG(3) << "pd_op.strided_slice must has starts,ends and strides input";
       return false;
     }
+    if (!pir::GetDefiningOpForInput(op, 1)
+             ->isa<paddle::dialect::FullIntArrayOp>() ||
+        !pir::GetDefiningOpForInput(op, 2)
+             ->isa<paddle::dialect::FullIntArrayOp>() ||
+        !pir::GetDefiningOpForInput(op, 3)
+             ->isa<paddle::dialect::FullIntArrayOp>()) {
+      VLOG(3) << "pd_op.strided_slice's starts/ends/strides input must be "
+                 "constant value";
+      return false;
+    }
     op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
     return true;
   }
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -741,121 +741,70 @@ def strided_slice_converter(network, paddle_op, inputs):
     ends_op = paddle_op.operands()[2].source().get_defining_op()
     strides_op = paddle_op.operands()[3].source().get_defining_op()
 
-    starts = (
-        starts_op.attrs()["value"]
-        if starts_op.name() == "pd_op.full_int_array"
-        else inputs[1]
-    )
-    ends = (
-        ends_op.attrs()["value"]
-        if ends_op.name() == "pd_op.full_int_array"
-        else inputs[2]
-    )
-    strides = (
-        strides_op.attrs()["value"]
-        if strides_op.name() == "pd_op.full_int_array"
-        else inputs[3]
-    )
+    if starts_op.name() == "pd_op.full_int_array":
+        starts = starts_op.attrs()["value"]
+
+    if ends_op.name() == "pd_op.full_int_array":
+        ends = ends_op.attrs()["value"]
+
+    if strides_op.name() == "pd_op.full_int_array":
+        strides = strides_op.attrs()["value"]
 
     input_shape = paddle_op.operands()[0].source().shape
     nchw_input_dims = len(input_shape)
 
     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_size_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):
-        if isinstance(starts, trt.ITensor):
-            start_tensor = get_shape_tensor_element(network, starts, i)
-        else:
-            start_tensor = add_1D_constant_layer(network, [starts[i]])
-
-        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 isinstance(strides, trt.ITensor):
-            stride_tensor = get_shape_tensor_element(network, strides, i)
-        else:
-            stride_tensor = add_1D_constant_layer(network, [strides[i]])
 
-        zero_tensor = add_1D_constant_layer(network, [0])
+    for i, trt_axis in enumerate(axes):
+        trt_start_dims[trt_axis] = starts[i]
+        trt_end_dims[trt_axis] = ends[i]
+        trt_step_dims[trt_axis] = strides[i]
+        if starts[i] < 0 or ends[i] < 0:
+            has_neg_indices = True
 
-        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:
-            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:
+    shape_tensor = trt_shape(network, input_tensor)
+    start_tensor = add_1D_constant_layer(network, trt_start_dims)
+    if has_neg_indices:
         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_vec_tensor = []
+    for i in range(len(trt_end_dims)):
+        end_vec_tensor.append(
+            get_shape_tensor_element(network, shape_tensor, i)
         )
-        end_tensor = network.add_select(
-            condition=has_neg_indices,
-            then_input=fixed_end_tensor,
-            else_input=end_tensor,
-        ).get_output(0)
 
-    # 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)
+    for i, trt_axis in enumerate(axes):
+        if ends[i] >= 0:
+            end_vec_tensor[trt_axis] = add_1D_constant_layer(network, ends[i])
+        else:
+            end_vec_tensor[trt_axis] = trt_sum(
+                network,
+                end_vec_tensor[trt_axis],
+                add_1D_constant_layer(network, ends[i]),
+            )
 
-    # floor_div_tensor computation
-    floor_div_tensor = trt_floor_div(network, size_tensor, step_tensor)
-    size_tensor = trt_sub(network, zero_tensor, floor_div_tensor)
+    size_tensor = trt_sub(
+        network,
+        start_tensor,
+        trt_min(network, trt_concat(network, end_vec_tensor), shape_tensor),
+    )
+    zero_t = add_1D_constant_layer(network, 0)
+    step_tensor = add_1D_constant_layer(network, trt_step_dims)
+    size_tensor = trt_sub(
+        network, zero_t, trt_floor_div(network, size_tensor, step_tensor)
+    )
 
-    # Create the slice layer
     layer = network.add_slice(
         input_tensor, trt_start_dims, trt_size_dims, trt_step_dims
     )
     layer.set_input(1, start_tensor)
     layer.set_input(2, size_tensor)
     layer.set_input(3, step_tensor)
-
     return layer.get_output(0)
 
 
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -448,7 +448,7 @@ class TestStrideSliceCase2TRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.strided_slice
         self.api_args = {
-            "x": np.random.random([3, 4, 10]).astype("int32"),
+            "x": np.random.random([3, 4, 10]).astype("int64"),
             "axes": [0, 1, 2],
             "starts": [1, 0, 2],
             "ends": [2, 3, 4],
@@ -484,15 +484,15 @@ 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"),
+            "x": np.random.random([1, 56, 56, 128]).astype("float32"),
+            "axes": [1, 2],
+            "starts": [0, 0],
+            "ends": [6, 6],
+            "strides": [2, 2],
         }
-        self.program_config = {"feed_list": ["x", "starts", "ends", "strides"]}
-        self.min_shape = {"x": [1, 5, 5]}
-        self.max_shape = {"x": [6, 5, 5]}
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 56, 56, 128]}
+        self.max_shape = {"x": [1, 56, 56, 128]}
 
     def test_trt_result(self):
         self.check_trt_result()
@@ -502,15 +502,18 @@ 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"),
+            "x": np.random.random([1, 56, 56, 128]).astype("float32"),
+            "axes": [1, 2],
+            "starts": [
+                1,
+                1,
+            ],
+            "ends": [10000, 10000],
+            "strides": [2, 2],
         }
-        self.program_config = {"feed_list": ["x", "starts", "ends", "strides"]}
-        self.min_shape = {"x": [1, 4, 10]}
-        self.max_shape = {"x": [5, 4, 10]}
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 56, 56, 128]}
+        self.max_shape = {"x": [1, 56, 56, 128]}
 
     def test_trt_result(self):
         self.check_trt_result()
