[Paddle Tensorrt] add pd_op.stack converter (#68839)

lizexu123 · web-flow · commit 22bd404b4843 · 2024-10-24T17:08:21.000+08:00
* kg

* merge develop

* stack

* fix

* 考虑int32类型的输入

* 使用getshapefromvalue
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -1400,6 +1400,43 @@ class NearestInterV2Pattern
   }
 };
 
+class StackOpPattern : public pir::OpRewritePattern<paddle::dialect::StackOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::StackOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::StackOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+
+    pir::Value x = op.operand_source(0);
+    int rank;
+    auto x_type = x.type();
+    if (x_type.isa<pir::VectorType>()) {
+      rank = x_type.dyn_cast<pir::VectorType>().size();
+    } else {
+      auto x_shape = pir::GetShapeFromValue(x);
+      rank = x_shape.size();
+    }
+
+    int axis;
+    if (op->HasAttribute("axis")) {
+      axis = op->attribute<pir::Int32Attribute>("axis").data();
+    } else {
+      axis = -1;
+    }
+    if (axis > rank || axis < -(rank + 1)) {
+      VLOG(3) << "Invalid axis value: " << axis
+              << ". Axis should be in range [-" << (rank + 1) << ", " << rank
+              << "], where rank is " << rank << ".";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 // Add ReduceCommonOpPattern base class to simplify code
 template <typename OpType>
 class ReduceCommonOpPattern : public pir::OpRewritePattern<OpType> {
@@ -1557,6 +1594,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<MinOpPattern>(context));
     ps.Add(std::make_unique<BilinearInterpV2Pattern>(context));
     ps.Add(std::make_unique<NearestInterV2Pattern>(context));
+    ps.Add(std::make_unique<StackOpPattern>(context));
     ps.Add(std::make_unique<TanhOpPattern>(context));
     return ps;
   }
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -654,3 +654,50 @@ def split_converter(network, paddle_op, inputs):
             ).get_output(0)
 
     return outputs
+
+
+@converter_registry.register("pd_op.stack", trt_version="8.x")
+def stack_converter(network, paddle_op, inputs):
+    input_tensors = inputs[0]
+    input_num = len(input_tensors)
+
+    inputs = []
+    for i in range(input_num):
+        inputs.append(input_tensors[i])
+
+    input_rank = len(input_tensors[0].shape)
+
+    output_rank = input_rank + 1
+    axis = paddle_op.attrs().get("axis")
+    if axis < 0:
+        axis += output_rank
+
+    shape_tensor = network.add_shape(input_tensors[0]).get_output(0)
+    shape_tensor_vec = []
+    for i in range(output_rank):
+        if i < axis:
+            shape_tensor_vec.append(
+                get_shape_tensor_element(network, shape_tensor, i)
+            )
+        elif i > axis:
+            shape_tensor_vec.append(
+                get_shape_tensor_element(network, shape_tensor, i - 1)
+            )
+        else:
+            shape_tensor_vec.append(add_1D_constant_layer(network, 1))
+
+    after_shape_tensor = network.add_concatenation(shape_tensor_vec).get_output(
+        0
+    )
+
+    for i in range(input_num):
+        shuffle_layer = network.add_shuffle(inputs[i])
+        shuffle_layer.set_input(1, after_shape_tensor)
+        reshaped_tensor = shuffle_layer.get_output(0)
+        inputs[i] = reshaped_tensor
+
+    concat_layer = network.add_concatenation(inputs)
+    concat_layer.axis = axis
+    output_tensor = concat_layer.get_output(0)
+
+    return output_tensor
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -343,5 +343,43 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestStackTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.stack
+        self.api_args = {
+            "x": [
+                np.array([[1.0, 2.0]]).astype("float32"),
+                np.array([[3.0, 4.0]]).astype("float32"),
+                np.array([[5.0, 6.0]]).astype("float32"),
+            ],
+            "axis": 0,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [[1, 2], [1, 2], [1, 2]]}
+        self.max_shape = {"x": [[3, 2], [3, 2], [3, 2]]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestStackCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.stack
+        self.api_args = {
+            "x": [
+                np.array([[1, 2]]).astype("int32"),
+                np.array([[3, 4]]).astype("int32"),
+                np.array([[5, 6]]).astype("int32"),
+            ],
+            "axis": -1,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [[1, 2], [1, 2], [1, 2]]}
+        self.max_shape = {"x": [[3, 2], [3, 2], [3, 2]]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()
