Updates plugin frontend function to convert data types between Tripy/TensorRT

Author: pranavm-nvidia

tripy/nvtripy/frontend/ops/plugin_qdp.py

@@ -14,12 +14,28 @@
 # limitations under the License.
 from typing import Any, Dict, List, Sequence, Union
 
+import tensorrt as trt
 from nvtripy import export, utils
+from nvtripy.common import datatype
 from nvtripy.common.exception import raise_error
 from nvtripy.frontend.ops import utils as op_utils
 from nvtripy.trace.ops.plugin import Plugin
 from nvtripy.utils.types import str_from_type_annotation, type_str_from_arg
 
+TRT_FROM_TRIPY_DTYPE = {
+    datatype.float32: trt.float32,
+    datatype.float16: trt.float16,
+    datatype.float8: trt.fp8,
+    datatype.bfloat16: trt.bfloat16,
+    datatype.int4: trt.int4,
+    datatype.int8: trt.int8,
+    datatype.int32: trt.int32,
+    datatype.int64: trt.int64,
+    datatype.bool: trt.bool,
+}
+
+TRIPY_FROM_TRT_DTYPE = {val: key for key, val in TRT_FROM_TRIPY_DTYPE.items()}
+
 
 # TODO (pranavm): Add link to custom layers guide once published
 @export.public_api(document_under="operations/functions")
@@ -93,11 +109,20 @@ def get_plugins_in_namespace(ns):
     input_descs = [None] * len(inputs)
     for i in range(len(inputs)):
         input_descs[i] = trtp._tensor.TensorDesc()
-        input_descs[i].dtype = inputs[i].dtype
+        input_descs[i].dtype = TRT_FROM_TRIPY_DTYPE[inputs[i].dtype]
         input_descs[i].shape_expr = trtp._tensor.ShapeExprs(inputs[i].rank, _is_dummy=True)
         input_descs[i]._immutable = True
+
     output_descs = utils.utils.make_tuple(trtp_op.register_func(*input_descs, attrs))
 
-    output_info = [(len(desc.shape_expr), desc.dtype) for desc in output_descs]
+    def tripy_from_trt_dtype(dtype):
+        if dtype not in TRIPY_FROM_TRT_DTYPE:
+            raise_error(
+                f"Unsupported TensorRT data type: '{dtype}'.",
+                details=[f"Supported types are: {list(TRIPY_FROM_TRT_DTYPE.keys())}."],
+            )
+        return TRIPY_FROM_TRT_DTYPE[dtype]
+
+    output_info = [(len(desc.shape_expr), tripy_from_trt_dtype(desc.dtype)) for desc in output_descs]
 
     return op_utils.create_op(Plugin, inputs, name, "1", namespace, output_info, kwargs)

tripy/tests/conftest.py

@@ -20,6 +20,7 @@
 
 import nvtripy as tp
 import pytest
+import tensorrt as trt
 import tensorrt.plugin as trtp
 import torch
 import triton
@@ -64,6 +65,8 @@ def add_kernel(x_ptr, n_elements, y_ptr, BLOCK_SIZE: tl.constexpr):
 
 @trtp.register("example::elemwise_add_plugin")
 def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> trtp.TensorDesc:
+    # QDPs should receive and return TRT data types
+    assert inp0.dtype == trt.float32
     return inp0.like()
 
 

tripy/tests/frontend/ops/test_plugin_qdp.py

@@ -14,9 +14,19 @@
 # limitations under the License.
 import nvtripy as tp
 import pytest
+import tensorrt as trt
+from nvtripy.common.datatype import DATA_TYPES
+from nvtripy.frontend.ops.plugin_qdp import TRIPY_FROM_TRT_DTYPE, TRT_FROM_TRIPY_DTYPE
 from tests import helper
 
 
+@pytest.mark.parametrize("dtype", DATA_TYPES.values())
+def test_dtype_trt_conversion(dtype):
+    assert dtype in TRT_FROM_TRIPY_DTYPE
+    assert isinstance(TRT_FROM_TRIPY_DTYPE[dtype], trt.DataType)
+    assert TRIPY_FROM_TRT_DTYPE[TRT_FROM_TRIPY_DTYPE[dtype]] == dtype
+
+
 class TestQuicklyDeployablePlugin:
     @pytest.mark.parametrize(
         "plugin_id, err",