设计与架构
导航
设计与架构#
本文档是为那些想了解 TVM 架构和/或积极开发项目的开发者准备的。本页的组织结构如下：
编译流程示例 概述了 TVM 将模型的高级别描述转化为可部署模块的步骤。要开始使用，请先阅读本节。
逻辑架构组件 部分描述了逻辑组件。之后的章节是专注于每个逻辑组件的具体指南，按照组件的名称组织。
设备/目标交互 页面描述了 TVM 如何与每个支持的物理设备和代码生成目标交互。
请随时查看 开发者指南，了解有用的开发技巧。
本指南提供了关于架构的几个互补性观点。首先，回顾了单一的端到端编译流程，并讨论了关键的数据结构和变换。这种基于运行时的观点着重于运行编译器时每个组件的相互作用。然后，将回顾代码库的逻辑模块和它们之间的关系。这一部分提供了一个静态的总体设计视图。
编译流程示例#
在本指南中，将研究编译器中编译流程的示例。下图显示了流程。在高层，它包含几个步骤：
Import：frontend 组件将模型摄取到 IRModule 中，IRModule 包含了一组内部表示（internally represent）模型的函数。
Transformation：编译器将 IRModule 变换成另一个在功能上等价或近似等价的（例如，在量化的情况下）IRModule。许多变换都是独立于目标（后端）的。我们还允许目标影响变换管道的配置。
Target Translation：编译器将 IRModule 翻译（codegen）为目标指定的可执行格式。目标翻译结果封装为
runtime.Module，可以导出、加载并在目标运行时环境中执行。Runtime Execution：用户加载回
runtime.Module，并在受支持的运行时环境中运行已编译的函数。
关键数据结构#
设计和理解复杂系统的最佳方法之一是识别操纵（变换）这些数据结构的关键数据结构和 API。一旦确定了关键数据结构，就可以将系统分解成逻辑组件，这些逻辑组件要么定义了关键数据结构的集合，要么定义了数据结构之间的变换。
IRModule 是整个堆栈中使用的主要数据结构。IRModule（中间表示模块）包含函数集合。目前，支持两种主要的函数变体。
relay::Function 是高级函数式程序表示。relay.Function 通常对应于端到端模型。您可以视 relay.Function 作为计算图，具有对控制流、递归和复杂数据结构的额外支持。
tir::PrimFunc 是低级的程序表示，它包含的元素有循环嵌套选择、多维加载/存储、线程化和矢量/张量指令。它通常用于表示在模型中执行（可能融合的）层的算子程序。
在编译过程中，relay 函数可能被降级为多个 tir::PrimFunc 函数和调用 tir::PrimFunc 函数的顶级函数。
变换#
既然已经介绍了关键的数据结构，现在讨论一下变换。每个变换都可以用于以下目的之一：
optimization：将程序转化为等价的，可能是更优化的版本。
lowering：将程序转换为更接近目标的较低层次表示。
relay/transform 包含优化模型的 passes 集合。优化包括常见的程序优化，如常数折叠和死代码消除，以及张量计算专用的 passes，如布局变换和缩放因子折叠。
在 relay 优化管道接近尾声的时候，将运行 pass（FuseOps）将端到端函数（例如 MobileNet）分解为子函数（例如：conv2d-relu）段。称这些为函数片段。这个过程帮助我们将原始问题划分为两个子问题：
对每个子函数进行编译和优化。
整体执行结构：需要对生成的子函数进行一系列调用，以执行整个模型。
使用低级 tir 阶段来编译和优化每个子函数。对于特定的目标，也可以直接进入目标变换阶段，并使用外部代码生成器。
有几种不同的方法（relay/backend）处理对整体执行问题的调用。对于形状已知且没有控制流的简单模型，可以使用计算图执行器将执行结构存储在计算图中。还支持动态执行的虚拟机后端。最后，计划支持提前编译，将高级执行结构编译成可执行文件和生成的原语函数。所有这些执行模式都封装在统一的 runtime.Module 接口，将在本指南的后半部分讨论。
tir/transform 包含 TIR 级函数的变换 pass。许多 tir pass 都是用来降级的。例如，有一些 pass 可以将多维访问扁平化为一维指针访问，将 intrinsic 扩展为特定于目标的内容，并装饰函数条目以满足运行时调用约定。当然，也有优化 pass，比如访问索引简化和死代码消除。
许多低级优化可以在目标阶段由 LLVM、CUDA C 和其他目标编译器处理。因此，将寄存器分配等低级优化留给下游编译器，只关注它们不涉及的优化。
搜索空间和基于学习的变换#
到目前为止，描述的转换 passes 都是确定的和基于规则的。TVM 堆栈的一个设计目标是支持针对不同硬件平台的高性能代码优化。为此，需要研究尽可能多的优化选择，包括但不限于多维张量访问、循环 tiling 行为、特殊的加速器内存层次结构和线程。
很难定义启发式来做出所有的选择。相反，将采取基于搜索和学习的方法。首先定义一组可以用来变换程序的行为。示例行为包括循环变换、内联、向量化。称这些行为为 调度原语 （scheduling primitive）。调度原语集定义了对程序进行优化的可能搜索空间。然后，系统在不同的可能调度序列中搜索出最优的调度组合。搜索过程通常由机器学习算法指导。
一旦搜索完成，可以为（可能融合的）算子记录最佳的调度序列。然后，编译器就可以查找最佳的调度序列，并将其应用到程序中。值得注意的是，这个调度应用阶段与基于规则的变换非常相似，使能够与传统 passes 共享相同的接口约定。
使用基于搜索的优化来处理初始 tir 函数生成问题。模块的这一部分称为 AutoTVM（auto_scheduler）。随着继续开发 TVM 堆栈，希望将基于学习的变换扩展到更多领域。
目标翻译#
目标翻译阶段将 IRModule 变换为相应的目标可执行格式。对于像 x86 和 ARM 这样的后端，使用 LLVM IRBuilder 来构建内存中的 LLVM IR。也可以生成源代码级的语言，比如 CUDA C 和 OpenCL。最后，支持通过外部代码生成器将 Relay 函数（子图）直接变换到特定目标。重要的是，最终的代码生成阶段要尽可能的轻量级。绝大多数的变换和降级应该在目标变换阶段之前进行。
还提供了 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(a)
print(a.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 生成主机端代码来计算启动参数（例如：线程组大小）然后从 CUDAModule 中调用另一个 PackedFunc，这个 CUDA 驱动 API 支持它。同样的机制也可以用于 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：包含 relay.Function 和 tir.PrimFunc
runtime.Module：包含 runtime.PackedFunc
编译的大部分是关键数据结构之间的变换。
relay/transform 和 tir/transform 是基于规则的确定变换。
auto_scheduler 和 autovm 包含基于搜索的变换
最后，编译流的例子只是 TVM 堆栈的典型用例。将这些关键数据结构和变换公开给 python 和 C++ API。因此，您可以像使用 numpy 一样使用 TVM，除了感兴趣的数据结构由 numpy.ndarray 转变为 tvm.IRModule。下面是一些示例用例：
使用 python API 直接构造 IRModule。
组成自定义的变换集（例如，自定义量化）。
使用 TVM 的 python API 直接操作 IR。
逻辑架构组件#
TVM 架构图#
上图显示了项目中的主要逻辑组件。请阅读以下部分，了解组件及其关系。
tvm/support#
support 模块包含了基础设施中最常用的工具，如通用 arena 分配器、套接字和日志记录。
tvm/runtime#
runtime 是 TVM 堆栈的基础。它提供了加载和执行已编译构件的机制。runtime 定义了一组稳定的标准 C API，用于与 Python 和 Rust 等前端语言进行接口翻译。
runtime::Object 是 TVM runtime 中除了 runtime::PackedFunc 之外的主要数据结构之一。它是引用计数的基类，具有类型索引以支持运行时类型检查和 downcasting。对象系统允许开发者在运行时引入新的数据结构，比如 Array、Map 和新的 IR 数据结构。
除了部署用例，编译器本身也大量使用了 TVM 的 runtime 机制。所有 IR 数据结构都是 runtime::Object 的子类，因此，它们可以直接从 Python 前端访问和操作。使用 PackedFunc 机制将各种 API 暴露给前端。
对不同硬件后端的运行时支持定义在 runtime 的子目录中（例如 runtime/opencl）。这些硬件特定的 runtime 模块为设备内存分配和设备函数序列化定义了 API。
runtime/rpc 实现了对 PackedFunc 的 RPC 支持。可以使用 RPC 机制将交叉编译的库发送到远程设备，并对执行性能进行基准测试。rpc 基础设施允许从广泛的硬件后端收集数据，以进行基于学习的优化。
tvm/node#
node 模块在 runtime::Object 的基础上为 IR 数据结构增加了额外的特性。其主要特性包括反射（reflection）、序列化、结构等效（structural equivalence）和哈希。
多亏了 node 模块，可以通过 Python 中的名称直接访问 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/relay 和 tvm/tir 共享，其中值得注意的包括
IRModule
Type
PassContext and Pass
Op
函数的不同变体（<例如 relay.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 是表示所有系统定义的原语 operator/intrinsics 的公共类。开发者可以注册新的 Ops 以及它们的附加属性（例如，不管 Op 是否为 elementwise）。
tvm/target#
target 模块包含所有将 IRModule 翻译为 target runtime.Module 的代码生成器。它还提供了通用的 Target 类来描述目标。
编译管道可以根据目标定制，查询目标中的属性信息和注册到每个目标 id 的内置信息（cuda, opencl）。
tvm/tir#
TIR 包含低级程序表示的定义。使用 tir::PrimFunc 来表示可以通过 TIR pass 进行变换的函数。除了 IR 数据结构外，tir 模块还通过公共 Op 注册表定义了一组内建 intrinsic 及其属性，以及在 tir/transform 中的变换。
tvm/arith#
这个模块与 TIR 紧密相连。低级代码生成中的一个关键问题是对索引的算术属性的分析——正定性、变量边界和描述迭代器空间的整数集。Arith 模块提供了一组进行（主要是整数）分析的工具。TIR pass 可以利用这些分析来简化和优化代码。
tvm/te#
名称 te 代表“张量表达式”（tensor expression）。这是领域专用语言（domain-specific language，简称 DSL）模块，它允许通过编写张量表达式来快速构建 tir::PrimFunc 变体。重要的是，张量表达式本身并不是可以存储在 IRModule 中的自包含函数。相反，它是 IR 的片段，可以把它拼接在一起来构建 IRModule。
te/schedule 提供了一组调度原语来控制生成的函数。未来，我们可能会把这些调度组件带到 tir::PrimFunc 本身。
tvm/topi#
虽然可以通过 TIR 或张量表达式（TE）为每个用例直接构造算子，但这样做很繁琐。topi （张量算子清单）提供了一组由 numpy 定义的预定义算子（在 TE 或 TIR 中），可在常见的深度学习工作负载中找到。还提供了一组通用的调度模板，以便在不同的目标平台上实现性能。
tvm/relay#
Relay 是用来表示完整模型的高级函数 IR。在 relay.transform 中定义了各种优化。Relay 编译器定义了多种方言，每种方言都被设计成支持特定类型的优化。值得注意的有 QNN（用于导入预量化模型）、VM（用于降级到动态虚拟机）、memory（用于内存优化）。
tvm/autotvm#
AutoTVM 和 AutoScheduler 都是基于程序优化的自动搜索组件。这一领域正在迅速发展，主要包括：
成本模型与特征提取。
用于存储成本模型构建所需的程序基准结果的记录格式。
程序变换上的一组搜索策略。
自动化程序优化仍然是一个活跃的研究领域。因此，尝试模块化设计，以便研究人员可以通过 Python 绑定快速修改组件或应用他们自己的算法，并从 Python 绑定定制搜索和拔插他们的算法。
前端#
前端从不同框架吸收模型到 TVM 堆栈。tvm.relay.frontend 是模型摄取 API 的命名空间。