【tvm官网教程01】设计和架构-1

最新推荐文章于 2024-08-16 18:29:18 发布

BesiseB

最新推荐文章于 2024-08-16 18:29:18 发布

阅读量1.6k

点赞数 2

分类专栏： tvm

本文链接：https://blog.csdn.net/b10030607/article/details/116039722

版权

tvm 专栏收录该内容

6 篇文章 8 订阅

订阅专栏

目的

如果你没听说过，是不会来学习tvm的，所以先不介绍其重要性。
如果你跟我一样陷入API，那只看到了皮毛，真正的开发者应该先了解架构，再了解源码，API只是表象。

注意：

并非要拽文，但是有些专业词汇我没办法找到合适、无歧义、公认的中文表达，所以迫不得已保留英文关键字，以防曲解。例如model、module、pass、operator、relay、lower、intrinsic、glue等。
一些专业词汇广为流传，形成共识，翻译成中文反而鸡肋，所以保留。例如：tensor、vector、layout、pipeline、access、build、run、runtime、launch、target、host等。

下文中：
原语–>Primitive
调度–>Schedule
算子–>Operator

【tvm官网教程01】设计和架构

【tvm官网教程02】tvm入门

【tvm官网教程03】张量表达与调度
【tvm官网教程04】TOPI:TVM算子清单

【tvm官网教程05】AutoTVM:基于模板的自动调优
【tvm官网教程06】AutoSchedule:无模板的自动调度

【tvm官网教程07】编译DL模型

【tvm官网教程08】优化张量算子
【tvm官网教程09】开发者教程

1. 编译流

The Example Compilation Flow gives an overview of the steps that TVM takes to turn a high level description of a model into a deployable module.
编译流示例展示TVM将model的高级描述调优成可部署module所实施的步骤。
在这里插入图片描述

step1: Import
前端组件把其他框架（tf/pytorch/caffe/onnx…）的模型提取、转换成IRModule。
step2. Transformation
主要是relay/transform目录的Relay优化pass，内部调用调度原语或者autotvm，将IRModule转换成功能相同或相近（例如量化场景）的IRModule。这些转换很多是与target无关的，但也支持target影响转换pipeline的配置。
step3. Target Transformation
主要是tir/transform目录的TIR下降pass，将IRModule转换/翻译/codegen成专用于指定target的可执行格式。转换的结果封装到runtime.Module类，以更加方便的在target的runtime环境加载、导出、执行。
step4. Runtime Execution
用户加载一个runtime.Module对象，在支持的runtime环境中run编译好的function。

1.1 关键数据结构

设计、理解大型复杂系统的最简单方法是找到关键的数据结构和对应接口。
一旦找到关键数据结构，就可以将复杂系统分解成若干逻辑组件，这些逻辑组件要么定义一系列关键数据结构，要么负责关键数据结构之间信息的流动。

贯穿整个调用栈的最核心数据结构是IRModule(intermediate representation module，中间表示模块)。
IRModule包含一个函数集合，这些函数可以分成两大变体：

relay::Function
high-level函数程序表达。一个relay.Function通常对应一个端到端model。可以将relay.Function看做额外支持控制流、递归和复杂数据结构的计算图（computational graph）。
tir::PrimFunc
low-level程序表达。其包含的元素有：循环、嵌套、选择，多维加载、存储，线程，vector/tensor指令，等等。它通常用于表示执行model中一个（融合）layer的operator程序。

在编译过程中，一个relay function可能下降成多个tir::PrimFunc函数和一个调用这些tir::PrimFunc函数的top-level函数。

1.2 Transformations

1.2.1 变换的分类

每个transformation的作用各不相同，但是可以完备地分为两大类：

optimization
将一个程序转换成等效、甚至更优的版本。
lowering
将一个程序转换成更靠近target的lower-level表示。

1.2.2 relay/transform

relay/transform包含一系列优化model的pass，例如：常量折叠、死代码消除等通用优化；又例如：layout转换、scale因子折叠等tensor计算特定pass。

1.2.2.1 子函数

在relay优化pipeline的末端，TVM run一个FuseOps的pass，将端到端函数（如MobileNet）切分成子函数（如conv2d-relu）段（segments）。这个过程将原始问题化解为两个子问题：

针对每个子函数段做编译和优化
TVM用low-level的tir阶段（phase）完成子函数的编译和优化。针对特定target，TVM也会直接在target translation阶段使用外部代码生成器（external code generators）。
Overall execution，组合调用子函数，完成整个model的执行。
relay/backend目录中有许多不同方法解决overall问题。

1.2.2.2 overall问题

针对已知shape且没有控制流的简单model来说，TVM将model lower成一个graph executor执行器，图执行器中用图存储执行结构。
TVM也支持用于动态execution的虚拟机后端。
TVM计划支持提前编译，将high-level执行结构提前编译为可执行的、已生成的原语（primitive）函数。

所有这些执行mode都由统一的runtime.Module 接口封装。

1.2.3 tir/transform

tir/transform包含用于TIR级别函数的转换pass。

大多tir pass的目的是lowering。例如：将多维access展开成1维指针access，将intrinsic展开成特定目标的函数，将函数入口装饰成runtime能调用的。

不过也有一些优化pass，例如简化访问下标和死代码消除。

例如LLVM,CUDA C等target编译器的target阶段已经可以处理许多low-level优化。所以，TVM把分配寄存器等low-level优化交给下游编译器啦，TVM聚焦于其他没被覆盖的优化。

1.3 搜索空间和基于学习的transformation

以上描述的变换pass都还是基于规则的、确定性的。TVM栈的一个设计目标是支持针对不同硬件平台的高性能代码优化。为此，TVM研究了尽可能多的优化选择，包括但不限于多维张量access，循环tiling，特殊的加速器内存结构和多线程。

很难定义一个启发式方法来做出所有选择。相反，TVM采用基于搜索和学习的方法。

定义可以用来转换程序的动作的集合
动作包括循环转换，内联，向量化。我们称这些动作为调度原语（scheduling primitives）。调度原语的集合定义了可以对程序进行的可能优化的搜索空间。
搜索不同的调度组合
搜索过程通常基于机器学习算法。
记录针对每个operator的最佳调度序列
编译器将最佳调度应用于程序
值得注意的是，此调度apply阶段与基于规则的转换完全一样，这样才能跟传统pass共用统一接口。

TVM使用基于搜索的优化，解决最初的tir函数生成问题。这部分module称为AutoTVM（auto_scheduler）。

1.4 Target Translation

target转换阶段将IRModule转换为专用于指定target的可执行格式。

对于x86和ARM等后端，TVM使用LLVM IRBuilder来构建内存中的LLVM IR。
在这里插入图片描述

TVM还可以生成诸如CUDA C和OpenCL之类的源代码级语言。

TVM支持通过外部代码生成器将Relay函数（sub-graph）直接转换为target专用版本。

最末端的代码生成阶段应尽可能轻巧，绝大部分的转换和lower都应在target转换阶段之前执行。

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

1.5 Runtime Execution

TVM运行时的主要目标是提供最小API，以使用选择的语言（包括Python，C ++，Rust，Go，Java和JavaScript）加载和执行已编译的工件。

例如Python伪码，处理一个简单的addone函数：

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.gpu(0))
fun: tvm.runtime.PackedFunc = mod["addone"]
fun(a)
print(a.asnumpy())

tvm.runtime.Module封装编译结果。runtime.Module包含根据name获取PackedFuncs的GetFunction方法。

tvm.runtime.PackedFunc是生成的函数的type-erased函数接口。runtime.PackedFunc的形参和返回值类型支持：POD类型（int，float），string，runtime.PackedFunc，runtime.Module，runtime.NDArray以及runtime.Object的其他子类。

tvm.runtime.Module和tvm.runtime.PackedFunc是将运行时模块化的强大机制。例如，要在CUDA上获得上述addone函数，可以使用LLVM生成host侧代码以计算launch参数（例如，线程组的大小），然后从由CUDAModule支持中调用CUDA driver API打包好的PackedFunc。 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 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()

TVM设计的主要优点是runtime.Module和runtime.PackedFunc足以封装算子级程序（如addone）以及端到端模型。

1.6 小结

编译流程中的关键数据结构为：

IRModule：包含relay.Function和tir.PrimFunc
runtime.Module：包含runtime.PackedFunc

编译的大部分内容是关键数据结构之间的transformation：

relay/transform和tir/transform是基于规则的确定性转换
auto_scheduler和autotvm是基于搜索/学习的转换

最后，TVM将关键数据结构和转换开放了python API和C ++ API。因此，除了将numpy.ndarray变为tvm.IRModule之外，您可以像使用numpy一样使用TVM。以下是一些用例示例：

使用python API直接构造IRModule
制定自定义的转换（例如，自定义量化）
使用python API直接操作IR

2. 逻辑组件

在这里插入图片描述

2.1 tvm/support

包含最常用的基础设施公用工具，例如分配器，socket和logging。

2.2 tvm/runtime

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

runtime :: Object是TVM运行时中除了runtime :: PackedFunc之外的主要数据结构之一。它是带引用计数的基类，它的type index支持运行时类型检查和downcast。对象系统允许开发人员向运行时引入新的数据结构，例如Array，Map和新的IR数据结构体。

除了部署用例之外，编译器本身还大量使用TVM的运行时机制。所有的IR数据结构都是runtime :: Object的子类，因此，可以从Python前端直接访问和操纵它们。TVM使用PackedFunc机制开放大量API给前端。

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

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

2.3 tvm/node

node模块在runtime::Object基础上为IR数据结构添加了其他功能。主要功能包括反射，序列化，结构等效和散列。

由于使用了node模块，可以在Python中的根据name直接访问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 node序列化为JSON格式，再将JSON加载回node。保存/存储和检查IR node的能力,是编译器更易于访问的基础。

2.4 tvm/ir

tvm/ir文件夹包含跨所有IR函数变体的统一的数据结构和接口。tvm/ir中的组件由tvm/relay和tvm/tir共享，包括:

IRModule
函数的不同变体（例如relay.Function和tir.PrimFunc）可以共存于IRModule中。
Type
尽管这些变体的内容表示形式不同，但它们使用相同的数据结构来表示type。因此，TVM使用相同的数据结构来表示这些变体的函数（类型）签名。一旦我们明确定义了调用约定，统一type系统就允许一个函数变体调用另一个函数。这为将来的跨函数变量优化（cross-function-variant optimizations）打开了大门。
PassContext、Pass
TVM还提供了一个统一的PassContext用于配置pass(https://tvm.apache.org/docs/dev/pass_infra.html)行为，并提供了通用的复合pass来执行pass pipeline。以下代码段给出了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是表示所有系统定义的原始operator/intrinsic的通用类。开发人员可以向系统注册新的Op以及它们的其他属性（例如Op是否是elementwise的）。

2.5 tvm/target

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

通过查询target中的属性信息和注册到每个target id（cuda，opencl）的内置信息，可以根据target定制编译pipeline。

2.6 tvm/tir

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

2.7 tvm/arith

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

2.8 tvm/te

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

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

InferBound Pass
Hybrid Frontend Developer Guide

2.9 tvm/topi

尽管可以针对每个用例直接通过TIR或张量表达式（TE）构造算子，但这样做很麻烦。 topi（张量算子清单）提供了一组由numpy定义并在常见深度学习工作负载中使用的预定义算子（在TE或TIR中）。TTVM还提供了一组公共调度模板，以在不同target平台上获得高性能的实现。

2.10 tvm/relay

Relay是用于表示完整模型的高级功能性IR。在relay.transform中定义了各种优化。Relay编译器定义了多种方言，每种方言旨在支持特定的优化样式。值得注意的是QNN（用于导入预量化模型），VM（用于lower至动态虚拟机），内存（用于内存优化）。
Introduction to Relay IR
Relay Operator Strategy
Convert Layout Pass