TVM设计与构架构建
本文档适用于希望了解TVM体系结构和/或在项目上进行积极开发的开发人员。该页面的组织如下：
- 实例编译流程Example Compilation Flow描述TVM把一个模型的高级描述到可部署模块的步骤。
- “逻辑体系结构组件” Logical Architecture Components部分描述了逻辑组件。针对每个逻辑组件的特定内容，按组件名称组织。
- 开发人员操作手册，以获取有用的开发技巧。
本文提供了一些体系结构的补充视图。首先，回顾一个端到端的编译流程，并讨论关键的数据结构和转换。这个基于运行时runtime的视图着重于运行编译器时每个组件的交互。然后，将回顾代码库的逻辑模块及其关系。这部分提供了设计的静态总体视图。
编译流程示例
本节将研究编译器中的示例编译流程。下图显示了流程。从高层次上讲，包含几个步骤：
- 导入：前端组件将模型提取到IRModule中，该模块包含内部表示模型的函数的集合。
- 转换：编译器将IRModule转换为另一个功能上等效或近似等效的（例如，在量化的情况下）IRModule。许多转换都是独立于目标（后端）的。还允许目标影响转换管道的配置。
- 目标翻译：编译器将IRModule转换（代码生成）为目标指定的可执行格式。目标翻译结果被封装为runtime.Module，可以导出，加载和执行对目标运行时runtime环境。
- 运行时runtime执行：用户负载背了runtime.Module在支持的运行环境和运行编译功能。
关键数据结构
设计和理解复杂系统的最佳方法之一就是，识别关键数据结构和操作（转换）。这些数据结构的API，一旦确定了关键数据结构，便可以将系统分解为逻辑组件，这些逻辑组件可以定义关键数据结构的集合或数据结构之间的转换。
IRModule是整个堆栈中使用的主要数据结构。IRModule（中间表示模块）包含函数的集合。当前，支持两种主要的功能变体。
- relay :: Function是高级功能代码表示。relay.Function通常对应于端到端模型。可以将relay作为计算图来查看，并额外支持控制流，递归和复杂的数据结构。
- tir :: PrimFunc是一种低级代码表示，其中包含一些元素，包括循环嵌套选择，多维加载/存储，线程和向量/张量指令。通常用于表示执行模型中（可能是融合的）层的算子代码。
在编译期间，可以将relay函数降低为多个tir :: PrimFunc函数和一个调用这些tir :: PrimFunc函数的顶级函数。
转换
已经涵盖了关键数据结构，来谈谈转换。每个转换可以满足以下目的之一：
- 优化：将代码转换为等效的，可能更优化的版本。
- 降维：将代码转换为更接近目标的较低层表示。
中继/转换包含优化模型的通道集合。这些优化包括常见的代码优化（例如恒定折叠和死代码消除）以及张量计算特定的遍历（例如布局转换和缩放因子折叠）。
在中继优化管道的末端附近，运行pass（FuseOps）将端到端功能（例如MobileNet）划分为子功能（例如conv2d-relu）段。称这些功能为段。此过程可将原始问题分为两个子问题：
- 每个子功能的编译和优化。
- 总体执行结构：需要对所生成的子函数进行一系列调用，以执行整个模型。
使用低级的Tir波段来编译和优化每个子功能。对于特定目标，也可以直接进入目标平移并使用外部代码生成器。
有几种不同的方法（在中继/后端）来处理对整个执行问题的调用。对于具有已知形状且无控制流的简单模型，可以降低到将执行结构存储在图中的图运行时runtime。还支持虚拟机后端进行动态执行。最后，计划支持提前编译，该编译将高级执行结构编译为可执行文件和生成的原始函数。所有这些执行模式都由统一的runtime.Module 接口封装 ，将在本文的后面部分中进行讨论。
tir / transform包含用于TIR级别功能的转换过程。许多Tir通道的目的是降维。例如，可以通过多种途径将多维接入到一维指针访问，将内在函数扩展为特定于目标的内在函数，并修饰函数条目，以满足运行时runtime调用约定。当然，也有一些优化过程，例如简化访问索引和消除无效代码。
LLVM，CUDA C和其它目标编译器，可以在目标处理许多低级优化。结果，将低级优化（例如寄存器分配）留给了下游编译器，只专注于未涵盖的优化。
搜索空间和基于学习的转换
到目前为止，描述的转换过程是确定性的且基于规则的。TVM堆栈的一个设计目标是支持针对不同硬件平台的高性能代码优化。为此，将需要研究尽可能多的优化选择，包括但不限于多维张量访问，循环切片行为，特殊的加速器内存层次结构和线程化。
很难定义做出所有选择的试探法。相反，将采用基于搜索和学习的方法。首先定义可以用来转换代码操作的集合。示例操作包括循环转换，内联，向量化。称这些操作为调度原语。调度原语的集合定义了，可以对代码进行的可能优化的搜索空间。然后，系统搜索不同的可能调度序列，以选择最佳调度组合。搜索过程通常以机器学习算法为原则。
搜索完成后，可以为（可能是融合的）算子记录最佳调度顺序。然后，编译器可以仅查找最佳调度序列，并将其应用于代码。值得注意的是，此调度应用代码布局完全类似于基于规则的变换，能够与传统流程共享相同的接口约定。
使用基于搜索的优化来处理最初的Tir函数生成问题。此模块部分称为AutoTVM（auto_scheduler）。随着继续开发TVM堆栈，希望将基于学习的转换扩展到更多领域。
目标转换
目标转换阶段将IRModule转换为相应的目标可执行格式。对于x86和ARM等后端，使用LLVM IRBuilder来构建内存中的LLVM IR。还可以生成诸如CUDA C和OpenCL之类的源代码级语言。最后，支持通过外部代码生成器将Relay函数（子图）直接转换为特定目标。重要的是，最终代码生成阶段应尽可能轻巧。绝大部分的转换和降维都应在目标转换之前进行。
还提供了一个Target结构来指定编译目标。目标翻译阶段之前的转换也可能受到目标的影响-例如，目标的向量长度会改变向量化行为。
运行时runtime执行
TVM运行时runtime的主要目标是提供一个最小的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], ctx=tvm.gpu(0))
fun: tvm.runtime.PackedFunc = mod["addone"]
fun(a)
print(a.asnumpy())
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是将运行时runtime模块化的强大机制。例如，要在CUDA上获得上述addone函数，可以使用LLVM生成主机端代码以计算启动参数（例如线程组的大小），然后从CUDAModule调用另一个由PackedFunc支持的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 gpu(0)
gmod: tvm.runtime.Module = factory["resnet18"](tvm.gpu(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).asnumpy()
主要优点是，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.ndarray更改为tvm.IRModule之外，可以像使用numpy，一样使用TVM。以下是一些用例示例：
- 使用python API直接构造IRModule。
- 组成一组自定义的转换（例如，自定义量化）。
- 使用TVM的python API直接操作IR。
逻辑架构组件
TVM体系结构图
上图显示了项目中的主要逻辑组件。阅读以下各节，以获取有关组件及其关系的信息。
tvm/support
支持模块包含最常用的基础构建实用代码，例如通用竞技场分配器，套接字和日志记录。
tvm/runtime
运行时runtime是TVM堆栈的基础。提供了加载和执行已编译工件的机制。运行时runtime定义了一组稳定的标准C API，以与诸如Python和Rust的前端语言进行接口。
除了runtime :: PackedFunc之外，runtime :: Object是TVM运行时runtime中的主要数据结构之一。它是带有类型索引的引用计数基类，以支持运行时runtime类型检查和向下转换。目标系统允许开发人员向运行时runtime引入新的数据结构，例如数组，映射和新的IR数据结构。
除了部署用例之外，编译器本身还大量使用TVM的运行时runtime机制。所有的IR数据结构都是runtime :: Object的子类，可以从Python前端直接访问和操作它们。使用PackedFunc机制将各种API公开给前端。
在运行时runtime的子目录（例如runtime / opencl）中定义了对不同硬件后端的运行时runtime支持。这些特定于硬件的运行时runtime模块定义，用于设备内存分配和设备功能序列化的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共享，包括
- ir模块
- 类型
- PassContext和Pass
- 算子
功能的不同变体（例如relay.Function和tir.PrimFunc），可以共存于IRModule中。尽管这些变体可能不具有相同的内容表示，但是它们使用相同的数据结构来表示类型。因此，使用相同的数据结构来表示这些变量的功能（类型）名称。一旦明确定义了调用约定，统一类型系统就允许一个函数变换调用另一个函数。这为将来的跨功能变量优化打开了大门。
还提供了一个统一的PassContext用于配置传递操作，并提供了通用的复合传递来执行传递管道。以下代码段给出了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是表示所有系统定义的原始算子/内部函数的通用类。开发人员可以向系统注册新的Op以及它们的其它属性（例如Op是否是元素化的）。
tvm/target
目标模块包含将IRModule转换为目标runtime.Module的所有代码生成器。还提供了描述目标的通用Target类。
通过查询目标中的属性信息和注册到每个目标id（cuda，opencl）的内置信息，可以根据目标定制编译管道。
tvm/ir
TIR包含低级代码表示的定义。使用tir :: PrimFunc表示可以通过TIR传递转换的函数。除IR数据结构外，tir模块还通过公共Op注册表，以及tir / transform中的转换，传递定义了一组内置的内在函数及其属性。
tvm/arith
此模块与TIR紧密相关。低级代码生成中的关键问题之一是分析索引的算术属性-正则性，变量边界以及描述迭代空间的整数集。arith模块提供了一组进行（主要是整数）分析的工具。TIR pass可以使用这些分析来简化和优化代码。
tvm/te
te代表“张量表达式”。这是一个特定领域的语言模块，允许通过编写张量表达式来快速构建tir :: PrimFunc变体。重要的是，张量表达式本身不是可以存储到IRModule中的自包含函数。相反，它是IR的一个片段，可以将其拼接在一起以构建IRModule。
te / schedule提供了一组调度原语，以控制所生成的功能。将来，可能会将其中的一些调度组件引入tir :: PrimFunc本身。
tvm/topi
可以针对每个用例直接通过TIR或张量表达式（TE）构造算子。 topi（张量算子清单）提供了一组预定义的算子（在TE或TIR中），由numpy定义并在常见的深度学习工作负载中找到。提供了一组公共时间表模板，以在不同目标平台上获得高性能的实现。
tvm/relay
继电器是用于表示完整模型的高级功能性IR。在relay.transform中定义了各种优化。Relay编译器定义了多种语言，每种语言旨在支持特定的优化样式。值得注意的是QNN（用于导入预量化模型），VM（用于降级为动态虚拟机），内存（用于内存优化）。
tvm/autotvm
AutoTVM和AutoScheduler都是自动进行基于搜索的代码优化的组件。主要包括：
- Cost models和特征提取。
- 一种记录格式，用于存储计划基准结果以进行cost model构建。
- 一组有关代码转换的搜索策略。
- Cost models and feature extraction.
- A record format for storing program benchmark results for cost model construction.
- A set of search policies over program transformations.
自动化代码优化仍然是活跃的研究领域。试图对设计进行模块化，以便研究人员可以通过Python绑定，快速修改组件或自己的应用算法，自定义搜索并从Python绑定中插入其算法。
前端
前端将来自不同框架的模型，存放在TVM堆栈中。 tvm.relay.frontend是模型提取API的命名空间。