设计与架构

本文档旨在帮助希望理解 Apache TVM 架构和/或积极参与项目开发的开发人员。本页面的组织结构如下

  • 总体流程 概述了 TVM 将模型的高级描述转换为可部署模块的步骤。要开始使用，请先阅读本节。

  • TVM 堆栈关键组件的简要介绍。也欢迎查看 TensorIR 深入探讨 和 Relax 深入探讨，了解有关 TVM 堆栈中两个主要组件的更多详细信息。

本指南提供了架构的一些补充视图。首先，我们回顾一个端到端的编译流程，并讨论关键数据结构和转换。这种基于运行时的视图侧重于编译器运行时每个组件的交互。然后，我们将回顾代码库的逻辑模块及其关系。这部分提供了设计的静态总体视图。

总体流程

在本指南中，我们将研究编译器中的示例编译流程。下图显示了该流程。从高层次来看，它包含几个步骤

  • 模型创建：创建要优化和编译的 IRModule，其中包含一系列函数，这些函数在内部表示模型。用户可以通过 NNModule、TVMScript 手动构建 IRModule，或从 Relax 前端导入预训练模型。

  • 转换：编译器将一个 IRModule 转换为另一个功能等效或近似等效（例如，在量化的情况下）的 IRModule。许多转换是目标（后端）独立的。我们也允许目标影响转换流水线的配置。

  • 目标转换：编译器将 IRModule 转换为目标指定的可执行格式。目标转换结果封装为 runtime.Module，可以导出、加载并在目标运行时环境中执行。

  • 运行时执行：用户加载回一个 runtime.Module，并在支持的运行时环境中运行编译后的函数。

../_static/downloads/tvm_overall_flow.svg

关键数据结构

设计和理解复杂系统的最佳方法之一是识别关键数据结构和操作（转换）这些数据结构的 API。一旦我们确定了关键数据结构，我们就可以将系统分解为定义关键数据结构集合或数据结构之间转换的逻辑组件。

IRModule 是整个堆栈中使用的主要数据结构。IRModule（中间表示模块）包含一系列函数。目前，我们支持两种主要类型的函数。

  • relax::Function 是一种高级函数式程序表示。relax.Function 表示高级图结构，通常对应于端到端模型或整个模型的子图。您可以将 relax.Function 视为具有控制流和复杂数据结构附加支持的计算图。

  • tir::PrimFunc 是一种低级程序表示，包含循环嵌套选择、多维加载/存储、线程和向量/张量指令等元素。它通常用于表示执行模型中（可能融合的）层的算子程序。

在编译和转换期间，所有 relax 算子都降级为 tir::PrimFunc 或 TVM PackedFunc，它们可以直接在目标设备上执行，而对 relax 算子的调用则降级为对低级函数的调用（例如，R.call_tir 或 R.call_dps）。

转换

现在我们已经介绍了关键数据结构，让我们谈谈转换。每个转换都可以服务于以下目的之一

  • 优化：将程序转换为等效的、可能更优化的版本。

  • 降级：将程序转换为更接近目标的低级表示。

relax 转换

relax 转换包含应用于 relax 函数的 pass 集合。优化包括常见的图级优化，例如算子的常量折叠和死代码消除，以及后端特定的优化，例如库分发。

tir 转换

tir 转换包含应用于 tir 函数的 pass 集合。主要有两种类型的转换

  • TensorIR 调度：TensorIR 调度旨在针对特定目标优化 TensorIR 函数，通过用户引导的指令和控制目标代码的生成方式。对于 CPU 目标，TIR PrimFunc 可以生成有效的代码并在目标设备上执行，而无需调度，但性能非常低。但是，对于 GPU 目标，调度对于生成具有线程绑定的有效代码至关重要。有关更多详细信息，请参阅 TensorIR 转换 部分。此外，我们提供 MetaSchedule 来自动化 TensorIR 调度的搜索。

  • 降级 Pass：这些 pass 通常在应用调度后执行，将 TIR PrimFunc 转换为另一个功能等效的 PrimFunc，但更接近于目标特定的表示。例如，有一些 pass 可以将多维访问展平为一维指针访问，将内部函数扩展为目标特定的函数，以及修饰函数入口以满足运行时调用约定。

许多低级优化可以通过 LLVM、CUDA C 和其他目标编译器在目标阶段处理。因此，我们将寄存器分配等低级优化

留给下游编译器，仅关注它们未涵盖的优化。

跨层级转换

Apache TVM 带来了一种统一的策略来优化端到端模型。由于 IRModule 同时包含 relax 和 tir 函数，因此跨层级转换旨在通过对这两种类型的函数应用不同的转换来改变 IRModule。

例如，relax.LegalizeOps pass 通过降级 relax 算子来改变 IRModule，将相应的 TIR PrimFunc 添加到 IRModule 中，并将 relax 算子替换为对降级后的 TIR PrimFunc 的调用。另一个例子是 relax 中的算子融合流水线（包括 relax.FuseOps 和 relax.FuseTIR），它将多个连续的张量运算融合为一个。与之前的实现不同，relax 融合流水线分析 TIR 函数的模式，并自动检测最佳融合规则，而不是人为定义的算子融合模式。

目标转换

目标转换阶段将 IRModule 转换为相应的目标可执行格式。对于 x86 和 ARM 等后端，我们使用 LLVM IRBuilder 构建内存中的 LLVM IR。我们还可以生成源代码级语言，例如 CUDA C 和 OpenCL。最后，我们支持通过外部代码生成器将 Relax 函数（子图）直接转换为特定目标。重要的是，最终代码生成阶段应尽可能轻量。绝大多数转换和降级应在目标转换阶段之前执行。

我们还提供了一个 Target 结构来指定编译目标。目标转换阶段之前的转换也可能受到目标的影响 — 例如，目标的向量长度会改变向量化行为。

运行时执行

TVM 运行时的主要目标是提供一个最小的 API，用于以他们选择的语言（包括 Python、C++、Rust、Go、Java 和 JavaScript）加载和执行编译后的工件。下面的代码片段显示了 Python 中的一个示例

import tvm
# Example runtime execution program in python, with type annotated
mod: tvm.runtime.Module = tvm.runtime.load_module("compiled_artifact.so")
arr: tvm.runtime.NDArray = tvm.nd.array([1, 2, 3], device=tvm.cuda(0))
fun: tvm.runtime.PackedFunc = mod["addone"]
fun(arr)
print(arr.numpy())

tvm.runtime.Module 封装了编译结果。runtime.Module 包含一个 GetFunction 方法，用于按名称获取 PackedFuncs。

tvm.runtime.PackedFunc 是为生成的函数提供的类型擦除函数接口。runtime.PackedFunc 可以接受以下类型的参数和返回值：POD 类型（int、float）、字符串、runtime.PackedFunc、runtime.Module、runtime.NDArray 以及 runtime.Object 的其他子类。

tvm.runtime.Module 和 tvm.runtime.PackedFunc 是模块化运行时的强大机制。例如，为了在 CUDA 上获得上述 addone 函数，我们可以使用 LLVM 生成主机端代码来计算启动参数（例如线程组的大小），然后从由 CUDA 驱动程序 API 支持的 CUDAModule 调用另一个 PackedFunc。相同的机制可以用于 OpenCL 内核。

上面的例子只处理了一个简单的 addone 函数。下面的代码片段给出了使用相同接口的端到端模型执行示例

import tvm
# Example runtime execution program in python, with types annotated
factory: tvm.runtime.Module = tvm.runtime.load_module("resnet18.so")
# Create a stateful graph execution module for resnet18 on cuda(0)
gmod: tvm.runtime.Module = factory["resnet18"](tvm.cuda(0))
data: tvm.runtime.NDArray = get_input_data()
# set input
gmod["set_input"](0, data)
# execute the model
gmod["run"]()
# get the output
result = gmod["get_output"](0).numpy()

主要的要点是 runtime.Module 和 runtime.PackedFunc 足以封装算子级程序（例如 addone）以及端到端模型。

总结与讨论

总而言之，编译流程中的关键数据结构是

  • IRModule：包含 relax.Function 和 tir.PrimFunc

  • runtime.Module：包含 runtime.PackedFunc

编译的大部分是关键数据结构之间的转换。

  • relax/transform 和 tir/transform 是确定性的基于规则的转换

  • meta-schedule 包含基于搜索的转换

最后，编译流程示例只是 TVM 堆栈的典型用例。我们将这些关键数据结构和转换暴露给 python 和 C++ API。因此，您可以像使用 numpy 一样使用 TVM，只是感兴趣的数据结构从 numpy.ndarray 变为 tvm.IRModule。以下是一些示例用例

  • 使用 python API 直接构造 IRModule。

  • 组合一组自定义转换（例如，自定义量化）。

  • 使用 TVM 的 python API 直接操作 IR。

tvm/support

support 模块包含基础设施最常用的实用程序，例如通用 arena 分配器、套接字和日志记录。

tvm/runtime

runtime 充当 TVM 堆栈的基础。它提供了加载和执行编译工件的机制。runtime 定义了一组稳定的标准 C API，用于与 Python 和 Rust 等前端语言交互。

runtime::Object 是 TVM runtime 中除了 runtime::PackedFunc 之外的主要数据结构之一。它是一个引用计数基类，带有类型索引，以支持运行时类型检查和向下转型。对象系统允许开发人员向运行时引入新的数据结构，例如 Array、Map 和新的 IR 数据结构。

除了部署用例之外，编译器本身也大量使用了 TVM 的运行时机制。所有 IR 数据结构都是 runtime::Object 的子类，因此，它们可以直接从 Python 前端访问和操作。我们使用 PackedFunc 机制将各种 API 暴露给前端。

对不同硬件后端的运行时支持在 runtime 的子目录（例如 runtime/opencl）中定义。这些特定于硬件的运行时模块定义了设备内存分配和设备函数序列化的 API。

runtime/rpc 实现了对 PackedFunc 的 RPC 支持。我们可以使用 RPC 机制将交叉编译的库发送到远程设备并基准测试执行性能。rpc 基础设施支持从各种硬件后端收集数据，用于基于学习的优化。

tvm/node

node 模块在 runtime::Object 之上为 IR 数据结构添加了附加功能。主要功能包括反射、序列化、结构等价和哈希。

得益于 node 模块，我们可以通过名称直接访问 TVM IRNode 的任何字段。

x = tvm.tir.Var("x", "int32")
y = tvm.tir.Add(x, x)
# a and b are fields of a tir.Add node
# we can directly use the field name to access the IR structures
assert y.a == x

我们还可以将任意 IR 节点序列化为 JSON 格式，并将其加载回来。保存/存储和检查 IR 节点的能力为使编译器更易于访问提供了基础。

tvm/ir

tvm/ir 文件夹包含跨所有 IR 函数变体的统一数据结构和接口。tvm/ir 中的组件由 tvm/relax 和 tvm/tir 共享，值得注意的包括

  • IRModule

  • 类型

  • PassContext 和 Pass

  • Op

不同变体的函数（例如 relax.Function 和 tir.PrimFunc）可以共存于 IRModule 中。虽然这些变体可能没有相同的内容表示，但它们使用相同的数据结构来表示类型。因此，我们使用相同的数据结构来表示这些变体的函数（类型）签名。统一的类型系统允许一个函数变体调用另一个函数，只要我们明确定义调用约定。这为未来的跨函数变体优化打开了大门。

我们还提供了一个统一的 PassContext 用于配置 pass 行为，以及通用的复合 pass 来执行 pass 流水线。以下代码片段给出了 PassContext 配置的示例。

# configure the behavior of the tir.UnrollLoop pass
with tvm.transform.PassContext(config={"tir.UnrollLoop": { "auto_max_step": 10 }}):
    # code affected by the pass context

Op 是表示所有系统定义的原始算子/内部函数的通用类。开发人员可以注册新的 Op 以及它们的附加属性（例如，Op 是否是元素级的）到系统中。

tvm/target

target 模块包含将 IRModule 转换为目标 runtime.Module 的所有代码生成器。它还提供了一个通用的 Target 类来描述目标。

编译流水线可以根据目标进行自定义，方法是查询目标中的属性信息以及注册到每个目标 id（cuda、opencl）的内置信息。

tvm/relax

Relax 是用于表示模型计算图的高级 IR。relax.transform 中定义了各种优化。请注意，Relax 通常与 TensorIR IRModule 紧密配合，大多数转换都应用于 IRModule 中的 Relax 和 TensorIR 函数。有关更多详细信息，请参阅 Relax 深入探讨。

tvm/tir

TIR 包含低级程序表示的定义。我们使用 tir::PrimFunc 来表示可以由 TIR pass 转换的函数。除了 IR 数据结构之外，tir 模块还包括

  • 一组用于控制 tir/schedule 中生成的代码的调度原语。

  • 一组 tir/tensor_intrin 中的内置内部函数。

  • 一组用于分析 tir/analysis 中的 TIR 函数的分析 pass。

  • 一组用于降级或优化 tir/transform 中的 TIR 函数的转换 pass。

有关更多详细信息，请参阅 TensorIR 深入探讨。

tvm/arith

此模块与 TIR 紧密相关。低级代码生成中的关键问题之一是索引的算术属性分析 — 正性、变量边界以及描述迭代器空间的整数集。arith 模块提供了一系列工具，用于执行（主要是整数）分析。TIR pass 可以使用这些分析来简化和优化代码。

tvm/te 和 tvm/topi

TE 代表张量表达式。TE 是一种特定领域语言 (DSL)，用于描述张量计算。重要的是，张量表达式本身不是一个可以存储到 IRModule 中的自包含函数。我们可以使用 te.create_prim_func 将张量表达式转换为 tir::PrimFunc，然后将其集成到 IRModule 中。

虽然可以直接通过 TIR 或张量表达式 (TE) 为每个用例构建算子，但这样做很繁琐。topi（张量算子清单）提供了一组由 numpy 定义并在常见深度学习工作负载中预定义的算子。

tvm/meta_schedule

MetaSchedule 是一个用于自动化基于搜索的程序优化的系统。它旨在成为 AutoTVM 和 AutoScheduler 的直接替代品，并且可以用于优化 TensorIR 调度。请注意，MetaSchedule 仅适用于静态形状工作负载。

tvm/dlight

DLight 是一组预定义的、易于使用且性能良好的 TIR 调度。DLight 的目标是

  • 完全支持 动态形状工作负载。

  • 轻量级。DLight 调度提供免调优或（极少次调优）调度，并具有合理的性能。

  • 鲁棒性。DLight 调度旨在对单个规则具有鲁棒性和通用性。如果该规则不适用，DLight 不会引发任何错误并自动切换到下一条规则。