[Paddle TensorRT No.19] Add pd_op.cumsum converter

paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc

@@ -1877,6 +1877,27 @@ class TopkOpPattern : public pir::OpRewritePattern<paddle::dialect::TopkOp> {
   }
 };
 
+class CumsumOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::CumsumOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::CumsumOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::CumsumOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+
+    if (!pir::GetDefiningOpForInput(op, 1)->isa<paddle::dialect::FullOp>()) {
+      VLOG(3) << "The 'axis' input of pd_op.cumsum must be an integer";
+      return false;
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 bool CheckSetValue(const pir::Operation *op, int starts_input_loc = 1) {
   paddle::dialect::FullIntArrayOp starts_defining_op =
       pir::GetDefiningOpForInput(op, starts_input_loc)
@@ -2228,6 +2249,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<FullWithTensorPattern>(context));
     ps.Add(std::make_unique<StridedSliceOpPattern>(context));
     ps.Add(std::make_unique<TopkOpPattern>(context));
+    ps.Add(std::make_unique<CumsumOpPattern>(context));
     ps.Add(std::make_unique<SetValueOpPattern>(context));
     ps.Add(std::make_unique<SetValue_OpPattern>(context));
     ps.Add(std::make_unique<SetValueWithTensorOpPattern>(context));

python/paddle/tensorrt/impls/math.py

@@ -16,6 +16,7 @@
 import tensorrt as trt
 
 from paddle.tensorrt.converter_utils import (
+    add_1D_constant_layer,
     add_cast_reduce_layer,
     add_elementwise_layer,
     add_reduce_layer,
@@ -227,6 +228,96 @@ def all_converter(network, paddle_op, inputs):
     )
 
 
+@converter_registry.register("pd_op.cumsum", trt_version="8.x")
+def cumsum_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    dtype = input_tensor.dtype
+    axis = paddle_op.operands()[1].source().get_defining_op().attrs()["value"]
+    input_shape = input_tensor.shape
+    rank = len(input_shape)
+
+    if axis < 0:
+        axis += rank
+    axis = int(axis)
+
+    # Obtain the number of cycles
+    if input_shape[axis] > 0:
+        axis_tensor = np.array(input_shape[axis], dtype=np.int32)
+        trip_limit = network.add_constant((), axis_tensor)
+    else:
+        dynamic_shape = network.add_shape(input_tensor).get_output(0)
+        axis_tensor = np.array(axis, dtype=np.int32)
+        index = network.add_constant((), axis_tensor).get_output(0)
+        trip_limit = network.add_gather(dynamic_shape, index, 0)
+
+    # Obtain the slice shape
+    shape_list = []
+    for i in range(rank):
+        if i == axis:
+            shape_list.append(add_1D_constant_layer(network, [1]))
+        elif input_shape[i] < 0:
+            dynamic_shape = network.add_shape(input_tensor).get_output(0)
+            index = network.add_constant(
+                (), np.array(i, dtype=np.int32)
+            ).get_output(0)
+            shape_index = network.add_gather(dynamic_shape, index, 0)
+            shuffle_layer = network.add_shuffle(shape_index.get_output(0))
+            shuffle_layer.reshape_dims = (1,)
+            shape_list.append(shuffle_layer.get_output(0))
+        else:
+            shape_list.append(add_1D_constant_layer(network, input_shape[i]))
+    slice_shape = network.add_concatenation(shape_list).get_output(0)
+
+    start = [0] * rank
+    size = [1] * rank
+    stride = [1] * rank
+    input_sliced = network.add_slice(input_tensor, start, size, stride)
+    input_sliced.set_input(2, slice_shape)
+
+    # squeeze axis
+    shape_list.pop(axis)
+    new_shape = network.add_concatenation(shape_list).get_output(0)
+    squeeze_layer = network.add_shuffle(input_sliced.get_output(0))
+    squeeze_layer.set_input(1, new_shape)
+
+    loop = network.add_loop()
+    loop.add_trip_limit(trip_limit.get_output(0), trt.TripLimit.COUNT)
+
+    iterator = loop.add_iterator(input_tensor, axis)
+    data = iterator.get_output(0)
+
+    # create zero tensor
+    zero_vec = np.array([0.0], dtype=np.float32)
+    zero = network.add_constant((1,), zero_vec).get_output(0)
+    lhs_val, rhs_val = broadcast(
+        network,
+        squeeze_layer.get_output(0),
+        zero,
+        squeeze_layer.get_output(0).name,
+        zero.name,
+    )
+    cast_tensor = trt_cast(network, rhs_val, dtype)
+    zero_tensor = network.add_elementwise(
+        lhs_val, cast_tensor, trt.ElementWiseOperation.PROD
+    ).get_output(0)
+
+    # Cycle and add according to the axis
+    running_sum = loop.add_recurrence(zero_tensor)
+    running_sum_tensor = running_sum.get_output(0)
+
+    cur_sum = network.add_elementwise(
+        data, running_sum_tensor, trt.ElementWiseOperation.SUM
+    ).get_output(0)
+
+    running_sum.set_input(1, cur_sum)
+
+    reverse_flag = trt.LoopOutput.CONCATENATE
+    loop_out = loop.add_loop_output(cur_sum, reverse_flag, axis)
+    loop_out.set_input(1, trip_limit.get_output(0))
+
+    return loop_out.get_output(0)
+
+
 @converter_registry.register("pd_op.floor_divide", trt_version="8.x")
 def floor_divide_converter(network, paddle_op, inputs):
     return add_elementwise_layer(

test/tensorrt/test_converter_math.py

@@ -252,6 +252,75 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestCumsumCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.cumsum
+        self.api_args = {
+            "x": np.random.randn(2, 2, 3).astype("float32"),
+            "axis": -1,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 2, 3]}
+        self.max_shape = {"x": [5, 2, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestCumsumCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.cumsum
+        self.api_args = {
+            "x": np.random.randn(2, 2, 3).astype("float32"),
+            "axis": 1,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 2, 3]}
+        self.max_shape = {"x": [5, 2, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestCumsumCase3TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.cumsum
+        self.api_args = {
+            "x": np.random.randn(2, 2, 3).astype("float32"),
+            "axis": 0,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 2, 3]}
+        self.max_shape = {"x": [5, 2, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestCumsumCase4TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.cumsum
+        self.api_args = {
+            "x": np.random.randn(2, 2, 3).astype("int64"),
+            "axis": 0,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 2, 3]}
+        self.max_shape = {"x": [5, 2, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 class TestFloorDivideFloatTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.floor_divide