论文阅读 [粗读]-APOLLO: AUTOMATIC PARTITION-BASED OPERATOR FUSION THROUGH LAYER BY LAYER OPTIMIZATION

今天来读一篇 MLSys 的文章，作者提出了 JIT 的 APOLLO 框架，可以同时考虑 memory- /compute-bound 的算子优化，比 XLA，TensorFlow 原生要快不少。

JIT (just-in-time)：和 AOT 相对，是指边运行边编译，不能支持某些编译优化方法。

作者来自郑州大学、华为、香港大学

Introduction

这部分说了目前 fusion 技术面临了三个问题：

现有的图编译器只考虑了 memory-bound op 的优化，对于 compute-bound op，只是分发到不同的 kernel，不做 fusion。因此 fusion 的搜索空间是不完全的
对于循环 loop fusion 来说：不同层次的编译器之间存才
stitching 融合技术，把没有依赖的不同算子放到同一个 kernel 来同时计算，提高硬件的并行性。这种技术对于自定义算子的支持非常差

本文提出 APOLLO 框架，把 sub-graph 进一步看成很多 micro-graph，在这个过程中同时考虑 memory- /compute-bound 的算子 fusion。划分的规则取消了上游图编译器的依赖，规则符合下游编译器的需要。可以符合一个多面体模型，同时也可以扩大规模。

APOLLO := Automatic Partition-based Operator fusion framework through Layer by Layer Optimization

首先，用一个 polyhedral loop fusion 算法让 micro-graph 内部的算子 fusion，并且让有依赖管的数据在更快的 local memory 里
接下里用一个 memory stitching 算法让 micro-graph 之间的算子 fusion，弥补前面 polyhedral 算法的劣势
最后，由一个 top layer 把前两步的结果组合在一起，更好的利用硬件的并行性。

APOLLO 可以把每一步中间生成的 IR 传给下一步。本文的贡献总结为：

扩大了 fusion 的搜索空间，考虑更多种 op
通过 operator-level 优化器反向反馈结果，可以进行更大规模的 fusion
同时考虑数据的局部性、并行性，效果更好
展现了 JIT compilation 的水平

ARCHITECTURE OF APOLLO

用一个例子来说明 apollo 结构，上图是一个计算图：

计算图可被分成上下两个子图 (相互是可以并行执行的)
Op5，op3 是计算瓶颈的节点
Op1,op6 是原子的节点
Op2,op4,op7 是可再分的节点

对于计算瓶颈节点 op5,op3, 已有编译器尝尝直接丢给 vendor library，并且把，每个子图 sub-graph 看成孤立的不同部分

Tensor compilers 面临上下游需求不同的问题，比如上游编译器希望 op7 编译成一个 kernel，但如果 op71 是 reduce op，下游编译器很难满足这个需求

另一个问题是当 batchsize 小的时候不能利用好硬件的并行性，现有方法经常把同一个子图中的比如 op1，op2 放到同一个 kernel，这和实际的 training scenarios 不符合。

本文提出的 APOLLO 框架如下：

首先把图分成不同子图 sub-graph，子图之间没有数据依赖 ( $F_{1}, F_{2}$ )
接下来，用一系列基于规则的方法把每个子图分出不同 micro-graph ( $G_{1}, G_{2} . . .$ )
对于每个 micro-graph，进行 fusion 生成 IR，用的方法是 polyhedral loop fusion heuristic
把上一层的结果传给下一层，layer 2 对结果进行 grouping，同时根据 reducing op 进行 fusion
最终 layer 3 接受上一层输入，进行子图之间的并行性的优化

polyhedral model

好多篇里都说了这个 polyhedral，到底是咋回事呢？这里拓展一下。

大概就是一个多重循环，里面的所有可能的取值 (i,j,k) 可以用一个矩阵来表示： $A_{3, 3} (\begin{matrix} i \\ j \\ k \end{matrix}) \geq (\begin{matrix} a \\ b \\ c \end{matrix})$ 之类的，总之这种方程在高位空间就是一个多面体，可以认为循环都是在多面体内执行。有什么作用呢？

循环间的数据依赖大概看成是多面体内内部有一些” 线 “，如果是水平的、垂直的，那么就不会有冲突，只有 “斜线 “会带来冲突，但如果我们对这个多面体进行仿射变换，就可以把斜线都变成直线了。
我们希望循环内可以更多利用 cache，同样利用仿射，可以把空间距离近的取值让时间距离也拉进。

大概就是这个原理，因为有了数学的表示，所有优化也就比较简单。

PARTITION PHASE

这一部分谈怎么变成 micro-graph

作者首先使用几条规则把图优化成符合多面体模型的图

接下来，要找出计算图中除了用户定义的 op 和 control flow ops 以外的 op，变成一些子图。

Opening Compound Operators

再接下来，作者把 compound op 拆解，比如说 logsoftmax： $S (t_{i}) = t_{i} - \ln \sum_{j = 1}^{N} e^{t_{j}}$

其对应蓝色的两个 op：减法，和一个循环计算的 op，这两个 op 有数据依赖。

如果把 op3 拆开分析，就能把 op31 和前面的 op1，op2 融合到一起，，增大 fusion 的搜索空间，这就是拆解 compound op 的好处。

Aggregating Primitive Operators

这一部分要分解出很多 micro-graph

首先把 primitive op 看做 micro-graph，然后用 aggregation rules 去逐步扩大 micro-graph。规则大概是下面这样定义的几条：

规则不用覆盖所有 op，因为有些 op 本身就不需要 fuse。

FUSION PHASE

fusion 的阶段符合自底向上的顺序

Layer I: Polyhedral Loop Fusion

这一部分讲了 micro-graph 内部是如何做 fusion 的。方法很多，很复杂，这里就不细讲了，感觉只有看代码才能深入了解，大概逻辑就是像下图这样：

Layer II: Memory Stitching

这一部分是说如何把 fusion 后的 micro-graph 再 fusion，再组装回 sub-graph

大概思路就是如果一个 micro-graph 是一个 reduce 型的 micro-graph，那么它也许可以和后面你的 micro-graph 的中间变量排布在一起，用到更快的 local memory。

Layer III: Parallelism Stitching

这一部分考虑子图之间的并行性，作者观察到并行性大多存在于：

branches of a multi-head/-tail op
两个子图没有数据的依赖关系。这种可以合并成上面那种

作者大概考虑了怎么优化 multi-head/-tail op，作者用一个 cost 模型逐步的选取并行效果最好的节点加入。

PUTTING IT ALL TOGETHER

Auto-Tuning: APOLLO 解决了 scalability 的问题，并且由于 peicewise 编译速度不慢，因此支持 auto-tuning
Piecewise Compilation：每一层的各个节点之间的编译是并行的，只需要在层交换的时候做一次不同
Code Generation：生成 A100 cuda 代码。

Result

大概就是说 APOLLO 处的代码效果更好

mindspore：是华为的开源 AI 框架，APOLLO 可以试做是在优化 mindspore 的编译

结论

polyhedral 优化的扩展性是一个 NP 问题，因此很难大范围的应用。本文通过把图拆解成子图，每一部分分别做，一定程度上解决了这个问题。
目前的规则很简单，不一定适合其他的、未来的算子
本文使用的 costmodel 很简单，只是单纯的在并行和同步之间做权衡