TVM：设计与架构

TVM：设计与架构

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

• 示例编译流程概述了 TVM 将模型的高层描述转换为可部署模块所采取的步骤。要开始使用，请先阅读本节。

• 逻辑架构组件部分描述了逻辑组件。后面的部分是针对每个逻辑组件的特定指南，按组件的名称组织。

• 设备/目标交互描述了 TVM 如何与每种支持的物理设备和代码生成目标进行交互。

• 请查看开发人员操作指南以获取有用的开发技巧。

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

示例编译流程

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

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

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

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

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

关键数据结构

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

IRModule 是全栈中使用的基本数据结构。 IRModule（中间表示模块）包含一组函数。目前，我们支持两种基本的函数变体。

• relay::Function 是一种高层程序表示。一个 relay.Function 通常对应一个端到端的模型。您可以将 relay::Function 视为一张计算图，它额外支持控制流、递归和复杂数据结构。

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

在编译过程中，一个中继函数可能会被 lowered 为多个 tir::PrimFunc 函数和一个调用这些 tir::PrimFunc 函数的高层函数。

转换

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

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

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

relay/transform 包含一组优化模型的 pass。优化包括通用的程序优化，例如常量折叠和无用代码消除，以及特定于张量计算的 pass，例如排布转换和缩放因子折叠。

在 relay 优化 pipeline 的最后，我们将通过 pass(FuseOps) 以将端到端function（例如 MobileNet）分解为sub-function（例如 conv2d-relu）碎片。我们称这些 function 碎片。这个过程帮助我们将原始问题分为两个子问题：

• 每个 sub-function 的编译和优化。

• 整体执行结构：我们需要对生成的 sub-function 进行一系列调用来执行整个模型。

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

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

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

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

搜索空间和基于学习的转换

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

很难定义一个启发式来做出所有选择。相反，我们将采用基于搜索和学习的方法。我们首先定义一组可以用来转换程序的操作。示例操作包括循环转换、内联、矢量化。我们将这些动作称为调度原语 （scheduling primitives）。调度原语的集合定义了我们可以对程序进行的可能优化的搜索空间。然后系统搜索不同的可能调度序列以选择最佳调度组合。搜索过程通常由机器学习算法来优化。

一旦搜索完成，我们可以记录（可能是经过融合的）算子的最佳调度顺序。然后编译器可以查找最佳调度序列并将其应用于程序。值得注意的是，这个调度应用阶段与基于规则的转换完全一样，使我们能够与传统 pass 共享相同的接口约定。

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

目标转换

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

我们还提供了一个 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 和 autotvm 包含基于搜索的转换

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

• 使用python API直接构建IRModule。

• 编写一组自定义转换（例如自定义量化）。

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

逻辑架构组件

上图显示了项目中的主要逻辑组件。 请阅读以下部分以获取有关组件及其关系的信息。

tvm/support

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

tvm/runtime

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

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

除了部署用例，编译器本身也大量使用 TVM 的运行时机制。所有 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

节点模块在 runtime::Object 之上为 IR 数据结构添加了额外的功能。 主要特征包括反射、序列化、结构等价和哈希。

有了 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 和 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 是表示所有系统定义的原始运算符/内在函数的通用类。 开发人员可以向系统注册新的 Ops 以及它们的附加属性（例如 Op 是否为 elementwise）。

• Pass Infrastructure

tvm/target

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

通过查询 target 中的属性信息和注册到每个 target id(cuda, opencl) 的 builtin 信息，可以根据target定制编译 pipeline。

• Device/Target Interactions

tvm/tir

TIR 包含低层程序表示的定义。 我们使用 tir::PrimFunc 来表示可以通过 TIR 通道转换的函数。 除了 IR 数据结构外，tir 模块还通过通用 Op 注册表定义了一组内置内在函数及其属性，以及 tir/transform 中的转换 pass。

tvm/arith

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

tvm/te

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

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

• InferBound Pass
• Hybrid Frontend Developer Guide

tvm/topi

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

tvm/relay

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

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

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