TVM设计与构架构建_tvm site:csdn.net

TVM设计与构架构建
本文档适用于希望了解TVM体系结构和/或在项目上进行积极开发的开发人员。该页面的组织如下：
• 实例编译流程Example Compilation Flow，描述TVM把一个模型的高级描述到可部署模块的步骤。
• “逻辑体系结构组件” Logical Architecture Components部分描述了逻辑组件。针对每个逻辑组件的特定内容，按组件名称组织。
• 开发人员操作手册，以获取有用的开发技巧。
本文提供了一些体系结构的补充视图。首先，回顾一个端到端的编译流程，讨论关键的数据结构和转换。这个基于运行时runtime的视图，着重于运行编译器时，每个组件的交互。回顾代码库的逻辑模块及其关系。这部分提供了设计的静态总体视图。
编译流程示例
本节将研究编译器中的示例编译流程。下图显示了流程。从高层次上，包含几个步骤：
• 导入：前端组件将模型提取到IRModule中，该模块包含内部表示模型的函数的集合。
• 转换：编译器将IRModule，转换为另一个功能上等效，或近似等效的（例如，在量化的情况下）IRModule。许多转换都是独立于目标（后端）的。允许目标影响转换管道的配置。
• 目标翻译：编译器将IRModule转换（代码生成）为目标指定的可执行格式。目标翻译结果，封装为runtime.Module，可以导出，加载和执行，对目标运行时runtime环境。
• 运行时runtime执行：用户load了runtime.Module，在支持的运行环境和运行编译功能。

关键数据结构
设计和理解复杂系统的最佳方法之一，识别关键数据结构和算子（转换）。这些数据结构的API，一旦确定了关键数据结构，便可以将系统分解为逻辑组件，这些逻辑组件可以定义关键数据结构的集合，或数据结构之间的转换。
IRModule是整个堆栈中使用的主要数据结构。IRModule（中间表示模块）包含函数的集合。支持两种主要的功能变体。
• relay :: Function是高级功能代码表示。relay.Function通常对应于端到端模型。可以将relay作为计算图，额外支持控制流，递归和复杂的数据结构。
• tir :: PrimFunc是一种低级代码表示，包含一些元素，包括循环嵌套选择，多维加载/存储，线程和向量/张量指令。通常用于表示执行模型中（可能是融合的）层的算子代码。
在编译期间，可以将relay函数降低为多个tir :: PrimFunc函数和一个调用这些tir :: PrimFunc函数的顶级函数。
转换
已经涵盖了关键数据结构，谈谈转换。每个转换可以满足以下目的之一：
• 优化：将代码转换为等效的，可能更优化的版本。
• 降维：将代码转换为更接近目标的，较低层表示。
Relay/转换包含优化模型的通道集合。这些优化包括常见的代码优化（如常量折叠和死代码消除），以及张量计算特定的遍历（如布局转换和缩放因子折叠）。
在Relay优化管道的末端附近，运行pass（FuseOps），将端到端功能（例如MobileNet），划分为子功能（例如conv2d-relu）段，称这些功能为段。可将原始问题分为两个子问题：
• 每个子功能的编译和优化。
• 总体执行结构：需要对所生成的子函数，进行一系列调用，执行整个模型。
使用低级的Tir，编译和优化每个子功能。对于特定目标，可以直接进入目标平移，使用外部代码生成器。
有几种不同的方法（在Relay/后端），处理对整个执行问题的调用。对于具有已知形状，无控制流的简单模型，可以降低到将执行结构，存储在图中的图运行时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"
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".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运行系统
• 调试器
• 将VM放入TVM：Relay虚拟机
• 模块序列化简介
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，用于配置传递算子，提供了通用的复合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是表示所有系统定义的原始算子/内部函数的通用类。开发人员可以向系统注册新的Op，以及其它属性（例如Op是否是元素化的）。
• Pass Infrastructure
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本身。
• InferBound Pass
• 混合前端开发人员指南
• InferBound Pass
• Hybrid Frontend Developer Guide
tvm/topi
可以针对每个用例直接通过TIR，或张量表达式（TE），构造算子。 topi（张量算子清单）提供了一组预定义的算子（在TE或TIR中），由numpy定义，在常见的深度学习工作负载中找到。提供了一组公共调度模板，在不同目标平台上，获得高性能的实现。
tvm/relay
Relay是用于表示完整模型的高级功能性IR。在relay.transform中定义了各种优化。Relay编译器定义了多种语言，每种语言旨在支持特定的优化样式。值得注意的是QNN（用于导入预量化模型），VM（用于降级为动态虚拟机），内存（用于内存优化）。
• Relay IR简介
• Relay算子策略
• 转换布局通道
• Introduction to Relay IR
• Relay Operator Strategy
• Convert Layout Pass
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绑定中，插入算法。
• 基准性能日志格式Benchmark Performance Log Format
前端
前端将来自不同框架的模型，存放在TVM堆栈中。 tvm.relay.frontend是模型提取API的命名空间。
• TensorFlow前端