pd_op.arange (#68757)

lizexu123 · web-flow · commit a7058b9a88a9 · 2024-10-18T10:35:58.000+08:00
diff --git a/python/paddle/tensorrt/converter_utils.py b/python/paddle/tensorrt/converter_utils.py
@@ -325,6 +325,19 @@ def build_size_tensor(
     return size_tensor
 
 
+# Reduce the given tensor in the TensorRT network to a scalar
+def trt_reduce_to_scalar(network, tensor):
+    if len(tensor.shape) == 0:
+        return tensor
+    axes = 0
+    for i in range(len(tensor.shape)):
+        axes |= 1 << i
+    reduce_layer = network.add_reduce(
+        tensor, trt.ReduceOperation.SUM, axes, keep_dims=False
+    )
+    return reduce_layer.get_output(0)
+
+
 def convert_conv2d(network, paddle_op, inputs):
     if (
         paddle_op.name() == "pd_op.conv2d"
diff --git a/python/paddle/tensorrt/impls/creation.py b/python/paddle/tensorrt/impls/creation.py
@@ -15,6 +15,15 @@
 import numpy as np
 import tensorrt as trt
 
+from paddle.tensorrt.converter_utils import (
+    add_1D_constant_layer,
+    cast_tensor,
+    trt_floor_div,
+    trt_max,
+    trt_reduce_to_scalar,
+    trt_reshape,
+    trt_sub,
+)
 from paddle.tensorrt.register import converter_registry
 
 
@@ -31,8 +40,36 @@ def full_int_array_converter(network, paddle_op, inputs):
 @converter_registry.register("pd_op.full", trt_version="8.x")
 def full_converter(network, paddle_op, inputs):
     shape = paddle_op.attrs()["shape"]
-    value = paddle_op.attrs().get("value", 1.0)  # 默认值为1.0
+    value = paddle_op.attrs().get("value", 1.0)
     full_layer = network.add_constant(
         shape, np.full(shape, value, dtype=np.float32)
     )
     return full_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
+    zero_tensor = add_1D_constant_layer(network, 0, np.int32)
+
+    delta = trt_sub(network, end, start)
+
+    f_quotient_tensor = trt_floor_div(network, delta, step)
+
+    if start.dtype == trt.DataType.FLOAT:
+        quotient_tensor = cast_tensor(network, f_quotient_tensor, trt.int32)
+    else:
+        quotient_tensor = f_quotient_tensor
+
+    number_tensor = trt_max(network, quotient_tensor, zero_tensor)
+
+    reshape_start_layer = trt_reshape(network, start, (1,))
+
+    start_tensor = trt_reduce_to_scalar(network, reshape_start_layer)
+
+    fill_layer = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
+    fill_layer.set_input(0, number_tensor)
+    fill_layer.set_input(1, start_tensor)
+    fill_layer.set_input(2, step)
+
+    return fill_layer.get_output(0)
diff --git a/test/tensorrt/test_converter_creation.py b/test/tensorrt/test_converter_creation.py
@@ -14,6 +14,7 @@
 
 import unittest
 
+import numpy as np
 from tensorrt_test_base import TensorRTBaseTest
 
 import paddle
@@ -31,5 +32,21 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestArangeTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.arange
+        self.api_args = {
+            "start": np.array([0]).astype("int32"),
+            "end": np.array([6]).astype("int32"),
+            "step": np.array([1]).astype("int32"),
+        }
+        self.program_config = {"feed_list": []}
+        self.min_shape = {}
+        self.max_shape = {}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == "__main__":
     unittest.main()
