feat: Automatically generate QDP plugins

Author: bowang007

examples/dynamo/auto_generate_plugin.py

@@ -0,0 +1,151 @@
+"""
+.. _auto_generate_converters:
+
+Automatically Generate a Plugin for a Custom Kernel
+===================================================================
+
+We are going to demonstrate how to automatically generate a plugin for a custom kernel using Torch-TensorRT using
+the new Python based plugin system in TensorRT 10.7.
+
+Torch-TensorRT supports falling back to PyTorch implementations of operations in the case that Torch-TensorRT
+does not know how to compile them in TensorRT. However, this comes at the cost of a graph break and will reduce the performance of the model.
+The easiest way to fix lack of support for ops is by adding a decomposition (see:
+`Writing lowering passes for the Dynamo frontend <https://pytorch.org/TensorRT/contributors/writing_dynamo_aten_lowering_passes.html>`_) - which defines the operator
+in terms of PyTorch ops that are supported in Torch-TensorRT or a converter (see:
+`Writing converters for the Dynamo frontend <https://pytorch.org/TensorRT/contributors/dynamo_converters.html>`_) - which defines the operator in terms of TensorRT operators.
+
+In some cases there isn't a great way to do either of these, perhaps because the operator is a custom kernel that is not part of standard PyTorch or
+TensorRT cannot support it natively.
+
+For these cases, it is possible to use a TensorRT plugin to replace the operator **inside** the TensorRT engine, thereby avoiding
+the performance and resource overhead from a graph break.
+
+Previously this involved a complex process in not only building a performant kernel but setting it up to run in TensorRT (see: `Using Custom Kernels within TensorRT Engines with Torch-TensorRT <https://pytorch.org/TensorRT/tutorials/_rendered_examples/dynamo/custom_kernel_plugins.html>`_).
+With TensorRT 10.7, there is a new Python native plugin system which greatly streamlines this process. This
+plugin system also allows Torch-TensorRT to automatically generate the necessary conversion code to convert the
+operation in PyTorch to TensorRT.
+"""
+
+# %%
+# Writing Custom Operators in PyTorch
+# -----------------------------------------
+#
+#  Pervious tutorials already cover creating custom operators in PyTorch which later get used with Torch-TensorRT.
+# Here we define a simple elementwise multiplication operator in Triton. This operator is then registered as a custom op in PyTorch.
+# with its host launch code as well as a "meta-kernel", A meta-kernel is a function that describes the shape and data type
+# transformations that the operator will perform. This meta-kernel is used by Dynamo and Torch-TensorRT, so it
+# is necessary to define.
+#
+
+from typing import Tuple
+
+import tensorrt_bindings.plugin as trtp
+import torch
+import torch_tensorrt
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def elementwise_scale_mul_kernel(X, Y, Z, a, b, BLOCK_SIZE: tl.constexpr):
+    pid = tl.program_id(0)
+    # Compute the range of elements that this thread block will work on
+    block_start = pid * BLOCK_SIZE
+    # Range of indices this thread will handle
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    # Load elements from the X and Y tensors
+    x_vals = tl.load(X + offsets)
+    y_vals = tl.load(Y + offsets)
+    # Perform the element-wise multiplication
+    z_vals = x_vals * y_vals * a + b
+    # Store the result in Z
+    tl.store(Z + offsets, z_vals)
+
+
+@torch.library.custom_op("torchtrt_ex::elementwise_scale_mul", mutates_args=())  # type: ignore[misc]
+def elementwise_scale_mul(
+    X: torch.Tensor, Y: torch.Tensor, b: float = 0.2, a: int = 2
+) -> torch.Tensor:
+    # Ensure the tensors are on the GPU
+    assert X.is_cuda and Y.is_cuda, "Tensors must be on CUDA device."
+    assert X.shape == Y.shape, "Tensors must have the same shape."
+
+    # Create output tensor
+    Z = torch.empty_like(X)
+
+    # Define block size
+    BLOCK_SIZE = 1024
+
+    # Grid of programs
+    grid = lambda meta: (X.numel() // meta["BLOCK_SIZE"],)
+
+    # Launch the kernel with parameters a and b
+    elementwise_scale_mul_kernel[grid](X, Y, Z, a, b, BLOCK_SIZE=BLOCK_SIZE)
+
+    return Z
+
+
+# %%
+# The meta kernel for an elementwise operation is just the shape and dtype of one of the inputs since we will not change the shape
+# in the course of the operation.
+
+
+@torch.library.register_fake("torchtrt_ex::elementwise_scale_mul")
+def _(x: torch.Tensor, y: torch.Tensor, b: float = 0.2, a: int = 2) -> torch.Tensor:
+    return x
+
+
+# %%
+# Here we use automatic plugin creation feature in Torch-TensorRT which enables plugin registration using
+# TensorRT QDP APIs
+torch_tensorrt.dynamo.conversion.plugins.generate_plugin(
+    "torchtrt_ex::elementwise_scale_mul"
+)
+
+
+# # %%
+# # Generating the Converter
+# # -------------------------------------------------------------------
+# # Given that we have defined the custom operator in PyTorch and TensorRT, we can now generate the converter for the operation.
+# # As long as the namespace and names match, the following function will automatically generate the converter for the operation.
+torch_tensorrt.dynamo.conversion.plugins.generate_plugin_converter(
+    "torchtrt_ex::elementwise_scale_mul", supports_dynamic_shapes=True
+)
+
+
+# # %%
+# # Above two commands can be replaced with the following single one line:
+# torch_tensorrt.dynamo.conversion.plugins.custom_op("torchtrt_ex::elementwise_scale_mul", supports_dynamic_shapes=True)
+
+
+# %%
+# Using our converter with a model
+# -------------------------------------------------------------------
+#
+# Now we can use our custom operator in a model and compile it with Torch-TensorRT.
+# We can see that the custom operator is used as one of the operations in the forward pass of the model.
+# The process of compiling the model at this point is identical to standard Torch-TensorRT usage.
+class MyModel(torch.nn.Module):  # type: ignore[misc]
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        z = torch.add(x, y)
+        res = torch.ops.torchtrt_ex.elementwise_scale_mul.default(x, z, b=0.5)
+
+        return res
+
+
+my_model = MyModel().to("cuda")
+m = torch.randint(0, 5, (64, 64), device="cuda", dtype=torch.float)
+n = torch.randint(0, 5, (64, 64), device="cuda", dtype=torch.float)
+
+with torch_tensorrt.logging.errors():
+    model_trt = torch_tensorrt.compile(
+        my_model, inputs=[m, n], debug=True, min_block_size=1
+    )
+    for i in range(300):
+        res = model_trt(m, n)
+        assert torch.allclose(res, my_model(m, n))
+
+print("Ran with custom plugin!")

py/torch_tensorrt/dynamo/conversion/plugins/__init__.py

@@ -1,3 +1,5 @@
+from torch_tensorrt.dynamo.conversion.plugins._custom_op import custom_op
+from torch_tensorrt.dynamo.conversion.plugins._generate_plugin import generate_plugin
 from torch_tensorrt.dynamo.conversion.plugins._generate_plugin_converter import (
     generate_plugin_converter,
 )

py/torch_tensorrt/dynamo/conversion/plugins/_custom_op.py

@@ -0,0 +1,21 @@
+from typing import Callable, Optional
+
+from torch.fx.node import Node
+from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo.conversion._ConverterRegistry import ConverterPriority
+from torch_tensorrt.dynamo.conversion.plugins._generate_plugin import generate_plugin
+from torch_tensorrt.dynamo.conversion.plugins._generate_plugin_converter import (
+    generate_plugin_converter,
+)
+
+
+def custom_op(
+    op_name: str,
+    capability_validator: Optional[Callable[[Node, CompilationSettings], bool]] = None,
+    priority: ConverterPriority = ConverterPriority.STANDARD,
+    supports_dynamic_shapes: bool = False,
+):
+    generate_plugin(op_name)
+    generate_plugin_converter(
+        op_name, capability_validator, priority, supports_dynamic_shapes
+    )

py/torch_tensorrt/dynamo/conversion/plugins/_generate_plugin.py

@@ -0,0 +1,195 @@
+import logging
+from types import FunctionType
+from typing import Tuple
+
+import tensorrt_bindings.plugin as trtp
+import torch
+from sympy import lambdify
+from torch._dynamo.source import LocalSource
+from torch._subclasses.fake_tensor import FakeTensorMode
+from torch.fx.experimental.symbolic_shapes import DimDynamic, ShapeEnv
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+def mksym(shape_env, value, source, dynamic_dim):
+    return shape_env.create_symintnode(
+        shape_env.create_symbol(
+            value,
+            source=source,
+            dynamic_dim=dynamic_dim,
+        ),
+        hint=value,
+        source=source,
+    )
+
+
+def generate_plugin(plugin_name: str):
+    namespace, name = plugin_name.split("::")
+
+    # retrieve the corresponding torch operation using the passed in string
+    torch_op = getattr(getattr(torch.ops, namespace), name)
+
+    # helper function that generates the required signature based on the torch operation
+    def generate_signature(torch_op):
+        schema = torch_op._schemas[""]
+
+        arg_list = []
+
+        register_func_annotation = {}
+        impl_func_annotation = {}
+
+        for arg in schema.arguments:
+            arg_list.append(arg.name)
+
+            # TODO: Torch types need to be converted to python primitive types here
+            # Some other types are not handled:
+            # - torch._C.ListType.ofT(<type>)
+            # - torch._C.TupleType.get()
+            # - torch._C.DictType.get(<key_type>, <value_type>)
+            # - torch._C.OptionalType.ofT(<type>)
+            # - torch._C.DeviceObjType.get()
+            # - torch._C.FunctionType.get()
+            # - torch._C.ClassType
+
+            if arg.type.isSubtypeOf(torch._C.TensorType.get()):
+                register_func_annotation[arg.name] = trtp.TensorDesc
+                impl_func_annotation[arg.name] = trtp.Tensor
+            elif arg.type.isSubtypeOf(torch._C.FloatType.get()):
+                register_func_annotation[arg.name] = float
+                impl_func_annotation[arg.name] = float
+            elif arg.type.isSubtypeOf(torch._C.IntType.get()):
+                register_func_annotation[arg.name] = int
+                impl_func_annotation[arg.name] = int
+            elif arg.type.isSubtypeOf(torch._C.Booltype.get()):
+                register_func_annotation[arg.name] = bool
+                impl_func_annotation[arg.name] = bool
+            elif arg.type.isSubtypeOf(torch._C.Stringtype.get()):
+                register_func_annotation[arg.name] = str
+                impl_func_annotation[arg.name] = str
+            else:
+                raise ValueError("arg type is not handled")
+
+        input_signature = ", ".join(arg_list)
+
+        plugin_signature = f"def add_plugin_desc({input_signature}):"
+
+        plugin_impl_arg_list = arg_list
+        plugin_impl_arg_list.append("outputs")
+        plugin_impl_arg_list.append("stream")
+        plugin_impl_input = ", ".join(plugin_impl_arg_list)
+        plugin_impl_signature = f"def add_plugin_impl({plugin_impl_input}):"
+
+        register_func_annotation["return"] = Tuple[trtp.TensorDesc]
+
+        impl_func_annotation["outputs"] = Tuple[trtp.Tensor]
+        impl_func_annotation["stream"] = int
+
+        return (
+            input_signature,
+            plugin_signature,
+            plugin_impl_signature,
+            register_func_annotation,
+            impl_func_annotation,
+        )
+
+    # Use the helper function to get the required signatures
+    (
+        input_signature,
+        plugin_signature,
+        plugin_impl_signature,
+        register_func_annotation,
+        impl_func_annotation,
+    ) = generate_signature(torch_op)
+
+    def _generic_plugin_desc(*args, **kwargs) -> Tuple[trtp.TensorDesc]:
+        shape_env = ShapeEnv()
+        fake_mode = FakeTensorMode(shape_env=shape_env)
+        syms_args = []
+        tensor_args = [elem for elem in args if isinstance(elem, trtp.TensorDesc)]
+
+        for tensor_arg in tensor_args:
+
+            sample = {f"{i}": 5 for i in range(tensor_arg.ndim)}
+            syms_arg = [
+                mksym(shape_env, v, LocalSource(k), DimDynamic.DYNAMIC)
+                for k, v in sample.items()
+            ]
+            syms_args.append(syms_arg)
+
+        with FakeTensorMode() as fake_mode:
+            fake_args = []
+            for syms_arg in syms_args:
+                fake_arg = torch.randn(syms_arg)
+                fake_args.append(fake_arg)
+
+            output = torch_op(*fake_args, **kwargs)
+
+        # We assume that number of dimensions are the same in torch op
+        shape_calc_fns = [None] * args[0].ndim
+        for i in range(args[0].ndim):
+            input_node_expr = [syms_arg[i].node.expr for syms_arg in syms_args]
+            shape_calc_fns[i] = lambdify(
+                tuple(input_node_expr), output.shape[i].node.expr, "math"
+            )
+
+        out_desc = tensor_args[0].like()
+        for i in range(out_desc.ndim):
+            input_shape_expr = [tensor_arg.shape_expr[i] for tensor_arg in tensor_args]
+            out_desc.shape_expr[i] = shape_calc_fns[i](*input_shape_expr)
+
+        return (out_desc,)
+
+    codegen_plugin = f"""
+{plugin_signature}
+    return _generic_plugin_desc({input_signature})
+    """
+
+    _LOGGER.warning(f"Plugin registration function: \n{codegen_plugin}")
+
+    plugin_code = compile(codegen_plugin, "<string>", "exec")
+
+    globals()["_generic_plugin_desc"] = _generic_plugin_desc
+
+    plugin = FunctionType(
+        plugin_code.co_consts[0],
+        globals(),
+        "plugin",
+    )
+
+    # Function annotation is required for dynamic function to work in TensorRT.Plugin
+    plugin.__annotations__ = register_func_annotation
+
+    trtp.register(plugin_name)(plugin)
+
+    def _generic_plugin_impl(outputs, stream, *args, **kwargs):
+        tensor_args = [elem for elem in args if isinstance(elem, trtp.Tensor)]
+        print(args)
+        non_tensor_args = [elem for elem in args if not isinstance(elem, trtp.Tensor)]
+        in_tensors = [torch.as_tensor(i, device="cuda") for i in tensor_args]
+
+        dest_tensors = [torch.as_tensor(o, device="cuda") for o in outputs]
+
+        stream = torch.cuda.ExternalStream(stream)
+        with torch.cuda.stream(stream):
+            out_tensors = torch_op(*in_tensors, *non_tensor_args, **kwargs)
+            if isinstance(out_tensors, torch.Tensor):
+                out_tensors = (out_tensors,)
+            [d.copy_(o) for (d, o) in zip(dest_tensors, out_tensors)]
+
+    plugin_impl_func = f"""
+{plugin_impl_signature}
+    _generic_plugin_impl(outputs, stream, {input_signature})
+    """
+
+    _LOGGER.warning(f"Plugin implementation function: \n{plugin_impl_func}")
+
+    plugin_impl_code = compile(plugin_impl_func, "<string>", "exec")
+
+    globals()["_generic_plugin_impl"] = _generic_plugin_impl
+
+    plugin_impl = FunctionType(plugin_impl_code.co_consts[0], globals(), "plugin_impl")
+
+    plugin_impl.__annotations__ = impl_func_annotation
+
+    trtp.impl(plugin_name)(plugin_impl)