[Paddle Tensorrt] add pd_op.assign,pd_op.assign_value_,pd_op.assign_out converter (#68775)

lizexu123 · web-flow · commit 8f6712867419 · 2024-10-22T16:27:59.000+08:00
* pd_op.assign,pd_op.assign_value_,pd_op.assign_out

* 修复split_with_num

* 消除冲突

* 修改assign_out

* 增加单测时间,修改仅为float32

* fix
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -76,6 +76,9 @@ DEFINE_GENERAL_PATTERN(Sigmoid, paddle::dialect::SigmoidOp)
 DEFINE_GENERAL_PATTERN(Sqrt, paddle::dialect::SqrtOp)
 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(AssignOut, paddle::dialect::AssignOut_Op)
 
 #undef DEFINE_GENERAL_PATTERN
 
@@ -1510,6 +1513,9 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Sqrt)
     ADD_PATTERN(Hardsigmoid)
     ADD_PATTERN(Hardswish)
+    ADD_PATTERN(AssignOut)
+    ADD_PATTERN(Assign)
+    ADD_PATTERN(AssignValue_)
 #if IS_TRT_VERSION_GE(8600)
     ADD_PATTERN(Layer_norm)
 #endif
diff --git a/python/paddle/tensorrt/impls/creation.py b/python/paddle/tensorrt/impls/creation.py
@@ -15,6 +15,7 @@
 import numpy as np
 import tensorrt as trt
 
+import paddle
 from paddle.tensorrt.converter_utils import (
     add_1D_constant_layer,
     cast_tensor,
@@ -47,6 +48,28 @@ def full_converter(network, paddle_op, inputs):
     return full_layer.get_output(0)
 
 
+@converter_registry.register("pd_op.assign", trt_version="8.x")
+@converter_registry.register("pd_op.assign_out_", trt_version="8.x")
+def assign_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    identity_layer = network.add_identity(input_tensor)
+    return identity_layer.get_output(0)
+
+
+@converter_registry.register("pd_op.assign_value_", trt_version="8.x")
+def assign_value_converter(network, paddle_op, inputs):
+    attrs = paddle_op.attrs()
+    shape = attrs['shape']
+    dtype = attrs['dtype']
+    values = attrs['values']
+
+    if dtype == paddle.float32:
+        dtype = np.float32
+
+    constant_layer = network.add_constant(shape, np.array(values, dtype=dtype))
+    return constant_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.arange", trt_version="8.x")
 def arange_converter(network, paddle_op, inputs):
     start, end, step = inputs
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -497,19 +497,19 @@ def split_with_num_converter(network, paddle_op, inputs):
     split_size = trt_floor_div(network, input_axis_size, num_splits_tensor)
 
     outputs = []
+    current_offset = add_1D_constant_layer(network, 0)
+
     for idx in range(num_splits):
         idx_tensor = add_1D_constant_layer(network, idx)
-        offset = trt_floor_div(network, idx_tensor, split_size)
-
+        # Calculate the slice start and size
         start_tensor = build_start_tensor(
-            network, input_shape_size, axis_tensor, offset
+            network, input_shape_size, axis_tensor, current_offset
         )
-
         size_tensor = build_size_tensor(
             network,
             input_shape_size,
             axis_tensor,
-            num_splits_tensor,
+            split_size,
             input_shape_tensor,
         )
 
@@ -525,6 +525,9 @@ def split_with_num_converter(network, paddle_op, inputs):
 
         outputs.append(slice_layer.get_output(0))
 
+        # Update current_offset for the next slice
+        current_offset = trt_sum(network, current_offset, split_size)
+
     return outputs
 
 
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -11,15 +11,15 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_model_bert PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_model_dummy PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_model_resnet50 PROPERTIES TIMEOUT "500")
-  set_tests_properties(test_converter_conv PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_conv PROPERTIES TIMEOUT "300")
   set_tests_properties(test_export PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_norm PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_ops PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_stat PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_math PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_activation PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_others PROPERTIES TIMEOUT "100")
-  set_tests_properties(test_converter_manipulation PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_manipulation PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_creation PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_attribute PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_common PROPERTIES TIMEOUT "300")
diff --git a/test/tensorrt/test_converter_creation.py b/test/tensorrt/test_converter_creation.py
@@ -32,6 +32,34 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestAssignTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.assign
+        self.api_args = {
+            "x": np.random.random([2, 2]).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 2]}
+        self.max_shape = {"x": [3, 2]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestAssignValueFloat32TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.assign
+        self.api_args = {
+            "x": np.array([[2.5, 2.5], [2.5, 2.5], [2.5, 2.5]], dtype='float32')
+        }
+        self.program_config = {"feed_list": []}
+        self.min_shape = {}
+        self.max_shape = {}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 class TestArangeTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.arange
@@ -48,5 +76,20 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestAssignOutTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.assign
+        self.api_args = {
+            "x": np.random.random([2, 2]).astype("float32"),
+            "output": np.zeros((2, 2), dtype="float32"),
+        }
+        self.program_config = {"feed_list": ["x", "output"]}
+        self.min_shape = {"x": [1, 2], "output": [1, 2]}
+        self.max_shape = {"x": [3, 2], "output": [3, 2]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -243,6 +243,22 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestSplitWithNumAllTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.split
+        self.api_args = {
+            "x": np.random.randn(1, 2).astype(np.float32),
+            "num_or_sections": 2,
+            "axis": np.array([1]).astype("int32"),
+        }
+        self.program_config = {"feed_list": ["x", "axis"]}
+        self.min_shape = {"x": [1, 2]}
+        self.max_shape = {"x": [1, 2]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 class TestSplitWithNumNegativeAxisTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.split
