【SCU】【Paddle TensorRT No.36】Add pd_op.flip converter (#69724)

PolaKuma · web-flow · commit a2d91524c9da · 2024-12-03T10:35:19.000+08:00
* add

* fix codestyle

* fix codestyle

* fix codestyle

* update

* fix codestyle

* add
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -86,6 +86,7 @@ DEFINE_GENERAL_PATTERN(Floor, paddle::dialect::FloorOp)
 DEFINE_GENERAL_PATTERN(Roll, paddle::dialect::RollOp)
 DEFINE_GENERAL_PATTERN(Softplus, paddle::dialect::SoftplusOp)
 DEFINE_GENERAL_PATTERN(ThresholdedRelu, paddle::dialect::ThresholdedReluOp)
+DEFINE_GENERAL_PATTERN(Flip, paddle::dialect::FlipOp)
 
 #undef DEFINE_GENERAL_PATTERN
 
@@ -2140,6 +2141,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Roll)
     ADD_PATTERN(Softplus)
     ADD_PATTERN(ThresholdedRelu)
+    ADD_PATTERN(Flip)
 #if IS_TRT_VERSION_GE(8600)
     ADD_PATTERN(Layer_norm)
 #endif
diff --git a/python/paddle/tensorrt/impls/linalg.py b/python/paddle/tensorrt/impls/linalg.py
@@ -16,7 +16,11 @@
 import tensorrt as trt
 
 from paddle.tensorrt.converter_utils import (
+    add_1D_constant_layer,
     broadcast,
+    get_shape_tensor_element,
+    trt_shape,
+    trt_sum,
 )
 from paddle.tensorrt.register import converter_registry
 
@@ -71,3 +75,42 @@ def bmm_converter(network, paddle_op, inputs):
         inputs[0], trt.MatrixOperation.NONE, inputs[1], trt.MatrixOperation.NONE
     )
     return out.get_output(0)
+
+
+@converter_registry.register("pd_op.flip", trt_version="8.x")
+def flip_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    input_dims = input_tensor.shape
+    rank = len(input_dims)
+    axis = paddle_op.attrs()["axis"]
+    axis = [a + rank if a < 0 else a for a in axis]
+    shape_tensor = trt_shape(network, input_tensor)
+
+    def get_axis_length(axis_idx):
+        dim_val = input_dims[axis_idx]
+        if dim_val >= 0:
+            return add_1D_constant_layer(network, [dim_val], is_scalar=True)
+        else:
+            return get_shape_tensor_element(
+                network, shape_tensor, axis_idx, is_scalar=True
+            )
+
+    for axis_idx in axis:
+        loop_layer = network.add_loop()
+        trip_limit = get_axis_length(axis_idx)
+        loop_layer.add_trip_limit(trip_limit, trt.TripLimit.COUNT)
+        iterator = loop_layer.add_iterator(input_tensor, axis_idx, reverse=True)
+        zero_tensor = add_1D_constant_layer(network, [0])
+        one_tensor = add_1D_constant_layer(network, [1])
+        iRec_layer = loop_layer.add_recurrence(zero_tensor)
+        iCur = iRec_layer.get_output(0)
+        iNext_layer = trt_sum(network, iCur, one_tensor)
+        iRec_layer.set_input(1, iNext_layer)
+        loop_out_layer = loop_layer.add_loop_output(
+            iterator.get_output(0), trt.LoopOutput.CONCATENATE, axis_idx
+        )
+        loop_out_layer.set_input(1, trip_limit)
+        input_tensor = loop_out_layer.get_output(0)
+
+    identity_layer = network.add_identity(input_tensor)
+    return identity_layer.get_output(0)
diff --git a/test/tensorrt/test_converter_linalg.py b/test/tensorrt/test_converter_linalg.py
@@ -67,5 +67,65 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestFlipTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.flip
+        self.api_args = {
+            "x": np.random.randn(2, 3, 4).astype("float32"),
+            "axis": [0, 2],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 4]}
+        self.max_shape = {"x": [5, 3, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestFlipNegAxisTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.flip
+        self.api_args = {
+            "x": np.random.randn(2, 3, 4).astype("float32"),
+            "axis": [-1, -3],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 4]}
+        self.max_shape = {"x": [5, 3, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestFlipIntTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.flip
+        self.api_args = {
+            "x": np.random.randn(2, 3, 4).astype("int64"),
+            "axis": [0, 2],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 4]}
+        self.max_shape = {"x": [5, 3, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestFlipIntNegAxisTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.flip
+        self.api_args = {
+            "x": np.random.randn(2, 3, 4).astype("int64"),
+            "axis": [-1, -3],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 4]}
+        self.max_shape = {"x": [5, 3, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()
