分区阶段¶

此阶段是可选的，由用户启用。它指示编译器将节点分为应在 PyTorch 中运行的节点和应在 TensorRT 中运行的节点。分离的标准包括：缺少转换器、用户显式将操作设置为在 PyTorch 中运行，或者节点具有由模块回退 Pass 通知分区应在 PyTorch 中运行的标志。

总的来说，Torch-TensorRT 分区阶段执行以下操作：

分段。按顺序遍历操作集，验证每个操作是否有对应的转换器。然后，大致将图分成 Torch-TensorRT 支持的部分和不支持的部分。
依赖分析。对于每个要编译的操作，都存在一个“完整依赖图”，这意味着每个输入都可以追溯到作为 Tensor 或 TensorList 的输入。在分段后遍历所有段，然后进行依赖分析，确保 TensorRT 段只有 Tensor/TensorList 输入和输出。
形状分析。对于每个段，根据用户提供的输入形状计算输入和输出形状。形状可以通过使用 JIT 运行图来计算。
转换。每个 TensorRT 段将被转换为 TensorRT 引擎。这部分在 compiler.cpp 中完成，但它仍然是我们分区过程中的一个阶段。
拼接。将所有 TensorRT 引擎与 PyTorch 节点拼接在一起。

以下是每个文件的功能的简要说明：

PartitonInfo.h/.cpp¶

core/partitioning/PartitionInfo.h

用于分区的自动回退 API。

SegmentedBlock.h/.cpp¶

core/partitioning/SegmentedBlock.h

用于维护分段后每个段信息的主要数据结构。

shape_analysis.h/.cpp¶

core/partitioning/shape_analysis.h

通过在 JIT 中运行来获取每个段形状的代码实现。

partitioning.h/.cpp¶

core/partitioning/partitioning.h

分区阶段的 API 和主要代码实现。

自动回退¶

要启用自动回退功能，您可以在 Python 中设置以下属性：

import torch
import torch_tensorrt as torchtrt

...
model = MyModel()
ts_model = torch.jit.script(model)
trt_model = torchtrt.ts.compile(model, **{
  ...
  "min_block_size" : 3,
  "torch_executed_ops": ["aten::add"],
  "torch_executed_modules": [],
})

enabled：默认情况下自动回退处于关闭状态。通过将其设置为 True 来启用。
min_block_size：必须满足被转换为 TensorRT 的连续操作的最小数量。例如，如果设置为 3，则必须有 3 个连续支持的操作，该段才会被转换。
forced_fallback_ops：一个字符串列表，包含用户明确希望在 PyTorch 节点中执行的操作名称。

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"

...
auto in = torch::randn({1, 3, 224, 224}, {torch::kCUDA});

auto mod = torch::jit::load("trt_ts_module.ts");
auto input_sizes =  std::vector<torchtrt::InputRange>{{in.sizes()}};
torchtrt::ts::CompileSpec cfg(input_sizes);
cfg.min_block_size = 2;
cfg.torch_executed_ops.push_back("aten::relu");
auto trt_mod = torchtrt::ts::compile(mod, cfg);
auto out = trt_mod.forward({in});

依赖感知分区¶

在分段过程中，Torch-TensorRT 使用输入 TorchScript 节点的依赖图来减少创建的段数量。考虑 tests/core/partitioning/test_segmentation.cpp 中 Partitioning.SegmentModelWithDependencyAwareness 测试中的这个例子：

graph(%x : Tensor, %y : Tensor):
    %3 : int = prim::Constant[value=0]()
    %20 : int = prim::Constant[value=1]()
    %add : Tensor = aten::add(%x, %y, %20)
    %x_lgamma : Tensor = aten::lgamma(%x)
    %mul : Tensor = aten::mul(%x, %y)
    %y_lgamma : Tensor = aten::lgamma(%y)
    %div : Tensor = aten::div(%x, %y)
    %div_lgamma : Tensor = aten::lgamma(%div)
    %27 : Tensor[] = prim::ListConstruct(%x_lgamma, %y_lgamma, %div_lgamma, %add, %mul)
    %12 : Tensor = aten::cat(%27, %3)
    return (%12)

在这个图中，aten::lgamma 不支持转换，必须被分区到 Torch 回退段中。如果 Torch-TensorRT 使用贪婪分段策略，按顺序遍历输入图中的节点，并将目标相同（TensorRT 或 Torch）的操作收集到一个段中，直到遇到目标不同的操作，则结果分区将包含 7 个段，其中许多段只有一个操作。

Segment Block @0:
    Target: TensorRT

    Graph: graph(%x : Tensor,
        %y : Tensor):
    %3 : int = prim::Constant[value=1]()
    %0 : Tensor = aten::add(%x, %y, %3)
    return ()

Segment Block @1:
    Target: Torch

    Graph: graph(%x : Tensor):
    %0 : Tensor = aten::lgamma(%x)
    return ()

Segment Block @2:
    Target: TensorRT

    Graph: graph(%x : Tensor,
        %y : Tensor):
    %0 : Tensor = aten::mul(%x, %y)
    return ()

Segment Block @3:
    Target: Torch

    Graph: graph(%y : Tensor):
    %0 : Tensor = aten::lgamma(%y)
    return ()

Segment Block @4:
    Target: TensorRT

    Graph: graph(%x : Tensor,
        %y : Tensor):
    %0 : Tensor = aten::div(%x, %y)
    return ()

Segment Block @5:
    Target: Torch

    Graph: graph(%1 : Tensor):
    %0 : Tensor = aten::lgamma(%1)
    return ()

Segment Block @6:
    Target: TensorRT

    Graph: graph(%1 : Tensor,
        %2 : Tensor,
        %3 : Tensor,
        %4 : Tensor,
        %5 : Tensor):
    %7 : int = prim::Constant[value=0]()
    %0 : Tensor[] = prim::ListConstruct(%1, %2, %3, %4, %5)
    %6 : Tensor = aten::cat(%0, %7)
    return ()

这种分区是有效的，但分段不是最优的。这些算术操作和 aten::lgamma 操作在我们在线性遍历图中在 Torch 和 TensorRT 目标之间切换时，各自被分割到自己的段中。

%add : Tensor = aten::add(%x, %y, %20)
%x_lgamma : Tensor = aten::lgamma(%x)
%mul : Tensor = aten::mul(%x, %y)
%y_lgamma : Tensor = aten::lgamma(%y)
%div : Tensor = aten::div(%x, %y)
%div_lgamma : Tensor = aten::lgamma(%div)

此段中的每个算术操作仅依赖于常量和输入 %x 和 %y。aten::lgamma 操作依赖于输入 %x、%y 和 aten::div 的输出。这意味着我们可以将输入图的这部分重写如下，而不会改变图的行为。使用上述贪婪分段方法，重排后的操作系列可以清晰地分区为仅 2 个段。

%add : Tensor = aten::add(%x, %y, %20)
%mul : Tensor = aten::mul(%x, %y)
%div : Tensor = aten::div(%x, %y)
%x_lgamma : Tensor = aten::lgamma(%x)
%y_lgamma : Tensor = aten::lgamma(%y)
%div_lgamma : Tensor = aten::lgamma(%div)

通过在基本贪婪分段方法中增加对操作之间依赖关系的感知，我们可以在不重写图的情况下实现相同分区。现在，我们在遍历图时将同时维护针对 Torch 和 TensorRT 的段。只有当我们遇到一个操作，该操作既依赖于段中的某个操作又具有不同的目标时，我们才会结束该段。这将允许分区通过跨段边界重新排序节点来创建更大的段，同时保证我们不会通过相对于其依赖关系重新排序节点来修改图的行为。在这个例子中，我们将算术操作收集到一个 TensorRT 段中，将 aten::lgamma 操作收集到一个 Torch 段中。当我们遇到 %div_lgamma : Tensor = aten::lgamma(%div) 操作时，我们可以看到它依赖于当前 TensorRT 段中的 %div : Tensor = aten::div(%x, %y)。这会触发包含 aten::div 操作的 TensorRT 段的结束，以确保它在最终分区中出现在其依赖项之前。包含 aten::lgamma 操作的 Torch 段在我们遇到针对 TensorRT 且依赖于 aten::lgamma 操作结果的 prim::ListConstruct 操作时结束。

Segment Block @0:
    Target: TensorRT

    Graph: graph(%x : Tensor,
        %y : Tensor):
    %3 : int = prim::Constant[value=1]()
    %0 : Tensor = aten::add(%x, %y, %3)
    %4 : Tensor = aten::mul(%x, %y)
    %5 : Tensor = aten::div(%x, %y)
    return ()

Segment Block @1:
    Target: Torch

    Graph: graph(%x : Tensor,
        %y : Tensor,
        %5 : Tensor):
    %0 : Tensor = aten::lgamma(%x)
    %2 : Tensor = aten::lgamma(%y)
    %4 : Tensor = aten::lgamma(%5)
    return ()

Segment Block @2:
    Target: TensorRT

    Graph: graph(%1 : Tensor,
        %2 : Tensor,
        %3 : Tensor,
        %4 : Tensor,
        %5 : Tensor):
    %7 : int = prim::Constant[value=0]()
    %0 : Tensor[] = prim::ListConstruct(%1, %2, %3, %4, %5)
    %6 : Tensor = aten::cat(%0, %7)
    return ()

在某些情况下，这种方法可能在分区中创建目标相同的相邻段。作为清理步骤，我们可以合并这些相邻段，以进一步减少最终分区中的段数量。合并段步骤识别图中相邻、目标相同且未标记为 do_not_merge 的段列表。来自这些段的节点将被合并到一个新的单个段中，该段将替换分区中已合并的段。do_not_merge 标记用于防止合并为条件节点和循环创建的段，这些段在图拼接中作为特殊情况处理，不应与相同类型的相邻段合并。