设计和架构

设计和架构#

本文档旨在为希望了解 Apache TVM 架构的开发者，以及积极参与该项目开发的人员提供参考。页面内容组织如下：

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

  • TVM 堆栈的关键组件简介。您也可以参考 TensorIR 深入解析 和 Relax 深入解析，以了解更多关于 TVM 堆栈中两个主要组件的详细信息。

本指南提供了对架构的几种补充视角。首先回顾一个单一的端到端的编译流程，并讨论关键的数据结构及其转换。这种基于运行时的观点着重于编译器运行时各组件之间的交互作用。接下来，将审视代码库的逻辑模块及其相互关系。这部分提供了静态的、全面的设计概览。

整体流程#

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

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

  • 变换：编译器将 IRModule 变换为另一个功能上等效或近似等效（例如，在量化的情况下）的 IRModule。许多变换是针对目标（target 后端）独立的。还允许目标影响变换管道的配置。

  • 目标翻译：编译器将 IRModule 编译（codegen）成由目标(target)指定的可执行格式。目标翻译的结果被封装为 runtime.Module，这个模块可以被导出、加载并在目标运行时环境中执行。

  • 运行时执行：用户重新加载 runtime.Module，并在支持的运行时环境中运行编译好的函数。

https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/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 函数的传递。这些优化包括常见的计算图级别优化，例如常量折叠和算子的死代码消除，以及特定于后端的优化，如库调度。

tir 变换#

tir 转换包含一系列适用于 tir 函数的变换。主要有两种类型的变换：

  • TensorIR 调度：TensorIR 调度旨在优化特定目标的 TensorIR 函数，通过用户指导的指令控制目标代码的生成方式。对于 CPU 目标，TIR PrimFunc 可以生成有效的代码并在目标设备上执行，无需调度但性能非常低。然而，对于 GPU 目标来说，调度对于生成带有线程绑定的有效代码至关重要。更多详情，请参阅：TensorIR 变换 部分。此外，也提供了 MetaSchedule 来自动化寻找 TensorIR 调度。

  • 降低阶段：这些阶段通常在调度应用后执行，将 TIR PrimFunc 变换为另一个功能上等价的 PrimFunc，但更接近目标特定的表示形式。例如，有些阶段会将多维访问展平成一维指针访问，将 intrinsics 扩展为目标特定的函数，并装饰函数入口点以符合运行时调用约定。

许多低级优化可以在目标阶段由 LLVM、CUDA C 以及其他目标编译器处理。因此，将诸如寄存器分配之类的低级优化留给它们来处理。

对于下游编译器而言，仅专注于那些未被它们覆盖的优化。

跨层次变换#

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

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

目标翻译#

目标翻译阶段将 IRModule 变换成对应的“目标可执行格式”。对于诸如 x86 和 ARM 这样的后端，使用 LLVM IRBuilder 来构建内存中的LLVM IR。还可以生成源代码级别的语言，例如 CUDA C 和 OpenCL。最后，支持通过外部代码生成器直接将 Relay 函数（子图）翻译成特定目标。确保最终的代码生成阶段尽可能轻量级是非常重要的。大部分变换和降低操作应该在目标翻译阶段之前完成。

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

运行时执行#

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：包含 relay.Function 和 tir.PrimFunc

  • runtime.Module: 包含 runtime.PackedFunc

该编译的大部分内容涉及关键数据结构之间的变换。

  • relay/transform 和 tir/transform 是基于确定性规则的变换。

  • auto_scheduler和 autotvm 包含了基于搜索的变换技术。

最后，提供的编译流程示例仅是 TVM 堆栈的典型使用案例之一。将这些关键的数据结构和变换暴露给 Python 和 C++ 的 API 接口。因此，您可以像使用 NumPy 一样使用TVM，只不过关注的不再是 numpy.ndarray 数据结构，而是 tvm.IRModule。以下是一些示例用例：

  • 直接使用 Python API 构建 IRModule。

  • 编写一组自定义的变换（例如，定制量化）。

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

tvm/support#

支持模块包含基础设施中最常见的实用工具，例如通用 arena 分配器、套接字和日志记录。

tvm/runtime#

运行时作为 TVM 栈的基础，它提供了加载和执行编译工件的机制。运行时定义了一套稳定的标准 C API 集合，以与前端语言（如 Python 和 Rust）进行接口交互。

runtime::Object 是 TVM 运行时中的主要数据结构之一，仅次于 runtime::PackedFunc。它是一个引用计数的基类，带有类型索引以支持运行时类型检查和向下转型(downcasting)。这个对象系统允许开发者向运行时引入新的数据结构，例如数组、映射以及新的中间表示（IR）数据结构。

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

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

runtime/rpc 实现了对 PackedFunc 的 RPC 支持。可以利用这一 RPC 机制将跨平台编译的库发送到远程设备，并对执行性能进行基准测试。该 RPC 基础设施能够从广泛的硬件后端收集数据，以支持基于学习的优化。

tvm/node#

这个节点模块在 runtime::Object 的基础上为 IR 数据结构增加了额外的功能。主要特性包括反射、序列化、结构等价性以及哈希处理。

由于节点模块的存在，我们可以直接通过名称在 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 和 Pass

  • Op

在 IRModule 中，不同的函数变体（例如 relay.Function 和 tir.PrimFunc）可以共存。尽管这些变体的“内容表示”可能不同，但它们使用相同的数据结构来表示类型。因此，采用同一数据结构来表示这些变体的函数（类型）签名。统一的类型系统允许一种函数变体在明确了调用约定后调用另一种函数。这为未来跨函数变体的优化开辟了道路。

还提供了统一的 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 是用于表示所有系统定义的 primitive operator/intrinsics 的通用类。开发者可以向系统中注册新的 Op 以及它们的附加属性（例如，Op 是否为逐元素运算）。

tvm/target#

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

根据目标的不同，编译流水线可以进行定制化设置。这可以通过“查询目标中的属性信息以及每个目标ID（如CUDA、OpenCL）注册的内置信息”来实现。

tvm/relax#

Relax 是一种高级的中间表示（IR），用于表达计算图。在 relax.transform 中定义了各种优化技术。需要注意的是，Relax 通常与 TensorIR 的 IRModule 紧密合作，大多数变换同时应用于 Relax 和 TensorIR 的功能上。更多详细信息请参见 Relax 深入探讨。

tvm/tir#

TIR 包含了低级程序表示的定义。使用 tir::PrimFunc 来表示可以通过 TIR 变换的函数。除了中间表示的数据结构，tir 模块还包括了：

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

  • 在 tir/tensor_intrin 中，有一组内建的 intrinsics。

  • 一组分析通过用于在 tir/analysis 中分析 TIR 功能。

  • 在 tir/transform 中，一系列变换被用于降低或优化 TIR 函数。

请参阅 TensorIR 深度解析 以获取更多详细信息。

tvm/arith#

这个模块与中间表示（TIR）密切相关。在低级代码生成中的关键问题是分析索引的算术属性——包括正性、变量界限以及描述迭代器空间的整数集。arith 模块提供了一套工具，主要进行整数分析。TIR 阶段可以使用这些分析来简化和优化代码。

tvm/te and 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 调度器提供无需调优或仅需极少量（shots tuning）调优的调度方案，并保持合理的性能。

  • 稳健性。DLight 调度旨在为单一规则提供稳健（鲁棒）且通用的调度，如果该规则不适用，DLight 不会引发任何错误，并会自动切换到下一条规则。