tvm学习笔记（五）：tvm工作原理

article/2025/10/3 19:46:24

一、总体流程：

TVM的工作流程：首先，将网络表示成统一的表示形式（Intermediate Representation），并进行一些可重用的图优化；然后，利用不同的后端生成对应设备代码，如图1所示。

图1 tvm 工作流程（摘自参考资料[1]）

首先，将不同的框架下的模型载入，并使用NNVM将模型转换成中间表示的计算图，并对图进行优化，如算子融合、减枝、图变换等；然后，TVM对张量运算进行优化，这里，TVM将代码的调度和计算分开（计算：定义需要进行的运算，调度：具体如何来进行运算）；最后，使用不同的后端，来生成对应设备代码，如图1所示，使用LLVM生成x86，ARM和Javescript/WASM系统代码，OpenCL、Metal和CUDA生成对应的GPU代码，通过这种中间堆栈（IR Stack）表示的方式，实现端到端的深度学习模型优化和部署，这种方式将实现op的复杂度转移到了编译规则的复杂度。

二、优化计算图

1、算子融合（operator Fusion）

算子融合，即将多个算子组合在一起放到同一个核中，通过算子融合的方式，我们不需要将中间结果保存到全局内存，进而减少执行所需要的时间，我们已知的算子融合分为四种，如图2所示（自己没翻译好，下面内容摘自参考资料[5]）：

图2 算子融合示意图（摘自参考资料[4]）

injective（单射性）：一到一的映射，如：add / sqrt / exp / sum 等操作算子（operator）；

reduction（简约）：多到少的映射，如：sum / max / min等操作操作算子（operator）；

complex-out-fusable：逐元素复用映射到输出，如：conv2d / bn / relu等操作算子（operator）；

opaque：不能被复用

像这种算子组合太多了，专门针对这些组合手写底层优化不太现实，这里就需要做一些自动代码生成。

2、数据布局变换：

当代计算架构中，从内存中载入数据的时间要远远大于进行一次浮点运算所耗费的时间，所以我们经常想要重复使用载入内存或寄存器中的数据。

首先我们看一下3x3的卷积操作，如图3所示：

图3 无tile的3x3卷积操作示意图（摘自参考资料[3]）

不采用tile的方式，每个线程载入一个3x3大小输入得到一个输出，16个线程需要进行16x9次数据载入，如果我们采用tile方式，如图4所示：

图4 有tile的数据载入（摘自参考资料[3]）

采用tile方式时，每个线程载入4x4大小输入得到2x2大小的输出，4个线程需要进行4x16次数据载入。

三、优化张量计算

张量表达语言（Tensor Expression Language）：直接描述每一个单元如何计算

这样的tensor表示（数学公式表达），可以涵盖几乎所有的高层算子，可以很容易做代码生成，因为对应的表达式已经确定了。然后就是将tensor expression映射到不同硬件上：

这里涉及到的问题有：算子张量化的问题、cache问题、数据类型问题（float32，float16,、int8）

解决方案： 将所有手工优化的可能（10亿级别的）总结起来，并将他们作为搜索空间的一部分，然后自动进行搜索，这里采用auto-tvm来自动进行搜素每个算子的最优实现。

为什么相信tvm的上限比手写优化做得更好？

如果是机器和人同时去解决一个问题的优化，人通过不断的去解决，可以做到比机器好一些，实际上，机器不一定要和人解决一样的问题，比如融合算子，其可能性太多，人可能没有力气去优化这些融合算子，机器通过去解决这些人没有解决的问题，进而达到更高的效率；反过来，当搜索空间越来越大，包含了人所有的搜索空间时，这时，哪怕直接和人的手写优化一一对应，机器也可以达到和人做的优化差不多，甚至更好都有可能。

总结起来就是：