自動為自訂核心產生轉換器¶

我們將示範如何使用 TensorRT 10.7 中新的基於 Python 的外掛系統，使用 Torch-TensorRT 自動產生自訂核心的轉換器。

在 Torch-TensorRT 不知道如何將運算子編譯到 TensorRT 中的情況下，Torch-TensorRT 支援回退到 PyTorch 運算子的實作。然而，這會以圖形中斷為代價，並降低模型的效能。解決缺少運算子支援的最簡單方法是新增分解（請參閱：為 Dynamo 前端撰寫 lowering passes）- 定義以 Torch-TensorRT 支援的 PyTorch 運算子表示的運算子，或轉換器（請參閱：為 Dynamo 前端撰寫轉換器）- 定義以 TensorRT 運算子表示的運算子。

在某些情況下，沒有很好的方法來做到這兩者，可能是因為運算子是自訂核心，而不是標準 PyTorch 的一部分，或者 TensorRT 無法原生支援它。

對於這些情況，可以使用 TensorRT 外掛程式來取代 TensorRT 引擎內部的運算子，從而避免圖形中斷造成的效能和資源開銷。

以前，這涉及一個複雜的過程，不僅要建構高效能的核心，還要設定它在 TensorRT 中執行（請參閱：在具有 Torch-TensorRT 的 TensorRT 引擎中使用自訂核心）。透過 TensorRT 10.7，有一個新的 Python 原生外掛系統，它極大地簡化了這個過程。這個外掛系統也允許 Torch-TensorRT 自動產生必要的轉換程式碼，以將 PyTorch 中的運算轉換為 TensorRT。

在 PyTorch 中撰寫自訂運算子¶

先前的教學已經涵蓋了在 PyTorch 中建立自訂運算子，這些運算子稍後將與 Torch-TensorRT 一起使用。

在這裡，我們在 Triton 中定義一個簡單的元素級乘法運算子。然後，此運算子在 PyTorch 中註冊為自訂運算子，並包含其主機啟動程式碼以及「meta-kernel」。 Meta-kernel 是一個描述運算子將執行的形狀和資料類型轉換的函式。 Dynamo 和 Torch-TensorRT 都使用此 meta-kernel，因此必須定義它。

from typing import Tuple

import tensorrt.plugin as trtp
import torch
import torch_tensorrt
import triton
import triton.language as tl


@triton.jit
def elementwise_mul_kernel(X, Y, Z, BLOCK_SIZE: tl.constexpr):
    # Program ID determines the block of data each thread will process
    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
    # Store the result in Z
    tl.store(Z + offsets, z_vals)


@torch.library.custom_op("torchtrt_ex::elementwise_mul", mutates_args=())  # type: ignore[misc]
def elementwise_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
    elementwise_mul_kernel[grid](X, Y, Z, BLOCK_SIZE=BLOCK_SIZE)

    return Z

元素級運算的 meta kernel 只是其中一個輸入的形狀和 dtype，因為我們在運算過程中不會更改形狀。

@torch.library.register_fake("torchtrt_ex::elementwise_mul")
def _(x: torch.Tensor, y: torch.Tensor, b: float = 0.2, a: int = 2) -> torch.Tensor:
    return x

使用快速部署外掛系統為 TensorRT 撰寫外掛程式¶

TensorRT 10.7 中的快速部署外掛系統允許使用更少的樣板程式碼在 Python 中建立自訂外掛程式。它使用類似 PyTorch 的系統，您可以在其中定義一個描述運算子將執行的形狀和資料類型轉換的函式，然後定義程式碼以啟動給定 GPU 記憶體控制碼的核心。

就像 PyTorch meta kernel 一樣，輸入和輸出之間在形狀或資料類型上沒有轉換，因此我們可以告訴 TensorRT 預期與我們輸入的形狀相同的形狀

@trtp.register("torchtrt_ex::elementwise_mul")
def _(
    x: trtp.TensorDesc, y: trtp.TensorDesc, b: float, a: int
) -> Tuple[trtp.TensorDesc]:
    return x.like()

在這裡，我們重複使用與 PyTorch 類似的主機啟動程式碼，但我們需要在啟動核心之前將 TensorRT 張量轉換為 PyTorch 張量。這些運算也是就地運算，因此結果必須放入 TensorRT 提供的輸出張量中。

@trtp.impl("torchtrt_ex::elementwise_mul")
def _(
    x: trtp.Tensor,
    y: trtp.Tensor,
    b: float,
    a: int,
    outputs: Tuple[trtp.Tensor],
    stream: int,
):
    # Define block size
    BLOCK_SIZE = 1024

    # Grid of programs
    grid = lambda meta: (x.numel() // meta["BLOCK_SIZE"],)

    x_t = torch.as_tensor(x, device="cuda")
    y_t = torch.as_tensor(y, device="cuda")
    z_t = torch.as_tensor(outputs[0], device="cuda")
    # Launch the kernel
    elementwise_mul_kernel[grid](x_t, y_t, z_t, BLOCK_SIZE=BLOCK_SIZE)

產生轉換器¶

鑑於我們已經在 PyTorch 和 TensorRT 中定義了自訂運算子，我們現在可以為運算產生轉換器。只要命名空間和名稱匹配，以下函式就會自動為運算產生轉換器。

torch_tensorrt.dynamo.conversion.plugins.generate_plugin_converter(
    "torchtrt_ex::elementwise_mul", supports_dynamic_shapes=True
)

將我們的轉換器與模型一起使用¶

現在我們可以在模型中使用我們的自訂運算子，並使用 Torch-TensorRT 編譯它。我們可以看到，自訂運算子被用作模型前向傳遞中的運算子之一。此時編譯模型的過程與標準 Torch-TensorRT 用法相同。

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_mul.default(x, z, a=1)

        return res


my_model = MyModel().to("cuda")
m = torch.full((64, 64), 2, device="cuda", dtype=torch.float)
n = torch.full((64, 64), 3, 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!")

腳本的總執行時間： ( 0 分鐘 0.000 秒)

由 Sphinx-Gallery 產生圖庫