整体流程#

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

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

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

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

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

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())

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()

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#

The tvm/ir folder contains the unified data structure and interfaces across for all IR function variants. The components in tvm/ir are shared by tvm/relax and tvm/tir, notable ones include

• IRModule

• Type

• PassContext 和 Pass

• Op

Different variants of functions(e.g. relax.Function and tir.PrimFunc) can co-exist in an IRModule. While these variants may not have the same content representation, they use the same data structure to represent types. As a consequence, we use the same data structure to represent function (type) signatures of these variants. The unified type system allows one function variant to call another function once we clearly define the calling convention. This opens doors for future cross-function-variant optimizations.

还提供了统一的 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 不会引发任何错误，并会自动切换到下一条规则。