tvm - Design And Architecture

发表于 | 分类于

本文翻译自 TVM 官方文档：Design and Architecture

设计与架构

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

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

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

示例编译流程

在本指南中，我们将研究编译器中的示例编译流程。下图显示了流程。在高层次上，它包含几个步骤：

  • 导入：前端组件将模型引入到 IRModule 中，IRModule 包含内部表示模型的函数集合。
  • 转换：编译器将一个 IRModule 转换为另一个功能等效或近似等效（例如在量化的情况下）IRModule。许多转换与目标（后端）无关。我们还允许目标影响转换管道的配置。
  • 目标翻译：编译器将 IRModule 翻译（代码生成）为目标指定的可执行格式。目标翻译结果被封装为一个runtime.Module，可以在目标运行时环境中导出、加载和执行。
  • 运行时执行：用户加载一个runtime.Module 并在支持的运行时环境中运行编译后的函数。

https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/tvm_dyn_workflow.svg

关键数据结构

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

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

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

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

转换

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

  • 优化：将程序转换为等效的、可能更优化的版本。
  • Lowering：将程序转换为更接近目标的 low-level 表示。

relay/transform 包含优化模型的通道集合。优化包括常见的程序优化，例如常量折叠和死代码消除，以及特定于张量计算的通道，例如布局转换和缩放因子折叠。

在 relay 优化管道接近尾声时，我们将运行 pass(FuseOps) 以将端到端功能（例如 MobileNet）分解为子功能（例如 conv2d-relu）段。我们称这些功能段。这个过程帮助我们将原始问题分为两个子问题：

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

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

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

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

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

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

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

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

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

我们使用基于搜索的优化来处理初始 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 中的这样一个示例：

1
2
3
4
5
6
7
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 函数。下面的代码片段给出了使用相同接口执行端到端模型的示例：

1
2
3
4
5
6
7
8
9
10
11
12
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。

逻辑架构组件

https://raw.githubusercontent.com/tlc-pack/web-data/main/images/design/tvm_static_overview.svg

TVM 架构图

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

tvm/support

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

tvm/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/node

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

由于 node 模块的存在，我们可以通过 Python 中的名称直接访问 TVM 的 IRNode 的任何字段。

1
2
3
4
5
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 用于配置传递行为，以及通用的复合传递来执行传递管道。以下代码片段给出了 PassContext 配置的示例。

1
2
3
# 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）。

tvm/target

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

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

tvm/tir

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

tvm/arith

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

tvm/te

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

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

tvm/topi

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

tvm/relay

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

tvm/autotvm

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

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

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

前端

前端将来自不同框架的模型摄取到 TVM 堆栈中。

安全

microTVM