AStitch: Enabling a New Multi-dimensional Optimization Space for Memory-Intensive ML Training and Inference on Modern SIMT Architectures

作者/机构:
Zhen Zheng, Xuanda Yang, Pengzhan Zhao, Guoping Long, Kai Zhu, Feiwen Zhu, Wenyi Zhao, Xiaoyong Liu, Jun Yang, Wei Lin (Alibaba Group China); Jidong Zhai (Tsinghua University China); Shuaiwen Leon Song (University of Sydney Australia)

A1 主要贡献

本文揭示了在最新的机器学习模型中，内存密集型计算正成为一个日益关键的性能影响因素。由于存在复杂的两级依赖关系和即时（just-in-time）需求的独特挑战，现有的机器学习（ML）优化编译器无法进行高效的算子融合。它们面临着两难的困境：要么执行融合但带来沉重的冗余计算，要么跳过融合导致产生大量的内核。此外，由于缺乏对具有不规则张量形状的真实世界生产负载的支持，这些编译器常常导致并行度低下。

核心问题与研究目标:
* 问题识别: 传统上，优化重点在于计算密集型操作（如GEMM、卷积），但许多新兴模型（如NLP领域的Attention）中，内存密集型操作（如element-wise和reduction）在执行时间和内核数量上占据了主导地位。如图1所示，在五个生产模型中，内存密集型计算平均占执行时间的63.2%和内核总数的89.6%。随着新一代GPU架构（如从NVIDIA V100到A100）的算力与内存带宽比急剧增加，这一问题变得更加严峻。
* 现有方案的局限: 现有的ML编译器（如XLA、TVM）虽然支持内核融合以减少开销，但在处理复杂的两级依赖（算子级和元素级）和不规则张量形状时存在以下关键挑战：
1. 复杂的两级依赖与JIT需求: 编译器在融合时面临两难选择：要么因无法有效处理一对多元素依赖而引入大量冗余计算，要么放弃融合而产生大量小内核，导致高昂的内核启动、上下文切换和内存访问开销。
2. 不规则张量形状: 生产环境中的张量形状通常不规则，导致编译器生成的线程映射方案并行度低下，硬件利用率不足。
* 研究目标: 提出一种名为AStitch的ML优化编译器，旨在为内存密集型ML计算开辟一个新的多维优化空间，通过支持对任意内存密集型子图的高效即时算子“缝合”（stitching），解决上述挑战。


图 1: 内存密集型计算的占比。该比例是内存密集型算子的指标占所有被研究GPU内核相应指标的比例。执行时间和内核数量的统计数据来自TensorFlow (v1.15)的执行。

创新点:
AStitch通过联合优化依赖特性、内存层次结构（局部性）和并行性，提出了以下创新技术：
1. 分层数据复用 (Hierarchical Data Reuse): 提出该技术以解决复杂的两级依赖问题，并扩大融合范围，避免了在“高计算开销的融合”与“不充分的融合”之间做出艰难选择。
2. 自适应线程映射 (Adaptive Thread Mapping): 提出该技术以适应不同的输入张量形状，生成合适的线程映射调度，从而最大化硬件利用率和并行性。
3. 多维度优化空间与自动编译器设计: 系统性地抽象出四种算子缝合方案，并设计了一个编译器，能够基于关键设计观察自动地、即时地为任意给定的ML模型（无论是训练还是推理）启用上述优化。如图2所示，AStitch将当前工作启用的小范围基础融合“缝合”成更大、更广泛的融合，从而在减少非计算开销的同时提升并行度。
4. 生产级实现与可移植性: 实现了一个生产就绪的ML编译器，可移植到任何版本的TensorFlow，用户无需修改源代码或调整代码生成调度即可使用。

a) 子图样本及XLA/TVM的融合计划。


图 2: AStitch如何优于XLA和TVM处理内存密集型子图的概念性说明。AStitch将大范围的内存密集型算子缝合在一起，以减少非计算开销并增加并行性。

A3 背景知识与当前挑战

本文的讨论聚焦于在目前最广泛使用的通用AI加速器GPU上，显著提升内存密集型ML应用的训练和推理效率。本文使用NVIDIA芯片及其术语作为验证平台，但所提出的技术具有通用性，可适用于其他GPU架构【索引编号1，AMD GPU-Powered Machine Learning Solutions】。

2.1 当前模型中基本的内存密集型算子

内存密集型算子在计算图中的角色。在现代框架中，ML工作负载通常表示为一个计算图。图中的大多数计算密集型算子是断开的，它们将计算图分割成一系列子图，每个子图包含数十甚至数百个内存密集型算子。现代ML模型中，覆盖了大部分内存密集型计算的两种广泛采用的算子是：element-wise（逐元素）算子和reduce（规约）算子。

内存密集型算子的分类。如图3所示，element-wise算子中的元素是独立地以逐元素方式处理的。这类算子可进一步分为轻量级（如add、sub）和重量级（如tanh、power、log）两种。Broadcast（广播）操作通常也被视为element-wise操作。Reduce算子接收一个张量作为输入，并对其一个或多个维度进行规约。如果规约的维度上元素在内存中是连续的，则称为行规约（row-reduce）；否则称为列规约（column-reduce）。行规约通常在一个GPU线程块内组织，相邻线程读取连续内存地址以提高效率。列规约用于处理不连续的元素，通常需要多个线程块和原子操作来保证正确性。


图 3: 典型的内存密集型操作。

现代模型中的复杂性。由于在现代ML计算图中reduce和broadcast操作的频繁使用，算子之间的张量形状变得日益多样化。例如，图1中的Transformer模型包含1666个reduce算子，约占总计算算子的10%。这些reduce算子可以形成任意的图拓扑结构，导致复杂的依赖关系。


图 4: Transformer模型中的一个典型子图。

2.2 内存密集型算子融合

现有融合技术的目标。诸如XLA【索引编号11，TensorFlow XLA】和TVM【索引编号18，{TVM}: An automated end-to-end optimizing compiler for deep learning, 2018, 13th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18)】等先进工作支持对通用的内存密集型算子进行内核融合，以减少片外内存访问、CPU-GPU上下文切换以及由频繁内核启动引起的框架级算子调度开销。

融合与代码生成能力。融合的一个最基本因素是代码生成能力。ML编译器通常根据其是否能生成高效代码来做出融合决策（例如，某些研究中的模式匹配过程）。例如，TVM/XLA的代码生成器通过将生产者的每个元素输入内联（inline）到消费者中来处理所有数据依赖关系，从而合并生产者和消费者。然而，我们在2.3节中发现这种代码生成方法存在一系列局限性。简单地修改融合决策逻辑（例如，扩大融合范围）可能会导致TVM/XLA的性能下降。

本文关注点。本文旨在解决基础的代码生成问题，而非表层的模式匹配问题。

2.3 当前技术的两大局限性

在生产环境中执行这些内存密集型ML模型时，出现了一系列独特的挑战，我们对此进行详细阐述。

2.3.1 挑战一：复杂的两级依赖与即时（Just-In-Time）需求加剧了训练/推理的低效性

问题描述。一个内存密集型子图通常由数十甚至数百个算子组成，存在两个层面的依赖关系：算子级依赖（描述算子间的连接，如图4中B和C依赖于A）和元素级依赖（指示张量内元素间的依赖，如图3中元素9依赖于元素1、4、4）。在算子层面，连接复杂；在元素层面，频繁的reduce和broadcast操作会引发许多一对多的依赖关系。此外，ML模型结构多样且定制化变体层出不穷，要求对任意给定的模型结构进行即时（JIT）优化，而非采用特定（ad-hoc）解决方案。这与传统领域（如HPC、数据库）中针对静态工作负载的内核融合方案不同，ML从业者需要通过频繁调整和执行来定制模型，这要求在每次运行时自动优化，而非每次试验前手动调优。因此，两级依赖与JIT需求的结合使得现代内存密集型ML模型的融合优化极具挑战性。

关键观察与低效根源。我们观察到，当前ML优化编译器（如XLA【索引编号11，TensorFlow XLA】、TVM【索引编号18，{TVM}: An automated end-to-end optimizing compiler for deep learning, 2018, 13th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 18)】）在这种两级依赖下无法执行高效融合。具体而言，对于具有复杂依赖的算子，这些编译器要么简单地跳过必要的融合，要么在融合后产生冗余计算。其根源在于，由于无法处理一对多元素级依赖，现有编译器无法高效融合两种常见的内存密集型模式：（1）reduce算子与其消费者（如图4中的橙色圆圈）；（2）昂贵的element-wise算子后跟broadcast算子（如图4中的蓝色和绿色圆圈）。当前ML编译器在对这两种模式进行融合时面临以下两难困境：

  • (i) 融合？导致严重的冗余计算。如果对这两种模式中的算子进行融合，我们观察到XLA和TVM都没有在线程间通信中间结果；它们仅利用每个线程的寄存器在编译器中融合算子。当存在一对多元素级依赖时（即生产者生成的一个元素被消费者的多个元素需要），消费者的每个线程都会独立计算这个共同的元素，导致显著的计算冗余。图5展示了TVM融合power<2> - broadcast<2,128> - add<2,128>时发生的冗余计算。在这种情况下，add的每128个元素都需要power产生的一个元素。然而，由于编译器无法通过其自动策略有效处理一对多依赖，power会在128个不同线程中为相同的值重复计算128次。Power是一个昂贵的element-wise算子，当张量形状很大时，会因冗余计算浪费大量GPU资源。此外，当融合reduce算子与其消费者时，冗余计算会更严重，因为消费者的每个线程都需要独立地重新计算整个规约过程。
    编译器自动优化与手动优化的区别：对于图5的情况，GPU专家可以手动将power的结果缓存在线程块内的共享内存中，以供add复用。然而，当给编译器一个具有频繁的一对多元素级依赖和分散的算子连接的子图时，要自动决定如何组织中间数据以供复用是非常棘手的。同时，编译器还需考虑算子的并行性，这进一步影响数据局部性。因此，现有ML编译器因缺乏对复杂依赖特性、并行性和局部性的联合自动优化考虑，在决定融合时会遇到严重的冗余计算问题。

    图 5: TVM在尝试用编译器自动融合power<2> - broadcast<2,128> - add<2,128>时出现的冗余计算。power的不同颜色代表处理输入张量中不同元素的线程。

  • (ii) 不融合？产生更多待执行的内核。为避免这些冗余计算，现有设计倾向于在遇到来自这两种模式的一对多依赖时跳过融合。例如，XLA在遇到模式(1)和(2)时会跳过融合，导致为内存密集型算子生成大量内核。对于Transformer模型，XLA为内存密集型计算生成的内核数量大约是计算密集型计算的3倍。而理想情况下，由于有效的融合应合并两个计算密集型算子之间的所有内存密集型算子，这两类计算的内核数量应大致相同。另一方面，TVM也会在reduce算子上跳过融合（即模式(1)），但会继续融合模式(2)，这虽然显著减少了生成的内核数量，但也引入了上述的大量冗余计算开销。

算子级冗余。在算子层面，一对多依赖（一个算子是多个算子的生产者）也可能导致冗余计算。例如，在图4中，算子B和C在现有ML编译器中会被分到两个不同的内核中，而算子A会冗余地内联在这两个内核里。

2.3.2 挑战二：真实世界生产工作负载中的不规则张量形状

问题描述。我们观察到生产工作负载中的张量维度常常表现出高度不均匀，我们称之为不规则张量形状。这使得编译器难以生成具有良好并行性的内核，因为它要求在预先不知道具体张量形状的情况下进行JIT优化。我们观察到，现有工作缺乏针对这一特性的自适应设计，导致性能不佳。

关键低效根源。不规则的张量形状导致在GPU上产生过多的小分区或过少的大分区。我们观察到生产工作负载中的张量形状通常与标准模型库中的基准测试不同。图6展示了两种由于生产环境中的张量形状导致的GPU并行性不佳的情况。
* 小块大小问题 (Small block size issue)：例如，在DIEN模型的一个真实案例中，处理从形状<750000,32>到<750000>的行规约时，XLA会产生750,000个GPU线程块，每个块大小仅为32，导致并行度低下。这是因为GPU能并发执行的线程块数量有上限；当线程块太小时，任意时刻的并发度也很低。
* 小块数量问题 (Small block count issue)：另一个案例是Transformer模型中的行规约，从形状<64,30000>到<64>。XLA自动生成64个大小为1024的线程块，而一块V100 GPU在同样块大小下可以并发调度160个线程块，导致严重的硬件资源未被充分利用。
因此，这些生产工作负载要求更好的编译器设计，以自动生成适应各种张量形状的线程映射。


图 6: 现有工作中典型的并行性不佳问题。

A2 方法细节

3. 关键设计方法论

高层目标。为应对之前讨论的挑战I和II，我们提出了AStitch，一个为内存密集型ML计算开启新的多维优化空间的机器学习优化编译器。它通过支持对任意内存密集型子图的高效即时（just-in-time）算子缝合来实现这一目标。本节将描述我们的核心思想：用于解决挑战I（第2.3.1节）的分层数据复用（第3.2节），以及用于解决挑战II（第2.3.2节）的自适应线程映射（第3.3节）。我们首先基于联合考量提出我们的算子缝合方案抽象（第3.1节），为依赖关系抽象的讨论做铺垫。

说明。本节重点介绍缝合的代码生成见解，这些见解既可以手动应用，也可以通过编译器优化实现。我们将在第4节基于本节的见解描述如何开发该优化编译器。

3.1 算子缝合方案抽象

四种缝合方案。如表1所示，我们抽象出四种类型的缝合方案，这些方案从依赖关系、内存层次结构和并行性的联合视角出发，覆盖了所有依赖场景。
* Independent (独立) 方案: 代表彼此独立的算子。
* Local (局部) 方案: 代表具有元素级一对一依赖关系的相邻算子（即element-wise方式），中间数据缓存在每个线程的寄存器中。这保证了线程级别的数据局部性。这两种方案已被现有设计【索引编号11, TensorFlow XLA; 索引编号18, {TVM}: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI'18; 索引编号55, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI'20】采用。
* Regional (区域) 方案: 表示存在一对多元素级依赖，中间数据缓存在GPU的共享内存中供其消费者使用。线程映射（并行性）应保证线程块级别的数据局部性。以图5中的例子为例，power的每个结果可以被缓存到共享内存中，并由执行add的多个线程复用。
* Global (全局) 方案: 用于应对任何复杂的依赖关系，中间数据缓存在全局内存中。这是一个面向并行性的方案，因此没有局部性要求（全局内存对所有线程可见）。与Regional方案类似，当前机器学习模型中最常见的情况仍是一对多元素级依赖。然而，有时优化块局部性会损害并行性，导致整体性能不佳。例如，块局部性要求将相应的线程组织到单个线程块中，这会妨碍通过将这些线程映射到更多线程块来获得更高的并行性。在这种情况下，AStitch会根据算子的特性（第4.3节）来权衡局部性与并行性（即Regional vs Global）。在第4节中，我们将讨论如何根据算子的线程映射自动确定最佳的缝合方案。

表 1: 联合考量下的缝合方案抽象。

3.2 分层数据复用阐述

分层数据复用打破融合困境。多种算子缝合方案促成了一种分层数据复用机制，借助该机制，我们可以高效地将任何算子缝合在一起，而不会产生大量的计算冗余。这有助于打破挑战I（第2.3.1节）中的困境。我们首先阐述对于每种算子缝合方案，算子间的中间数据是如何被维护和复用的。

两级数据复用。基于表1中的缝合方案，AStitch在内存层次结构（即寄存器、共享内存和全局内存）中实现了两个级别的数据复用，以消除融合困境：
* 元素级数据复用。在AStitch中，对于一对多的元素级依赖，生产者只处理每个数据项一次，不会产生冗余计算。其结果可以被保存在GPU共享/全局内存缓冲区中，供其消费者在regional/global缝合方案中复用。
* 算子级数据复用。对于一对多的算子级依赖，AStitch只处理生产者一次，并将其结果缓冲起来供多个消费者复用。对于local缝合方案，待复用的数据保存在寄存器中；而对于regional/global方案，数据则保存在共享/全局内存中。需要注意的是，在当前的ML编译器中，多个消费者（如图4中的算子B和C）可能会被分割到不同的内核中，导致生产者（如图4中的算子A）的冗余计算。

内核形态示例。在图7-(a)中，我们为一个示例性的子图展示了采用分层数据复用后生成的内核是什么样的。这里我们假设算子的缝合方案为：reduce.1采用regional方案，power.1和reduce.2采用global方案，multiply.1采用independent方案，其他算子采用local方案。图7-(b)和(c)展示了XLA和AStitch为该子图生成内核的不同方式。AStitch生成一个GPU内核，而XLA则生成4个内核，分别以reduce.1、power.1、reduce.2和multiply.1结尾。AStitch通过细粒度的数据管理和多级线程屏障，消除了3次内核启动，从而减少了CPU-GPU上下文切换和框架调度开销（在一个内核中使用额外的轻量级线程屏障）。reduce.1的输出无需刷新到片外内存，它被缓存在片上供其消费者读取（即元素级数据复用）。parameter.2和broadcast.2的值只需从片外内存加载一次并缓存在寄存器中（即算子级数据复用），而XLA则在不同的内核中从全局内存加载它们两次。此外，TVM会为这种情况生成3个内核，其中power.1和reduce.2被合并到一个内核中，这导致了power.1的冗余计算。AStitch避免了这类冗余计算。关于自动线程映射和缝合策略的详细设计优化将在第4.3节讨论。


a) 从真实工作负载中简化的内存密集型子图


图 7: 内存密集型子图的执行方案。AStitch通过分层数据复用减少了内核启动和片外内存访问。Pr: parameter。A: add。B: broadcast。R: reduce。D: divide。Pw: power。M: multiply。

全局屏障 (Global Barrier)。Global缝合方案需要一个作用于所有GPU线程的全局屏障，并满足以下约束：GPU线程块的总数不应超过每波（wave）在GPU上可同时执行的允许数量【索引编号50，Inter-block GPU communication via fast barrier synchronization, 2010, IPDPS】。否则，它将导致活动和非活动线程块之间的死锁。这一约束在第3.3节中得到满足，该节的方法有助于限制线程块数量，同时仍保持高并行性。与分离的内核相比，global缝合将内核调用之间隐式的全局线程屏障内联到单个内核中。

3.3 自适应线程映射

应对不规则张量形状。如第2.3.2节所述，现有ML编译器缺乏支持生产工作负载中不规则张量形状的设计（挑战II），导致了显著的性能问题。为解决此问题，我们提出了一种基于GPU执行的SIMT特性的任务打包和拆分方法，以自适应地处理各种张量形状。需要注意的是，受不规则张量形状影响的主要是reduce算子，它们是内存密集型计算中最耗时的算子。

任务打包 (Task packing)。任务打包包括水平和垂直两个维度。
* 水平打包：将多个处理单行规约的小块打包成一个大的线程块。这解决了图6-(a)中展示的小块大小问题。例如，对于从形状<750000,32>到<750000>的行规约，我们将32个大小为32的小块打包成一个大小为1024的线程块以增加并行性，线程块中每32个线程处理一行。
* 垂直打包：将多个线程块的任务打包成一个，以减少块数。这有助于将多波线程块打包成一波，以满足全局屏障的要求。在垂直打包中，块大小不变，每个线程按顺序处理来自多个任务的元素。图8-(a)描述了二维任务打包如何工作以增加GPU利用率，同时有效限制每波的块数。

任务拆分 (Task splitting)。另一方面，任务拆分是将一个线程块内的任务拆分成几个线程块，以在因块数过少导致利用率不足的情况下增加块数（图6-(b)）。它要求在拆分的块之间进行跨块规约以完成行规约。图8-(b)描述了如何通过跨块原子操作将一个块内的行规约任务拆分成两个块来增加并行性。


图 8: GPU上行规约的任务打包和拆分优化。现有工作中的任务映射见图6。

4. 编译器设计与优化

挑战与观察。要为一个给定的复杂内存密集型子图自动应用第3节中描述的优化是具有挑战性的。需要为所有算子即时确定缝合方案和线程映射。然而，枚举每个算子的方案会导致组合爆炸。幸运的是，根据缝合方案的抽象，我们识别出内存密集型算子的两个基本特性。我们有一个关键观察：AStitch只需要确定几个关键算子的缝合方案和相应的线程映射，这些信息将传播到子图中的所有其他算子。接下来我们描述AStitch如何自动实现所提出的优化。

4.1 缝合范围识别

识别待缝合的算子。首先，我们描述AStitch如何在一个给定的机器学习模型中识别出哪些算子需要缝合在一起。为了最小化CPU-GPU上下文切换、框架调度开销并减少片外内存访问，AStitch会尽可能大地将内存密集型算子缝合在一起（在资源限制下）。值得注意的是，AStitch能够有效管理硬件资源，以防止较大内核的资源爆炸（第4.4节，第4.5节）。

缝合操作的形成。给定一个ML计算图，AStitch首先应用BFS算法自动识别内存密集型子图。然后，它将每个子图替换为一个名为stitch op的新算子。为了进一步扩大缝合范围，AStitch将不相连的stitch op组合在一起，形成一个更大的stitch op，我们称之为远程缝合（remote stitching）。远程缝合的过程是遍历所有现有的stitch op，如果它们之间没有数据依赖，就将两个合并。形成stitch op的一个约束是不允许出现循环依赖。如果合并会导致循环，AStitch则不会合并这些计算。每个stitch op最终将被编译成一个CUDA内核。

4.2 关键设计观察

关于缝合方案抽象的观察。我们对缝合方案的抽象（第3.1节）有以下观察：
* 观察-A: 如果一个算子是local方案，其线程映射可以通过从其消费者的线程映射传播来确定。这是因为local方案表示元素级一对一的依赖关系。给定消费者算子的调度，其生产者的调度可以通过逐元素索引传播直接推导出来。同理，如果一个算子及其消费者都是local方案，该算子的线程映射调度可以传播给其消费者。
* 观察-B: reduce和昂贵的element-wise算子后跟broadcast算子的模式需要由regional或global方案支持，因为这两种模式会引发复杂的一对多元素级依赖（见表1）。

4.3 自动编译器优化设计

三步法自动优化。根据这两个关键设计观察，我们提出了一个即时编译器设计，它能自动为每个算子确定缝合方案和线程映射。以图7-(a)中的图拓扑为例，图9展示了编译器如何分三步工作。


图 9: 对图7-(a)所示图拓扑的调度传播和数据管理规划。

第1步：主导算子识别与算子分组 (dominant identification and op grouping)。根据第4.2节的观察，AStitch只需确定几个关键算子（我们称之为主导算子，dominant ops）的线程映射，然后将它们传播到所有其他算子。AStitch首先识别几个主导算子的候选者，并通过主导合并最终确定。
* 候选主导算子识别：根据观察-B，轻量级element-wise算子和不后跟broadcast的昂贵element-wise算子默认为local方案。非local方案的算子将被视为主导算子的候选者。这意味着reduce和后跟broadcast的昂贵element-wise算子是候选者。一个特例是stitch op的输出（如图9中的M.1），它也被视为主导算子候选者。根据此规则，在图7-(a)中，reduce.1、power.1和reduce.2是主导算子的候选者。
* 主导合并与最终确定：如果两个候选主导算子通过仅local方案的算子连接，AStitch会选择其中一个作为最终主导算子，并将另一个视为次主导算子（sub-dominant）。这样我们只需确定主导算子的线程映射调度，即可通过传播得到次主导算子的调度。在合并时，AStitch倾向于选择reduction作为最终主导算子，因为使用像reduction这样耗时的算子来生成整体调度通常会带来更好的性能。
* 算子分组：最后，AStitch会为每个主导算子形成一个组，包含所有仅通过local方案与该主导算子连接的算子。这样，不同的组仅通过主导和次主导算子相互连接。之后，AStitch将在每个组内传播线程映射，并通过确定主导和次主导算子的缝合方案来明确不同组如何通信。在图9-1中，reduce.1和reduce.2是最终主导。Power.1（和multiply.1）可以通过仅local方案的算子连接到reduce.1（和reduce.2）。因此，Power.1和multiply.1被视作次主导。

第2步：自适应线程映射与调度传播 (adaptive thread mapping and schedule propagation)。AStitch根据第3.3节为每个主导算子生成并行代码，并在相应的组内传播线程映射调度。通过这种方式，我们得到所有算子的线程映射调度。
* 张量形状自适应：AStitch根据张量形状和硬件资源自动对主导算子应用任务打包和拆分。以行规约为例，如果要规约的行数小于每波允许的块数，且每行包含大量数据项（即大于1024），AStitch会拆分行以增加并行性。否则，AStitch将应用任务打包以形成足够大的块，并限制块数小于每波的阈值，以满足全局屏障的要求。
* 主导合并的好处：值得强调的是，第1步中的主导合并可以实现更多的算子级数据复用（第3.2.1节）。例如，在图7-(a)中，如果不将reduce.2和multiply.1合并到一个组中，broadcast.2将出现在两个组中。这可能导致broadcast.2在两个组中有不同的线程映射调度；因此，由于调度不兼容，该算子的每个线程加载的值无法在两个组之间复用。在我们的主导合并下，由于整个组合并后由相同的线程映射调度主导，broadcast.2现在可以被缓存到寄存器中以供复用。

第3步：最终确定 (Finalization)。AStitch在最后一步确定主导和次主导算子的缝合方案。注意，每个组内的其他算子均为local方案（在第1步中已讨论）。不同组通过主导和次主导算子连接，其缝合方案为regional或global。
* 局部性检查与方案选择：Regional方案要求块级数据局部性。对于一个输出为regional方案的算子，当该算子在一个块内产生一定范围的数据时，其消费者应在同一个块内读取相同范围的数据。否则，该方案应回退到global。我们应用局部性检查来在regional和global之间确定缝合方案。
* 被动块局部性检查：以图7-(a)为例，我们计算reduce.1、power.1和reduce.2每个块产生的连续元素数量，以及它们的消费者每个块需要的连续元素数量。在这种情况下，只有reduce.1匹配块局部性，并被赋予regional方案，数据缓存在GPU共享内存上；其他两个被赋予global方案，数据缓存在全局内存上。
* 主动块局部性适应：对于一个只包含element-wise算子的算子组，AStitch会主动调整该组的线程映射调度，以匹配其生产者组的块局部性。根据该组的生产者每个线程块生成多少连续元素，可以为该组以逐元素的方式确定线程映射调度。
* 总结：Reduce算子优先考虑并行性，因为它们需要更多计算。因此，由reduce主导的组将执行被动的块局部性检查，而不调整其线程映射。另一方面，element-wise算子由于其计算成本低，通常优先考虑局部性。因此，由它们主导的组会执行主动的块局部性检查，以实现更多的块级局部性。

4.4 内存使用优化

内存资源管理。适度使用GPU共享/全局内存至关重要，例如，高共享内存使用会损害内核并行性。AStitch在考虑内存资源限制及其对性能影响的前提下运行。例如，AStitch会尽可能复用先前分配的内存，以减少不必要的内存分配请求。AStitch使用支配树算法（dominance tree algorithm）【索引编号19，A simple, fast dominance algorithm, 2001, Software Practice & Experience】进行内存数据流分析，以最大化算子的内存复用。

方案回退机制。此外，regional方案需要额外的共享内存。如果一个共享内存请求超过了线程块的硬件限制，AStitch会逐一将一个stitch op内的主导和次主导算子的缝合方案从regional更改为global，直到内存使用量低于限制。

4.5 资源感知的启动配置

全局屏障约束与启动配置。AStitch需要一个合适的GPU内核启动维度，其中线程块数不超过每波最大块数（$max\_blocks\_per\_sm$），以满足全局屏障的要求。然而，AStitch在进行GPU二进制编译的线程映射规划时需要该信息，但该信息只有在编译后才能获得。我们设计了一种“假设-放宽-应用”（assume-relax-apply）方法来解决这个问题。基本思想是：在优化前假设一个目标$max\_blocks\_per\_sm$，然后得到一个放宽的寄存器使用量，最后通过CUDA编译器注解来应用寄存器使用限制以达到目标。

具体步骤。
1. 假设 (Assume): 首先，我们假设寄存器使用上限是一个非常小的值（即32）。
2. 放宽 (Relax): 其次，我们尝试在可能的情况下放宽这个上限。其洞见是，并行性可能受限于共享内存使用而非寄存器使用，此时我们可以放宽寄存器使用上限。我们根据假设的寄存器使用上限、规划的共享内存使用（第4.4节）和指定的块大小【索引编号4，CUDA Occupancy Calculator】来计算$max\_blocks\_per\_sm$。在此过程中，我们将块大小指定为CUDA允许的上限（例如1024），因为更大的块大小会导致更小的$max\_blocks\_per\_sm$，进而导致更小的全局屏障开销。然后，我们根据$max\_blocks\_per\_sm$推导出允许的最大寄存器使用量，这个最大允许值就是放宽后的上限。
3. 应用 (Apply): 最后，AStitch在将线程映射调度降级到GPU IR时，添加注解信息以应用最大寄存器上限。在我们的评估中，使用此方法没有观察到寄存器溢出问题。

5. 实现

系统概览。图10展示了AStitch的系统概览。AStitch的设计和见解可以直接应用于所有当前的机器学习框架和ML优化编译器，以增强其功能。


图 10: AStitch 系统概览。

与TensorFlow/XLA集成。目前，我们将AStitch实现为TensorFlow的附加组件。AStitch利用TensorFlow内置的XLA引擎进行系统实现，保留了XLA的所有优化，除了其融合策略和代码生成通道。AStitch接受XLA表示的计算图，但替换了XLA的融合和代码生成过程。

可移植性与易用性。我们将AStitch构建为一个独立的编译器引擎，可移植到任何版本的TensorFlow。它利用TensorFlow的自定义图优化通道API（custom graph pass API）来重写计算图并为stitch op生成GPU代码。然后，它利用TensorFlow的自定义算子API（custom op API）将stitch op插入到计算图中算子的运行时序列中。用户只需指定AStitch引擎的路径和一个环境变量来初始化它，无需修改模型脚本中的任何内容，使AStitch成为一个非常易用的机器学习编译器。

A4 实验

本节我们使用单块NVIDIA V100 GPU（16GB设备内存）、CUDA工具包10.0和cuDNN 7.6进行详细评估。

实验环境

  • 数据集与模型: 评估使用了一系列具有代表性的内存密集型ML应用，包括：
    • BERT【索引编号21，BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, 2018, arXiv】 和 Transformer【索引编号44，Attention is all you need, 2017, arXiv】 用于自然语言处理（NLP）。
    • DIEN【索引编号58，Deep interest evolution network for click-through rate prediction, 2019, AAAI】 用于推荐系统。
    • ASR【索引编号47，ESPnet: End-toEnd Speech Processing Toolkit, 2018, Interspeech】 用于自动语音识别。
    • CRNN【索引编号40，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】 用于光学字符识别。
      这些模型在生产环境中被广泛使用，其构件包括感知机、注意力机制、卷积、RNN以及广泛的内存密集型算子。所有工作负载的批量大小（batch size）均选自真实世界的生产配置，具体见表2。
      表 2: 评估使用的工作负载。
  • 硬件配置:
    • GPU: 单块 NVIDIA V100，16GB HBM2 内存。
    • 在6.1.1节中也使用了NVIDIA T4 GPU进行推理评估。
  • 软件配置:
    • CUDA Toolkit 10.0, cuDNN 7.6。
    • 基线系统: TensorFlow (v1.15), TensorFlow XLA (v1.15), TensorRT (v7.0)。
    • TVM Ansor【索引编号55，Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI'20】 也被用于特定案例研究。
  • 评估指标: 评估推理时间或单次迭代的训练时间，重复10次取平均性能。

实验结果

6.1.1 总体结果

性能加速比。如图11所示，AStitch在所有工作负载上均优于其他对比技术。
* 推理任务 (图11a):
* 与TensorFlow相比，AStitch最高实现4.06倍加速，平均2.37倍。
* 与XLA相比，最高实现2.73倍加速，平均1.84倍。
* 与TensorRT相比，最高实现4.46倍加速，平均2.47倍。
* 训练任务 (图11b):
* 与TensorFlow相比，AStitch平均加速1.34倍。
* 与XLA相比，平均加速1.30倍。
* 其他平台与优化: 在NVIDIA T4 GPU上以及与自动混合精度（AMP）【索引编号2，Automatic Mixed Precision for Deep Learning】优化结合使用时，AStitch也表现出与图11相似的加速效果（如图12所示），证明了其通用性和兼容性。


图 11: 端到端性能加速比。


图 12: 推理加速比，基线和AStitch均开启AMP优化。

6.1.2 性能分解

时间开销分解。如图13所示，AStitch显著减少了内存密集型算子执行时间（MEM）和非计算开销（OVERHEAD）。例如，对于Transformer模型，AStitch节省了约2/3的OVERHEAD时间和1/4的MEM时间。OVERHEAD的减少主要得益于内核调用次数的减少，而MEM的性能提升则源于并行度的增加。


图 13: 性能分解，未显示计算密集型算子的时间。

内核调用减少。如表3所示，得益于彻底的缝合，AStitch平均减少了65.7%的内存密集型计算内核调用，大幅降低了上下文切换导致的非计算开销。此外，AStitch平均减少了43.2%的CUDA memcpy/memset活动。

表 3: 内核数量。MEM: 内存密集型算子的内核。CPY: CUDA memcpy/memset调用。

并行度提升。如图14所示，我们使用nvprof工具收集了排名前80%的内存密集型内核的两个性能计数器：achieved_occupancy（占用率，指示并行度）和sm_efficiency（SM效率，指示GPU利用率）。得益于自适应线程映射，AStitch全面提升了并行度和GPU利用率。


图 14: 排名前80%的内存密集型计算的平均并行度。AS: AStitch。Occu: 占用率。Effi: SM效率。

6.1.3 CRNN案例综合研究

消融研究。如表4所示，在CRNN模型上，我们逐一启用AStitch的优化。
1. 在XLA的融合范围上启用自适应线程映射（ATM），性能提升8.9%。
2. 在ATM基础上应用分层数据管理的彻底缝合（HDM），额外提升8.2%。
3. 在HDM基础上加入主导算子合并（完整的AStitch），最终再提升18.7%。

表 4: CRNN的消融研究。

性能计数器分析。
* 并行性: 如图15所示，AStitch生成的核心内核相比XLA具有更高的占用率和SM效率。
* 内存流量: 如表5所示，AStitch减少了GPU全局内存的读写事务，这得益于分层数据管理将大量中间值缓存在片上。
* 指令数: 如表5所示，由于减少了冗余计算，AStitch减少了fp32指令数。


图 15: CRNN的占用率和SM效率趋势。X轴表示按执行时间降序排列的内存密集型算子。注意AStitch的算子数量更少。


表 5: CRNN中所有内存密集型算子的总性能计数器。DR_transactions: dram读事务。DW_transactions: dram写事务。

6.2 与TVM Ansor的案例比较

BERT推理案例。在BERT推理任务上，AStitch耗时31.75ms，而经过2000次试验调优的Ansor耗时42.02ms，AStitch实现了1.3倍的加速。AStitch为内存密集型算子生成的GPU内核比Ansor少53%，并且在所有内存密集型计算上实现了1.4倍的加速。如图16所示，AStitch生成的核心内核比Ansor具有更高的并行度和硬件利用率，同时减少了近40%的片外内存事务。


图 16: BERT的占用率和SM效率趋势。坐标轴含义与图15相同。注意AStitch的算子数量更少。

6.3 生产环境评估

AStitch已部署在一个生产集群中，并在一周内为70,000个任务节省了约20,000个GPU小时，证明了其鲁棒性。AStitch有效优化的工作负载主要包括基于Transformer的模型、推荐模型和RNN模型。

6.4 开销分析
  • 优化开销: 对于包含5000到10000个节点的计算图，AStitch的优化开销平均为90秒，而原生XLA为30秒。这是一次性开销，相比于搜索和调优类优化仍然高效。
  • 全局屏障开销: 如表6所示，在V100上，全局屏障的开销不超过2.72微秒，低于10微秒量级的内核启动开销。在CRNN模型上的实验表明，移除全局屏障（结果不正确）并未观察到明显性能提升，说明全局屏障不是瓶颈。

表 6: 内联全局屏障的开销。

A5 结论

本文揭示了内存密集型计算是近期机器学习模型中一个日益严峻的性能关键因素。我们提出了分层数据复用技术来处理复杂的依赖关系以扩大融合范围，减少非计算开销。我们提出了自适应线程映射技术来解决不规则张量形状的问题。我们开发了一个名为AStitch的JIT编译器，集成了这些优化，并具有很高的易用性。实验结果表明，AStitch相比于最先进的编译器最高可实现2.73倍的加速。我们相信，AStitch填补了机器学习编译器领域一个长期被忽视的空白。

A6 附录

A.1 摘要

本构件（artifact）包含用于验证AStitch论文中主要结果的必要软件组件。我们提供一个docker镜像以简化环境设置。该镜像包含AStitch的已编译二进制文件、评估推理和训练性能的脚本，以及绘制图表的脚本。它需要一个在配备NVIDIA V100 GPU的x86_64机器上运行的Linux系统，且该系统需装有能够运行CUDA 10.0的NVIDIA驱动。启动docker容器后，用户可以运行一个脚本来收集所有性能数据。这需要一些手动操作，将性能数据填入几个python脚本中，以绘制论文中最重要的图表，展示AStitch的加速比和分解信息。

A.2 构件清单（元信息）

  • 二进制文件: AStitch的docker镜像。
  • 运行时环境: 带有NVIDIA驱动（能运行CUDA 10.0）的Linux系统。
  • 硬件: NVIDIA V100 GPU。
  • 输出: 性能结果和显示加速比的图表（需要一些手动收尾工作）。
  • 实验: 对于论文中评估的所有推理工作负载，它包含了对原生TensorFlow、XLA、TensorRT和AStitch的评估。对于BERT和Transformer训练，它包含了对原生TensorFlow、XLA和AStitch的评估。
  • 所需磁盘空间（大约）: 12GB。
  • 准备工作流程所需时间（大约）: 30分钟。
  • 完成实验所需时间（大约）: 下载docker镜像需要数十分钟。然后可以一次性运行一个脚本收集所有性能结果，执行过程约2.5小时，期间可以做其他事情。最后，需要约20分钟手动整理以绘制加速比和分解图。
  • 是否公开: 是。docker镜像公开，其中包含已编译的二进制文件。源代码正在开源过程中，我们将在2022年初发布。
  • 代码许可证: Apache-2.0。
  • 数据许可证: Apache-2.0。
  • 存档DOI: 10.5281/zenodo.5733989

A.3 描述

A.3.1 如何访问

我们同时在dockerhub和zenodo上提供了docker镜像。
Docker-hub URL: https://hub.docker.com/r/jamesthez/astitch/tags
Zenodo URL: https://zenodo.org/record/5733989

A.3.2 硬件依赖

配备NVIDIA V100 GPU的x86_64机器。

A.3.3 软件依赖

能够运行CUDA 10.0的带NVIDIA驱动的Linux系统。

A.4 安装

您只需拉取docker镜像并启动一个容器：

docker pull \
    jamesthez/astitch:astitch_asplos_ae
docker run \
    --gpus all --net=host --pid=host -it \
    --name <your-container-name> \
    jamesthez/astitch:astitch_asplos_ae bash

或者，您可以下载docker镜像的tar文件并导入。下载URL是https://http://zenodo.org/record/5733989 ：

gzip -d astitch_asplos_ae.tar.gz
docker import - astitch_asplos_ae < \
    astitch_asplos_ae.tar
docker run \
    --gpus all --net=host --pid=host -it \
    astitch_asplos_ae bash

如有需要，请使用sudo运行docker。

A.5 评估与预期结果

我们在提供的docker镜像中提供了一个read-me文件，说明如何复现关键结果。您可以在启动docker容器后在/root目录下找到它。