add pd_op.cast converter for PIR TRT (#68216)

zhink · web-flow · commit 042eb6d23b94 · 2024-10-09T10:30:38.000+08:00
* add cast op trt layer

* codestyle

* codestyle

* with add_identity

* code style
diff --git a/python/paddle/tensorrt/impls/manipulation.py b/python/paddle/tensorrt/impls/manipulation.py
@@ -214,6 +214,34 @@ def squeeze_converter(network, paddle_op, inputs):
     return layer.get_output(0)
 
 
+@converter_registry.register("pd_op.cast", trt_version="8.x")
+@converter_registry.register("pd_op.cast_", trt_version="8.x")
+def cast_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    out_dtype = int(paddle_op.attrs().get("dtype"))
+    # Reference paddle/phi/common/data_type.h enum DataType
+    if out_dtype == 1:
+        out_dtype = trt.bool
+    elif out_dtype == 7:
+        out_dtype = trt.int32
+    elif out_dtype == 9:
+        out_dtype = trt.int32
+    elif out_dtype == 10:
+        out_dtype = trt.float32
+    elif out_dtype == 11:
+        out_dtype = trt.float32
+    elif out_dtype == 15:
+        out_dtype = trt.float16
+    else:
+        raise RuntimeError(
+            f"cast converter currently doesn't support dtype: {out_dtype}"
+        )
+    cast_layer = network.add_identity(input_tensor)
+    cast_layer.set_output_type(0, out_dtype)
+    cast_layer.get_output(0).dtype = out_dtype
+    return cast_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.slice", trt_version="8.x")
 def slice_converter(network, paddle_op, inputs):
     input_tensor = inputs[0]
diff --git a/test/tensorrt/test_converter_manipulation.py b/test/tensorrt/test_converter_manipulation.py
@@ -21,6 +21,42 @@
 from paddle import _C_ops
 
 
+class TestCast0TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.cast
+        self.api_args = {
+            "x": np.random.randn(7, 3).astype(np.float32),
+            "out_dtype": np.bool_,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [3, 3]}
+        self.max_shape = {"x": [10, 3]}
+
+
+class TestCast1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.cast
+        self.api_args = {
+            "x": np.random.randn(7, 3).astype(np.float16),
+            "out_dtype": np.int32,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [3, 3]}
+        self.max_shape = {"x": [10, 3]}
+
+
+class TestCast2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.cast
+        self.api_args = {
+            "x": np.random.randn(7, 3).astype(np.float32),
+            "out_dtype": np.int64,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [3, 3]}
+        self.max_shape = {"x": [10, 3]}
+
+
 class TestConcatTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.concat
