tensorrt.plugin.impl

tensorrt.plugin.impl(plugin_id: str) → Callable

包装一个函数，为已通过 trt.plugin.register 注册的插件定义实现。

此 API 仅旨在用作装饰器。被装饰的函数不需要对输入参数或返回值进行类型提示；但是，指定的任何类型提示都将针对 trt.plugin.register 签名进行一致性验证。

该函数的模式如下

(inp0: Tensor, inp1: Tensor, ..., attr0: SupportedAttrType, attr1: SupportedAttrType, outputs: Tuple[Tensor], stream: int, tactic: Optional[int]) -> None

  • 首先传递输入张量，每个张量都由 Tensor 描述。

  • 接下来声明插件属性。

    • 并非 trt.plugin.register 中包含的所有属性都必须在此处指定 - 它们可以是子集。

    • 包含的属性将被序列化到 TRT 引擎。因此，应仅包含插件实际需要执行推理的属性（在 trt.plugin.impl 的主体中）。

  • tactic 是一个可选参数。如果插件使用自定义策略，则必须指定它以接收用于当前插件执行的策略值。

参数:

plugin_id – 插件的 ID，格式为“{namespace}::{name}”，必须与 trt.plugin.register 期间使用的 ID 匹配

使用 OpenAI Triton 内核实现 elementwise 插件
 1import tensorrt.plugin as trtp
 2import triton
 3import triton.language as tl
 4
 5@triton.jit
 6def add_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
 7    pid = tl.program_id(0)
 8    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
 9    mask = offsets < n_elements
10    x = tl.load(x_ptr + offsets, mask=mask)
11    tl.store(y_ptr + offsets, x + 1, mask=mask)
12
13@trtp.register("my::add_plugin")
14def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> Tuple[trtp.TensorDesc]:
15    return inp0.like()
16
17@trtp.impl("my::add_plugin")
18def add_plugin_impl(inp0: trtp.Tensor, block_size: int, outputs: Tuple[trtp.Tensor], stream: int) -> None:
19
20    n = inp0.numel()
21    inp0_t = torch.as_tensor(inp0, device="cuda")
22    out_t = torch.as_tensor(outputs[0], device="cuda")
23
24    add_kernel[(triton.cdiv(n, block_size),)](inp0_t, out_t, n, BLOCK_SIZE = block_size)