pd_op.strided_slice (#69015)

lizexu123 · web-flow · commit 6643ab4284b6 · 2024-10-29T11:06:09.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
@@ -1,4 +1,5 @@
 
+
 // Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -971,6 +972,7 @@ class SplitOpPattern : public pir::OpRewritePattern<paddle::dialect::SplitOp> {
     return true;
   }
 };
+
 class SplitWithNumOpPattern
     : public pir::OpRewritePattern<paddle::dialect::SplitWithNumOp> {
  public:
@@ -1038,6 +1040,7 @@ class SplitWithNumOpPattern
     return true;
   }
 };
+
 class GreaterEqualOpPattern
     : public pir::OpRewritePattern<paddle::dialect::GreaterEqualOp> {
  public:
@@ -1532,6 +1535,32 @@ class FullWithTensorPattern
   }
 };
 
+class StridedSliceOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::StridedSliceOp> {
+ public:
+  using pir::OpRewritePattern<
+      paddle::dialect::StridedSliceOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::StridedSliceOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    if (!op->HasAttribute("axes")) {
+      VLOG(3) << "The necessary attribute of the pd_op.strided_slice operator "
+                 "axes are missing.";
+      return false;
+    }
+    if (!op.operand_source(1) || !op.operand_source(2) ||
+        !op.operand_source(3)) {
+      VLOG(3) << "pd_op.strided_slice must has starts,ends and strides input";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class TrtOpMarkerPass : public pir::PatternRewritePass {
  public:
   TrtOpMarkerPass() : pir::PatternRewritePass("trt_op_marker_pass", 2) {}
diff --git a/python/paddle/tensorrt/converter_utils.py b/python/paddle/tensorrt/converter_utils.py
@@ -502,3 +502,15 @@ def add_cast_reduce_layer(network, paddle_op, inputs, op_type):
     layer.set_output_type(0, trt.bool)
     layer.get_output(0).dtype = cast_layer.get_output(0).dtype
     return layer.get_output(0)
+
+
+def fix_negative_indices(network, input_shape, indices):
+    rank = len(input_shape.shape)
+    zero_tensor = add_1D_constant_layer(network, [0] * rank)
+    minus_one_tensor = add_1D_constant_layer(network, [-1] * rank)
+
+    min_indices_zero = trt_min(network, indices, zero_tensor)
+    sign = trt_max(network, min_indices_zero, minus_one_tensor)
+    sub = trt_mul(network, sign, input_shape)
+    fixed_indices = trt_sub(network, indices, sub)
+    return fixed_indices
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -21,6 +21,7 @@
     build_size_tensor,
     build_start_tensor,
     cast_tensor,
+    fix_negative_indices,
     get_axes_for_reduce_op,
     get_positive_dim,
     get_shape_tensor_element,
@@ -701,3 +702,91 @@ def stack_converter(network, paddle_op, inputs):
     output_tensor = concat_layer.get_output(0)
 
     return output_tensor
+
+
+@converter_registry.register("pd_op.strided_slice", trt_version="8.x")
+def strided_slice_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    axes = paddle_op.attrs()["axes"]
+
+    starts_op = paddle_op.operands()[1].source().get_defining_op()
+    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]
+    )
+
+    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_step_dims = [1] * nchw_input_dims
+
+    has_neg_indices = False
+    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
+
+    shape_tensor = network.add_shape(input_tensor).get_output(0)
+    start_tensor = add_1D_constant_layer(network, trt_start_dims)
+
+    if has_neg_indices:
+        start_tensor = fix_negative_indices(network, shape_tensor, start_tensor)
+
+    end_vec_tensor = []
+    for i in range(len(trt_end_dims)):
+        end_vec_tensor.append(
+            get_shape_tensor_element(network, shape_tensor, i)
+        )
+
+    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)
+            ).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
+            )
+
+    concat_end_tensor = network.add_concatenation(end_vec_tensor).get_output(0)
+    min_tensor = trt_min(network, concat_end_tensor, shape_tensor)
+    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 = trt_floor_div(network, size_tensor, step_tensor)
+    size_tensor = trt_sub(network, zero_t, floor_div_tensor)
+
+    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
@@ -381,5 +381,59 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+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"),
+            "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]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+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"),
+            "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]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestStrideSliceCase3TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.strided_slice
+        self.api_args = {
+            "x": np.random.random([3, 4, 10]).astype("bool"),
+            "axes": [0, 1, 2],
+            "starts": [0, -1, 0],
+            "ends": [2, -3, 5],
+            "strides": [1, -1, 1],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        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()
