APOLLO: AUTOMATIC PARTITION-BASED OPERATOR FUSION THROUGHLAYER BY LAYER OPTIMIZATION

文章标题: APOLLO: 通过逐层优化的基于自动分区的算子融合
作者/机构: Jie Zhao 1, Xiong Gao 2, Ruijie Xia 2, Zhaochuang Zhang 2, Deshi Chen 2, Lei Chen 3, Renwei Zhang 2, Zhen Geng 2, Bin Cheng 2, Xuefeng Jin 2

A1 主要贡献

本文介绍了一个名为APOLLO（Automatic Partition-based Operator fusion framework through Layer by Layer Optimization）的即时编译（JIT）框架，旨在解决深度神经网络（DNN）中的算子融合问题。

核心问题与挑战:
1. 融合搜索空间不完整（挑战1）: 许多图编译器（如Google XLA, nGraph）只对内存密集型算子（memory-bound ops）进行节点分组，而忽略了计算密集型算子（compute-bound ops），如卷积（conv）和矩阵乘法（matmul），因为它们假设这些算子的内核启动开销与执行时间相比微不足道。这导致融合机会的错失和融合空间的受限。
2. 下游循环优化器的可伸缩性问题（挑战2）: 现有的张量编译器（如TVM, Tiramisu）采用“节点分组后进行循环融合”的常规流程，这种自上而下的强制方式给下游的循环优化器带来了巨大压力，挑战了循环融合启发式算法的可伸缩性。例如，上游图编译器生成的子图可能强制要求下游使用一个大的循环移位因子，这会显著增加编译时间。
3. 对自定义算子和训练场景支持不足（挑战3）: 现有的利用独立算子间并行性的方法（如TASO, Rammer）通常依赖于手动调度模板，采用预先编译（ahead of time）的方式，这使得它们难以支持训练场景中的自定义算子。

研究目标与创新点:
APOLLO旨在通过一个自动化的、分层的优化框架来应对上述挑战。其核心贡献如下：
* 扩展融合搜索空间: APOLLO同时考虑内存密集型和计算密集型算子进行融合，通过打破算子边界（op boundaries），生成了更多原本被阻碍的有价值的跨层调度方案。
* 解决融合的可伸缩性问题: APOLLO允许来自下游算子级优化器的反向反馈。通过一个基于规则的分区算法，将子图划分为对下游多面体优化器友好的“微图”（micro-graphs），从而解决了可伸缩性问题，实现了一个全自动的融合框架。
* 同时利用数据局部性和并行性: APOLLO通过一个多层融合策略，不仅优化了数据局部性（通过循环融合和内存拼接），还利用了独立算子间的并行性（通过并行拼接），从而生成了比现有技术更高效的代码。
* 高效的JIT编译: APOLLO采用分段编译（piecewise compilation）策略，显著降低了编译开销，使其能够胜任即时编译器的角色，并同时适用于训练和推理负载。

实验结果表明，在训练负载上，APOLLO在单GPU上的性能优于TensorFlow和XLA 1.86倍和1.37倍，在多GPU上优于1.96倍和1.18倍。在特定领域的加速器上，APOLLO比厂商提供的DNN框架性能提升了19.7%。

A3 框架概览

DNN计算图示例。图1展示了一个典型的DNN计算图，其中包含两种不相交的子图。op3和op5是计算密集型算子。op2、op4和op7是复合算子（compound ops），而op1和op6是原始算子（primitive ops）。复合算子是由多个原始算子组成的函数，例如SoftMax。图中的op2由两个蓝色原始算子（op21, op22）组成，op4由两个灰色原始算子（op41, op42）组成，op7由三个黄色原始算子（op71, op72, op73）组成。

图1：一个说明性的DNN计算图。
图1：一个说明性的DNN计算图。

现有图编译器的局限性。许多图编译器【【17, Xla: Optimizing compiler for machine learning, 2017, Google】；【55, Dlvm: A modern compiler infrastructure for deep learning systems, 2018】】不考虑融合op3或op5这类计算密集型算子，而是依赖厂商调优的库来处理它们，这自然地将每个子图分割成多个部分。虽然可以使用动态规划策略【【15, Ios: Inter-operator scheduler for cnn acceleration, 2021, MLSys】】来评估每种融合可能性，但随着子图复杂度的增加，其编译开销可能成为一个问题。

现有张量编译器的可伸缩性问题。张量编译器【【7, Tvm: An automated end-to-end optimizing compiler for deep learning, 2018, USENIX Conference on Operating Systems Design and Implementation】；【8, Learning to optimize tensor programs, 2018, NIPS】；【51, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated gpu kernels, automatically, 2019, ACM Trans. Archit. Code Optim.】；【5, Tiramisu: A polyhedral compiler for expressing fast and portable code, 2019, CGO】】将融合与分块（tiling）【【23, Supernode partitioning, 1988, POPL'88】】结合起来，但它们的融合启发式算法受到上游图编译器施加的约束，从而遭受可伸缩性问题【【38, Revisiting loop fusion in the polyhedral framework, 2014, PPoPP’14】；【60, Optimizing the memory hierarchy by compositing automatic transformations on computations and data, 2020, MICRO53】】。例如，上游图编译器期望为op7生成单个核函数实现，但如果op71是一个归约（reduction）算子，下游的循环优化器可能无法满足此要求。

现有并行性优化方法的不足。近期的一些工作【【25, Taso: Optimizing deep learning computation with automatic generation of graph substitutions, 2019, SOSP’19】；【36, Rammer: Enabling holistic deep learning compiler optimizations with rtasks, 2020, OSDI 20】；【63, Fusionstitching: Boosting memory intensive computations for deep learning workloads, 2020】】通过将多个独立分支（如op1和op2）打包到单个核函数中来利用硬件并行性，但它们很少考虑训练场景或专用芯片。

APOLLO框架设计。为了解决上述问题，我们设计了APOLLO框架，其架构如图2所示。其分区阶段（§3）首先提取最大的子图集合P，然后将其分割成m个独立的子图$F_x$（$1 \le x \le m$）。与XLA不同，APOLLO在执行融合时会考虑计算密集型算子。接着，这m个子图被一个基于规则的算法分割成n个微图（micro-graphs）$G_y$（$1 \le y \le n$），该算法充分考虑了融合阶段的可伸缩性挑战，这与TVM的图引擎不同，后者在分区时没有意识到其循环优化器的需求，可能导致可伸缩性问题。

图2：APOLLO的架构。
图2：APOLLO的架构。

三层融合阶段。APOLLO的融合阶段（§4）包含三个层次。第一层（§4.1）对每个$G_y$执行循环融合。它通过在多面体模型【【53, Scheduling for ppcg, 2017】】中使用多面体循环融合启发式算法，为每个$G_y$始终生成单个融合组，从而解决了可伸缩性问题。第二层（§4.2）通过聚合第一层的输出来实现节点分组，以实现归约算子与其后续算子之间的融合，这是先前方法【【61, Akg: Automatic kernel generation for neural processing units using polyhedral transformations, 2021, PLDI 2021】；【51, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated gpu kernels, automatically, 2019, ACM Trans. Archit. Code Optim.】】未研究过的。第三层（§4.3）负责利用独立算子间的并行性。与TASO【【25, Taso: Optimizing deep learning computation with automatic generation of graph substitutions, 2019, SOSP’19】】不同，我们考虑了不同类型算子之间的并行拼接，并支持为GPU和专用加速器生成代码。APOLLO最终生成一个或多个由我们的自动调优器优化的核函数，并且第一层和第二层的执行是并行的，以适应JIT编译。

A2 方法细节

3 分区阶段

图预处理优化。在DNN计算图进入分区阶段之前，我们会执行一些预处理优化来简化图。这些优化包括：
* ♣ 代数简化：利用结合律、交换律和分配律，类似于函数内联，也被Google【【17, Xla: Optimizing compiler for machine learning, 2017, Google】】采用。
* ♣ 数据流优化：执行公共子表达式消除和常量折叠，nGraph【【12, Intel ngraph: An intermediate representation, compiler, and executor for deep learning, 2018】】也考虑了这一点。
* ♣ 控制流优化：通过消除无用的生产者算子和分支来简化图，这可能是代数简化的结果。
* ♣ 数据布局转换：根据需求改变张量的存储方式，以保持生产者和消费者数据空间之间的维度对齐。

这些准备工作简化了DNN计算图，减少了编译开销，并满足了多面体模型的“静态仿射控制”要求【【52, Isl: An integer set library for the polyhedral model, 2010, ICMS’10】】，尽管可能不完全。

3.1 提取子图集群

筛选符合条件的节点。APOLLO首先从计算图中提取出符合条件的节点集合。有两类算子不被APOLLO考虑。第一类是用户自定义和/或具有复杂计算逻辑的特殊算子，一个典型的例子是用于训练语音识别的all-reduce【【3, Deep speech 2 : End-to-end speech recognition in english and mandarin, 2016, PMLR】】，部署这类算子最合适的方法是封装一个高度优化的库，如Cho等人【【11, Blueconnect: Decomposing all-reduce for deep learning on heterogeneous network hierarchy, 2019, MLSys】】的工作。第二类是控制流算子，如TensorFlow的RefSwitch，也应被排除。剩下的节点类型可以构成一组与其他节点不相连的子图P。我们使用融合来最小化每个子图内算子之间的生产者-消费者距离。

3.2 打开复合算子

复合算子的融合挑战。虽然可以在每个子图$F_x$内部应用基于依赖的融合算法，但算子之间的依赖模式通常很复杂，导致无法为$F_x$甚至单个复合算子生成单一的核函数。

激活函数带来的问题。激活函数的使用是主要原因之一。一个激活函数通常涉及许多算术运算，可以分解为多个原始算子。我们以SoftMax的一个稳定变体LogSoftMax为例来说明这个问题。其公式可以表示为：
$S(t_i) = t_i - \ln\left(\sum_{j=1}^{N} e^{t_j}\right)$
其中$t_i$是长度为N的输入向量的第i个元素，e是自然常数。公式(1)需要两个操作：一个计算所有向量元素归约后的对数，另一个执行减法。这两个操作必须通过循环分块分解成多个并行执行单元。分块后的减法必须等待归约操作所有并行执行的分块完成，这阻止了这两个分块操作之间的融合。因此，一个结合了循环分块的循环融合启发式算法无法为这个复合算子生成单个核函数。因此，LogSoftMax不是一个可融合的候选者。

算子边界的限制。在图1中，复合算子用蓝色圆角矩形表示。这种表示方式形成了一个我们称之为“算子边界”的约束，它限制了循环融合启发式算法可探索的融合模式。算子边界通过将复合算子内部的原始算子与外部的原始算子隔离开来，导致了次优的融合模式。我们分解每个复合算子以消除所有算子边界。例如，在图2的每个$F_x$中，这些边界已经被移除。

示例1。假设图3(a)中的op3是一个由归约算子op31和减法算子op32组成的LogSoftMax函数，其前面有两个原始算子。当存在算子边界时，融合模式如图3(b)所示，每个融合组用红色块表示。如果移除该边界，则可能出现如图3(c)所示的另一种融合模式。

图3：算子边界的影响。(a) LogSoftMax；以及当算子边界(b)存在或(c)不存在时的融合模式。
图3：算子边界的影响。(a) LogSoftMax；以及当算子边界(b)存在或(c)不存在时的融合模式。

3.3 聚合原始算子

自底向上的微图构建。最后一步是将每个$F_x$划分为微图。我们不直接分割$F_x$，而是为$F_x$中的每个原始算子实例化一个独立的微图G，然后使用模仿循环融合启发式算法所能容忍的融合行为的聚合规则，逐步合并这些不相交的微图。这种方法使我们能够明确列出所有可以被多面体启发式算法解决的融合模式，从而避免了可能导致第一层（§4.1）出现可伸缩性问题的未知场景，代价是可能在第一层失去一些融合机会。这种副作用要么在第二层（§4.2）被克服，要么通过离线学习来更新分区阶段的规则。

算子类型定义。我们的经验表明，可伸缩性问题与待融合的原始算子类型密切相关。一个原始算子以p个多维张量为输入，并将结果输出到另一个张量。换言之，每个原始算子维护p个输入数据空间$I_k \in V_Z^d$（$1 \le k \le p$）和一个输出数据空间$O \in V_Z^d$，其中$V_Z^d$（$d \ge 1$）表示正整数向量空间。可以在第k个输入和输出张量之间建立一个数据流关系$I_k \to O$；或者第k个输入张量可以被丢弃。

数据流关系分类。可以安全地假设每个$I_k$与O具有相同的维度d，因为我们可以通过张量广播来转换两个张量以使其形状兼容。因此，一个数据流关系可以写成d个一维函数的合取：
$$\bigwedge_{1<l<d} f_l := I_k^l \to O^l \quad (I_k^l, O^l \in V_{\mathbb{Z}})$$<br /> 其中整数$e_i, e_o \ge 1$表示$I_k^l$和$O^l$的循环次数。$f_l$被假定为双射（bijective, $e_i=e_o$）、单射（injective, $e_i=1 \land e_o \neq 1$）、满射（surjective, $e_i \neq 1 \land e_o = 1$）或一般函数。因此，一个数据流关系$I_k \to O$被认为是：
1. 双射关系：如果每个$f_l$都是双射的。
2. 单射关系：如果d个$f_l$中有u个（$u \ge 1$）是单射的，其余都是双射的。
3. 满射关系：如果其d个一维函数中有u个（$u \ge 1$）是满射的，其余d-u个$f_l$是双射或单射函数。
4. 一般关系：如果至少有一个$f_l$是一般函数。

原始算子类型。给定一个原始算子，我们可以使用以下定义来确定其类型：
* 定义1：若每个$I_k \to O$（$1 \le k \le p$）都是双射数据流关系，则该算子是按元素（element-wise）算子。
* 定义2：若p个数据流关系中有u个（$u \ge 1$）是单射的，其余都是双射函数，则该算子被视为广播（broadcast）算子。
* 定义3：若其p个数据流关系中至少有一个是满射数据流关系，其余是双射或单射关系，则该算子类型为归约（reduction）。
* 定义4：若算子至少有一个一般数据流关系，则称之为不透明（opaque）算子。

聚合规则。我们的定义将重塑（reshaping）操作、(批量)矩阵乘法和卷积归类为不透明算子。一旦确定了原始算子的类型，就可以定义两个微图聚合结果的类型。我们在表1中总结了合并两个微图（$G_p$, $G_c$）的规则。

表1：聚合规则。$G_p$和$G_c$具有生产者-消费者关系；$G_a$是合并后的微图。
表1：聚合规则。$G_p$和$G_c$具有生产者-消费者关系；$G_a$是合并后的微图。

规则的不完备性与反馈机制。这些规则不需要覆盖所有的算子组合模式，因为某些算子对不应该被融合。例如，我们最初为重塑算子后跟按元素或广播算子定义了一个聚合规则，但发现多面体循环优化器无法在合理时间内完成调度过程。因此，根据下游优化器的反馈（§4.1中介绍），该规则从表1中移除。这种移除导致表1中的规则集不完整，这在实践中可能会生成更多的子图，但避免了融合的可伸缩性问题。这对JIT编译工具至关重要。

规则应用。最右侧列指示聚合结果的类型，而前面两列枚举了待聚合的微图对的每种可能的类型组合。这种定义使得这些规则可以递归使用，并在遇到未定义的组合模式时保证算法终止。一个微图在用原始算子初始化时只有一个输出张量，但在合并过程中可能会创建多个输出张量。附录A中的算法1描述了如何聚合微图。

示例2。算法1通过应用规则①和④获得了图3(c)的第一个融合组。由单个算子组成的第二个融合组与其余部分分开，因为表1中没有定义将归约算子与后续算子聚合的规则。

4 融合阶段

自底向上融合策略。融合阶段以自底向上的方式执行融合。第一层还执行循环分块。选择的分块大小决定了每个微图占用的内存量，这向上层的第二层和第三层融合启发式算法暴露了已使用的快速内存量。否则，后两层将无法感知这些信息。这种自底向上的融合策略在支持训练场景中的自定义算子方面更具灵活性：当分区阶段产生的任何$G_y$不可接受或无法生成单个融合组时，融合阶段可以向分区阶段反馈。

4.1 第一层：多面体循环融合

使用AKG进行循环融合和分块。当一个$G_y$被下沉到这一层时，它已经被转换成一系列循环嵌套。我们使用一个多面体优化器AKG（Automatic Kernel Generator）【【61, Akg: Automatic kernel generation for neural processing units using polyhedral transformations, 2021, PLDI 2021】】来执行循环融合和分块。

多面体模型的优势。多面体模型通过求解一个整数线性规划（ILP）问题【【6, A practical automatic polyhedral parallelizer and locality optimizer, 2008, PLDI’08】】来确定每个循环嵌套的并行性和可分块性，并引入辅助变换（如循环交换、移位和缩放）以确保必要时循环嵌套之间的对齐。在多面体模型中，这些循环变换用仿射关系【【53, Scheduling for ppcg, 2017】】表示，其中一些变换在手动调度方法中是无法获得的，但它们有时对于最小化生产者-消费者距离至关重要。即使可以补充更多的调度原语，编写手动调度模板也可能无法最小化循环嵌套之间的生产者-消费者距离，因为辅助循环变换必须与循环分块和融合一起使用【【59, Flextended tiles: A flexible extension of overlapped tiles for polyhedral compilation, 2019, ACM Trans. Archit. Code Optim.】；【24, An effective fusion and tile size model for polymage, 2020, ACM Trans. Program. Lang. Syst.】】，这需要对建模的变换进行系统性的组合。列表1展示了一个需要系统性组合循环缩放和融合的例子；多面体模型可以将其转换为列表2所示的形式。

for i in [0,M) for j in [0,N) a(i,j)=a(i,j)+bias; //S1   
for i in [0,M/2) for j in [0,N/2) pool(i,j)=max(a(2*i,2*j), a(2*i,2*j+1), a(2*i+1,2*j), a(2*i+1,2*j+1)); //S2

列表1：原始循环嵌套。

for i in [0,M) for j in [0,N){ 
    a(i,j)=a(i,j)+bias; 
    if(i+1) mod 2 == 0 and (j+1) mod 2 == 0 
        pool((i-1)/2,(j-1)/2)= max(a(i-1,j-1),a(i,j-1), a(i-1,j),a(i,j)); //S2 
}

列表2：融合后的代码。

解决AKG的可伸缩性问题。不幸的是，AKG仍然遵循TVM制定的常规变换编排，因为它继承了后者的图引擎。TVM用于执行图级节点分组的规则不允许来自下游AKG的反馈，其循环融合启发式算法有时可能会生成非常低效的融合模式，或者无法在合理的时间内终止。其他多面体编译器【【51, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated gpu kernels, automatically, 2019, ACM Trans. Archit. Code Optim.】；【54, Polyhedral parallel code generation for cuda, 2013, ACM Trans. Archit. Code Optim.】】也面临类似的可伸缩性问题。

示例3。仍然考虑图3(a)中的例子。如果没有我们的聚合规则，所有四个算子都将被TVM的图引擎传递给AKG，后者会为这四个算子构建一个单一的融合组。如§3.2所述，需要一个非常大的循环移位因子来保证分块算子的执行顺序。我们观察到，不仅编译时间增加，而且借助大循环移位因子获得的单一融合组在实践中也降低了执行性能。

聚合规则的作用。我们的聚合规则通过始终生成其循环嵌套组合可预测的微图来克服这一弱点。因此，多面体模型的融合启发式算法永远不会受到可伸缩性问题的挑战。一旦多面体调度器无法在合理时间内完成，其算子组合模式将被反馈给开发人员，然后开发人员可以利用这些信息来更新表1中的聚合规则。

归约算子的优化。我们使用PANAMERA【【62, Parallelizing neural network models effectively on gpu by implementing reductions atomically, 2022, PACT’22】】来优化归约算子的融合，它可以有效地将归约算子与其前面的按元素/广播算子融合。这得益于一个维度扁平化优化，该优化总是将归约算子的循环嵌套合并为三种规范形式，从而简化了调度过程并降低了多面体复杂性。手动调度方法或厂商库【【10, cudnn: Efficient primitives for deep learning, 2014】】没有考虑这种优化，但它使APOLLO能够在实践中找到更好的融合模式。

示例4。列表3展示了一个归约算子前面有一个按元素算子。APOLLO首先通过循环交换将其转换为三种规范形式之一（称为y-reduce），如列表4所示，这使得APOLLO能够获得列表5中的融合结果。张量下标由多面体模型自动更新。

for i in [0,M) and j in [0,N) and k in [0,P) and l in [0,Q) 
    a(i,j,k,l) = a(i,j,k,l) + bias   
for i in [0,M) and j in [0,N) and k in [0,P) and l in [0,Q) 
    b(i,k) += a(i,j,k,l)

列表3：原始循环嵌套。

for x in [0,M*P) and y in [0,N*Q) 
    a(x/P,y/Q,x%P,y%Q) = a(x/P,y/Q,x%P,y%Q) + bias 
for x in [0,M*P) and y in [0,N*Q) 
    b(x/P,x%P) += a(x/P,y/Q,x%P,y%Q)

列表4：规范形式。

for x in [0,M*P) and y in [0,N*Q){ 
    a(x/P,y/Q,x%P,y%Q) = ... 
    b(x/P,x%P) += ... 
}

列表5：融合结果。

传递到第二层。一旦每个微图被多面体模型融合，我们就可以生成优化的中间表示（IR），并将其传递给第二层，以探索微图之间的融合可能性。

4.2 第二层：内存拼接

补充聚合规则。我们定义了补充的聚合规则来利用微图之间的融合可能性。这一层考虑的微图类型组合总结在表2中。它允许递归地应用其每个规则，贪婪地搜索微图之间的融合可能性。这些规则只考虑以归约算子开始的微图类型组合，从而在遇到其他情况时保证贪婪融合方法的终止。

表2：补充聚合规则。
表2：补充聚合规则。

内存拼接的挑战。尽管如此，由于DNN模型中归约场景的多样性，为第二层设计融合算法仍然不简单。归约可以沿着算子循环嵌套的任何一个或多个维度进行。例如，列表3所示的循环嵌套中，归约是沿着第二和第四维度进行的。此外，第一层的循环分块将归约微图的各个循环维度绑定到不同的硬件维度。第二层必须检查并确保循环维度和硬件参数都匹配。

循环维度匹配。两个微图的循环维度之间的匹配由PANAMERA生成的三种规范形式来保证，如§4.1所述。APOLLO总是将一个归约微图转换为all-reduce形式（其中归约微图的循环嵌套被扁平化为一个单一的归约循环），或x-/y-reduce形式（其中融合微图的多维循环嵌套已被合并为二维循环嵌套，并且归约仅沿着外/内循环维度进行）。如果两个微图已被扁平化为相同的规范形式，我们就可以融合它们。

硬件参数匹配。硬件设置在第一层指定，并与特定循环变换的参数密切相关。例如，循环优化器找到的分块大小对目标为GPU时的网格/块维度配置有显著影响。我们主要关注前一层找到的调优分块大小，它决定了连接两个微图的张量将被分配到哪里。当（$G_p$, $G_c$）的分块大小和硬件参数设置都相同时，它们可以被融合。附录A中的算法2描述了内存拼接方法。

示例5。在图5中，每个用红色虚线框表示的微图都是从一个真实的DNN工作负载中提取的。相同颜色的原始算子属于同一个微图。算法2使用规则⑦融合这些微图，每个蓝色边携带的中间张量可以分配到本地内存上。

4.3 第三层：并行拼接

利用算子间并行性。第一层和第二层没有考虑独立算子/微图之间的内在并行性。这也是许多现有张量编译器【【7, Tvm: An automated end-to-end optimizing compiler for deep learning, 2018, USENIX Conference on Operating Systems Design and Implementation】；【51, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated gpu kernels, automatically, 2019, ACM Trans. Archit. Code Optim.】】的弱点，因此它们通常用于为单个设备生成代码。

寻找并行候选。确定每对算子之间的独立性并非易事，可能导致组合爆炸问题。我们观察到这种并行性主要存在于多头/多尾算子的分支之间。另一个应考虑的场景是图2中$F_1$和$F_2$之间的独立性，这两个子图之间没有共同的生产者/消费者。可以通过引入一个虚拟的共同消费者/生产者将这种情况转换为多头/多尾情况，而不会违反语义。因此，我们接下来只讨论多头/多尾场景的处理。

遍历算法。为了在一个分支上挑选出可以与另一分支上的算子并行执行的合适算子，我们沿着一个分支向后/向前遍历，直到到达另一个多头/多尾算子。在遍历过程中不考虑共同的头/尾。遍历方向是相对于算子之间生产者-消费者依赖关系的方向。一个多头算子如果是向后遍历的终点，则被认为可以与其他分支上的算子并行执行；如果它出现在向前路径的末端，则被排除。相反，一个多尾终点在向前路径中被收集，或在向后遍历中被排除。此外，计算密集型算子也应被排除，因为其庞大的数据量通常会耗尽硬件资源。

示例6。图6(b)右下角的mul算子的两个分支首先被遍历。它们被认为是可并行的，用蓝色框表示。然后通过引入一个虚拟尾部来评估这两个独立分支，可并行的部分用红色框表示。

代价模型。APOLLO现在可以尝试组合收集到的可并行候选来利用算子间的并行性。同一分支上的算子不能并行执行，因为它们相互依赖。我们重复地从每个分支中选择一个算子，并考虑它们的组合可能性。理想情况下，所有这些提取的算子都可以同时执行，但并行化受到可用硬件资源的限制。将更多数量的可并行算子调度到有限的硬件资源（如GPU流式多处理器）上并不意味着更高的性能。我们使用一个代价模型来评估并行化的潜在性能增益：
$$gain = \sum_{op=m}^{k} cost_{op} - \max_{m \leq op \leq k} (cost_{op})$$
其中k是被评估的算子数量，m是这k个已排序算子的起始编号。$cost_{op}$可以看作是一个算子的执行时间，可以根据其张量的形状和分配的硬件资源来估计。求和部分计算所有被评估算子的顺序执行时间。最大值部分表示这些算子中的最大执行时间，它决定了并行化后的执行时间。它们的差值就是性能增益。并行拼接在附录A的算法3中有正式描述。

示例7。假设我们有五个分支，从每个分支提取的算子按成本降序排列为{op3, op1, op5, op2, op4}。算法3尝试组合所有算子（m=1, k=5），但性能增益不是正的。它更新k（在第8行）并继续。我们假设它仍然无法组合{op3, op1, op5, op2}，并使用{op3, op1, op5}重试。如果算法3成功，则将评估{op2, op4}，最终的融合组可能是{op3, op1, op5}和{op2, op4}。

平台无关性。请注意，第三层执行的优化是平台无关的，因为多面体模型可以管理传统和新兴的多向内存系统。相反，现有的框架【【25, Taso: Optimizing deep learning computation with automatic generation of graph substitutions, 2019, SOSP’19】；【63, Fusionstitching: Boosting memory intensive computations for deep learning workloads, 2020】】仅在典型平台上进行评估，因为它们严重依赖于只能管理GPU或类似目标的传统内存层次金字塔的TVM。

5 整体实现

适用性目标。为了使APOLLO适用于训练和推理场景，我们必须在提供有前景的性能的同时，保证轻量级的编译开销。为实现这一目标，我们通过以下步骤来完善APOLLO。

自动调优（Auto-Tuning）。APOLLO循环变换的多功能性在于其多面体引擎。尽管AKG中仍然不可避免地存在较长的编译时间，但APOLLO解决了其可伸缩性问题，因此总能在几秒钟（最坏情况下为几分钟）内产生一个有效的解决方案，该方案通过稍后介绍的分段编译策略进一步优化。我们还提供了一个自动调优策略。在调优过程中，APOLLO可能会捕获到被代价模型(3)阻止并行的算子组合模式。开发人员可以利用这些信息来设计新的（补充）规则。

分段编译（Piecewise Compilation）。分段编译实现了多个微图/子图的更快编译。当APOLLO的复杂性成为问题时，它被设计用来加速编译开销。分段编译沿着图2中所示的红色/紫色箭头进行。它通过让多面体模型在更短的时间内解决ILP问题来进一步减少编译开销。每个编译过程都可以就特定目标的优化权衡做出自己的决定。分段编译的同步发生在图2中的绿色点划线上。

代码生成（Code Generation）。APOLLO在NVIDIA V100 GPU上生成CUDA代码，在华为昇腾910（Ascend 910）【【35, Ascend: a scalable and unified architecture for ubiquitous deep neural network computing : Industry track paper, 2021, HPCA】】上生成所谓的可执行CCE代码。在多面体模型下生成CUDA代码已有近十年的历史：可以借鉴PPCG【【54, Polyhedral parallel code generation for cuda, 2013, ACM Trans. Archit. Code Optim.】】和TC【【51, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated gpu kernels, automatically, 2019, ACM Trans. Archit. Code Optim.】】的后端来支持GPU设备。通过isl（整数集库）【【52, Isl: An integer set library for the polyhedral model, 2010, ICMS’10】】，可以保证对GPU上传统的内存层次金字塔和Ascend 910芯片上新兴的多级、多向内存层次进行自动内存管理。我们还集成了一些补充性的代码优化策略，包括利用SIMD硬件内在函数和对发出指令之间同步的底层优化，这些都超出了多面体模型的能力范围，从而进一步提高了生成代码的性能。

示例8。我们最后回顾图1来说明APOLLO的端到端编译过程。分区阶段首先从计算图中提取P，并获得两个子图$F_1$和$F_2$。$F_1$被分割成三个微图，其中第三个微图$G_3$包含一个被考虑用于融合的计算密集型算子，这是图编译器未曾考虑的。类似地，$F_2$也被分区阶段分成了两个微图。由于op71是一个原始归约算子，它必须从其原始的复合算子op7中分离出来。我们在表1中的规则允许它被认为可与其前面的按元素算子（op6）融合。结果，我们得到了$G_4$和$G_5$。然后，融合阶段的第一层用于在短时间内搜索每个$G_y$内的融合计划。输出随后被传递到第二层，该层分别在$G_1/G_2$和$G_3$之间，以及$G_4$和$G_5$之间执行内存拼接。最后，第三层利用$G_1$和$G_3$之间的并行性，同时也研究了$F_1$和$F_2$之间的并行性。编译期间发生的自动调优过程会在需要时触发反馈。

A4 实验环境与结果

实验环境

  • 实现：APOLLO在MindSpore【【22, Mindspore, 2020, Huawei】】框架内实现，包含14.4k行C++代码和2k行Python代码。
  • 硬件配置：
    • GPU平台：NVIDIA V100 GPU。
    • 专用加速器平台：华为Ascend 910芯片。
  • 软件配置：
    • GPU代码：使用CUDA Toolkit 10.1编译，开启-O3优化。
    • Ascend代码：使用Ascend芯片的原生编译器编译。
  • 模型与数据集：
    • 主要对比模型：BERT【【14, BERT: Pre-training of deep bidirectional transformers for language understanding, 2019, ACL】】, Transformer【【52, Attention is all you need, 2017, NIPS】】, Wide&Deep【【9, Wide & deep learning for recommender systems, 2016, DLRS】】, Yolo-v3的DarkNet部分【【47, You only look once: Unified, real-time object detection, 2016, CVPR】】, 和DeepFM【【18, Deepfm: A factorization-machine based neural network for ctr prediction, 2017, IJCAI-17】】。
    • 模型精度：BERT使用混合精度（FP16和FP32），其余四个工作负载使用单精度（FP32）进行训练。
    • 对比基线：TensorFlow 1.15（Wide&Deep使用2.0），XLA，以及原生的MindSpore。
    • 其他模型：MindSpore模型库中的多个模型（图7），用于验证通用性；PanGu-α【【57, Pangu-α: Large-scale autoregressive pretrained chinese language models with auto-parallel computation, 2021】】用于Ascend芯片评估；EPP-MVSNet用于推理场景评估。

实验结果

子图案例研究 (GPU)
* 可伸缩性问题解决: 如图4所示，一个包含reshape算子的子图导致AKG的多面体优化器在30分钟内无法找到融合方案。APOLLO通过其分区规则，将reshape视为不透明算子并且没有定义相应的聚合规则，从而避免了将这个难题交给融合阶段，并在1秒内返回了一个可行的分区方案（图中红色框）。这证明了APOLLO的反馈机制和分区策略能有效解决可伸缩性问题。
图4：一个关于可伸缩性问题的案例研究。
* 第二层（内存拼接）效果: 如图5所示，多个由归约算子分隔的微图（红色虚线框）通过算法2进行融合。表3的数据显示，第二层通过将中间张量提升到GPU共享内存和寄存器，为图5中的子图带来了平均69%的性能提升。
图5：内存拼接效果示例。
* 第三层（并行拼接）效果: 如图6所示，多个独立分支（用相同颜色的框标出）被识别并并行化。表3的数据显示，第三层为图6中的子图带来了平均92%的加速。并行拼接带来的收益随着可并行候选数量的增加而增长。
图6：并行拼接效果示例。
表3：子图的执行时间（单位：µs）。

单GPU训练性能 (表4)
* APOLLO vs. MindSpore: APOLLO平均性能提升2.23倍。第一层（循环融合）贡献最大（1.91倍），第二、三层进一步提升性能。
* APOLLO vs. TensorFlow/XLA: 在排除DeepFM（其在不同框架中实现差异巨大）后，APOLLO使MindSpore的性能平均比TensorFlow高1.86倍，比XLA高1.37倍。这主要得益于APOLLO有效地融合了XLA未考虑的计算密集型算子，且没有引入可伸缩性问题。
表4：单GPU上的吞吐量。

MindSpore模型库通用性验证 (图7)
* APOLLO在MindSpore模型库的多个模型上进行了测试，涵盖计算机视觉、语音识别、NLP等领域，平均性能比原生MindSpore提升了29.6%，展示了其广泛的适用性。对于MobileNet-v2和LSTM等模型，由于其结构特性（可融合算子较少），性能提升有限。
图7：MindSpore模型库的执行时间（y轴：对数尺度时间，单位ms；越低越好）。

多GPU训练性能 (表5)
* 多GPU上的性能趋势与单GPU一致。APOLLO使MindSpore的性能平均比TensorFlow高1.96倍，比XLA高1.18倍。
* 对于BERT-large，XLA性能略微领先2%，这归因于MindSpore引入的inplace assign算子导致了性能下降。
表5：多GPU上的吞吞吐量。

Ascend 910芯片性能 (图8)
* 在Ascend 910上对BERT和PanGu-α进行测试，APOLLO的性能平均比依赖厂商库的MindSpore高19.7%。
* APOLLO的优势在于能更好地管理Ascend复杂的多级、多向内存系统，并利用其矢量单元的并行性。
图8：BERT和PanGu-α在Ascend上的吞吐量（单位：examples/s）。括号内为批大小。越高越好。

编译开销 (表6)
* 实验数据显示，APOLLO能在几秒到几分钟内完成代码生成，验证了其作为JIT编译器在训练场景中的高效性。
表6：编译开销（单位：秒）。

推理工作负载性能 (表7)
* 在Wide&Deep、Yolo-v3和EPPMVSNet等推理任务上，APOOLLO同样展现出优越的性能，证明了其框架的通用性。
表7：GPU上推理工作负载的吞吐量。

A5 结论

本文提出了APOLLO框架，它统一了循环融合、内存拼接和并行拼接三种优化技术。通过打破算子边界扩展融合空间，并同时利用数据局部性和并行性，APOLLO在广泛的实验中超越了现有工作。其分段编译策略使其能够作为高效的JIT编译器，同时适用于训练和推理场景。

相较于其底层的代码生成器AKG，APOLLO做出了显著贡献：
1. 解决了AKG的可伸缩性问题：通过分区阶段和分层融合，避免了AKG在处理包含分块归约算子的复杂子图时遇到的长时间编译问题。
2. 增强了优化能力：在AKG利用数据局部性的基础上，APOLLO增加了第三层的并行拼接，进一步提升了性能。
3. 提供了反馈机制：APOLLO支持的向上反馈机制使得可以根据下游编译器的表现来更新上游的分区规则，从而更好地管理多面体模型这一“黑盒”变换，增强了用户与框架的交互能力。

局限性与未来工作：
* 规则集不完备：当前聚合规则（表1和表2）虽然能处理所有测试过的工作负载，但并非完备，未来可以进一步补充。
* 代价模型简单：用于并行拼接的代价模型（公式3）较为简单，未能精确建模并行与通信/同步之间的权衡。
尽管存在这些局限，APOLLO的多层融合方法仍比先前工作取得了进步，为DNN编译器的设计与实现提供了有价值的见解。

A6 附录

A 算法

算法1：聚合算法
* 输入: 子图$F_x$
* 流程:
1. 初始化: 将$F_x$中的每个原始算子op实例化为一个微图G，并根据op的类型设置G的类型，存入集合M。
2. 更新$F_x$: 用微图集合M替换$F_x$的内容。
3. 规则优先级: 定义规则的迭代顺序（Rules集合）。
4. 循环聚合:
* 对Rules中的每个规则r进行遍历。
* 在一个while循环中，持续扫描$F_x$中所有通过边e连接的生产者-消费者微图对($G_p$, $G_c$) 。
* 如果$G_p$和$G_c$的类型与规则r匹配，并且合并它们不会引入环路（通过acyclic函数检查），则执行聚合。
* 聚合操作aggregate(Gp, Gc)生成新的微图$G_a$，其类型由规则r指定。
* 从$F_x$中移除$G_p$和$G_c$，并加入$G_a$。
* 设置changed为true以继续while循环，直到当前规则下没有更多可聚合的对。

算法2：内存拼接算法
* 输入: 由第一层为每个$G_y$生成的IR集合$IR_{F_x}$
* 流程:
1. 规则集: 定义补充聚合规则Rules（表2中的规则7和8）。
2. 遍历融合:
* 对每个规则r和$F_x$中所有生产者-消费者微图对($G_p$, $G_c$)进行遍历。
* 检查条件: 只有当以下所有条件都满足时才进行融合：
* $G_p$和$G_c$的类型与规则r匹配，且合并不会引入环路。
* 两个微图的（并行和顺序）循环维度数量和硬件参数一致（consist_dim_and_param）。
* 如果规则为8（归约-归约融合），两个归约微图具有相同的规范形式（have_identical_canonical_form）。
* 执行融合:
* 聚合$G_p$和$G_c$为$G_a$，并更新其类型。
* 连接它们的IR，生成$IR_{G_a}$。
* 更新$F_x$和$IR_{F_x}$。
* 内存分配:
* 检查连接边e上的张量e.tensors的大小。
* 如果大小不超过本地内存容量，则分配到本地内存（如GPU共享内存）。
* 否则，分配到全局内存。
* 输出: 更新后的$IR_{F_x}$

算法3：并行拼接算法
* 输入: IR集合$IR_P = \{IR_{F_x}\}$
* 流程:
1. 模式优先级: 定义并行模式的检查顺序Patterns（多头、多尾、独立）。
2. 提取候选集: 对每种模式p，提取所有可并行的候选集Sets。
3. 遍历候选集: 对每个候选集s进行处理。
* 初始化: 将s中所有算子标记为“未访问”。
* 循环分组: 当s中仍有至少两个分支包含未访问的算子时，循环执行：
* 从每个分支中各提取一个算子，组成一个group。
* 按算子成本对group进行降序排序。
* 贪心组合: 使用一个内部while循环和代价模型(3)来寻找最佳的并行组合。
* 初始化m=1, k=n（group中算子数量）。
* 如果从第m个到第m+k-1个算子的并行化有正收益（is_positive），则执行并行拼接，并更新m和k以处理剩余算子。
* 如果没有收益，则减少k（k = k - 1），缩小组合规模再试。
* 将group中所有算子标记为“已访问”。
* 输出: 更新后的$IR_P$

方法细节中的参考文献引用分析

  • 引用 [3] Amodei, D., et al. Deep speech 2 : End-to-end speech recognition in english and mandarin. PMLR, 2016.

    • 引用位置: §3.1 提取子图集群, 第 1 段
    • 原文描述: "A typical example is all-reduce used for training speech recognition (Amodei et al., 2016)." (一个典型的例子是用于训练语音识别的all-reduce)

  • 引用 [5] Baghdadi, R., et al. Tiramisu: A polyhedral compiler for expressing fast and portable code. CGO, 2019.

    • 引用位置: §2 APOLLO架构, 第 3 段
    • 原文描述: "Tensor compilers (Chen et al., 2018a;b; Vasilache et al., 2019; Baghdadi et al., 2019) perform fusion together with tiling..." (张量编译器...将融合与分块一起执行)

  • 引用 [6] Bondhugula, U., et al. A practical automatic polyhedral parallelizer and locality optimizer. PLDI’08, 2008.

    • 引用位置: §4.1 第一层：多面体循环融合, 第 2 段
    • 原文描述: "The polyhedral model determines the parallelism and tilability of each loop nest by solving an integer linear programming (ILP) problem (Bondhugula et al., 2008)..." (多面体模型通过求解一个整数线性规划问题来确定每个循环嵌套的并行性和可分块性)

  • 引用 [7] Chen, T., et al. Tvm: An automated end-to-end optimizing compiler for deep learning. OSDI’18, 2018a.

    • 引用位置: §2 APOLLO架构, 第 3 段；§2 APOLLO架构, 第 5 段；§4.3 第三层：并行拼接, 第 1 段
    • 原文描述: "Tensor compilers (Chen et al., 2018a;b; ...)"；"...different from the graph engine of TVM (Chen et al., 2018a) that partitions a sub-graph without the awareness of its loop optimizer’s requirements..."；"This is also the weakness of many existing tensor compilers (Chen et al., 2018a; Vasilache et al., 2019)..."

  • 引用 [8] Chen, T., et al. Learning to optimize tensor programs. NIPS’18, 2018b.

    • 引用位置: §1 引言, 第 5 段；§2 APOLLO架构, 第 3 段
    • 原文描述: "...goes beyond prior work (Ragan-Kelley et al., 2013; Chen et al., 2018a;b; Vasilache et al., 2019)..."； "Tensor compilers (Chen et al., 2018a;b; ...)"

  • 引用 [10] Chetlur, S., et al. cudnn: Efficient primitives for deep learning, 2014.

    • 引用位置: §4.1 第一层：多面体循环融合, 第 5 段
    • 原文描述: "This optimization was not considered by manual scheduling approaches or vendor libraries (Chetlur et al., 2014)..." (这种优化没有被手动调度方法或厂商库所考虑)

  • 引用 [11] Cho, M., et al. Blueconnect: Decomposing all-reduce for deep learning on heterogeneous network hierarchy. MLSys, 2019.

    • 引用位置: §3.1 提取子图集群, 第 1 段
    • 原文描述: "The most appropriate solution to deploy such ops is wrapping a highly-crafted library like that of Cho et al. (2019)." (部署这类算子最合适的方法是封装一个高度优化的库，如Cho等人的工作)

  • 引用 [12] Cyphers, S., et al. Intel ngraph: An intermediate representation, compiler, and executor for deep learning, 2018.

    • 引用位置: §3 分区阶段, 第 1 段
    • 原文描述: "Data-flow optimization performs common subexpression elimination and constant folding, which is also considered by nGraph (Cyphers et al., 2018)." (数据流优化执行公共子表达式消除和常量折叠，nGraph也考虑了这一点)

  • 引用 [15] Ding, Y., et al. Ios: Inter-operator scheduler for cnn acceleration. MLSys, 2021.

    • 引用位置: §2 APOLLO架构, 第 2 段
    • 原文描述: "A dynamic programming strategy (Ding et al., 2021) is used to evaluate each fusion possibility..." (使用一种动态规划策略来评估每种融合可能性)

  • 引用 [17] Google. Xla: Optimizing compiler for machine learning, 2017.

    • 引用位置: §1 引言, 第 2 段；§2 APOLLO架构, 第 2 段；§3 分区阶段, 第 1 段
    • 原文描述: "Many graph compilers (Google, 2017; Wei et al., 2018; Rotem et al., 2019) only perform node grouping..."；"...as also adopted by Google (2017)."

  • 引用 [23] Irigoin, F. and Triolet, R. Supernode partitioning. POPL’88, 1988.

    • 引用位置: §2 APOLLO架构, 第 3 段
    • 原文描述: "Tensor compilers ... perform fusion together with tiling (Irigoin & Triolet, 1988)..." (张量编译器将融合与分块一起执行)

  • 引用 [24] Jangda, A. and Bondhugula, U. An effective fusion and tile size model for polymage. ACM Trans. Program. Lang. Syst., 2020.

    • 引用位置: §1 引言, 第 2 段；§4.1 第一层：多面体循环融合, 第 2 段
    • 原文描述: "...node grouping (Jia et al., 2019a; Jangda & Bondhugula, 2020)..."；"...auxiliary loop transformations have to be used together with loop tiling and fusion (Zhao & Cohen, 2019; Jangda & Bondhugula, 2020)..."

  • 引用 [25] Jia, Z., et al. Taso: Optimizing deep learning computation with automatic generation of graph substitutions. SOSP’19, 2019a.

    • 引用位置: §1 引言, 第 2 段；§1 引言, 第 3 段；§2 APOLLO架构, 第 4 段；§2 APOLLO架构, 第 6 段；§4.3 第三层：并行拼接, 第 7 段
    • 原文描述: 多处引用TASO来对比其在节点分组、并行性利用方面的做法和局限性。

  • 引用 [35] Liao, H., et al. Ascend: a scalable and unified architecture for ubiquitous deep neural network computing : Industry track paper. HPCA, 2021.

    • 引用位置: §5 整体实现, 第 4 段
    • 原文描述: "...and so-called CCE code executable on Huawei Ascend 910 (Liao et al., 2021)." (以及在华为Ascend 910上可执行的所谓CCE代码)

  • 引用 [36] Ma, L., et al. Rammer: Enabling holistic deep learning compiler optimizations with rtasks. OSDI 20, 2020.

    • 引用位置: §1 引言, 第 3 段；§2 APOLLO架构, 第 4 段
    • 原文描述: "TASO (Jia et al., 2019a) and Rammer (Ma et al., 2020) also pack independent ops..."

  • 引用 [38] Mehta, S., et al. Revisiting loop fusion in the polyhedral framework. PPoPP’14, 2014.

    • 引用位置: §2 APOLLO架构, 第 3 段
    • 原文描述: "...and thus suffer from the scalability issue (Mehta et al., 2014; Zhao & Di, 2020)." (因此遭受可伸缩性问题的困扰)

  • 引用 [45] Ragan-Kelley, J., et al. Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines. PLDI’13, 2013.

    • 引用位置: §1 引言, 第 5 段
    • 原文描述: "The coupled implementation of multi-layer fusion goes beyond prior work (Ragan-Kelley et al., 2013; ...)" (多层融合的耦合实现超越了先前的工作)

  • 引用 [51] Vasilache, N., et al. The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated gpu kernels, automatically. ACM Trans. Archit. Code Optim., 2019.

    • 引用位置: 多处引用，作为现有张量编译器的代表，讨论其在循环融合、可伸缩性问题和代码生成方面的特点。

  • 引用 [52] Verdoolaege, S. Isl: An integer set library for the polyhedral model. ICMS’10, 2010.

    • 引用位置: §3 分区阶段, 第 1 段；§5 整体实现, 第 4 段
    • 原文描述: "...meets the 'static affine control' requirement (Verdoolaege, 2010) of the polyhedral model..."；"...is guaranteed by isl (integer set library) (Verdoolaege, 2010)."

  • 引用 [53] Verdoolaege, S. and Janssens, G. Scheduling for ppcg. Report CW, 706, 2017.

    • 引用位置: §2 APOLLO架构, 第 6 段；§4.1 第一层：多面体循环融合, 第 2 段
    • 原文描述: "...in the polyhedral model (Verdoolaege & Janssens, 2017)."； "Expressed using affine relations (Verdoolaege & Janssens, 2017) in the polyhedral model..."

  • 引用 [54] Verdoolaege, S., et al. Polyhedral parallel code generation for cuda. ACM Trans. Archit. Code Optim., 2013.

    • 引用位置: §4.1 第一层：多面体循环融合, 第 3 段；§5 整体实现, 第 4 段
    • 原文描述: "Similar scalability issue is also faced by other polyhedral compilers (Vasilache et al., 2019; Verdoolaege et al., 2013)."；"...the backends of PPCG (Verdoolaege et al., 2013) and TC ... can be borrowed..."

  • 引用 [55] Wei, R., et al. Dlvm: A modern compiler infrastructure for deep learning systems, 2018.

    • 引用位置: §1 引言, 第 2 段；§2 APOLLO架构, 第 2 段
    • 原文描述: "Many graph compilers (Google, 2017; Wei et al., 2018; ...)"

  • 引用 [59] Zhao, J. and Cohen, A. Flextended tiles: A flexible extension of overlapped tiles for polyhedral compilation. ACM Trans. Archit. Code Optim., 2019.

    • 引用位置: §4.1 第一层：多面体循环融合, 第 2 段
    • 原文描述: "...auxiliary loop transformations have to be used together with loop tiling and fusion (Zhao & Cohen, 2019; ...)"

  • 引用 [60] Zhao, J. and Di, P. Optimizing the memory hierarchy by compositing automatic transformations on computations and data. MICRO53, 2020.

    • 引用位置: §1 引言, 第 2 段；§2 APOLLO架构, 第 3 段
    • 原文描述: "An example that demonstrates the profit of such fusion strategies was described by Zhao & Di (2020)."； "...suffer from the scalability issue (Mehta et al., 2014; Zhao & Di, 2020)."

  • 引用 [61] Zhao, J., et al. Akg: Automatic kernel generation for neural processing units using polyhedral transformations. PLDI 2021, 2021.

    • 引用位置: §2 APOLLO架构, 第 6 段；§4.1 第一层：多面体循环融合, 第 1 段
    • 原文描述: "...which was not studied by prior approaches (Zhao et al., 2021; ...)"； "...we use AKG (Automatic Kernel Generator) (Zhao et al., 2021), a polyhedral optimizer..."

  • 引用 [62] Zhao, J., et al. Parallelizing neural network models effectively on gpu by implementing reductions atomically. PACT’22, 2022.

    • 引用位置: §4.1 第一层：多面体循环融合, 第 5 段
    • 原文描述: "We optimize the fusion of reduction ops using PANAMERA (Zhao et al., 2022)..." (我们使用PANAMERA优化归约算子的融合)

  • 引用 [63] Zheng, Z., et al. Fusionstitching: Boosting memory intensive computations for deep learning workloads, 2020.

    • 引用位置: §1 引言, 第 3 段；§2 APOLLO架构, 第 4 段；§4.3 第三层：并行拼接, 第 7 段
    • 原文描述: 多处引用，将stitching作为利用独立算子并行性的一种技术进行讨论。