FusionStitching: Boosting Memory Intensive Computations for Deep Learning Workloads

文章标题：FusionStitching：为深度学习工作负载加速内存密集型计算
作者/机构：Zhen Zheng, Pengzhan Zhao, Guoping Long, Feiwen Zhu, Kai Zhu, Wenyi Zhao, Lansong Diao, Jun Yang, Wei Lin (Alibaba Group)

A1 主要贡献

本文指出，在众多深度学习模型中，内存密集型计算因片外内存访问和CPU-GPU上下文切换开销而导致严重的性能问题。针对此问题，现有的即时（JIT）内核融合和代码生成技术存在局限性，例如融合计划探索策略粗糙、代码生成能力有限。为此，本文提出了FusionStitching，一个能够将具有不同数据依赖和非同构并行性的内存密集型算子融合成大型GPU内核的深度学习编译器，以自动减少全局内存访问和上下文切换开销。FusionStitching通过引入中间值的数据复用，拓宽了融合技术可处理的算子组合范围，超越了以往的JIT工作。它探索了巨大的融合空间，以决定最优的融合计划，同时考虑了内存访问成本、内核调用和资源使用限制。FusionStitching利用一个领域特定的成本模型高效地调整最优的拼接方案。实验结果显示，与现有最先进技术相比，FusionStitching最高可实现2.21倍的加速，平均加速1.45倍。该方法已被集成到编译器产品中，并部署在拥有数千个GPU的生产集群上，在4个多月的时间里，为每月约30,000个任务平均节省了7,000个GPU小时。

核心贡献总结如下：
* 揭示了内存密集型算子的重要性，并通过引入待融合算子的数据复用，拓宽了JIT融合所能处理的算子组合范围。
* 为机器学习工作负载中的内存密集型算子提供了一套融合方案抽象，并提出了一种在给定非常复杂的融合模式下生成高效GPU内核的方法。 同时，提出了一种在扩展的算子组合搜索空间中即时探索良好融合计划的技术，并辅以一个双层GPU内核成本模型。
* 提供了一个对用户透明的工业级实现，并在生产集群上用各种现代机器学习模型进行了评估。

A3 背景知识/关键Observation/设计原则

2.1 动机

现有JIT融合技术的局限性。 JIT融合是减少机器学习工作负载中密集内存操作的上下文切换和内存访问开销的最先进技术。作为最先进的融合引擎，XLA仅支持融合中的线程局部数据传输，这依赖于索引分析和重计算来提高线程局部性。一个糟糕的案例是在融合模式的中间放置一个reduction操作。如果不同的线程需要相同的reduction结果值，每个线程都需要独立且冗余地重计算该reduction。为了避免这种情况，XLA不允许融合导致中间出现reduction的操作。因此，这种策略错过了许多优化探索空间。

FusionStitching与XLA在层归一化（Layer Normalization）融合上的差异。 图1展示了XLA和FusionStitching如何对层归一化进行融合。XLA形成了4个融合，其中fusion.3和fusion.7以reduction结束，fusion.2以昂贵操作结束。XLA为了避免重计算开销而没有进行其他融合，其缺点是带来了CPU-GPU上下文切换和片外内存访问开销。而FusionStitching只生成一个内核，所有中间结果都驻留在片上内存中，并且不涉及CPU-GPU切换。FusionStitching中没有引入重计算。

图1：XLA和FusionStitching在层归一化融合模式上的差异。

现有库和编译器方案的不足。 静态库无法解决此问题，因为单个内存密集型操作很轻量，而它们的组合在不同模型中变化多样。TVM 【11, {TVM}: An automated end-to-end optimizing compiler for deep learning, 2018, 13th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18)】【55, Ansor: Generating high-performance tensor programs for deep learning, 2020, arXiv preprint arXiv:2006.06762】主要关注计算密集型操作而非融合优化，而XLA仅为内存密集型操作提供简单的基于规则的融合策略。此外，TVM依赖于预定义的调度进行代码生成，并且通常需要很长的调优时间。

2.2 观察

本文通过分析多种机器学习工作负载（第7节），得出了两个主要观察。

严重的上下文切换开销。 CPU和GPU之间的上下文切换开销是一个严重的性能问题。表2显示了多种机器学习应用的分解分析。TensorFlow实现的GPU内核调用次数（#）可达10,406次，导致巨大的内核启动开销。即使使用XLA进行融合，内核调用次数仍可高达6,842次。因此，对于某些模型（如BERT推理、DIEN、ASR和CRNN），由机器学习框架带来的CPU上的调度和前/后处理时间占据了主导地位。

内存密集型操作占比巨大。 同样在表2中，对于某些模型，内存密集型操作的执行时间可占总时间的40%。片外内存流量时间通常占据内存密集型操作总执行时间的很大一部分。通过融合大量操作，我们的系统能够减少CPU-GPU上下文切换开销，并利用高速的片上内存来复用操作间的中间数据。

2.3 挑战

尽管数据复用是融合大量操作以避免冗余重计算的潜在机会，但在给定融合模式下应用复用以及决定操作融合策略方面存在几个主要挑战。

在给定融合模式下应用复用的挑战：
* 如何评估复用收益。 复用并不总是优于重计算。复用会导致额外的线程间数据通信，而重计算引入了冗余的计算开销但没有线程间通信。此外，通过共享内存引入复用可能会损害内核占用率，从而损害并行性。
* 如何应用复用。 数据复用有两种类型：warp内复用和block内复用。Warp内/block内复用意味着一个warp/block内的线程复用同一warp/block内其他线程产生的中间数据。其权衡在于，warp内复用内存访问开销较小，而block内复用并行性更好。

决定融合哪些操作的挑战。 对于给定的机器学习模型，挑战在于决定哪些操作应该被融合在一起。XLA也试图解决这个挑战，但其探索空间有限。一个主要问题是，通过数据复用形成一个融合模式并不总是比分离的内核更好。复用要求线程块或warp内的数据局部性，这可能会限制并行性。以reduction为例，复用reduction结果需要在块或warp内计算reduction。然而，某些reduction操作通过多个块协同工作表现更好，具有更高的并行性。不将这样的reduction与其消费者融合可能会带来更好的性能，即使存在上下文切换和片外内存访问开销。块内复用可能进一步损害并行性，因为它需要额外的共享内存。另一个典型问题是，当操作A可以与操作B或操作C融合，而B和C由于某些限制不能一起融合时，决定A与谁融合并不总是容易的。像XLA这样的基于规则的方法，无法为不同模型找到有效的融合计划。

A2 方法细节

3 概述

3.1 数据复用

通过数据复用扩展融合可能性。 如前所述，FusionStitching通过引入数据复用拓宽了融合的可能性，这在当前最先进的JIT融合技术中很少使用。当处理操作B的不同线程需要由操作A生成的相同值时，可以应用数据复用。我们假设B是A的消费者。需要相同值的线程可以通过线程间通信读取相同的数据，而不必像在XLA中那样冗余地重计算它。我们观察到，由于诸如reduce、broadcast、gather和slice等操作，一系列内存密集型操作的张量形状总是频繁地收缩和扩展。因此，探索数据复用的机会很大。

warp内与block内复用。 在我们的工作中，我们同时探索了warp内复用和block内复用，对应于图3中的warp组合和block组合。以reduction为例，warp内复用意味着每个warp为一行数据执行reduction，并将结果存储在warp的第一个通道的寄存器中。reduction的消费者通过寄存器混洗（register-shuffle）从第一个通道读取数据。Block内复用则使用块内的所有线程为该行数据执行reduction，并将结果存储在共享内存中。reduction的消费者从共享内存中读取数据。这种方法能够在避免重计算的同时，利用片上内存拼接操作。为了正确地进行数据复用，必须保证数据局部性。Warp内复用需要warp级别的数据局部性，而block内复用需要block级别的数据局部性。

图2：FusionStitching概览。

3.2 FusionStitching 系统

两阶段优化过程。 我们在FusionStitching中将优化过程分为两个阶段：融合探索和代码生成。这两个阶段在概念上是独立的，但高度相关。

系统组件。 如图2所示，FusionStitching由融合探索器和代码生成器组成。融合探索器生成可能利用数据复用的融合模式。然后，它应用波束搜索（beam search）来利用这些融合模式形成几个候选融合计划，并使用成本模型从候选计划中选择最佳的融合计划。代码生成器为融合探索器产生的每个融合模式生成GPU内核。我们为每种算子预定义了一组调度（即模板）。一个调度描述了关于一个操作如何并行运行的代码生成方式。同时，它指明了一个操作的线程是否与其他线程复用数据以及复用范围，以及该操作是否为消费者的复用产生数据。FusionStitching将一个融合模式的操作分成几个组，并枚举组中调度的所有有效组合。通过一个估计每个调度和启动维度设置性能的成本模型，FusionStitching选择融合模式的最佳配置并生成GPU内核代码。

双层成本模型。 我们在FusionStitching中设计了一个双层成本模型。融合探索器需要在大的搜索空间中搜索，并应用delta-evaluator（5.4节），它速度快但精度较低。代码生成器在合并的GPU内核上操作，需要更准确的性能估计，因此我们应用latency-evaluator（4.3节），它更准确但速度较慢。

章节安排。 我们将首先在第4节描述代码生成器的工作原理，然后在第5节描述融合探索器。

4 代码生成

代码生成器的任务。 代码生成器接收一个融合模式作为输入，并产生一个实现融合算子的GPU内核。由于各种依赖场景和并行性不兼容问题，将多个算子融合成一个高性能的GPU内核并非易事。

基本算子类型与调度。 机器学习工作负载中内存密集型算子的组合模式数量众多，但基本类型的内存密集型算子是有限的。我们进行了三种分类：轻量级element-wise（大多数elem-wise操作）、昂贵element-wise（tan、exponential等）和reduction操作。我们为每种内存密集型算子预定义了一组调度。代码生成的剩余问题是如何将不同的算子拼接成一个内核，以及每个独立算子应用何种调度。

自动生成方案。 我们首先系统地研究了四种内核组合方案（4.1节），这些方案涵盖了内存密集型操作的常见执行模式。利用这些组合方案，我们使用基于性能建模的自动生成解决方案（4.3节）来为融合模式找到好的调度并生成代码（4.2节）。

4.1 内核组合方案

四种内核组合方案。 我们研究了四种内核组合方案，它们揭示了常见内存密集型操作的主要行为。不同的方案表示待融合内核的不同数据依赖和并行行为，范围从无依赖到复杂的跨线程依赖，以及从统一并行到非同构计算。图3展示了这四种组合方案。

图3：将算子A和B拼接在一起的四种组合方案。Gmem：全局内存（不同的方块表示不同的内存地址）。Reg：寄存器。Shmem：共享内存。

• 内核打包（Kernel Packing）：将没有数据依赖的算子计算打包在一起。该方案有助于减少内核启动和框架调度的上下文切换开销，同时在指令级别减少循环控制开销。为了减少控制流开销，当多个element-wise算子具有相同的并行维度时，我们执行积极的循环融合【5, Optimizing compilers for modern architectures: a dependence-based approach, 2002】【18, Collective loop fusion for array contraction, 1992, International Workshop on Languages and Compilers for Parallel Computing】将它们合并到单个循环结构中。
• 线程组合（Thread Composition）：融合相互依赖的算子，并通过本地线程上下文内的寄存器传递中间结果。这是XLA采用的融合方案，可能会在需要相同值的线程之间引入冗余计算。
• Warp组合（Warp Composition）：融合相互依赖的算子，并应用warp内复用。该方案利用寄存器混洗（register shuffle）在warp内的线程间进行通信。一个典型案例是warp reduction及其消费者，它应用warp级别的reduction并利用寄存器将中间结果传递给其消费者的线程。
• 块组合（Block Composition）：应用块内复用，释放了将非同构计算组合成大型融合内核的潜力，只要这些计算可以在块级别内通信。它利用共享内存来传递中间结果以实现数据复用。

Warp和Block组合的灵活性。 Warp组合和块组合是灵活的方案，因为它们允许不同的算子在融合内核中以独立的调度执行。唯一的要求是在生产者和消费者之间保持warp/block的局部性。这两种方案对于高效地组合具有各种并行特性和依赖关系的广泛算子种类至关重要。

FusionStitching的创新。 据我们所知，FusionStitching是第一个为机器学习工作负载的内存密集型操作进行即时编译而全面研究上述所有组合技术的工作。XLA仅实现了内核打包和线程组合。我们不拼接涉及块间通信的算子，因为它会导致全局内存级别的同步并引入高昂的开销。

4.2 内核生成

基于成本模型的调优方法。 由于我们上面讨论的组合方案的权衡，FusionStitching中的代码生成不像XLA那样直接。我们在FusionStitching中使用基于成本模型的调优方法，并通过算子分组来简化过程。

调度预定义与枚举。 根据四种组合方案，我们为每种算子（轻量级element-wise、昂贵element-wise和reduction）预定义了一组调度。具体来说，我们为轻量级element-wise算子定义了一个单一模板，代表内核打包和线程组合（我们发现这两种可以用相同调度描述）。我们为昂贵element-wise算子和reduction算子定义了三种调度：第一种代表内核打包和线程组合；第二种将中间结果存入每个warp第一个通道的寄存器，代表warp组合；第三种将中间结果存入共享内存，代表block组合。

通过算子分组简化枚举。 FusionStitching枚举所有算子的调度组合，并根据成本模型（4.3节）评估这些组合。为了简化枚举过程，我们将融合模式中的算子分为几个组。每个组应用一种调度，不同组之间通过warp/block级别的复用传递中间数据。

分组策略。 我们根据算子类型对算子进行分组。给定一个融合模式，我们枚举一小组分组策略。我们称一个组的输出算子为子根（sub-root），融合的输出为根（root）。我们首先识别子根，并生成以子根和根为根的组。Reduce算子总是被视作子根。昂贵的element-wise算子被枚举为子根和非子根。其他算子均不为子根。这导致了一系列以reduction、昂贵element-wise和根算子为根的组。

分组方法的洞见。 分组方法的洞见在于，非子根的调度可以通过张量索引传播由子根的调度确定。因此，我们只需要枚举子根和根算子的调度，就可以通过传播得到所有算子的调度。

生成最佳代码策略。 FusionStitching枚举分组策略，并模拟每个子根/根算子的调度以及融合内核的启动维度。由于数据复用要求在复用范围内有正确的数据局部性，不满足数据局部性要求的调度将被丢弃。在使用latency-evaluator评估每次枚举的性能后，FusionStitching选择具有最佳估计性能的代码生成策略。

4.3 内核评估：Latency-Evaluator

性能评估需求。 FusionStitching需要对内核性能进行准确估计以选择代码生成策略。幸运的是，由于代码生成的搜索空间不是很大，我们可以容忍一个相对较慢的成本模型。

Latency-Evaluator公式。 我们构建了如公式1所示的latency-evaluator，其中L是融合内核的估计执行周期。
$$L=N_{\text{wave}} \times L_{\text{warp}}$$
$$N_{\text{wave}} = \frac{N_{\text{warp}}}{\text{Occupancy}}$$
$$L_{\text{warp}} = N_{\text{instruction}} \times CPI$$
其中$N_{wave}$表示一个GPU将处理多少波次的warp，大型GPU内核的warp将被分成几波在GPU上执行，同一波中的warp并发执行。$L_{warp}$表示单个warp在融合内核中花费的延迟（周期）。$N_{wave}$和$L_{warp}$的乘积代表执行融合内核的总周期。

$N_{wave}$的估计。 我们用要发出的总warp数（$N_{warp}$）和融合内核的占用率（Occupancy）来估计$N_{wave}$。$N_{warp}$根据启动维度和张量形状计算。我们根据启动维度、共享内存使用量（见4.4节）和估计的寄存器使用量来计算Occupancy。我们通过分析每个中间值的生命周期来估计寄存器使用量，从而得到一个线程的最大寄存器使用量。同样，我们通过生命周期分析来估计共享内存使用量。这种方法对于我们计算Occupancy足够准确。

$L_{warp}$的估计。 对于$L_{warp}$，我们使用已报告的不同类型操作的CPI（每指令周期数）数据【21, Dissecting the nvidia turing t4 gpu via microbenchmarking, 2019, arXiv preprint arXiv:1903.07486】【22, Dissecting the nvidia volta gpu architecture via microbenchmarking, 2018, arXiv preprint arXiv:1804.06826】，并将其与我们估计的总指令数（$N_{instruction}$）相乘。

4.4 共享内存优化

优化共享内存使用的重要性。 适度使用共享内存至关重要，因为大量的共享内存使用会损害内核并行性，特别是对于大粒度的组合。为了在不损害并行性的前提下尽可能多地使用共享内存，我们探索了一种基于数据流的共享内存共享技术。其核心思想是，FusionStitching尽可能重用先前分配的共享内存，以减少不必要的共享内存分配。

基于支配树的数据流分析。 我们使用支配树算法【12, A simple, fast dominance algorithm, 2001】进行共享内存数据流分析。该方法以计算图和共享内存请求为输入，输出一个分配图。为了优化共享空间，我们按拓扑顺序遍历计算图中的算子。当一个算子不需要共享空间时，先前的分配信息将向前传播。如果一个算子需要共享空间，我们合并其所有操作数的分配信息，测试支配关系以检查是否可以为当前算子共享任何先前分配的空间，如果可能则重用该空间。

4.5 计算复用优化

线程局部冗余计算的优化。 除了通过共享内存和寄存器混洗在生产者和消费者之间复用中间结果外，FusionStitching还减少了线程局部的冗余计算。

冗余计算的来源与解决。 线程局部的冗余计算主要来自内存访问索引计算和线程局部中间值。这是因为融合内核中的不同部分可能使用不同的调度，而索引和一些中间值在每个调度中是独立生成的，即使在同一个线程内也是如此。在生成代码之前，FusionStitching首先分析整体的索引和中间值特性，然后组织输出代码以尽可能地复用计算和数据。

5 融合探索

基于成本的搜索方法。 如前所述，XLA中基于规则的方法不足以处理复杂的融合决策。我们使用一种基于成本的搜索方法来寻找好的融合计划。我们提出一种近似动态规划方法来搜索一组有希望的融合模式（5.2节），这些模式可能相互重叠。搜索过程由一个轻量级的领域特定成本模型（5.4节）指导。FusionStitching最终通过选择这些融合模式来生成整体融合计划（5.3节）。

5.1 融合问题定义

将融合探索定义为子图搜索问题。 我们将融合探索形式化为一个子图搜索问题。对于计算图 $G = (V, E)$，其中V和E分别是顶点和边的集合。我们将一个融合模式 $P_i = (V_i, E_i)$ 定义为G的一个子图，其中 $V_i \subseteq V, E_i \subseteq E$。一个融合计划是一组不相交的融合模式 $S = \{P_0, \dots, P_{k-1}\}$。$P_i$ 的得分函数记为 $f(P_i)$。性能越高，$f(P_i)$ 越大。因此，计算融合问题的目标是找到一个融合计划S，使得 $\sum_{i=1}^{k} f(P_i)$ 最大化。

5.2 探索融合模式

近似动态规划方法。 暴力枚举所有融合模式的复杂度在没有剪枝的情况下高达 $O(2^V)$。我们提出了一种复杂度为 $O(V+E)$ 的近似动态规划方法来寻找好的融合模式。

探索的基本思想。 融合探索的基本思想是，我们从图中每个顶点按后序开始生成候选模式，并用这些候选模式来选择和组合最终的融合计划。从顶点 $V_i$ 开始的候选模式，是其生产者节点为 $V_i$ 的融合模式集合。本节描述我们如何为每个顶点生成候选模式，5.3节描述如何组合最终的融合计划。

生成候选模式的过程。 给定一个计算图G，我们得到一个拓扑排序列表。我们按后序为顶点生成候选模式，从最后一个顶点到第一个顶点。这是为了保证每次搜索都得到一个以当前顶点为生产者的融合模式。

近似动态规划公式。 公式2展示了近似动态规划过程。$C_i$ 是 $V_i$ 的所有消费者的集合。PatternReduction是根据其所有消费者的候选模式来获取 $V_i$ 候选模式的方法。

Pi = PatternReduction(Ci)

PatternReduction示例。 我们通过图4中的一个例子来描述PatternReduction方法，其中V0是产生整个图输出的根顶点。假设在V8之前的顶点的候选模式已经生成。（我们根据得分函数f为每个顶点只探索前k个模式，本例中k设为3。）我们将利用V8的消费者的信息来为V8生成候选模式。一个朴素的方法是将V8附加到其消费者候选模式集合中所有模式的可能组合上，并从这些附加后的组合中选出前3个模式。然而，当消费者数量和k值较大时，组合数会非常庞大，而JIT方法要求及时优化。因此，我们设计了一个近似的分治过程PatternReduction来以有限的复杂度找到前3个模式。

图4：为V8生成候选模式的PatternReduction案例。

分治过程。 我们首先将V8的消费者分成几个组，独立地考虑这些组为V8寻找候选模式，最后通过归约所有组的结果来组合最终的候选模式。本例中我们假设组数为2（组{V5, V6}和组{V7}）。对于组{V5, V6}，我们枚举V5和V6候选模式中所有可能的模式组合。具体来说，V5和V6之间有15种可能的组合，包括空集。我们将V8附加到这15种组合中的每一种，并选择前3个模式作为与组{V5, V6}相关的临时候选。同样，我们考虑组{7}得到另外3个临时候选模式。最后，我们从上述所有6个临时候选模式中选出最终的前3个作为V8的候选模式。注意，我们根据得分函数f来验证顶尖模式。

递归的PatternReduction。 如果一个顶点的消费者数量非常大，上述PatternReduction过程是递归的。我们首先将消费者分为两组。要获取一个组的候选模式，我们进一步划分该组，直到它足够小。

图5：远程模式的融合合并。

图6：循环依赖：在将节点A和C融合后出现循环依赖。

融合模式的约束。 融合模式的一个约束是不允许出现循环依赖。图6展示了一个融合后出现循环依赖的例子。FusionStitching在搜索过程中会丢弃带有循环依赖的模式。同时，FusionStitching只探索代码生成器能够处理的融合模式，它不会形成需要跨块通信的融合模式。

远程融合（Remote Fusion）。 为了进一步减少CPU和GPU之间的上下文切换开销，在上述过程之后，我们尝试合并图中不相邻的融合模式。如图5所示，我们为所有顶点添加一个虚拟顶点h作为生产者，并应用PatternReduction。我们最终得到Vh的候选模式，其中包括了不相邻的远程模式的融合。远程模式融合有助于减少生成的内核数量，从而减少上下文切换开销。远程融合在代码生成阶段对应于内核打包。

5.3 生成整体融合计划

从候选模式生成融合计划。 所有顶点的所有候选模式（可能相互重叠）形成一个新的集合E。形成良好融合计划的洞见是从E中选择得分f高的不重叠模式。注意，f表示该模式节省的成本（见5.4节）。

使用波束搜索（Beam Search）。 FusionStitching使用波束搜索【17, Limited discrepancy beam search, 2005, IJCAI】生成前3个（波束宽度）候选融合计划，并最终使用latency-evaluator（见4.3节）从这3个候选计划中选择最佳计划。

搜索过程。 具体来说，FusionStitching维护3个缓冲区集合来存储候选融合计划。它从操作图G中的生产者顶点遍历到消费者顶点，并依次尝试将每个顶点的每个候选模式附加到每个缓冲区集合中，前提是不引入重叠。在每次迭代中附加后可能会产生多于3个临时集合，此时将选择累积f最高的3个作为新的缓冲区集合。

选择最终计划。 遍历完所有顶点后，每个缓冲区集合就形成一个候选融合计划。FusionStitching将使用更准确的成本模型latency-evaluator来评估哪个候选计划性能最佳。最佳者即为最终的融合计划。

5.4 融合评估：Delta-Evaluator

Delta-Evaluator的洞见。 FusionStitching应用delta-evaluator来构成得分函数f。一个洞见是，我们只需要估计形成一个融合模式时的性能增益或损失，而不需要准确估计整体执行时间。基于此，得分函数f仅代表一个融合模式的性能增益或损失。

得分函数构成。 delta-evaluator中有三个主要因素：减少的内存访问延迟、减少的CPU-GPU上下文切换开销以及内核融合的性能惩罚。得分函数f是这三部分的总结，如公式3所示。

$f = T_{reduced\_mem} + T_{reduced\_calls} - T_{penalty}$
$f = T_{reduced\_mem} + T_{reduced\_calls} - T_{penalty}$

各部分估计方法。
* 我们用两个因素估计减少的内存访问延迟（$T_{reduced\_mem}$）。第一个是待融合算子间的内存流量，可以根据输入输出张量的形状和类型计算。第二个是存储操作间中间值的内存类型的变化。我们建立了一个回归模型来预测当内存类型从全局内存变为寄存器或共享内存时，给定内存流量大小所减少的内存访问延迟。该回归模型基于我们在各种GPU供应商上针对不同数据流量离线收集的延迟数据。
* $T_{reduced\_calls}$可以简单地通过将融合的内核数量乘以我们收集的代表平均CPU-GPU上下文切换时间的固定值来估计。
* 性能惩罚（$T_{penalty}$）通过一个简化版的latency-evaluator（见4.3节）来估计。由于融合计划探索面临非常大的搜索空间，直接使用latency-evaluator进行JIT优化耗时太长。latency-evaluator中最耗时的部分是Occupancy的估计。为了快速估计Occupancy，我们使用一个固定数值（16）作为寄存器数量，并使用融合模式内任何算子内部和之间最大的共享内存使用量作为整体共享内存使用量。在delta-evaluator中，寄存器和共享内存的生命周期分析被舍弃了。

6 实现

作为独立插件实现。 我们研究的融合探索和代码生成技术需要大量的实现工作，如果留给用户将是一个负担。同时，TensorFlow发展迅速，将FusionStitching的实现从一个TensorFlow版本移植到另一个版本工作量很大。为了简化FusionStitching的使用并使其可移植到任何TensorFlow版本，我们将其构建为一个独立的插件。用户无需更改任何模型脚本即可利用所有技术，只需指定FusionStitching路径并设置环境以在TensorFlow XLA之外启用FusionStitching。

与XLA的集成。 具体来说，我们将FusionStitching添加为XLA框架之上的一个额外pass。在XLA完成其所有优化（包括其基本的融合pass）后，我们将XLA的融合结果输入到我们的融合探索器中，以获得更大范围的融合。我们没有禁用基本的XLA融合pass，因为在朴素融合的基础上应用FusionStitching没有冲突，而且基本融合可以减少FusionStitching的优化时间。

代码生成流程。 我们利用FusionStitching中的代码生成器技术实现了一个GPU IR发射器模块。由融合探索器生成的融合计划会通过这个IR发射器。IR发射器将融合模式依次编译成LLVM IR、GPU PTX IR，最终生成GPU二进制文件。

未被融合的内核处理。 FusionStitching的输入是XLA的基本融合结果，并非每个基本融合最终都会被合并到一个FusionStitching的融合模式中。那些未被FusionStitching合并成更大融合的基本融合，最终将通过XLA的基本编译pass处理。

实现为TensorFlow的钩子。 所有上述实现都作为一个TensorFlow的钩子实现，如果加载了FusionStitching并设置了环境，它就会被调用。

异步编译模式。 为了隐藏优化调优时间，我们在FusionStitching库中开发了异步编译模式。对于训练，如果用户将环境设置为异步优化，前几次迭代将不会执行由FusionStitching生成的二进制文件，而FusionStitching仍在后台异步运行。在FusionStitching为模型生成优化内核后，后续迭代中未优化的算子将被替换为FusionStitching的优化版本。

A4 实验环境

实验应用：我们使用了一系列来自不同领域的机器学习应用来评估FusionStitching，其特征总结在表1中。这些应用涵盖了图像（CRNN 【37, An end-to-end trainable neural network for image-based sequence recognition and its application to scene text recognition, 2016, IEEE transactions on pattern analysis and machine intelligence】）、语音（ASR 【52, AUTOMATIC SPEECH RECOGNITION, 2016, Springer】）、自然语言处理（Transformer 【47, Attention is all you need, 2017, Advances in neural information processing systems】、BERT 【15, Bert: Pre-training of deep bidirectional transformers for language understanding, 2018, arXiv preprint arXiv:1810.04805】）以及互联网规模的电子商务搜索和推荐系统（DIEN 【57, Deep interest evolution network for click-through rate prediction, 2019, Proceedings of the AAAI conference on artificial intelligence】）。这些工作负载的构建模块包括感知机、注意力机制、卷积、RNN以及大量的内存密集型算子。

模型详情：
表1：用于评估的工作负载。

硬件配置：
* GPU：NVIDIA V100 GPU (16 GB设备内存)，部分测试在NVIDIA T4 GPU上进行。
* 服务器平台：Red Hat Enterprise Linux 7.2。

软件配置：
* 框架与基线：默认的TensorFlow实现和XLA（社区最新功能）。
* 依赖库：CUDA toolkit 10.2 和 cuDNN 7.6。

A4 实验结果

7.2 概述

评估方法：我们通过比较在相同批次大小下，TF、XLA和FusionStitching的推理成本或单次迭代的训练时间来评估FusionStitching的加速效果。测试期间，训练中每次迭代的准确率和推理结果与TF和XLA相同。我们重复10次实验并取平均性能来验证加速比。对于训练工作负载，我们收集第11到第20次迭代的执行时间，以避免早期训练迭代的初始化开销。

整体性能：如图7所示，我们将TensorFlow的执行时间归一化为1。与TensorFlow相比，我们的方法最高实现了2.42倍的加速，平均为1.66倍。与XLA相比，我们的方法最高实现了2.21倍的加速，平均为1.45倍。值得注意的是，XLA在DIEN（训练和推理）上表现出性能下降，而FusionStitching在所有案例中均未出现负优化。在NVIDIA T4 GPU上的推理工作负载测试也获得了类似的加速效果。

图7：FusionStitching的性能加速比。

生产环境效益：FusionStitching已在生产环境中应用，并测量了其性能效益。结果显示，在一个月内，FusionStitching为约30,000个任务节省了约7,000个GPU小时。集群上的性能结果表明FusionStitching是稳健的。需要注意的是，作为最先进的JIT引擎，XLA因在许多情况下存在负优化而无法默认启用。

性能提升来源：总的来说，FusionStitching的性能提升来自于减少的CPU-GPU上下文切换和片外内存访问。

7.3 性能分解

性能分解指标：表2展示了内核的分解信息，包括内存密集型操作（Mem）、计算密集型操作（Math）的执行时间（T），CPU时间（CPU，包括内核启动和框架调度），CUDA memcpy/memset活动（Cpy）以及内核调用次数（#）。注意，分解分析过程与端到端性能测量过程不同，因为分析会引入开销。

XLA和FusionStitching对GEMM行为的影响：需要指出的是，XLA会影响矩阵运算（GEMM和GEMV）的行为，它倾向于将GEMV融合成GEMM，以及在有共享输入的矩阵时将GEMM融合成更大的GEMM。其他代数变换和循环不变量代码移动也减少了GEMM的数量。FusionStitching建立在XLA之上，表现出相同的行为。

表2：内核执行分解。TF：原生TensorFlow。FS：FusionStitching。CPU：CPU上的调度和前/后处理指标。Math：计算密集型内核。Mem：内存密集型内核。Cpy：CUDA memcpy和memset调用。E2E：单次迭代的端到端时间（毫秒）。T：执行时间。#：内核调用次数。

观察结果分析：
* 减少了上下文切换开销：FusionStitching有效减少了所有工作负载的内存密集型内核调用，从而降低了内核启动和框架调度开销。如表2所示，使用FusionStitching的内存密集型内核调用次数平均只有XLA的38.0%（范围从27.8%到48.4%）。同时，FusionStitching比XLA减少了CUDA memcpy/memset活动（Cpy时间），平均减少了34.3%。这是因为我们比XLA融合了更多的操作到更大的内核中，许多CUDA memcpy/memset被合并了。表2中的CPU时间差异表明了因内核调用和CUDA memcpy/memset活动减少而节省的时间。与XLA相比，FusionStitching在CPU时间上最高节省61.0%，平均节省41.0%。
* 以DIEN-train为例：使用FusionStitching，内存密集型操作的内核调用数为2109，而XLA为6842。同时，CUDA memcpy/memset活动减少到1395次，而XLA为1996次。最终，FusionStitching的CPU时间显著少于TF和XLA。
* 减少了内存密集型操作的执行时间：FusionStitching减少了内存密集型操作的总执行时间。与XLA相比，所有工作负载的内存密集型操作平均加速1.39倍，最高可达1.74倍。性能提升主要来自于减少的全局内存访问。通过积极融合内存密集型操作，中间值可以被缓存在寄存器和共享内存中。
* 以CRNN为例：使用XLA时，它读取了667.6 MB的全局内存，而FusionStitching将流量减少到225.8 MB。内存密集型操作的全局内存访问减少了约66%，其执行时间因此比XLA快了1.48倍。

结论：总的来说，FusionStitching支持比XLA更复杂的融合模式和有效的内核生成，放宽了融合条件，从而减少了最终的内核数量和中间全局内存事务。这两个因素对于内存密集型操作的性能至关重要。

7.4 融合模式分析

Layer Normalization案例分析：我们再次以图1中的Layer Normalization为例来分析XLA和FusionStitching的融合结果。XLA形成了四个不同的融合内核，而FusionStitching能够将它们全部融合并生成一个更高效的内核。我们收集了两种版本所有内核的性能。FusionStitching生成的单个内核相比XLA的4个内核总和实现了1.23倍的加速。此测试未计算XLA版本中4个内核的上下文切换开销，这会进一步损害XLA的性能。

XLA无法进行更大粒度融合的原因：
1. Reduce操作的限制：xla-fusion.7和xla-fusion.3中的reduce操作是原因之一。如前所述，XLA不探索对生产者操作中间结果的复用。在内核中间融合reduce操作会导致不同线程对相同值进行冗余计算，从而损害性能。
2. 处理小张量的昂贵操作：xla-fusion.2中处理小张量的昂贵element-wise操作是另一个原因。XLA不倾向于在内核中间融合处理小张量的昂贵操作，因为这种融合会引入昂贵指令的冗余计算。

FusionStitching的优势：相反，FusionStitching在这种情况下发现了更大粒度融合的显著潜力，并将所有操作融合到一个内核中。它应用数据复用以防止冗余计算，从而避免了中间的全局内存事务和CPU-GPU上下文切换。

7.5 开销分析

一次调优多次运行的场景：与XLA类似，FusionStitching专为“一次调优，多次运行”的场景设计，这是深度学习工作负载的基本特征。对于可能长达数天的训练，FusionStitching只需在第一次训练迭代中运行。对于推理任务，可执行内核可以一次性准备好，并在未来多次运行。

JIT编译开销：我们测量了FusionStitching相较于XLA在训练中引入的一次性开销。该值为FusionStitching的JIT编译时间减去XLA的JIT编译时间。结果显示，对于本文评估的工作负载，额外开销少于30分钟，远低于总训练时间。

成本模型的有效性：我们尝试在delta-evaluator中使用完整的latency-evaluator来估计$T_{penalty}$。结果显示调优时间长得多，但调优结果的性能并未更好。这表明简化的成本模型在融合探索方面工作良好。

局限性：FusionStitching和XLA的一个问题是，它们无法以较低的调优开销处理动态形状，这在一些深度学习工作负载中会出现。原因在于XLA服务框架的设计对动态形状不友好，而FusionStitching是基于XLA服务框架实现的。这个实现问题不影响FusionStitching所展示的洞见。

A7 补充细节

8 相关工作

内存密集型操作优化：当前许多AI编译优化工作主要关注计算密集型操作【11, {TVM}: An automated end-to-end optimizing compiler for deep learning, 2018, 13th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18)】【13, Automatic generation of highperformance quantized machine learning kernels, 2020, Proceedings of the 18th ACM/IEEE International Symposium on Code Generation and Optimization】【24, A code generator for high-performance tensor contractions on gpus, 2019, 2019 IEEE/ACM International Symposium on Code Generation and Optimization (CGO)】【29, Rammer: Enabling holistic deep learning compiler optimizations with rtasks, 2020, 14th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 20)】【30, Boda: A holistic approach for implementing neural network computations, 2017, Proceedings of the Computing Frontiers Conference】【34, Halide: decoupling algorithms from schedules for high-performance image processing, 2017, Communications of the ACM】【45, Latte: a language, compiler, and runtime for elegant and efficient deep neural networks, 2016, Proceedings of the 37th ACM SIGPLAN Conference on Programming Language Design and Implementation】，但对内存密集型操作关注不足。XLA【26, Xla: Tensorflow, compiled, 2017】是深入优化内存密集型操作的一项工作，它提供了一个即时融合引擎来减少上下文切换和片外内存访问开销。然而，XLA缺乏对中间值数据复用的探索，限制了融合探索空间，并且它使用简单的规则来探索融合策略，无法适应多样的操作和张量形状组合。FusionStitching揭示了JIT融合的数据复用机会，并提出了基于成本的方法进行融合搜索和代码生成调优。

静态模式的融合优化：一些工作研究了机器学习工作负载中静态模式的融合优化。Li等人【27, Automatic horizontal fusion for gpu kernels, 2020, arXiv preprint arXiv:2007.01277】探索了GPU内核的水平融合以增加线程级并行性。Appleyard等人【7, Optimizing performance of recurrent neural networks on gpus, 2016, arXiv preprint arXiv:1604.01946】和Diamos等人【16, Persistent rnns: Stashing recurrent weights on-chip, 2016, International Conference on Machine Learning】研究了加速RNN工作负载的内核融合技术。这些工作不研究处理未知融合模式的即时融合问题，也没有探索操作组合策略。Abdolrashidi等人【3, Learning to fuse, 2019, NeurIPS】使用近端策略优化（Proximal Policy Optimization）【36, Proximal policy optimization algorithms, 2017, arXiv preprint arXiv:1707.06347】算法探索操作组合策略，这是一种提前调优方法，在与XLA类似的简单融合规则中搜索，不探索复杂模式的融合。

具备融合能力的计算密集型优化工作：一些专注于计算密集型操作的工作也具备操作融合能力。近期的TVM【11, {TVM}: An automated end-to-end optimizing compiler for deep learning, 2018, 13th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18)】和Ansor【55, Ansor: Generating high-performance tensor programs for deep learning, 2020, arXiv preprint arXiv:2006.06762】的实现具有基于简单规则处理简单模式的融合能力，但通常需要对给定的静态计算描述进行长时间的调优，并非为及时处理多样的操作组合而设计。Tiramisu【9, Tiramisu: A polyhedral compiler with a scheduling language for targeting high performance systems】采用了与TVM类似的融合方法。Tensor Comprehensions【46, Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions, 2018】提供了一个基于多面体的JIT编译器，能够融合操作，关注并行性与数据局部性的权衡。Glow【35, Glow: Graph lowering compiler techniques for neural networks, 2018, arXiv preprint arXiv:1805.00907】和Latte【45, Latte: a language, compiler, and runtime for elegant and efficient deep neural networks, 2016, Proceedings of the 37th ACM SIGPLAN Conference on Programming Language Design and Implementation】支持基本融合，不探索复杂的组合场景。Astra【40, Astra: Exploiting predictability to optimize deep learning, 2019, Proceedings of the TwentyFourth International Conference on Architectural Support for Programming Languages and Operating Systems】和Ashari等人的工作【8, On optimizing machine learning workloads via kernel fusion, 2015, ACM New York, NY, USA】支持GEMM和基本element-wise操作的融合。这些工作都没有即时地通过共享内存和寄存器混洗来揭示待融合操作间中间数据复用的融合机会，也缺乏对即时融合策略探索的深入研究。

其他领域的融合方法：融合方法也被应用于广泛的领域，如HPC【48, Scalable kernel fusion for memory-bound gpu applications, 2014, SC’14: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】、数据库【50, Kernel weaver: Automatically fusing database primitives for efficient gpu computation, 2012, 2012 45th Annual IEEE/ACM International Symposium on Microarchitecture】、图像处理【32, From loop fusion to kernel fusion: a domain-specific approach to locality optimization, 2019, 2019 IEEE/ACM International Symposium on Code Generation and Optimization (CGO)】【33, Automatic kernel fusion for image processing dsls, 2018, Proceedings of the 21st International Workshop on Software and Compilers for Embedded Systems】，但它们面临的挑战与AI工作负载的即时编译不同。具体来说，它们不融合事先未知的内核，也不面临在巨大搜索空间中选择融合哪些操作的问题。

相关技术：CUDA Graph【19, Getting started with cuda graphs, 2019, URL: https://devblogs. nvidia. com/cudagraphs】减少了内核启动开销，但由于图创建【31, Dag-based scheduling with resource sharing for multi-task applications in a polyglot gpu runtime, 2020, arXiv preprint arXiv:2012.09646】而面临更严重的初始化开销和大量的额外GPU内存使用，且不减少片外内存流量。GPU和其他加速器的性能模型是另一个相关的研究课题【14, An accurate gpu performance model for effective control flow divergence optimization, 2012, 2012 IEEE 26th International Parallel and Distributed Processing Symposium】【23, A learned performance model for the tensor processing unit, 2020, arXiv preprint arXiv:2008.01040】【28, Delta: Gpu performance model for deep learning applications with indepth memory system traffic analysis, 2019, 2019 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS)】【51, A performance model for gpu architectures that considers on-chip resources: application to medical image registration, 2019, IEEE Transactions on Parallel and Distributed Systems】【53, Understanding the gpu microarchitecture to achieve bare-metal performance tuning, 2017, Proceedings of the 22nd ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming】【54, A quantitative performance analysis model for gpu architectures, 2011, 2011 IEEE 17th international symposium on high performance computer architecture】。然而，我们为融合和代码生成需求设计了一个领域特定的成本模型系统。

A5 结论

本文解决了机器学习工作负载中优化内存密集型算子的问题。我们证明了内存密集型算子对各种深度学习模型的端到端性能至关重要。我们提出了FusionStitching，它支持即时融合具有复杂依赖和非同构并行性的算子，以减少内存访问和上下文切换开销。FusionStitching通过引入复用方案，拓宽了超越现有最先进JIT技术的融合可能性。FusionStitching由融合探索器和代码生成器组成。融合探索器以适当的计算复杂度从巨大的搜索空间中选择候选融合模式，并产生一个预期性能良好的融合计划。我们提供了一套内核组合方案，代码生成器利用这些方案来拼接算子并尽可能使用片上内存复用中间值，并调整调度以为给定的融合模式生成高性能的GPU代码。一个双层成本模型辅助FusionStitching的搜索和调优过程。结果显示，FusionStitching的性能优于最先进的JIT技术，最高加速比为2.21倍，平均为1.45倍。