TVM：设计与架构

本文档适用于想要了解 TVM 架构和/或积极开发项目的开发人员。页面组织如下：

• 示例编译流程概述了 TVM 将模型的高层描述转换为可部署模块所采取的步骤。要开始使用，请先阅读本节。
• 逻辑架构组件部分描述了逻辑组件。后面的部分是针对每个逻辑组件的特定指南，按组件的名称组织。
• 设备/目标交互描述了 TVM 如何与每种支持的物理设备和代码生成目标进行交互。
• 请查看开发人员操作指南以获取有用的开发技巧。

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

示例编译流程

在部分，我们将研究编译器中的示例编译流程。 下图展示了该流程。 总体来说，它包含几个步骤：

• 导入(Import)：前端组件将模型引入到 IRModule 中，IRModule 包含内部表示模型的函数集合。

• 转换(Transformation)：编译器将一个 IRModule 转换为另一个功能等效或近似等效（例如在量化的情况下）IRModule。 许多转换与目标（后端）无关。目标转换pipeline的配置还可以根据具体设备目标改变。

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

• 运行时执行(Runtime Execution)：用户加载一个 runtime.Module 并在支持的运行时环境中运行编译好的函数。

关键数据结构

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

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

• relay::Function 是高级功能程序表示。 一个relay.Function 通常对应一个端到端的模型。 您可以将 relay.Function 视为计算图，并额外支持控制流、递归和复杂数据结构。
• tir::PrimFunc 是一种低级程序表示，它包含包括循环嵌套选择、多维加载/存储、线程和向量/张量指令在内的元素。 它通常用于表示在模型中执行（可能是融合的）层的算子程序。

在编译期间，一个relay::Function可能会降低为多个 tir::PrimFunc 函数和一个调用这些 tir::PrimFunc 函数的顶级函数。

---

---

TVM的复杂是由其目标决定的：

• 向上，兼容所有不同的深度学习框架（也叫 frontend），例如 pytorch、TensorFlow、onnx；
• 向下，兼容所有不同的底层硬件和推理框架（也叫 backend），同时性能最大化：不同的底层硬件，例如 x86 cpu、arm cpu、mali gpu、nvidia gpu同样的底层硬件，往往有不同的推理框架，例如 nvidia gpu，就支持 TensorRT、CUDA、CUDNN

有的推理框架，是 NN 算子级别的，例如 TensorRT、CUDNN，集成了常用的卷积、池化；但是支持的算子却不完全一致，例如 NPU 通常都加入了卷积和激活的 Composite OP，但在 CPU 和 GPU 的推理框架中不存在的。

有的推理框架，是运算级别的，例如 CUDA，需要用户开发加减乘除等基本运算，考虑缓存、并行等机制来实现 NN 算子加速。

除了上面两种情况，TVM 还可以实现不同深度学习框架的模型转换，例如 pytorch → TensorFlow，onnx → pytorch。

因此，为了覆盖上述的全部场景，TVM 中引入了两个 IR（Intermediate Representation，即中间表示）：

• 为了向上兼容的 Relay IR（简称 Relay），基于不同深度学习平台训练好的模型，在进入 TVM 后都会首先被转换为 Relay 的表示，消除表示差异；
• 为了向下兼容的 Tensor IR（简称 TIR），所有模型在编译为指定硬件上的源代码之前，都要先 lower 为 TIR。

---

---

转换(Transformation)

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

• 优化(optimization)：将程序转换为等效的，可能更优化的版本。
• 降低(lowering)：将程序转换为更接近目标的较低级别表示。

relay/transform 包含一组优化模型的passes。 优化包括常见的程序优化，例如常量折叠和死代码消除，以及特定于张量计算的通道，例如布局转换和缩放因子折叠。
在 relay 优化 pipeline 的最后，我们将通过 pass(FuseOps) 以将端到端function（例如 MobileNet）分解为sub-function（例如 conv2d-relu）碎片。调用这些function。 这个过程帮助我们将原始问题分为两个子问题：

• 每个子功能的编译和优化。
• 整体执行结构：我们需要对生成的子函数进行一系列调用来执行整个模型。

我们使用low-level tir 阶段来编译和优化每个 sub-function。 对于具体的目标，我们也可以直接进入目标翻译阶段并使用外部代码生成器。

在relay/backend有几种不同的方法去处理对整体执行问题的调用。
对于具有已知形状且没有控制流的简单模型，我们可以降低为将执行结构存储在图中的图执行器。 我们还支持用于动态执行的虚拟机后端。 最后，我们计划支持提前编译，将高级执行结构编译成可执行和生成的原始函数。
所有这些执行模式都被一个统一的runtime.Module接口封装，我们将在指南的后半部分讨论。

tir/transform 包含 TIR 级函数的转换过程。
许多 tir passes的目的是lowering。 例如，将多维访问扁平化到一维指针访问，将内在函数(intrinsics)扩展为特定于目标的内在函数，以及修饰函数入口以满足runtime调用约束。 当然，也有优化passes，比如访问索引简化和死代码消除。

LLVM、CUDA C 和其他目标编译器可以在目标阶段处理许多低级优化。 因此，我们将寄存器分配等低级优化留给下游编译器，只关注它们未涵盖的优化。

搜索空间和基于学习的转换(Search-space and Learning-based Transformations)

到目前为止，我们描述的转换过程是确定性的和基于规则的。** TVM 栈的一个设计目标是支持针对不同硬件平台的高性能代码优化**。 为此，我们需要研究尽可能多的优化选择，包括但不限于多维张量访问、循环平铺行为、特殊加速器内存层次结构和线程。

很难定义一个启发式来做出所有选择。 相反，我们将采用基于搜索和学习的方法。
我们首先定义一组可以用来转换程序的操作。 示例操作包括循环转换、内联、矢量化。 我们将这些动作称为调度原语(scheduling primitives.)。 调度原语的集合定义了我们可以对程序进行的可能优化的搜索空间。 然后系统搜索不同的可能调度序列以选择最佳调度组合。** 搜索过程通常由机器学习算法指导。**
一旦搜索完成，我们可以记录（可能是经过融合的）算子的最佳调度顺序。然后编译器可以查找最佳调度序列并将其应用于程序。值得注意的是，这个调度应用阶段与基于规则的转换完全一样，使我们能够与传统 pass 共享相同的接口约定。

我们使用基于搜索的优化来处理初始 tir 函数生成问题。这部分模块称为 AutoTVM(auto_scheduler)。随着我们继续开发 TVM ，我们希望将基于学习的转换扩展到更多领域。

目标转换(target Translation)

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

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

运行时执行(Runtime Execution)

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 是两个生成函数的类型擦除函数( type-erased function )接口。 runtime.PackedFunc 可以采用以下类型的参数和返回值：POD 类型（int，float）、字符串、runtime.PackedFunc、runtime.Module、runtime.NDArray 和 runtime.Object 的其他子类。
tvm.runtime.Module 和 tvm.runtime.PackedFunc 是模块化runtime的强大机制。 例如，要在 CUDA 上获取上述 addone 函数，我们可以使用 LLVM 生成主机端代码来计算启动参数（例如线程组的大小），然后从由支持的 CUDAModule 调用另一个 PackedFunc CUDA 驱动程序 API。 相同的机制可用于 OpenCL 内核。

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

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

support模块包含基础设施最常用的实用程序，例如generic arena allocator, socket, and logging.

tvm/runtime

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

runtime::Object 是 TVM runtime 中除 runtime::PackedFunc 之外的主要数据结构之一。 它是一个具有类型索引的引用计数基类，以支持运行时类型检查和向下转换。 object system 允许开发人员向运行时引入新的数据结构，例如 Array、Map 和新的 IR 数据结构。
除了部署用例，编译器本身也大量使用 TVM 的runtime机制。 所有 IR 数据结构都是 runtime::Object 的子类，因此可以从 Python 前端直接访问和操作它们。 我们使用 PackedFunc 机制向前端公开各种 API。

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

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

• TVM Runtime System
• Runtime-Specific Information
• Debugger
• Putting the VM in TVM: The Relay Virtual Machine
• Introduction to Module Serialization
• Device/Target Interactions

tvm/node

node模块在 runtime::Object 之上为 IR 数据结构添加了额外的功能。 主要特征包括reflection, serialization, structural equivalence, and hashing.

得益于 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 是表示所有system-defined primitive operator/intrinsics 的通用类。 开发人员可以向系统注册新的 Ops 以及它们的附加属性（例如 Op 是否为 elementwise）。

• Pass Infrastructure

tvm/target

target模块包含将 IRModule 转换为目标 runtime.Module 的所有代码生成器。 它还提供了一个描述目标的通用 Target 类。
通过查询target中的属性信息和注册到每个target id(cuda, opencl)的builtin信息，可以据此定制编译过程。

• Device/Target Interactions

tvm/tir

TIR 包含低级程序表示的定义。 我们使用 tir::PrimFunc 来表示可以通过 TIR pass转换的函数。 除了 IR 数据结构外，tir 模块还通过通用 Op 注册表定义了一组内置内在函数及其属性，以及tir/transform中的transformation passes
TIR 的所有类型和函数，都定义在 tvm.tir 当中。

---

虽然 Relay 和 TIR 所表达的层次不同，但是两者有一些共同点。
首先，Relay 和 TIR 都是一个有向无环图 DAG：

• TIR 的 AST 是树结构，树是特殊的 DAG
• Relay 表示的神经网络自然也是 DAG

其次，Relay 和 TIR 都设计到 lower，即需要翻译成更低一级别的表示；Relay 需要翻译为 TIR，而 TIR 需要被翻译为目标源代码。这个过程就涉及到对 Relay 和 TIR 两种 DAG 的遍历。

因此，在 TVM 中所有设计到 DAG 的基本要素和功能都封装到了 tvm.ir 当中。tvm.tir 和 tvm.relay 的类型和函数，几乎都是对 tvm.ir 的继承和调用。

---

tvm/arith

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

tvm/te

te 代表“张量表达式”。 这是一个特定领域的语言模块，它允许我们通过编写张量表达式来快速构造 tir::PrimFunc 变体。
重要的是，张量表达式本身并不是一个可以存储到 IRModule 中的自包含函数。 相反，它是 IR 的一个片段，我们可以拼接起来构建一个 IRModule。

te/schedule 提供了一组调度原语来控制正在生成的函数。 将来，我们可能会将其中一些调度组件带到 tir::PrimFunc 自身中去。

• InferBound Pass
• Hybrid Frontend Develop Guide

tvm/topi

虽然可以通过 TIR 或张量表达式 (TE) 为每个用例直接构造算子，但这样做很乏味。 topi（张量算子清单）提供一组由 numpy 定义的预定义算子（在 TE 或 TIR 中），并在常见的深度学习工作负载中找到。 我们还提供了一组常见的调度模板，以获得跨不同目标平台的高性能实现。

tvm/relay

Relay 是用于表示完整模型的高级功能 IR。 在 relay.transform 中定义了各种优化。 Relay 编译器定义了多种dialects，每种dialect都旨在支持特定的优化风格。 值得注意的包括 QNN（用于导入预量化模型）、VM（用于降低到动态虚拟机）、memory（用于内存优化）。

• Introduction to Relay IR
• Relay Operator Strategy
• Convert Layout Pass

Relay 就是神经网络的一种统一中间表示。但是和同为统一表示的 ONNX 不同，Relay 不包含神经网络算子的任何具体计算逻辑，单纯是算子的占位符，无法直接进行推理计算。

Relay 中包含如下信息：

• 算子的类型
• 算子的结构参数，例如 shape、group、layout、data type
• 算子的输入输出连接，即如何构成神经网络的 Graph

Relay 中不包含如下信息：

• 算子的具体参数，例如卷积的 kernel 和 bias，或者 BN 的 mean 和 std
• 算子如何进行计算。例如给定 input，经过卷积后 output 的具体数值

Relay 要进行推理，必须先指定硬件平台（即 Target），再对其进行编译。整个过程将涉及到将 Relay 翻译为 TIR（这一步被称为 lower），然后再 CodeGen 成指定的源代码，最后调用指定 Compiler 编译出可执行程序。

Relay 的所有类型和函数，都定义在 tvm.relay 当中。

tvm/autotvm

AutoTVM 和 AutoScheduler 都是自动化基于搜索的程序优化的组件。 这正在迅速发展，主要包括：

• 成本模型和特征提取。
• 用于存储成本模型构建的程序基准结果的记录格式。
• 一组关于程序转换的搜索策略。

自动化程序优化仍然是一个活跃的研究领域。 因此，我们尝试将设计模块化，以便研究人员可以通过 Python 绑定快速修改组件或应用他们自己的算法，并从 Python 绑定自定义搜索和插件他们的算法。

• Benchmark Performance Log Format

Frontends

前端将来自不同框架的模型读取到 TVM 栈中。 tvm.relay.frontend 是模型摄取 API 的命名空间。

• TensorFlow Frontend

Security

• Security Guide

microTVM

• microTVM Design Document
• microTVM Project API
• Model Library Format

参考：https://zhuanlan.zhihu.com/p/532873577
https://tvm.apache.org/docs/arch/index.html#example-compilation-flow