编写Dynamo转换器¶

Torch-TensorRT中的dynamo转换器库位于TensorRT/py/torch_tensorrt/dynamo/conversion。

转换器实现¶

注册¶

转换器是一个用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: 转换器实现的节点目标（例如，torch.ops.aten.leaky_relu.default）

enabled: 转换器是否应作为可以在转换器注册表中使用的转换器启用

capability_validator: 一个可以接受torch.fx.Node并确定转换器是否可以正确处理此节点的lambda函数。如果验证器返回False，子图分区器将确保在编译图中在PyTorch中运行此节点。

priority: 允许开发者在转换器注册表中覆盖现有的转换器

转换器所需的只是密钥。

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

ctx : 编译器的当前状态。转换器主要会操作ctx.net，这是正在构建的tensorrt.INetworkDefinition。此结构体中还包含用户提供的设置等额外元数据。

target: 上面call_module或call_function中的目标键。例如：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中使用，匹配正在转换的操作的输出签名。

能力验证¶

有一些转换器需要考虑特殊情况。在这些情况下，应该使用capability_validators来注册转换器，使用@dynamo_tensorrt_converter。我们通过torch.ops.aten.embedding.default来说明这一点。它有参数scale_grad_by_freq和sparse，这些参数目前不受实现支持。在这种情况下，我们可以编写验证器embedding_param_validator，它实现了给定这些参数时不支持转换器，并通过以下方式注册转换器：

类型合约¶

该函数应遵循签名所建立的类型契约。这包括接受有效的PyTorch类型和numpy数组的联合，用于常量张量和TensorRT ITensors。如果转换器仅支持类型的子集，您还可以添加torch_tensorrt.dynamo.conversion.converter_utils.enforce_tensor_types，它允许您指定输入位置和这些输入可以接受的类型之间的字典映射。在可能的情况下，装饰器将转换输入以匹配这些类型，优先使用提供的顺序。字典中的int键将引用args中的位置参数。str键将引用kwargs中的关键字参数。

示例: `卷积`¶

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

@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。

操作符分解¶

在PyTorch中，有一些转换器可以分解为子操作，不需要单独的转换器注册。这些转换器可以通过分解来实现

示例: `addmm`¶

分解通过register_torch_trt_decomposition装饰器注册我们定义了addmm_replacement并用torch操作替换它，这些操作将调用它们相应的转换器。

@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 降低操作。

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