最近粗略的看完了天奇大佬的MLC课程（顺便修了一些语法和拼写错误，也算是做了微弱的贡献hh），对TVM的近期发展有了一些新的认识。之前天奇大佬在《新一代深度学习编译技术变革和展望》一文中（链接：https://zhuanlan.zhihu.com/p/）讲解了TVM Unify也即统一多层抽象的概念。这里的统一多层抽象具体包括AutoTensorization用来解决硬件指令声明和张量程序对接，TVM FFI（PackedFunc）机制使得我们可以灵活地引入任意的算子库和运行库函数并且在各个编译模块和自定义模块里面相互调用。TensorIR负责张量级别程序和硬件张量指令的整合。Relax (Relax Next) 引入relay的进一步迭代，直接引入first class symbolic shape的支持 （摘抄自《新一代深度学习编译技术变革和展望》一文）。然后这些抽象可以相互交互和联合优化来构造深度学习模型对应的最终部署形式。
我个人感觉TVM Unify类似于MLIR的Dialect，但是这几个抽象的直接交互能力相比于MLIR的逐级lower我感觉是更直观方便的，毕竟是Python First（这个只是我最近看MLC课程的一个感觉）。对这部分内容感兴趣的读者请查看天奇大佬的TVM Unify介绍原文以及MLC课程。
这篇文章我将结合TVM Unify相关的抽象以及之前的一些积累重新梳理一下TVM的整体流程。我会从前端，中端（图优化Pass机制），代码生成（Schedule），Runtime，开发工具几个角度来介绍一遍。我对TVM的代码并没有做到精细的阅读，所以本文将尽量避免涉及到底层C++代码的细枝末节，而是从较为宏观的视角来讲清楚目前TVM的架构。
本篇文章的所有参考资料以及idea主要来自我维护的这个仓库（https://github.com/BBuf/tvm_mlir_learn）里面搜集的TVM的相关资料，TVM官方doc以及源码，MLC课程。上面这个仓库基本收集了TVM中文社区里面的大部分高质量博客或者专题，对TVM感兴趣的小伙伴可以自行下载或者收藏，更欢迎点个star。写作不易，这篇文章对你有用的话也请点个赞👍。文章有错误也请指出，我动态修改。之后的计划应该会学习TVM如何和硬件的指令对接。
TVM为了向上兼容所有的机器学习框架如PyTorch，TensorFlow，ONNX等引入了Relay IR，机器学习模型在进入TVM之后首先会被转换为Relay IR。同时TVM为了向下兼容所有的硬件，引入了Tensor IR简称TIR，模型在被编译为指定硬件的源代码之前都会被Lower为TIR。另外，TVM社区正在开发新一代中间表示Relax（也被称为下一代Relay，目前还没有upstream主分支：https://github.com/tlc-pack/relax/tree/relax/python/tvm/relax），Relax是实现前言里面提到的TVM Unify关键的一环。TVM前端的架构可以粗略的表示为：
接下来我们分别介绍一下 Relay，TIR，Relax这几种不同的前端表示。
由于无论是Relay还是新一代的Relax中间表示，它们最后都会被Lower到TIR（离硬件最近的IR），所以我们这里先介绍一下TIR。TIR的代码被封装在中，一个TIR可以被编译成目标硬件的源代码或者中间表示例如C++源码，CUDA源码，LLVM IR等等。那么TIR是如何被编译为目标硬件的代码呢？这是因为TIR的数据结构其实是一个AST（抽象语法树），然后这个语法树可以表示变量的声明，初始化，变量的计算，函数调用以及控制流（如if-else条件判断，循环等等）等等。所以只要我们遍历一下TIR对应的AST就可以实现一对一的将其翻译到目标硬件了。可以借助这个图来理解：
在上图中有几个细节需要解释。首先是IRModule，IRModule 是在机器学习编译中保存元张量函数（也即PrimFunc）集合的容器对象，它是TVM进行编译的最小完整单元。TVM不同的前端表示最终都会被封装到IRModule中进行编译，在Linux下IRModule就是一个.so动态链接库。然后PrimFunc叫作元张量函数，它内部封装了一个完整的TIR AST。当IRModule被编译之后，每个PrimFunc都对应了这个动态库的一个函数入口，因此一个IRModule可以有很多个PrimFunc。然后上面的Codegen实际上就是对TIR AST进行中序遍历然后一对一的将AST Node翻译为相应的TIR Node对应的数据结构并发送给回调函数VisitExpr_ 和 VisitStmt。VisitExpr_ 用于处理 Expression Node，而 VisitStmt 用于处理 Statement Node。后续在介绍Codegen的时候我们再仔细探索一下这个转换流程。
这里还需要说明的一点是，在0.8之前的TVM要声明一个TIR AST依赖于对Tensor Expression的编译。现在TVM基于Python AST实现了一种新的特定领域的方言让我们可以直接使用Python来编写TIR AST。我们这里举一个例子：
它实现的功能对应的numpy代码为：
其中，表示被修饰的MyModule是一个待编译的IRModule，而表示被修饰的main函数是元张量函数（PrimFunc），这个函数内部定义的就是TIR AST。
继续讲Relay IR以及Relax之前我们先了解一下这个抽象，无论是TIR还是Relay/Relax IR它们都对应了IRModule这个统一的最小编译单元，同时它们也对应的有一套共用的IR基础设置，具体实现在和目录下。
对于IR来说，Type和Expr是尤为关键的两个概念。Type包含基础的数据类型如Int，Float，Double等等，也包含一些自定义的复杂类型比如函数类型，Tensor类型等。而对于Expr来说，既包含可以直接映射到Low-level IR的PrimExpr，又包含RelayExpr。我们可以在中看到对PrimTypeNode的定义：
可以看到PrimType可以直接对应到Low-level IR的基础数据类型。我们还可以找到FuncTypeNode的定义：
从注释可以看到FuncType类似C++的模板函数，记录了函数的参数类型和返回值类型以及模板参数，约束等信息。然后我们还可以关注一下和深度学习模型结合得很紧密的TensorTypeNode类型。
我们从TensorTypeNode的定义可以看到shape也是TensorType的一部分，所以TVM在做类型推断的时候也包含了Shape的推断。也正是因为在IR中Shape是Type的一部分（比如和是不同的Type）导致TVM对动态Shape的支持非常困难，因为Expr的类型推断是不支持动态Shape的。这里需要提一下，Relax通过引入一个新的Type叫作DynTensor较好的解决了动态Shape的表示问题，DynTensor包含的信息是Dtype和Shape的纬度，但Shape本身的表达式是独立存储的。也就是和都是同一个Type， 但是和是不同的Type，这样就从原生上支持了动态Shape。我们从这里可以看到DynTensor的定义：
我们紧接着看一下Expr的定义（），Expr分成PrimExpr以及RelayExpr。其中PrimExpr保存了一个runtime时候的Dtype，然后
例如表示一个整数的Expr就可以通过继承PrimExprNode来实现，IntImm表示的是整数字面值表达式，所以它记录了一个int类型的value成员。
RelayExpr的定义如下：
总的来说，无论是高级别的Relay，Relax还是低级别的TIR，它们最终都是由这里的Expr和Type为基础来表达的。因为对于Relay和TIR来讲，它们的op定义都是继承自RelayExprNode：。除了对Op名字，类型以及参数，属性等定义外还有一个特殊的参数，从注释上看应该是用来解释当前Op的等级，值越小表示这种Op类型等级越高（暂不清楚具体的作用）。
最后我们看一下IRModule的定义，。我们说过IRModule是TVM编译的最小单元，我们可以从它的定义中发现它就是一系列BaseFunc（在下一节Relay的介绍中我们会讲到它的实现）的映射。
其中type_definitions是对ADT的定义，本文不关注Relay中函数式编程的概念，所以不展开ADT以及Let Binding部分的概念和源码，感兴趣的朋友可以参考张伟大佬的这篇文章或者官方文档对Relay的介绍学习一下：https://zhuanlan.zhihu.com/p/ 。后面在介绍Relax IR的时候我们会看到，实际上Relax相比于Relay就类似于TensorFlow的静态图到PyTorch动态图的过度，更加强调数据流图的概念而非函数式编程的概念，我个人感觉也是为了易用性考虑吧。
接下来我们简单介绍一下Relay IR。首先Relay IR目前仍然是TVM和其它深度学习框架对接的主要方式，我之前在《【从零开始学TVM】三，基于ONNX模型结构了解TVM的前端》文章中以ONNX为例介绍了模型是如何转换为Relay IR的，然后这个Relay IR会被进一步封装为IRModule给TVM编译。从源码角度来看，Relay的基类Expr就是tvm.ir基础设施中定义的RelayIR（）。
然后Relay还定义了ConstantExpr，TupleExpr，VarExpr，CallNodeExpr，LetNodeExpr，IfNodeExpr等多种Expr。我们可以看一下ConstantExprNode的定义，类定义中声明了数据data并定义了tensor_type方法返回data的类型，然后is_scalar函数用来判断这个常量是否为标量。
然后我们再看一下VarNode的定义，Var就是Relay里面的变量，它的定义如下：
首先Id vid表示的就是变量的名称，可以理解为一个字符串，比如我们在可视化Relay IR时看到的以@开头的全局变量以及%开头的局部变量。这里的type_annotation表示变量的类型注释，这个字段是可选的。接下来我们再看一个FunctionNode的定义，FunctionNode就是IRModule中的BaseFunc在Relay里面的具体实现了：
FunctionNode的定义中有函数参数，函数体以及返回值类型和参数类型。其它类型的Relay表达式定义我们就不看了，感兴趣的读者可以直接在阅读。接下来我们解析一下Relay中的Op定义，上一节tvm.ir基础设施中我们已经提到无论是Relay还是TIR的Op都定义为一种RelayExpr，也就是OpNode的定义。我们这里看一个Relay定义的bias_add Op的例子来加深理解。首先，我们为BiasAdd Op定一个属性类型记录它所有的属性，，属性定义时我们还可以给属性设置描述和默认值：
第二步，我们给Biass Add Op定义类型推断函数（）：
假设这里指定的操作是 c = nn.bias_add(a , b)，这里的逻辑就是根据输入a的类型推断b和c的类型并重写（Assign）。第三步，我们把nn.BiasAdd Op注册到全局表中（）：
注意到这里的op name/describe/num_inputs/arguments/support_level是对应了OpNode类的成员，然后OpNode还有一个attrs_type_key和attrs_type_index成员对应的就是BiasAddAttrs了。然后我们再看一下这个FTVMCompute这个用来描述Op计算逻辑的额外属性，它使用Op的输入，属性参数以及输出类型来确定这个Op的计算逻辑。到这里可能你还有一个疑问，我们知道TVM的核心是计算和调度分离，Relay Op的调度逻辑是怎么注册的呢？TVM没有为每个Relay OP注册compute和schedule，而是为其注册fcompute和fschedule，然后根据输入和属性参数，输出类型等生成对应的compute和schedul，这种compute和schedule的组合对应了OpImplementation（）。
从OpImplementation类的实现我们看出，它的Compute和Schedule就是根据fcompute和fschedule来生成的。
然后由于特定的OpImplementation需要特定的条件，所以又按照这个条件（condition）进行分组，每一组被叫作OpSpecialization（）
最后使用一个OpStrategy类来记录这个Relay Op的所有OpImplementation。（）
其中，AddImplementation函数通过FFI机制在Python层也可以调用，大多数的Relay Op都是在Python端注册它的Strategy。我们以Relay的nn.Softmax Op为例看一下，它的Strategy（包含fcompute+fschedule）注册在 和 。
然后在将实现的Strategy注册到nn.softmax op。
其实Relay Op除了Strategy属性之外，还又一些其它的属性，比如我们在 这里可以看到Op还可以有FInferCorrectLayout和TOpPattern属性用于后续优化（比如算符融合Pass就依赖了TOpPattern属性，Ansor的data layerout transform依赖FInferCorrectLayout属性）。
Relay就暂时讲到这里，Relay IR做为函数式风格的IR目前是TVM和其它深度学习框架交互的桥梁并且也经历了多年的维护完备性是比较好的（支持TensorFlow，PyTorch，Paddle，OneFlow各种主流深度学习框架）。但Relay的缺点在于由于共用了TVM的 tvm.ir 基础设施没办法支持Dynamic Shape导致Relay IR也无法支持Dynamic Shape，并且Relay IR这种函数式编程的风格相比于数据流图形式的计算图来说不是太直观。
由于Relax这个前端还没有正式upstream到apache tvm主分支，所以我这里就不从源码的角度来看。我们可以从Relax的wiki发现它不仅原生的支持动态Shape（通过提供DynTensor的抽象并将Shape从Tensor的type中分离出来实现的）还做了一个TVM Unify抽象，也就是天奇在《新一代深度学习编译技术变革和展望》一文中提到的，这个特点可以让不同的抽象之间相互交互和联合优化。这里提到的抽象包含AutoTensorization用来解决硬件指令声明和张量程序对接，TVM FFI（PackedFunc）机制使得我们可以灵活地引入任意的算子库和运行库函数并且在各个编译模块和自定义模块里面相互调用。TensorIR负责张量级别程序和硬件张量指令的整合。还有这里的Relax (Relax Next)。我们可以从下面的例子体会：
注意这里展示的代码片段是Relax wiki提供的，由于没有upstream主分支，它的用法也许会有微小变化。我们从这个代码中可以看到，Relax把Relax Function和TIR Function放到了同一个IRModule（最小的编译单元）也就是说在任意时刻我们都可以同时拿到这两个不同层次的IR进行修改（或者说联合优化）这就摆脱了编译器范式里因为Lower导致丢失高层语义信息无法联合优化的问题。知乎上思远指出了一个很经典的例子，我这里附上他回答链接（）并截图来说明一下：接下来我们翻译一下Relax的设计关键点来进一步体会Relax相比于Relay的变化（中间插了一些个人理解）。
D0：数据流块作为第一优先级的构造
大部分的relax_func都封装在with R.dataflow()构造里面。数据流块下的所有操作都是没有副作用的，并且不包含高级的控制流（比如if-then-else）或者嵌套区域。一个数据流块可以有效地视为嵌入在程序里面的计算图。请注意，数据流块里面的大多数绑定变量（上面Relax脚本中的lv0, lv1, lv2, lv3）是local的，这意味着它们仅是块内可见的。这些变量可以被视为计算图的“内部节点”。我们可以将变量标记为输出（gv0），在这种情况下，该变量将在程序的后面部分可见。这些输出变量可以被视为计算图中的输出节点。请注意， 在数据流块之外。数据流块之外的所有内容都可能产生副作用。因此，除非我们进行更仔细的分析，否则我们无法执行优化，例如根据拓扑顺序重新排序这些绑定。我们预计大多数优化将发生在数据流块级别。这些优化可以由熟悉计算图概念的 ML 工程师完成。隔离和表示有效组件的能力还为需要它们的地方提供了更高级别的优化机会。
D1：形状推导作为第一优先级的计算
形状推导对于动态模型工作负载至关重要。在动态形状设置下，我们通常需要在运行计算之前计算中间张量的形状。此外，我们还需要处理形状本身依赖于数据（例如unique op）的情况。最后，大多数动态形状工作负载仍然包含大量（部分）静态形状，理想情况下，我们希望利用这些静态形状信息进行优化。
上述程序涵盖了形状推断的典型场景（在注释中标记）。重要的是，形状现在与张量值一起成为计算的一部分。这反映了形状的计算可以在运行时发生的事实。而文本格式类型注释 显示了每个Shape的值。这只是一个语法糖，从 IR 的角度来看，Shape字段 不是 的一部分。lv0 的类型是 ，Shape是附加到每个 Expr 的特殊值字段。我们做出这个显式的选择是为了简化类型推断，这样我们就不需要进入完全依赖类型的领域。有两个与符号Shape计算相关的关键结构：
D1a: match_shape
形状匹配构造接受一个 lhs 值和pattern（整型符号表达式）。它有两个重载语义：
-
当 lhs 为 Tensor 时，将 lhs.shape 匹配到 pattern 中，如果第一次出现在 pattern 中，则填充对应的整型符号变量，然后返回一个与 lhs 相同但 shape 字段更新为 pattern 的 Tensor。
-
lhs 也可以是直接匹配 pattern 的 Shape。当我们想要分离出不对应于任何张量值的 Shape 函数时，这很有用。
比如：
特别注意这里lv2的Shape就被设置为(n, m)，并且match_shape的lhs是一个Shape表达式，而不是Tensor。
D1b. 从符号整数元组构造Shape
在我们得到 n 和 m 等符号化整数之后。我们可以将它们重新组合在一起以形成一个 Expr。任何符号整数表达式的元组都可以在 Relax 中被识别为Shape 值。比如 (n, m) 就是一个表示 Shape 的值。
Shape传播的方法
重要的是，现在Shape是计算过程中值的一部分。编译时Shape推断可以被看作是对发生在Shape上的操作的常量折叠，程序有几种Shape计算的方法：
-
方法1: 符号化的形状传播。可以将Shape分解为符号整数比如上个脚本中的n和m，然后我们可以使用符号整数的表达式来代表Shape的计算比如。值得注意的是，静态形状是符号整数的一种特殊情况，然后我们可以重新组合符号整数来构造一个新的Shape如。
-
方法2: 不透明的Shape函数调用。我们还可以实现不透明的Shape函数比如（看上上个Relax脚本），这些不透明的Shape函数是快速激活成功教程运行时Shape函数的有用fallback（这里应该是说加上手工干预的形状推导？）。
-
方法3：对于数据相关的Shape（如Unique），我们将简单地推迟到一个运行时的调用 它接收一个输入张量，分配并返回输出张量。然后我们可以通过match_shape构造从Tensor值中获得lv5的形状。（看上上个Relax脚本）
Implications for pass writing
很多优化Pass都需要知道Shape信息。既然很多Shape可以是符号化的比如 (n, 4)，那么理想的优化Pass将需要更泛化一点以利用符号信息。比如在上述脚本中，我们知道所有的都对应同一个值。这种约束很有用。因为符号化的整数（我们之前讲过对应 ）动态的执行常量折叠，当输入是静态Shape时计算的结果也应该动态的折叠为整形常数，保留我们执行静态Shape优化依赖的属性。因为我们现在可以在元组(n, 4)表示混合的静态符号Shape，所以我们可以尝试利用静态信息进行额外的优化。
D2：与 TensorIR 和 PackedFunc 直接交互
我们做出的最后一个关键设计决策是允许高层 IR 能够直接交互并调用低层 TensorIR 和 PackedFunc。TensorIR 函数和许多外部库采用目标传递约定（我们需要显式分配输出并作为参数传入函数）。我们使用 dps(destination passing) 来表示这个约定。dps 在低级 ML 优化中非常重要，因为它允许我们在可能的情况下一次性全局分配中间存储，并在没有主动内存分配的情况下执行计算。调用 dps 函数意味着在调用之后，结果通过函数参数（例如，下面示例中的结果）而不是函数的返回值传回。
dps 风格在本质上意味着突变（输出）。我们需要一种将调用桥接到Relax Dataflow的方法(可以观察一下Relax这一节开头那部分的脚本)，以便我们可以对一系列 tir 调用执行计算图样式的重写。
D2a. call_tir
是将调用桥接到Relax Dataflow的内嵌函数。它的命名含义是：“调用一个tir转换”
call_tir 接受输出形状，lowlevel_func(can be packed func, tir PrimFunc) 和一个输入元组。call_tir 的语义可以通过上面的代码来演示。值得注意的是，当我们lower 时，我们不需要选择单独的分配输出张量。编译器可以选择创建中间张量的内存计划，并将它们联系在一起以实现有效重用。值得注意的是，call_tir 内嵌函数的 参数可以是不透明的形状值、符号整数元组或常量形状（支持动态Shape）。 可以是任何带有签名的函数：最常见的两种情况包括：(1) TIR 函数 (2) 不透明的packed func
实现笔记
call_tir 可以实现为特殊的内嵌函数 (Op)，以最大限度地减少对 IR 更改的影响（而不是独立的 IR 节点）。从 AST 的角度来看，这变为：
这也将允许 call_tir 的未来迭代而不改变 IR 本身，这可能在特定时间点需要：
-
在同一个数组上启用多个突变序列（在 concat 相关操作的情况下）
-
启用将符号化的Shape提示传递给融合操作。
对整合的影响
D2 使我们能够将较低级别的抽象直接嵌入到高级抽象（R.function）中。这释放了很多机会，包括但不限于：
-
使用不同的策略逐步lower程序的不同部分。
-
我们可以将call_tir节点作为AST的一部分进行优化，然后将一些关键信息比如data layerout信息带回到high level的IR获得更好的优化结果。
-
将 BYOC 流作为转换的自然部分（通过将图的一部分转换为不透明打包函数的调用）。
这里的第二点实际上对应了Ansor引入的weight layout rewrite，即在算子auto-tuning之后，我们去分析最高效的weight layout，并且在编译时改写，来提高运行时的效率。那么没有Relax之前是怎么完成这个工作的呢？一个op 更适合的weight layout是要在tuning之后才能够知道的，而这个时候图IR已经被lower，不能修改了。所以Ansor用了一个非常tricky的方法，先lower一遍把tuning做好，再带着这些信息重新lower一遍。所以Relax通过消除lower的边界隔阂可以较好的解决这一问题。
D2b. Packed function calls
我们使用 来指示对Packed Func的调用。从 AST 的角度来看，我们不需要引入额外的调用节点，而是可以引入一个 ExternFunc 构造，它表示我们可以调用的打包函数。
仅用作表示上述 AST 节点的语法糖。这使我们能够统一所有调用。值得注意的是，它还允许我们在必要时混合打包函数和 call_tir。
对应于下面的 AST。
当我们想要将低级库（例如 cudnn）直接集成到高级而不调用内存分配时，外部打包函数上的 CallTIR 会很有用。关于这一点在MLC课程中也有演示，通过dlpack调用PyTorch的Op来做优化，感兴趣的读者可以看一下，链接：https://mlc.ai/zh/chapter_end_to_end/index.html。这里简单做一个总结，Relax作为下一代Relay不仅原生支持动态Shape且使用体验更加靠近PyTorch这种数据流图的编程方式。尤其重要的是Relax在为TVM Unify而服务，通过和TensorIR抽象，TVMFFI（Packed Func）的交互（通过MLC教程可以知道，也可以和Auto Schedule交互）使得TVM Unify的目标得到实现。当然我也要说一下我目前看到的Relax的不完善的地方，那就是Relax目前和其它深度学习框架对接还不够完善，如果能实现Relay到Relax的自动转换那将是一个振奋人心的消息，可以最小化我们的迁移成本。
让我们回到开头的这个图：
我们可以发现Relay要到TIR有2条路径，第一条就是直接到TIR比如PrimExpr派生的节点比如一个IntImmNode可以直接映射到TIR，另外一条就是Relay里面类似Conv的Op的计算逻辑是用TOPI来表达的，TOPI是TVM自己的一个算子库，这些算子可以通过TE来进行表达。除此之外，我们在前端介绍Relax的时候已经可以看到要直接编写TIR AST，一种方法是使用TVMScript来表示抽象的计算逻辑，另外一种方法就是要通过TE，TE的代码无法被直接编译成目标硬件的代码，而是需要先Lower为TIR的元张量函数才可以进行编译。其实我之前写过一些Schedule相关的文章比如《【TVM 三代优化巡礼】在X86上将普通的矩阵乘法算子提速90倍》，也都是基于TE的。由此可见，TE不仅提供了另外一种编写TIR AST的方法，还提供了一系列变换TIR AST的Schedule。在0x5节我们会提一下Schedule。我们先看一下给予TVM Script写的这个向量加法的例子：
输出：
然后我们使用TE DSL来表达这个向量加法：
输出：
从两个输出中我们可以看到，最后创建的IRModule其实是完全一样的。然后这个IRModule可以被编译为目标硬件上可以执行的代码。如果你想更加深入的了解TE是如何被编译成TIR的，可以看一下 《TVM 自底向上（三）：TE 的概念和编译原理》 这篇文章，我这里借一下作者文章中的核心图简要说明一下：
我们从上往下看，这里的List[PrimExpr]就是这个lambda表达式中的PrimExpr集合，第一个PrimExpr是A(*i)，第二个PrimExpr是1.0，然后+对应了TIR中的ExprOp（），Expr作用在1个或者多个PrimExpr上得到的结果仍然是PrimExpr。实际上，这里的List[PrimExpr]就对应了这个lambda表达式的AST表示。接下来我们看一下te.compute的代码（）：
在compute的实现中最后返回的是TensorComputeOp对象的output()成员（也是一个tvm.te.Tensor）， 同时这个tvm.te.Tensor包含这个TensorComputeOp对象（通过来访问，在可以看到）。最后这行代码完成了TE到TIR的转换。这个api对应的c++实现在这个文件，感兴趣的读者可以自行查看。基本流程就是将所有Operation对应的PrimExpr AST连在一起构成一个AST Graph，然后使用Post-DFS算法遍历这个AST Graph分别处理每一个Operation创建对应的TIR节点，最后构造一个完整的TIR PrimFunc。TE除了可以构造TIR之外，另外一个重要的点就是它支持Schedule（），我在【TVM 三代优化巡礼】在X86上将普通的矩阵乘法算子提速90倍 文章中对GEMM优化的介绍就是基于TE Schedule来做变换进行优化计算的。
现在我们把目光转向图优化的Pass。首先，我们看一下TVM Pass的基类定义（）：
从operator()的定义可知，Pass做的主要是IRModule到IRModule的变换，另外这里的PassInfo和PassContext分别表示每个Pass的关键信息和多个Pass执行过程中的共同上下文信息。我们分别看一下它们的定义()：
这里需要注意的是在PassContextNode定义中出现了一个类，这个类是为开发者设计的一个工具，开发者可以实现一些函数运行在每个Pass执行前或者执行后（）：
我们可以在这个测试文件中找到PassInstrument机制的示例用法， 这个功能可以方便的让我们观察每一个IRModule经过一个Pass之后变成新的IRModule之后有什么变化，方便debug或者可视化。然后TVM为了方便实现了3个级别的Pass，即Module-Level的Pass直接操作IRModule，以及Function-Level的Pass遍历Module 中的Function进行处理，还有Sequential Pass包含一堆顺序执行的Pass（对比PyTorch的nn.Sequential）。感兴趣的读者可以自行阅读源码。接下来我们讲一讲图优化Pass遍历以及重写AST节点的原理。注意，我们这里讲的Pass是TVM内置的作用于TIR AST上的Pass，我们知道TIR AST是由一系列PrimExpr和RelayExpr（非PrimExpr）来表示的，它们都继承了TVM的Expr基础类。所以TVM针对TIR AST的遍历专门做了一个工具类ExprFunctor来做，它定义在 ：
从类的定义可以看到ExprFunctor主要提供了VisitExpr函数接口，并根据Expr的具体类型转发到对应的 VisitExpr_ 。VisitExpr_ 则由派生类负责实现，当然从代码也可以看出，VisitExpr 本身也可以被重载。有了这个转发机制之后，就可以很容易的实现一个遍历所有类型Expr的类了，在TVM中叫作ExprVisitor（）：
比如对于中的这个类，就继承了，并通过，访问数据。的成员函数实现如下()：
可以看到这个类实际上调用的是父类()的，而的的实现如下：
可以看到设置了虚函数，在解析时会回到来解析节点，而这个类继承了，这样我们只需要在类中重写各个Expr节点类型的函数就可以了。在的实现中有一个宏，这个宏的定义如下：
这里的即为的的实现中的，而指向。又因为方法调用的是的函数，所以这里的指向的是实例。以为例子，看看的实现。由于指向的是实例，最后会在实例中生成的列表。
是在中定义的一个，来标记在遍历AST时某种Expr是否出现，同时记录下出现的次数来保证每个Expr都只会被访问一次。
显然，如果AST很复杂，这样递归可能会导致Stack Overflow. 为了解决这个问题，TVM 提供了 MixedModeVisitor 来实现和 ExprVisitor 一样的功能，但是避免了 Stack Overflow。我们上面提到对于AST除了遍历，还有改写的需求，所以TVM提供了一个ExprMutator ，同样继承了 ExprFunctor。类的定义如下：
注意 Mutate 只是 VisitExpr 的别名。ExprMutator 的 VisitExpr 会返回一个修改后的新 Expr, 看一下 VisitExpr 的实现：
可以看到存储了图中的各个节点。参考IfNode的实现：
如果的子节点都没有被修改，那么就返回这个节点本身。否则创建新的节点并返回。这里构造新节点的类If的定义和实现分别在和中：
TVM的Pass里面有一个经典的算符融合Pass，之前在【从零开始学深度学习编译器】八，TVM的算符融合以及如何使用TVM Pass Infra自定义Pass 这里讲过，感兴趣的小伙伴可以看一下。
我认为TVM的Schedule主要分为三个部分：TE Schedule，TIR Schedule以及Auto Schedule。由于精力有限我还没有探索Schedule在TVM的源码实现，不过最近TVM圈子的这篇来自Kord大佬的《TVM 自底向上（四）：TE/TIR Schedule 的原理》文章为我们理清了TE/TIR Schedule的原理，个人推荐大家去阅读。链接：https://zhuanlan.zhihu.com/p/ 。
基础概念1: PackedFunc
为了便于Python和C++混合编程，TVM使用了统一的PackedFunc机制。PackedFunc可以将C++的函数打包成统一的函数接口并导出到Python端供用户使用，同时也支持从Python中注册一个函数，并伪装成PackedFunc在C++和Python中调用。这里推荐一篇讲解PackedFunc原理的优质博客：https://hjchen2.github.io/2020/01/10/TVM-PackedFunc%E5%AE%9E%E7%8E%B0%E6%9C%BA%E5%88%B6/ 。
基础概念2: tvm.runtime.Module
tvm.runtime.Module是tvm编译的结果（这一节之后简称Module）。Module中包含一系列可以运行的PackedFunc（所以这里的Module可以看作<name, PackedFunc>的哈希表），并且Module可以import另一个Module，从而访问其它Module的PackedFunc。我们看一下Module的接口定义（）：
然后Module的具体实现由ModuleNode负责，并且不同的target对应不同的ModuleNode实现。我们来看一下CUDAModuldeNode的定义（）, 请注意看下面的注释:
我们看一下核心的GetFunction实现（https://github.com/apache/tvm/blob/main/src/runtime/cuda/cuda_module.cc#L244-L257）：
这里首先根据函数的名称找到描述这个函数的FunctionInfo，而FunctionInfo里面包含了launch_param_tags成员，这个成员中存储了CUDA Kernel Launch时需要的gridDim/blockDim/SharedMemorySize，然后将上下文打包到CUDAWrappedFunc中并包装为一个PackFunc返回。然后我们可以看一下CUDAWrappedFunc是怎么执行的（https://github.com/apache/tvm/blob/main/src/runtime/cuda/cuda_module.cc#L164-L203）。
从这里可以看到CUDAWrappedFunc会根据func_name在CUDAModuleNode中找到CUfunction然后根据launch_param_config_进行Kernel Launch。这里的fcache_[device_id]是用来缓存当前device上的CUFunction的，避免重复查找带来的额外开销。另外在CUDAModuleNode::GetFunction的定义中提到如果name是tvm_prepare_global_barrier，则将CUDAPrepGlobalBarrier包成一个PackedFunc返回，在CUDA 9.0之前是不支持Global Barrier的，所以这里TVM通过类似spin lock的方式，自旋地检查一个全局变量的值来block 线程执行，从而实现Global Barrier。核心实现见：
除了CUDAModuleNode之外，其它的硬件抽象都实现了一个对应的ModuleNode比如OpenCLModuleNode，ROCMModuleNode等等。借助Module和PackFunc我们可以将不同device生成的代码打包成统一的形式。但如果想要执行这些生成的代码，我们需要做内存管理，同步等一系列操作，TVM将这些操作抽象为DeviceAPI。
TVM通过DeviceAPI 类来对硬件的能力进行抽象，形成了几个统一的接口（在OneFlow中有一个硬件抽象模块EP和这个类似）。只要为每一种device重载了这些统一的接口，那么执行器（runtime）就可以通过访问这些统一的接口使用device的某种能力，比如查询参数，内存分配，数据拷贝，同步等等。DeviceAPI的定义在：。这里有一些通用的接口比如SetDevice，GetAttr，GetTargetProperty，AllocDataSpace等等，然后对于不同的device比如cpu，cuda，hexagon，metal，rocm，vulkan，opencl都会基于各自的runtime api重写这些接口。这些接口对于TVM的执行引擎非常重要。Module，PackFunc，DeviceAPI分别从不同的角度对硬件的功能进行了封装，比如Module封装了加载device Module（比如CUModule），加载Kernel，统一打包设备代码等功能，DeviceAPI封装了内存分配释放，数据拷贝等功能，但这些功能必须要有一个执行引擎凑到一起才可以run起来。TVM提供了2种执行引擎。
Graph Executor
GraphExecutor是TVM为静态模型设计的执行引擎（不支持动态Shape和Control Flow）。我们先看一个GraphExecutor执行一个Relay Function的示例（https://github.com/BBuf/tvm_mlir_learn/blob/main/relay/simplenet.ipynb）：
这里首先创建了一个GraphExecutor对象并使用Relay Function的编译结果对其进行初始化，RelayFunction的编译结果包含序列化图结构（对应executor_config）、kernel（对应mod）、weight（对应params）。
接下来为GraphExecutor对象设置输入数据，然后调用run子函数来执行kernel，最后get_output获取输出结果。GraphExecutor的实现主要有2个函数，第一个函数就是Init（https://github.com/apache/tvm/blob/main/src/runtime/graph_executor/graph_executor.cc#L77）。
这个函数中主要包含json参数解析。为每一个算子的input/output edge准备对应的memory（对应SetupStorage） 以及为每一个算子准备一个可调用的kernel function用来做实际的计算（对应SetupOpExecs）。
json就是计算图的表示，表示了node之间的连接关系，输入、输出node、输入shape等信息，上面的代码中Load(Read)会提取json中的信息，存储在graph_executor成员变量中。
Virtual Machine
目前我基本没有使用过这种运行时，并且了解也比较少，所以这里就留坑不展开了。VM是TVM中更加灵活的一种运行时，它可以支持动态模型（也就是带动态Shape和Control Flow的）的执行。其实，从MLC的课件也可以看到Relax在处理动态Shape程序时也用到了这个运行时。一位Intel的工程师在《TVM Runtime System 概述》介绍了TVM的Relay Virtual Machine运行时，感兴趣的小伙伴可以去阅读一下：https://zhuanlan.zhihu.com/p/ 。
之前提到IRModule是编译的最小单元，然后当我们执行类似于 可以将IRModule编译为，这里的target参数就是用来选择使用哪一个CodeGen来编译TIR AST的。比如我们要编译CPU可以执行的代码，那么target参数可以选择'c'或者'llvm'。如果要编译成CUDA代码，那么参数设置为'cuda'或者“llvm”。然后tvm.build会根据target参数找已经注册的build函数，在TVM中使用TVM_REGISTER_GLOBAL宏注册build函数。例如：https://github.com/apache/tvm/blob/main/src/target/source/codegen_c_host.cc#L466 这里的 以及 https://github.com/apache/tvm/blob/main/src/target/opt/build_cuda_on.cc#L165 这里的 。我们这里以生成c代码为例介绍一下Codegen的原理。当target='c'时，tvm.build调用的是提前注册的target.build.c的全局函数（）。代码实现如下：
上面代码中的核心是CodeGenCHost这个类，这个类定义在 https://github.com/apache/tvm/blob/main/src/target/source/codegen_c_host.h#L40 。这个类又继承自CodegenC类，https://github.com/apache/tvm/blob/main/src/target/source/codegen_c.h#L59 。我们看一下CodegenC类的定义（简化了代码）：
CodegenC类的定义中重载了VisitExpr_和VisitStmt_两种函数分别处理TIR AST中的Expression节点（表达式） 和 Statement节点（语句）。Expression（表达式）中包含了常见的变量声明、运算、判断、函数调用，而 Statement（语句）中包含了控制流（if-else，Loop 等）、内存管理、赋值等操作。在https://github.com/apache/tvm/blob/main/src/target/source/codegen_c.cc 中对每一种AST节点进行对应的代码生成（定向到一个文件输出流中），比如：
其它类型的Codegen比如CUDA，LLVM IR等的原理都是一样的，只不过根据target的不同AST Node翻译的目标代码语句的语法又一点区别而已。
这一节为大家介绍2个有用的工具。第一个工具是《FFI Navigator: 跨语言调用跳转IDE插件》原文见：https://zhuanlan.zhihu.com/p/ 。这个工具的作用就是支持tvm项目中从c++和python之间的函数调用跳转以及类型object定义的跳转。除了tvm最近小伙伴还加入了对pytorch，mxnet，dgl的支持，有兴趣的同学也可以尝试一下。可以在vscode中直接配置使用。工具的github链接：https://github.com/tqchen/ffi-navigator/第二个工具是《Relay IR可视化》，应该也可以用到Relax上，这个工具来自一个TVM的PR（https://github.com/apache/tvm/pull/3259/files），这个PR提供了一个python/tvm/relay/visualize.py文件，我们可以稍加修改进行使用。修改后的脚本如下（注意要放到python/tvm/relay/visualize.py这个路径）：
然后我们在tvm_learn/tmp/tvm/python/tvm/relay/init.py把这个visualize注册一下，添加 。还需要安装一下pydot和graphviz可视化包：
最后我们就可以使用这个模块来做Relay IR的可视化了，还是以第6节的那个例子：
在当前目录会生成可视化的png图片，预览一下：
我们知道TIR AST是由一系列PrimExpr和RelayExpr（非PrimExpr）来表示的，它们都继承了TVM的Expr基础类。所以TVM针对TIR AST的遍历专门做了一个工具类ExprFunctor。而这可视化个工具就是通过继承ExprFunctor来遍历计算图并自定义可视化效果。
这篇文章就是对TVM的重新梳理，从前端到图优化以及后端，比较宏观的叙述了TVM整个架构，希望对入门TVM的读者有帮助。
参考
-
其它博客精选（TVM&MLIR 相关）
-
深度学习编译器 TVM 代码串讲
-
TVM Overview
-
TVM - Relay IR计算图可视化
-
TVM - 代码生成流程
-
TVM/VTA代码生成流程
-
tvm算子优化schedule（一）--CPU篇
-
tvm算子优化schedule（二）--GPU篇
-
TVM Runtime System 概述
-
TVM PackedFunc实现机制
-
向外借力：Pluto助力MLIR编译器的多面体优化
-
TVM 自底向上（一）：基本框架和概念
-
TVM 自底向上（二）：TIR 的概念和编译原理
-
TVM 自底向上（三）：TE 的概念和编译原理
-
TVM 自底向上（四）：TE/TIR Schedule 的原理
-
陈天奇 MLC课程
-
深度学习编译器学习笔记和实践体会
-
FFI Navigator: 跨语言调用跳转IDE插件
版权声明：
本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如若内容造成侵权、违法违规、事实不符，请将相关资料发送至xkadmin@xkablog.com进行投诉反馈，一经查实，立即处理！
转载请注明出处，原文链接：https://www.xkablog.com/cjjbc/51598.html