系统概述

Torch-TensorRT 主要是一个 C++ 库，并计划提供 Python API。我们使用 Bazel 作为构建系统，目前主要针对 Linux x86_64 和 Linux aarch64（仅原生支持）。我们使用的编译器是 GCC 7.5.0，该库尚未在早于此版本的编译器上进行测试，因此如果您尝试使用旧版本编译器，可能会出现编译错误。

仓库结构如下

• core: 主要编译器源代码

• cpp: C++ API

• tests: C++ API、核心和 Converter 的测试

• py: Python API

• notebooks: 使用 Torch-TensorRT 构建的示例应用

• docs: 文档

• docsrc: 文档源文件

• third_party: 依赖库的 BUILD 文件

• toolchains: 不同平台的 Toolchain

C++ API 尚不稳定，在库成熟之前可能会发生变化，但大部分工作都在核心内部进行。

核心包含几个主要部分：顶层编译器接口，负责协调模块的摄取、Lowering、转换以及生成新模块并将其返回给用户。编译器有三个主要阶段：Lowering 阶段、转换阶段和执行阶段。

编译器阶段

Lowering

Lowering 阶段

Lowering 阶段由一系列 Pass 组成（部分来自 PyTorch，部分特定于 Torch-TensorRT），这些 Pass 在图 IR 上运行，将庞大的 PyTorch opset 映射到更容易转换为 TensorRT 的精简 opset。

Partitioning

Partitioning 阶段

此阶段是可选的，由用户启用。它指示编译器将节点分离为应在 PyTorch 中运行的节点和应在 TensorRT 中运行的节点。分离的标准包括：缺少 Converter、用户明确将 operator 设置为在 PyTorch 中运行，或者节点具有通过模块回退 Pass 告知 Partitioning 在 PyTorch 中运行的标志。

转换

转换阶段

在转换阶段，我们遍历 Lowering 后的图并构建等效的 TensorRT 图。转换阶段由三个主要组件组成：用于管理编译时数据的上下文、用于执行可在编译时解析的操作的 Evaluator 库，以及用于将 JIT 的 op 映射到 TensorRT 的 Converter 库。

编译与运行时

部署 Torch-TensorRT 程序

最终的编译阶段构建一个 TorchScript 程序来运行转换后的 TensorRT Engine。它获取一个序列化的 Engine，并在 Engine Manager 中实例化它，然后编译器将构建一个引用此 Engine 并将其封装在 Module 中的 JIT 图，最后将其返回给用户。当用户执行该 Module 时，JIT 程序将在由 Torch-TensorRT 扩展的 JIT 运行时中运行，并使用用户提供的数据。