TVM学习笔记（二）编译流程

本文是基于TVM官方文档对其编译流程的梳理

TVM编译流程

上图^[1]展示了TVM的编译流程

关键步骤

① Import ：将模型读取至IRModule中的前端组件，其中包含一系列函数用来在内部表示模型。

② Transformation ：TVM编译器将一个 IRModule 转换为另一个功能等效或近似等效的（例如在量化的情况下）IRModule。许多转换与目标（后端设备）无关。我们还允许转换目标影响转换pipline的配置。

③ Target Translation ：将 IRModule 翻译（代码生成）为目标（后端设备）指定的可执行格式。目标翻译结果被封装成一个runtime.Module，并可以在目标运行环境上导出、加载和执行。

④ Runtime Execution ：用户重新加载一个 runtime.Module 并在支持的运行时环境中运行已编译的函数。

关键数据结构

通过研究TVM的关键数据结构和API对数据结构的操作，可以对TVM进行深入的了解。据此，可以将TVM拆解为一组关键数据结构的定义或者在数据结构之间进行转换（Transformation）的逻辑组件（logical components）。

① IRModule（intermediate representation module，中间表达模块）是整个 TVM stack 的主要数据结构（primary data structure），其包含了一组函数，目前TVM支持以下两种基础函数。

② relay::Function是一种高级函数程序表达。一个 relay.Function 通常对应一个端到端模型，可以将其视为一个具有额外控制流、递归和复杂数据结构支持的计算图。

③ tir::PrimFunc是一种低级程序表达。包括循环嵌套选择（loop-nest choices）、多维加载/存储（multi-dimensional load/store）、线程和向量/张量指令（vector/tensor instructions）。

编译期间，一个relay函数可以 lower 为多个tir::PrimFunc函数和一个调用这些tir::PrimFunc函数的顶级（top-level）函数。

转换 Transformations

每个转换可以用于以下目的之一：

① optimization 优化：将程序转换为等效的、可能更优化的版本。
② lowing 低层化：将程序转换为更低层次表达方式，使其更接近后端设备。

relay/transform：包含一组 pass 用来进行模型优化，这些优化包含常用的程序优化（constant folding and dead-code elimination）和张量计算优化（layout transformation and scaling factor folding）

靠近relay优化流程末端的地方，TVM会调用 FuseOps 将端到端函数（如 MobileNet ）分解为子函数段（sub-function segment）（如 conv2d-relu），这一过程将原始问题划分为两个子问题：

对每个子函数进行优化。Compilation and optimization for each sub-function.
整体执行结构：通过对一系列子函数的调用，完成整个模型的执行。Overall execution structure: we need to do a sequence of calls into the generated sub-functions to execute the whole model.

再此之后，TVM使用低级的TIR段对每个子函数进行编译和优化。对于一些特定的Target，TVM会直接使用外部代码生成器进入目标翻译阶段。

对于形状固定没有控制流的简单模型，TVM可以将整体执行的调用降低为一个图执行器（graph executor）将执行结构保存至 graph 中。TVM同样支持用于动态执行的 virtual machine backend。

tir/transform：包含TIR级别函数的转换 pass，其中许多 pass 的目标是底层化（lowering）（例：其中有 pass 的目标是将多维度的数据访问平整为单维度指针访问，以此将 intrinsics 扩展为特定于目标的函数，同时装饰函数入口以满足运行时调用要求。）

许多低级优化可以由LLVM或者NVCC等后端设备编译器完成，因此，TVM中仅关注于这些编译器未覆盖的地方。

基于搜索空间和学习的转换 Search-space and Learning-based Transformations

TVM stack 的设计目标是针对不同硬件平台实现高性能代码优化。为此，其需要尽可能多的调查优化选项，包括但不限于：multi-dimensional tensor access、 loop tiling behavior、 special accelerator memory hierarchy 以及 threading。

这很难定义一个后发式的方法来实现所有选择。与之相反，TVM采用了基于搜索和学习的方法。在TVM中，定义了一组可以用来进行转换的操作（如：loop transformations, inlining, vectorization），称之为调度原语（scheduling primitives）。系统会搜索所有可能调度（scheduling ）序列并选取一组最佳的调度（ scheduling）组合。搜索过程通常由机器学习算法主导完成。

一旦搜索完成，系统会记录最佳算子调度序列（the best schedule sequence for an (possibly-fused) operator），之后编译器仅查找最佳调度序列并将其应用到程序当中。值得注意的是，调度的应用阶段与基于规则的转换（ rule-based transformations）是完全一致的，这使得TVM可以使用与传统 pass 相同的接口规约。

TVM使用了基于搜索的优化方法，用来处理TIR函数的生成问题，这一部分被称为AutoTVM。

目标翻译 Target Translation

如前文所述，目标翻译主要作用是将IRModule翻译成后端设备指定的可执行格式。如X86、ARM以及CUDA、OpenCL。代码生成阶段当尽可能的轻量化，绝大多数转换和低层化都应在目标翻译之前进行。

通过指定编译目标，其同样会影响在此之前的转换过程，如：目标设备的向量长度会改变向量化过程。

运行时执行 Runtime Execution

TVM运行时的主要目标是提供一个最小化API，用于加载和执行编译结果。

以下为TVM官方文档中提供的Python示例 ^[1]:

import tvm
# Example runtime execution program in python, with type annotated
# tvm.runtime.Module encapsulates the result of compilation.
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())

Compilation Flow https://tvm.apache.org/docs/dev/index.html#example-compilation-flow ↩ ↩