编写 Dynamo Converter

Torch-TensorRT 中的 dynamo converter 库位于 TensorRT/py/torch_tensorrt/dynamo/conversion。

Converter 实现

注册

Converter 是一个使用 torch_tensorrt.dynamo.dynamo_tensorrt_converter 装饰器装饰的函数，遵循以下函数签名

@torch_tensorrt.dynamo.conversion.dynamo_tensorrt_converter(torch.ops.aten.leaky_relu.default)
def leaky_relu_converter(
    ctx: torch_tensorrt.dynamo.conversion.ConversionCtx,
    target: Target,
    args: Tuple[Argument, ...],
    kwargs: Dict[str, Argument],
    name: str,
) -> Union[tensorrt.ITensor, Sequence[tensorrt.ITensor]]:

该装饰器接受多个参数

• key: 实现 converter 的节点目标（例如，torch.ops.aten.leaky_relu.default）

• enabled: 此 converter 是否应被启用，以便在 converter 注册表中可用

• capability_validator: 一个 lambda 函数，它接受一个 torch.fx.Node 并判断此 converter 是否能正确处理此节点。如果验证器返回 False，子图分区器将确保此节点在编译图中于 PyTorch 中运行。

• priority: 允许开发者覆盖 converter 注册表中现有的 converter

Converter 所需的全部内容只有 key。

函数体负责获取网络的当前状态，并添加下一个子图以使用 TensorRT 操作执行装饰器中指定的操作。函数提供的参数与原生 PyTorch 操作提供的参数相同，但增加了用于冻结张量属性的 numpy 数组或 TensorRT ITensor（它们是前一个节点的输出张量，对应于图中的中间操作的边/输出张量）。要确定预期的类型以及 converter 的返回类型，请查看要转换的操作的定义。对于 aten 操作，此文件是事实来源：https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml 由于开发者可能编写的许多 converter 是较低级别操作的组合，因此无需在原始 TensorRT 中实现 converter，torch_tensorrt.dynamo.conversion.impl 子包包含许多可以链式组合以创建 TensorRT 子图的操作实现。

• ctx : 编译器的当前状态。Converter 主要操作 ctx.net，它是正在构建的 tensorrt.INetworkDefinition。此结构中还包含其他元数据，包括用户提供的设置。

• target: 上述 call_module 或 call_function 中的目标 key。例如：torch.ops.aten_.leaky_relu.default。请注意，torch.ops.aten._leaky_relu 是 OpOverloadPacket，而 torch.ops.aten_.leaky_relu.default 是 OpOverload。

• args: 传递给特定节点的参数（由 torch_tensorrt.dynamo.conversion.TRTInterpreter 收集）。这些参数连同 kwargs 用于构建一个特定的 TensorRT 子图，表示 INetworkDefinition 中的当前节点。

• kwargs: 传递给特定节点的参数（由 torch_tensorrt.dynamo.conversion.TRTInterpreter 收集）。

• name: 包含目标名称的字符串

函数预期返回 tensorrt.ITensor 或一些 tensorrt.ITensor 的集合，供 torch_tensorrt.dynamo.conversion.TRTInterpreter 使用，匹配被转换操作的输出签名

能力验证

有些 converter 需要考虑特殊情况。在这种情况下，应该使用 capability_validators 通过 @dynamo_tensorrt_converter 注册 converter。我们通过 torch.ops.aten.embedding.default 来阐述这一点。它有两个参数 - scale_grad_by_freq 和 sparse，这些参数目前不支持。在这种情况下，我们可以编写 validator embedding_param_validator，它实现的是，如果给定这些参数，则不支持此 converter，并通过以下方式注册 converter

类型契约

函数应遵循签名建立的类型契约。这包括接受有效的 PyTorch 类型与常量张量的 numpy 数组以及 TensorRT ITensor 的并集。如果 converter 中仅支持部分类型，您还可以添加 torch_tensorrt.dynamo.conversion.converter_utils.enforce_tensor_types，它允许您指定一个字典，映射输入位置与其可以接受的类型。在可能的情况下，装饰器将转换输入以匹配这些类型，优先考虑提供的顺序。字典中的 int key 将引用 args 中的位置参数。 str key 将引用 kwargs 中的关键字参数。

示例: Convolution

默认的卷积 converter 同时使用能力验证器和类型强制来防止在不支持的情况下运行。能力验证器在分区期间运行，以确定特定的卷积节点是否可以转换为 TensorRT 或需要在 PyTorch 中运行。在此，验证器确保卷积不超过 3D。类型强制器会在调用 converter 之前进行自动类型转换，将输入转换为 converter 中支持的类型，从而限制作者必须处理的情况数量。

@dynamo_tensorrt_converter(
    torch.ops.aten.convolution.default, capability_validator=lambda conv_node: conv_node.args[7] in ([0], [0, 0], [0, 0, 0])
)  # type: ignore[misc]
@enforce_tensor_types(
    {
        0: (TRTTensor,),
        1: (np.ndarray, torch.Tensor, TRTTensor),
        2: (np.ndarray, torch.Tensor, TRTTensor),
    }
)  # type: ignore[misc]
def aten_ops_convolution(
    ctx: ConversionContext,
    target: Target,
    args: Tuple[Argument, ...],
    kwargs: Dict[str, Argument],
    name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:

评估器

有些操作不会产生 TensorRT 子图作为副作用。这些操作被称为评估器。

示例: operator.getitem

评估器之所以这样分类，是因为它们不对图进行任何修改。这在 py/torch_tensorrt/dynamo/conversion/op_evaluators.py 中实现，并附带相应的 capbility_validator。操作码是 operator.getitem。

操作符分解

有些 converter 可以分解为 PyTorch 中的子操作，无需单独注册 converter。此类 converter 可以通过分解实现

示例: addmm

分解通过 register_torch_trt_decomposition 装饰器注册。我们定义 addmm_replacement 并将其替换为 torch ops，这将调用其相应的 converter。

@torch_tensorrt.dynamo.lowering.register_torch_trt_decomposition(torch.ops.aten.addmm)
def addmm_replacement(
    input_: torch.Tensor, mat1: torch.Tensor, mat2: torch.Tensor, *, beta=1, alpha=1
) -> torch.Tensor:
    return torch.add(
        torch.mul(input_, beta), torch.mul(torch.matmul(mat1, mat2), alpha)
    )

您可以通过编辑 torch_tensorrt.dynamo.lowering.torch_enabled_decompositions 和 torch_tensorrt.dynamo.lowering.torch_disabled_decompositions 来修改运行的分解

注意：torch_tensorrt.dynamo.lowering.torch_enabled_decompositions 和 torch_tensorrt.dynamo.lowering.torch_disabled_decompositions 必须是不相交的集合，并且已经在 torch_tensorrt.dynamo.lowering 中定义的分解将优先于 torch 下层操作。

在大多数情况下，这比实现 converter 要容易得多。因此，在可能的情况下，应首先尝试这种方法。