【TRT_Converter】add tile & share_data op TRT_Converter (#69086)

chang-wenbin · web-flow · commit a5cf82a86b49 · 2024-11-08T18:14:52.000+08:00
* add pir_converter_tile

* update tile op_converter

* update converter  tile

* add shared_data

* update tile&amp; share_data converter
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -79,6 +79,8 @@ DEFINE_GENERAL_PATTERN(Hardsigmoid, paddle::dialect::HardsigmoidOp)
 DEFINE_GENERAL_PATTERN(Hardswish, paddle::dialect::HardswishOp)
 DEFINE_GENERAL_PATTERN(Assign, paddle::dialect::AssignOp)
 DEFINE_GENERAL_PATTERN(AssignValue_, paddle::dialect::AssignValue_Op)
+DEFINE_GENERAL_PATTERN(Tile, paddle::dialect::TileOp)
+DEFINE_GENERAL_PATTERN(Share_Data, paddle::dialect::ShareDataOp)
 DEFINE_GENERAL_PATTERN(AssignOut, paddle::dialect::AssignOut_Op)
 DEFINE_GENERAL_PATTERN(Roll, paddle::dialect::RollOp)
 
@@ -1613,6 +1615,8 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(AssignOut)
     ADD_PATTERN(Assign)
     ADD_PATTERN(AssignValue_)
+    ADD_PATTERN(Tile)
+    ADD_PATTERN(Share_Data)
     ADD_PATTERN(Roll)
 #if IS_TRT_VERSION_GE(8600)
     ADD_PATTERN(Layer_norm)
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -39,6 +39,8 @@
 )
 from paddle.tensorrt.register import converter_registry
 
+from ..util import get_trt_version_list
+
 
 @converter_registry.register("pd_op.reshape", trt_version="8.x")
 def reshape_converter(network, paddle_op, inputs):
@@ -704,6 +706,61 @@ def stack_converter(network, paddle_op, inputs):
     return output_tensor
 
 
+@converter_registry.register("pd_op.tile", trt_version="8.x")
+def tile_converter(network, paddle_op, inputs):
+    input = inputs[0]
+    input_shape = paddle_op.operands()[0].source().shape
+    input_shape_tensor = network.add_shape(input).get_output(0)
+    rank = len(input_shape)
+
+    repeat_times_op = paddle_op.operands()[1].source().get_defining_op()
+    if repeat_times_op.name() == "pd_op.full_int_array":
+        repeat_times = repeat_times_op.attrs()["value"]
+        repeat_tensor = add_1D_constant_layer(network, repeat_times)
+        repeat_rank = len(repeat_times)
+    else:
+        repeat_tensor = inputs[1]
+        repeat_shape = paddle_op.operands()[1].source().shape
+        repeat_rank = repeat_shape[0]
+
+    if rank > repeat_rank:
+        one_rank_tensor = add_1D_constant_layer(
+            network, [1] * (rank - repeat_rank)
+        )
+        repeat_expand_tensor = trt_concat(
+            network, [one_rank_tensor, repeat_tensor]
+        )
+    elif rank < repeat_rank:
+        one_rank_tensor = add_1D_constant_layer(
+            network, [1] * (repeat_rank - rank)
+        )
+        input_shape_tensor = trt_concat(
+            network, [one_rank_tensor, input_shape_tensor]
+        )
+        input = trt_reshape(network, input, input_shape_tensor, "", True)
+        repeat_expand_tensor = repeat_tensor
+    else:
+        repeat_expand_tensor = repeat_tensor
+
+    start = [0] * max(rank, repeat_rank)
+    stride = [1] * max(rank, repeat_rank)
+    output_shape = [0] * max(rank, repeat_rank)
+    output_shape_tensor = trt_mul(
+        network, input_shape_tensor, repeat_expand_tensor
+    )
+
+    slice_layer = network.add_slice(input, start, output_shape, stride)
+    slice_layer.set_input(2, output_shape_tensor)
+
+    version_list = get_trt_version_list()
+    if version_list >= [8, 6, 0]:
+        slice_layer.mode = trt.SampleMode.WRAP
+    else:
+        slice_layer.mode = trt.SliceMode.WRAP
+
+    return slice_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.strided_slice", trt_version="8.x")
 def strided_slice_converter(network, paddle_op, inputs):
     input_tensor = inputs[0]
diff --git a/python/paddle/tensorrt/impls/others.py b/python/paddle/tensorrt/impls/others.py
@@ -136,3 +136,12 @@ def multiclass_nms3_converter(network, paddle_op, inputs):
         constant_layer.get_output(0),
         batch_nms_layer.get_output(0),
     )
+
+
+@converter_registry.register("pd_op.share_data", trt_version="8.x")
+def share_data_converter(network, paddle_op, inputs):
+    x = inputs[0]
+
+    identity_layer = network.add_identity(x)
+
+    return identity_layer.get_output(0)
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -381,6 +381,51 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestTileTRTPatternCase0(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.tile
+        self.api_args = {
+            "x": np.random.randn(1, 2, 3).astype("float32"),
+            "repeat_times": (2, 4),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 2, 3]}
+        self.max_shape = {"x": [2, 2, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestTileTRTPatternCase1(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.tile
+        self.api_args = {
+            "x": np.random.randn(1, 2, 3).astype("int32"),
+            "repeat_times": np.array([1, 2, 3, 4]).astype("int32"),
+        }
+        self.program_config = {"feed_list": ["x", "repeat_times"]}
+        self.min_shape = {"x": [1, 2, 3]}
+        self.max_shape = {"x": [2, 2, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestTileTRTPatternCase2(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.tile
+        self.api_args = {
+            "x": np.random.randn(1, 2, 3).astype("float32"),
+            "repeat_times": [1, 2, 3],
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 2, 3]}
+        self.max_shape = {"x": [2, 2, 3]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 class TestStrideSliceCase1TRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.strided_slice
diff --git a/test/tensorrt/test_converter_others.py b/test/tensorrt/test_converter_others.py
@@ -17,9 +17,14 @@
 import numpy as np
 from tensorrt_test_base import TensorRTBaseTest
 
+import paddle
 from paddle import _C_ops
 
 
+def api_wrapper(x):
+    return paddle._C_ops.share_data(x)
+
+
 def multiclass_nms3(
     bboxes,
     scores,
@@ -81,5 +86,19 @@ def test_trt_result(self):
         self.check_marker(expected_result=False)
 
 
+class TestShareDataTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = api_wrapper
+        self.api_args = {
+            "x": np.random.rand(4, 3, 5).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [4, 3, 5]}
+        self.max_shape = {"x": [6, 3, 5]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()
