|
| 1 | +""" |
| 2 | +.. _auto_generate_converters: |
| 3 | +
|
| 4 | +Automatically Generate a Converter for a Custom Kernel |
| 5 | +=================================================================== |
| 6 | +
|
| 7 | +We are going to demonstrate how to automatically generate a converter for a custom kernel using Torch-TensorRT using |
| 8 | +the new Python based plugin system in TensorRT 10.7. |
| 9 | +
|
| 10 | +Torch-TensorRT supports falling back to PyTorch implementations of operations in the case that Torch-TensorRT |
| 11 | +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. |
| 12 | +The easiest way to fix lack of support for ops is by adding a decomposition (see: |
| 13 | +`Writing lowering passes for the Dynamo frontend <https://pytorch.org/TensorRT/contributors/writing_dynamo_aten_lowering_passes.html>`_) - which defines the operator |
| 14 | +in terms of PyTorch ops that are supported in Torch-TensorRT or a converter (see: |
| 15 | +`Writing converters for the Dynamo frontend <https://pytorch.org/TensorRT/contributors/dynamo_converters.html>`_) - which defines the operator in terms of TensorRT operators. |
| 16 | +
|
| 17 | +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 |
| 18 | +TensorRT cannot support it natively. |
| 19 | +
|
| 20 | +For these cases, it is possible to use a TensorRT plugin to replace the operator **inside** the TensorRT engine, thereby avoiding |
| 21 | +the performance and resource overhead from a graph break. |
| 22 | +
|
| 23 | +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>`_). |
| 24 | +With TensorRT 10.7, there is a new Python native plugin system which greatly streamlines this process. This |
| 25 | +plugin system also allows Torch-TensorRT to automatically generate the necessary conversion code to convert the |
| 26 | +operation in PyTorch to TensorRT. |
| 27 | +""" |
| 28 | + |
| 29 | +# %% |
| 30 | +# Writing Custom Operators in PyTorch |
| 31 | +# ----------------------------------------- |
| 32 | +# |
| 33 | +# Pervious tutorials already cover creating custom operators in PyTorch which later get used with Torch-TensorRT. |
| 34 | +# Here we define a simple elementwise multiplication operator in Triton. This operator is then registered as a custom op in PyTorch. |
| 35 | +# with its host launch code as well as a "meta-kernel", A meta-kernel is a function that describes the shape and data type |
| 36 | +# transformations that the operator will perform. This meta-kernel is used by Dynamo and Torch-TensorRT, so it |
| 37 | +# is necessary to define. |
| 38 | +# |
| 39 | + |
| 40 | +from typing import Tuple |
| 41 | + |
| 42 | +import tensorrt_bindings.plugin as trtp |
| 43 | +import torch |
| 44 | +import torch_tensorrt |
| 45 | +import triton |
| 46 | +import triton.language as tl |
| 47 | + |
| 48 | + |
| 49 | +@triton.jit |
| 50 | +def elementwise_mul_kernel(X, Y, Z, BLOCK_SIZE: tl.constexpr): |
| 51 | + # Program ID determines the block of data each thread will process |
| 52 | + pid = tl.program_id(0) |
| 53 | + # Compute the range of elements that this thread block will work on |
| 54 | + block_start = pid * BLOCK_SIZE |
| 55 | + # Range of indices this thread will handle |
| 56 | + offsets = block_start + tl.arange(0, BLOCK_SIZE) |
| 57 | + # Load elements from the X and Y tensors |
| 58 | + x_vals = tl.load(X + offsets) |
| 59 | + y_vals = tl.load(Y + offsets) |
| 60 | + # Perform the element-wise multiplication |
| 61 | + z_vals = x_vals * y_vals |
| 62 | + # Store the result in Z |
| 63 | + tl.store(Z + offsets, z_vals) |
| 64 | + |
| 65 | + |
| 66 | +@torch.library.custom_op("torchtrt_ex::elementwise_mul", mutates_args=()) # type: ignore[misc] |
| 67 | +def elementwise_mul( |
| 68 | + X: torch.Tensor, Y: torch.Tensor, b: float = 0.2, a: int = 2 |
| 69 | +) -> torch.Tensor: |
| 70 | + # Ensure the tensors are on the GPU |
| 71 | + assert X.is_cuda and Y.is_cuda, "Tensors must be on CUDA device." |
| 72 | + assert X.shape == Y.shape, "Tensors must have the same shape." |
| 73 | + |
| 74 | + # Create output tensor |
| 75 | + Z = torch.empty_like(X) |
| 76 | + |
| 77 | + # Define block size |
| 78 | + BLOCK_SIZE = 1024 |
| 79 | + |
| 80 | + # Grid of programs |
| 81 | + grid = lambda meta: (X.numel() // meta["BLOCK_SIZE"],) |
| 82 | + |
| 83 | + # Launch the kernel |
| 84 | + elementwise_mul_kernel[grid](X, Y, Z, BLOCK_SIZE=BLOCK_SIZE) |
| 85 | + |
| 86 | + return Z |
| 87 | + |
| 88 | + |
| 89 | +# %% |
| 90 | +# 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 |
| 91 | +# in the course of the operation. |
| 92 | + |
| 93 | + |
| 94 | +@torch.library.register_fake("torchtrt_ex::elementwise_mul") |
| 95 | +def _(x: torch.Tensor, y: torch.Tensor, b: float = 0.2, a: int = 2) -> torch.Tensor: |
| 96 | + return x |
| 97 | + |
| 98 | + |
| 99 | +# %% |
| 100 | +# Writing Plugins for TensorRT using the Quick Deploy Plugin system |
| 101 | +# ------------------------------------------------------------------- |
| 102 | +# The quick deployment plugin system in TensorRT 10.7 allows for the creation of custom plugins in Python with significantly |
| 103 | +# less boilerplate. It uses a similar system PyTorch where you define a function that describes the shape and data type transformations |
| 104 | +# that the operator will perform and then define the code to launch the kernel given GPU memory handles. |
| 105 | +# |
| 106 | + |
| 107 | + |
| 108 | +# %% |
| 109 | +# Just like the PyTorch meta kernel, there is no transformation in shape or data type between the input and output so |
| 110 | +# we can just tell TensorRT to expect the same shape as we get in |
| 111 | +# |
| 112 | +@trtp.register("torchtrt_ex::elementwise_mul") |
| 113 | +def _( |
| 114 | + x: trtp.TensorDesc, y: trtp.TensorDesc, b: float, a: int |
| 115 | +) -> Tuple[trtp.TensorDesc]: |
| 116 | + return x.like() |
| 117 | + |
| 118 | + |
| 119 | +# %% |
| 120 | +# Here we reuse similar host launch code as PyTorch but we need to convert the TensorRT tensors into PyTorch tensors prior to launching the kernel |
| 121 | +# These operations are also in-place, so the result must be put in the the output tensors provided by TensorRT. |
| 122 | +@trtp.impl("torchtrt_ex::elementwise_mul") |
| 123 | +def _( |
| 124 | + x: trtp.Tensor, |
| 125 | + y: trtp.Tensor, |
| 126 | + b: float, |
| 127 | + a: int, |
| 128 | + outputs: Tuple[trtp.Tensor], |
| 129 | + stream: int, |
| 130 | +): |
| 131 | + # Define block size |
| 132 | + BLOCK_SIZE = 1024 |
| 133 | + |
| 134 | + # Grid of programs |
| 135 | + grid = lambda meta: (x.numel() // meta["BLOCK_SIZE"],) |
| 136 | + |
| 137 | + x_t = torch.as_tensor(x, device="cuda") |
| 138 | + y_t = torch.as_tensor(y, device="cuda") |
| 139 | + z_t = torch.as_tensor(outputs[0], device="cuda") |
| 140 | + # Launch the kernel |
| 141 | + elementwise_mul_kernel[grid](x_t, y_t, z_t, BLOCK_SIZE=BLOCK_SIZE) |
| 142 | + |
| 143 | + |
| 144 | +# %% |
| 145 | +# Generating the Converter |
| 146 | +# ------------------------------------------------------------------- |
| 147 | +# Given that we have defined the custom operator in PyTorch and TensorRT, we can now generate the converter for the operation. |
| 148 | +# As long as the namespace and names match, the following function will automatically generate the converter for the operation. |
| 149 | +torch_tensorrt.dynamo.conversion.plugins.generate_plugin_converter( |
| 150 | + "torchtrt_ex::elementwise_mul", supports_dynamic_shapes=True |
| 151 | +) |
| 152 | + |
| 153 | + |
| 154 | +# %% |
| 155 | +# Using our converter with a model |
| 156 | +# ------------------------------------------------------------------- |
| 157 | +# |
| 158 | +# Now we can use our custom operator in a model and compile it with Torch-TensorRT. |
| 159 | +# We can see that the custom operator is used as one of the operations in the forward pass of the model. |
| 160 | +# The process of compiling the model at this point is identical to standard Torch-TensorRT usage. |
| 161 | +class MyModel(torch.nn.Module): # type: ignore[misc] |
| 162 | + def __init__(self): |
| 163 | + super().__init__() |
| 164 | + |
| 165 | + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 166 | + z = torch.add(x, y) |
| 167 | + res = torch.ops.torchtrt_ex.elementwise_mul.default(x, z, a=1) |
| 168 | + |
| 169 | + return res |
| 170 | + |
| 171 | + |
| 172 | +my_model = MyModel().to("cuda") |
| 173 | +m = torch.full((64, 64), 2, device="cuda", dtype=torch.float) |
| 174 | +n = torch.full((64, 64), 3, device="cuda", dtype=torch.float) |
| 175 | + |
| 176 | +with torch_tensorrt.logging.errors(): |
| 177 | + model_trt = torch_tensorrt.compile( |
| 178 | + my_model, inputs=[m, n], debug=True, min_block_size=1 |
| 179 | + ) |
| 180 | + for i in range(300): |
| 181 | + res = model_trt(m, n) |
| 182 | + assert torch.allclose(res, my_model(m, n)) |
| 183 | + |
| 184 | +print("Ran with custom plugin!") |
0 commit comments