ALCOP: Automatic Load-COmpute Pipelining in Deep Learning Compiler for AI-GPUs

文章标题:ALCOP:面向AI-GPU的深度学习编译器中自动加载-计算流水线技术
作者/机构:Guyue Huang (UC Santa Barbara), Yang Bai (Chinese University of Hong Kong), Liu Liu (Rensselaer Polytechnic Institute), Yuke Wang (UC Santa Barbara), Bei Yu (Chinese University of Hong Kong), Yufei Ding (UC Santa Barbara), Yuan Xie (UC Santa Barbara, Alibaba DAMO Academy)

A1 主要贡献

核心问题:在GPU上,数据加载和计算之间的流水线化是关键的张量程序优化。为了充分发挥最新GPU的高性能,必须在GPU的多级缓冲层次结构中进行多阶段流水线的协同优化。然而,积极的大规模分块(Tiling)会限制分块(tile)的数量,从而阻碍分块间的并行性,而这正是实现高利用率的关键GPU机制。因此,恢复因积极分块而损失的并行性成为一项重要任务。流水线化(Pipelining)——即数据加载与计算的重叠——是释放分块内并行性的理想机制。

现有方法局限性:现有的深度学习框架依赖于cuBLAS等手写库来进行流水线优化,这种方法对新算子不具备可扩展性,也无法与张量编译器先前的优化(如自动融合和自动分块)相结合。一些研究工作仅探讨了双缓冲(double-buffering),这只是多阶段、多级流水线设计空间中一个简单特例(两阶段、单级别),限制了性能提升的潜力。

本文方法与创新
本文提出了ALCOP(Automatic Load-COmpute Pipelining),这是第一个编译器原生且完全支持多阶段、多级流水线的框架。ALCOP的架构如下图所示:

Fig. 4: ALCOP概览。

ALCOP通过利用深度学习编译器的渐进式降低结构以及在每个级别暴露的信息,将自动流水线化这一复杂问题分解为三个解耦且协作的编译模块,以应对工作负载复杂性硬件复杂性设计空间复杂性这三大挑战:
1. 流水线缓冲区检测:在调度阶段进行,此时整个数据流可见,能够处理多样化的深度学习算子。
2. 流水线程序变换:在程序变换阶段,对复杂的for循环结构和数据移动进行修改,利用前一模块的安全检查来执行稳健的变换。
3. 分析模型指导的设计空间搜索:在自动调优阶段,将流水线和其他技术协同优化,并设计了一个分析硬件模型以加速设计空间搜索。

主要贡献
1. 设计了检查每个缓冲区以应用流水线的方法,包括流水线和其他调度变换的排序,以避免相互干扰。(章节 II)
2. 设计了一个程序变换过程,处理索引操作、同步注入和序言(prologue)注入等变换。(章节 III)
3. 提出了一个感知流水线的分析性能模型。将其与现有的基于机器学习(ML)的调优算法相结合,显著提高了调度调优的效率。(章节 IV)

性能优势
- 动机示例:图1展示了流水线对于提升大规模分块下GPU利用率的重要性。仅使用分块会导致次优性能,而流水线能释放分块内并行性,从而提升性能。
(a) 分块、分块间并行和流水线并行概念。

(b) 激励示例:在NVIDIA A100上测试2048 × 2048 × 2048矩阵乘法,采用不同分块和流水线选择的性能。

Fig. 1: 自动流水线的动机。(a-3)解释了流水线(即数据加载与计算重叠)的概念。(b)给出了一个激励示例。仅使用分块时,性能总是次优。流水线释放了块内并行性,并提高了大分块下的性能。
- 多阶段与多级流水线概念:图2和图3说明了多阶段和多级流水线的重要性。双缓冲(两阶段)不足以隐藏数据加载延迟,而多阶段流水线可以实现计算单元的完全利用。类似地,在GPU的多级内存(共享内存、寄存器文件)上进行多级流水线并结合内部流水线融合,可以获得最佳性能。

Fig. 2: 多阶段流水线概念。(a) 两阶段流水线(或称双缓冲)不足以隐藏数据加载延迟。(b) 四阶段流水线可以隐藏数据加载延迟并实现计算单元的完全利用。ALCOP支持多阶段流水线。
(c) 无内部流水线融合的多级流水线。

(d) 带内部流水线融合的多级流水线。

Fig. 3: 多级流水线和内部流水线融合的概念。(a)展示了GPU内存层次结构,有两级缓冲区:共享内存和寄存器文件。(b)展示了单级(仅共享内存)流水线的执行时间线。(c)通过对内层循环(寄存器加载和计算)进行流水线化,改进了(b)。(d)通过内部流水线融合改进了(c),它将重复的内层循环视为一个整体循环并对其进行流水线化。ALCOP支持(d)中的优化,其性能在(b)-(d)中最佳。

A2 方法细节

II. 调度变换

自动流水线化的起点。自动流水线化的第一步是识别潜在的流水线机会。我们通过在编译器中实现一个调度变换过程来实现这一点,该过程将流水线原语附加到程序中的缓冲变量上。流水线可以应用于一个“加载-使用”循环,其中加载步骤将数据复制到缓冲区,使用步骤从该缓冲区读取数据。因此,调度变换的目的是识别并记录程序中的“加载-使用”结构。该过程会将这类加载-使用循环内的缓冲变量标记为流水线化缓冲区。随后,在第三节中描述的程序变换过程会将该加载-使用结构转换为其流水线化版本。

需要解决的关键问题。必须解决两个重要问题:首先,我们必须确定应用什么规则来识别可以进行流水线化的缓冲区。其次是确定流水线相对于其他调度变换(如分块)的顺序,并意识到它们之间的相互影响。

A. 用于流水线化的缓冲区识别

流水线化的约束条件。流水线化的约束不仅来自算法层面,即缓冲区如何被使用,也来自硬件能力,即何种形式的内存复制可以被异步执行。对于每个缓冲变量,我们评估以下三个规则来确定是否可以应用流水线。

规则一:由异步内存复制产生。我们不 对不是由异步内存复制产生的缓冲区进行流水线化。异步内存复制意味着内存复制是非阻塞的,因此我们可以提前为未来的循环迭代发起内存复制,同时继续当前迭代的计算。只有当遇到显式的同步指令时,程序才会阻塞以等待内存复制的完成。如果缓冲区中的数据不是由直接内存复制产生,而是由某些计算操作产生,则该缓冲区不满足此条件。

规则二:在顺序循环内产生。我们不对在顺序循环之外产生的缓冲区进行流水线化。流水线化的目的是将未来迭代的数据加载操作与当前迭代的计算重叠。这个加载-使用循环必须是顺序的,不能被并行化(绑定到并行线程)或展开。这个条件通常在那些使用分块来增加输入张量重用,但缓冲区只被填充和使用一次的模板(stencil)算法中被违反,这与在顺序循环中生成缓冲区的情况相反。因此,流水线化方法不能应用于这些缓冲区。

规则三:同步点的匹配。最后一条规则是关于流水线的同步:如果硬件平台仅支持基于作用域(scope-based)的同步,我们会检查同一作用域内的所有缓冲区,如果它们的同步位置不匹配,则拒绝为它们进行流水线化。同步流水线需要特殊的内存屏障,以等待特定的加载指令(例如,在4阶段流水线中,倒数第四次迭代发出的指令)。在NVIDIA Ampere GPU上,这种内存屏障是为共享内存作用域提供的。因此,如果两个缓冲区都在共享内存作用域内,但它们的屏障必须插入在程序中的不同位置,硬件将无法解决这种冲突。如果发生这种冲突,我们的调度变换会拒绝为这些缓冲区进行流水线化。

B. 调度变换的顺序

与其他变换的交互。流水线化适用于三种已有的调度变换:缓存读取(cache-reading)、分块(tiling)和融合(fusion)。我们将简要介绍这些变换,然后确定流水线应该在它们之前还是之后应用。

缓存读取(Cache-reading)应先于流水线。缓存读取是指为一个张量输入插入一个读取缓冲区。给定一个算法和从张量S1计算出的张量S2,应用缓存读取意味着插入一个新张量S1_buf,它是S1的一个相同副本,但具有一个缓冲区作用域。缓存读取应该在流水线之前应用,因为流水线需要应用于前者生成的缓冲区。

分块(Tiling)应先于流水线。分块是将输出张量划分为块的过程。与缓存读取相结合,它可以在缓冲区内缓存数据以提高数据重用。分块也应该在流水线之前执行。一个缓冲区是否有资格进行流水线的第二个条件,即是否存在一个顺序的加载-使用循环,必须基于分块后的for循环结构来检查。

内联融合(Inlining)应后于流水线。融合意味着避免在两个算子之间写回中间数据。内联是一种特定类型的融合,它应该在流水线之后进行。内联一个张量意味着在它被使用的地方精确地产生该张量的值;这项技术常用于轻量级的逐元素操作,如数据类型转换。图5展示了一个例子,其中S2最初是通过对S1应用逐元素函数f(·)产生的,并且通过缓存读取在S2之后注入了一个缓冲张量S2_buf。内联S2相当于先应用f(·),然后将数据直接复制到S2_buf中,而不将数据写回内存。根据我们在前一小节中概述的第一个规则,一个流水线化的内存缓冲区应该由异步内存复制产生。然而,对于这里的S2_buf,产生它的操作不再是异步的,因为显式的f(·)强制程序停顿,等待数据加载。由于缓冲区S2_buf不是异步产生的,它不能被流水线化。这里在案例1中,内联阻碍了流水线化的机会。然而,如果流水线在内联之前应用,如案例2所示,S2的内联仍然可以应用,但方式不同:我们不是将S2内联到S2_buf中,而是缓存读取S1,并将计算f(·)融合到S3的产生过程中。因此,我们确保了两方面都得到满足:缓冲区通过异步复制产生并且可以被流水线化,同时计算f(·)被融合,我们避免了显式生成一个中间张量。


Fig. 5: 内联和流水线优化顺序的有效性研究。在案例1中,内联后,S2_buf不再能被流水线化,因为它不再由异步内存复制产生。在案例2中,流水线化后,内联仍然可以应用。

III. 程序变换

本节介绍自动流水线化的第二个组成部分:变换程序中间表示(IR)以实现流水线。在第二节概述的调度变换之后,程序被降低到其IR形式,由for循环和加载/存储/计算操作组成。图7给出了流水线过程的示例输入和变换后的IR。图6也描绘了变换步骤。


Fig. 6: 流水线程序变换的工作流程及示例输入输出

A. 分析

第一步:收集流水线提示。给定一个程序IR,流水线化的第一步是收集由调度变换插入的流水线提示,包括要进行流水线化的缓冲区和每个缓冲区的阶段数。

第二步:重构生产者和消费者。给定一组我们想要应用流水线化的缓冲区,第二个分析任务是重构这些缓冲区的生产者张量和消费者张量。然后我们可以通过判断一个流水线缓冲区的生产者是否也是一个流水线化缓冲区来推断是否存在多级缓冲区。由于流水线化缓冲区总是通过异步内存复制产生的,要确定生产者张量,只需检索它从哪个张量复制而来即可。消费者的确定发生在IR遍历遇到从该缓冲区加载数据的操作时。

第三步:确定顺序的加载-使用循环。这一步是为每个流水线化缓冲区确定其顺序的加载-使用循环。这标识了要进行流水线化的迭代变量,并且是变换步骤中所有索引移位操作所必需的。顺序循环可以如下确定:从将数据复制到缓冲区的指令开始,从内到外遍历所有的for循环,找到第一个其迭代变量不用于索引此缓冲区内部的顺序循环。这意味着该缓冲区在该循环的每次迭代中都被重用,这就是我们想要进行流水线化的循环。以图7为例,A_shared的流水线循环是迭代变量为ko的循环,而A_reg的流水线循环是迭代变量为ki的循环。

第四步:记录加载和使用代码片段。我们应该记录加载和使用该缓冲区的代码片段。这些信息是注入同步原语和序言(prologue)所必需的。在图7的输入IR中,A_shared的“加载”部分是第8行,“使用”部分是ki循环及其内部的所有内容。A_reg的加载部分是第13行,“使用”部分是第15行。

第五步:确定序言注入位置。我们还需要决定在何处注入序言。因为我们将程序转换为在为当前迭代进行计算的同时为未来迭代发出内存复制,我们需要将内存复制的前几个阶段移到主加载-使用循环开始之前。这个预先放置的加载代码块就是一个序言。通常,序言可以简单地注入在流水线循环之前。然而,当出现多级流水线时,内部流水线的序言必须注入到最外层流水线的顺序循环中,以便构建一个整体的流水线,而不是如图3c所示的递归流水线。

B. 变换

将一个加载-使用循环变换为一个流水线循环需要五个步骤。图7中的变换后IR(Transformed IR)显示了同一图中输入IR(Input IR)的变换版本。


Fig. 7: 一个示例,说明如何将原始的Tensor-IR(左)转换为其流水线版本(右)。

第一步:增加缓冲区大小。这一步将内存缓冲区的大小增加流水线阶段数的倍数。相关的变换后代码以浅黄色高亮显示。

第二步:移动内存访问索引。这一步移动内存访问中使用的索引。相关代码以蓝色高亮显示。在每个加载-使用迭代中,我们为未来的迭代而不是当前迭代发出异步内存复制。因此,我们需要在内存访问索引中增加流水线循环变量。例如,如果是一个3阶段的流水线,我们应该提前2个迭代加载数据。

第三步:处理缓冲区滚动和越界回绕的索引。相关代码以绿色高亮显示。我们需要在两种情况下对索引进行回绕:首先,当我们使用流水线变量来索引缓冲区的一个块时,我们应该使用流水线迭代变量除以流水线阶段数的模。其次,由于我们增加了流水线变量,可能会索引到其生产者张量的边界之外。我们必须取流水线变量除以其自身范围的模,以避免索引越界。一个复杂的情况是,在多级流水线中,内部流水线的溢出导致外部流水线变量的增加。变换后IR的第26行处理了这种情况。

第四步:注入序言原语。序言的内容是前n_stage - 1个数据块的内存复制,其中n_stage是流水线阶段数。我们在前述分析过程中记录的位置注入序言。

第五步:注入同步原语。最后一步是注入同步原语。流水线由四个原语保护:producer_acquireproducer_commitconsumer_waitconsumer_releaseproducer_commit提交一批异步加载操作。consumer_wait阻塞直到前一批加载完成。当流水线已满时,producer_acquire会阻塞直到consumer_release被调用。生产者/消费者原语对分别放置在缓冲区的加载/使用部分周围,如变换后IR的第15、17、22和30行所示。

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

IV. 静态分析引导的调优

本节介绍我们如何将静态分析性能模型与现有的基于机器学习(ML)的自动调优【【索引编号19,Learning to optimize tensor programs+2018+arXiv】】相结合来选择调度参数。关键组成部分是一个新颖的性能模型,它能够感知流水线及其与其他优化的相互作用,如图8所示。

Fig. 8: 性能模型的高层视图。与先前工作【【索引编号20,Delta: Gpu performance model for deep learning applications with in-depth memory system traffic analysis+2019+ISPASS】】相比,我们的模型考虑了流水线、分块和空间并行性之间的约束和权衡。

A. 顶层模型

内核延迟模型。我们的分析模型如表I所示。在顶层,线程块被分组为线程块批次(threadblk batch),一个线程块批次一次性占用所有流式多处理器(SM)。由于所有线程块执行相同的程序,一个内核的延迟等于线程块延迟乘以批次数。内核中的线程块批次数取决于GPU的调度策略,我们通过性能剖析来学习。每个SM的最大线程块数受限于每个SM可提供的共享内存和寄存器文件的大小,以及线程块的请求。我们的模拟GPU调度策略会考虑所有这些因素来决定$N_{threadblk\_batch}$。

TABLE I: 分析性能模型
$T_{kernel} = T_{threadblk} \times N_{threadblk\_batch}$
Input: $T_{load}, T_{use}, N_{loop}, N_{pipe}, N_{mplx}$
Output: $T_{load\_use\_loop}$
If $T_{load} \le (N_{pipe} \times N_{mplx} - 1) \times T_{use}$: $T_{load\_use\_loop} = T_{use} \times N_{loop}$
Else: $T_{load\_use\_loop} = (T_{load} + T_{use}) \times N_{loop} \div N_{pipe}$
$T_{threadblk} = T_{init} + T_{main\_loop} + T_{epilogue}$
$T_{init} = T_{smem\_load} + T_{reg\_load}$
$T_{main\_loop} = PipelineLatencyModel(T_{smem\_load}, T_{smem\_use}, N_{smem\_loop}, N_{smem\_pipe\_stage}, N_{threadblk\_per\_SM})$
$T_{smem\_use} = PipelineLatencyModel(T_{reg\_load}, T_{compute}, N_{reg\_loop}, N_{reg\_pipe\_stage}, N_{warp\_per\_threadblk})$
$T_{compute} = \frac{FLOPS_{one\_reg\_loop}}{Throughput_{SM} \times Util(N_{warp\_per\_threadblk}, N_{threadblk\_per\_SM})}$
$T_{smem\_load} = MAX(T_{LLC\_load}, T_{DRAM\_load})$
$T_{LLC\_load} = LAT_{LLC\_read} + \frac{Bytes_{one\_smem\_loop} \times N_{threadblk\_per\_threadblk\_batch}}{BW_{LLC}}$
$T_{DRAM\_load} = LAT_{DRAM\_read} + \frac{Bytes_{threadblk\_batch\_workset}}{BW_{DRAM}}$
$T_{epilogue} = LAT_{DRAM\_write} + \frac{Bytes_{output\_tile} \times N_{threadblk\_per\_threadblk\_batch}}{BW_{DRAM\_write}}$

线程块延迟模型。在线程块级别,我们通过将三个阶段的延迟相加来估计其最终性能:(1)初始阶段$T_{init}$,其中请求第一个数据块,流水线等待其到达;(2)主循环$T_{main\_loop}$,其中加载-使用流水线以稳定速率前进;(3)结尾阶段$T_{epilogue}$,其中最终结果被写回到全局内存。$T_{epilogue}$使用DELTA【【索引编号20,Delta: Gpu performance model for deep learning applications with in-depth memory system traffic analysis+2019+ISPASS】】中提出的结尾模型方程确定。

主循环延迟与嵌套流水线。让我们考虑$T_{main\_loop}$。它描述了共享内存级别的加载-使用循环,包括将数据从设备内存复制到共享内存,将数据读入寄存器,以及使用张量核进行计算。我们采用下一小节描述的流水线延迟模型来计算该循环的延迟。该模型考虑了流水线和多路复用因子,$N_{pipe}$和$N_{mplx}$,即流水线拥有的阶段数,以及可以被多路复用以隐藏内存复制延迟的并行工作单元数。在共享内存级别,这两个参数分别等于外部加载-使用循环的阶段数$N_{smem\_pipe\_stage}$和SM中的并行线程块数$N_{threadblk\_per\_SM}$。

内部流水线融合。计算$T_{main\_loop}$仍然需要该循环中使用阶段的延迟。然而,使用阶段是另一个将数据加载到寄存器文件并使用张量核进行计算的流水线。我们可以通过估计内部流水线在稳定状态下的延迟(通过内部流水线融合)来计算使用阶段的延迟。对于这个内部加载-使用循环,使用延迟指的是在一个循环内使用张量核执行算术运算的延迟。流水线和多路复用因子由这个内部加载-使用循环的阶段数和线程块中的并行warp数决定。

B. 获取详细延迟

流水线延迟模型。现在我们解决估计加载-使用循环在稳定状态下延迟的核心问题。直观上,预测应该根据瓶颈是加载还是使用而有所不同。表I中流水线延迟模型的第3行是确定瓶颈的标准。图9说明了计算或加载是瓶颈的两种情况。其直觉是,在加载一个数据块期间,计算单元可以用于计算该流水线中的其他数据块($N_{pipe}$),或用于其他并行工作单元($N_{mplx}$)。如果数据加载的延迟超过了所有可以与之重叠的计算的延迟,那么加载就成为瓶颈,使得循环延迟等于单次加载-使用迭代的延迟除以重叠流的数量,即$(T_{load} + T_{use})/N_{pipe}$。

Fig. 9: 流水线延迟模型的解释。一次加载可以被其他线程块中的计算,或同一线程块的其他阶段的计算所重叠。

计算和内存延迟模型。为了获得计算延迟,我们可以简单地将一个循环内执行的浮点运算次数除以一个SM中的张量核吞吐量。在确定内存复制的延迟时,必须考虑四个参数:传输的数据量、可用带宽、与之共享带宽的并行工作单元(线程块或warp)数量,以及一个恒定的往返延迟$LAT$。需要注意的是,GPU的LLC被所有SM共享。因此,DRAM流量不能通过所有线程块加载的数据之和来计算,因为数据可能在LLC中命中。我们通过确定一个线程块批次的工作集来建模DRAM流量。

C. 模型引导的自动调优

结合分析模型与机器学习模型。现在,我们将讨论如何使用分析性能模型进行调度调优。自动调优的工作流程由一个用于从调度预测性能的成本模型和一个用于提出新试验的采样方法组成。与我们开发的分析模型不同,TVM不使用分析模型,而是使用一个基于机器学习(ML)的成本模型,该模型仅从剖析的性能结果中学习。分析模型和基于ML的调优提供了互补的优势:分析模型不需要编译和运行采样调度的复杂性,但不能非常准确,因为它难以完全捕捉内存系统等硬件因素。基于ML的调优从包含这些复杂因素的实测性能中学习成本模型,但它需要大量的采样数据,导致调优过程漫长。

预训练ML模型。最终,我们利用分析性能模型的预测来预训练基于ML的模型,使得ML模型能够获得先验知识,同时仍然利用剖析数据。表II比较了我们的方法(Model-Assisted XGB)与其他可用的自动调优方法。
TABLE II: 编译器搜索方法比较。

A4 实验环境

  • 硬件配置
    • GPU:NVIDIA A100-SMX4,配备40GB设备内存。选择该GPU是因为其支持异步内存复制硬件特性,这对流水线至关重要。
  • 软件配置
    • CUDA版本:v11.4。
    • 基础框架:实现基于TVM v0.8【【索引编号10,TVM: An automated End-to-End optimizing compiler for deep learning+2018+OSDI 18】】。
    • 基线(Baselines)
      1. TVM:原版TVM。
      2. TVM DB:手动为TVM插入双缓冲原语的版本。
      3. XLA【【索引编号18,Xla: Optimizing compiler for machine learning+2021+https://www.tensorflow.org/xla】】 :Tensorflow v2.9.1集成的编译器,用于端到端模型比较。
      4. Vendor Libraries:cuBLAS【【索引编号16,cuBLAS+N/A+https://docs.nvidia.com/cuda/cublas/index.html】】和cuDNN【【索引编号23,cuDNN+N/A+https://docs.nvidia.com/deeplearning/cudnn/developer-guide/index.html】】 。
    • 消融研究版本
      1. ALCOP without ML:仅在共享内存级别进行流水线,无多级优化。
      2. ALCOP without MS, ML:仅允许两阶段流水线,且无多级优化。
  • 数据集/算子
    • 单算子性能测试:从真实DNN工作负载中提取的四种算子(矩阵乘法MatMul、批量矩阵乘法BMM、二维卷积Conv2D、三维卷积Conv3D),具有多种形状。所有算子均使用半精度并在Tensor Cores上运行。
    • 端到端模型测试:评估了六个流行的深度学习模型:
      • NLP模型:BERT, BERT-Large【【索引编号2,Bert: Pre-training of deep bidirectional transformers for language understanding+2018+arXiv】】, GPT-2【【索引编号21,Language models are unsupervised multitask learners+2019+OpenAI blog】】。
      • 视觉模型:ResNet-18, ResNet-50【【索引编号1,Deep residual learning for image recognition+2016+CVPR】】, VGG-16【【索引编号22,Very deep convolutional networks for large-scale image recognition+2014+arXiv】】。

A4 实验结果

A. 单算子性能

  • 实验内容:在NVIDIA A100上,将ALCOP生成的算子性能与TVM、TVM DB以及ALCOP的消融版本进行比较。所有版本都穷举搜索以找到最佳调度。
  • 实验结果:如图10所示,ALCOP生成的算子性能比TVM平均快1.23倍,最高可达1.73倍。

    Fig. 10: A100上单算子性能相对于TVM的归一化结果。
  • 分析结论
    • 流水线适用场景:对于输出形状较小但归约轴较长的问题(如MM RN50 FC),流水线效果尤其显著,因为它能发掘额外的并行性。对于输出形状大(已有充足的空间并行性)或归约维度小(无法摊销初始加载延迟)的问题,流水线带来的收益有限。
    • 消融研究:多级(Multi-level)和多阶段(Multi-stage)流水线对最终加速都至关重要。TVM DB相比TVM无明显提升。没有多级流水线,ALCOP平均加速仅为1.13倍;既没有多级又没有多阶段,则仅为1.01倍。

B. 端到端性能

  • 实验内容:在六个流行的深度学习模型上,比较ALCOP与TVM和XLA的端到端推理性能。
  • 实验结果:如表III所示,ALCOP相比TVM实现了1.02-1.18倍的端到端加速,相比XLA实现了1.01-1.64倍的加速。
    TABLE III: 流水线带来的模型加速

C. 与库的比较

  • 实验内容:将ALCOP生成的核性能与供应商库(cuBLAS/cuDNN)中高度优化的核进行比较。
  • 实验结果:如图11所示,ALCOP能够达到与库核相当的性能,平均为库性能的93%。对于某些算子(如BMM BERT QK),由于编译器能够搜索整个调度空间,甚至生成了比cuBLAS更快的核。

    Fig. 11: 单算子性能与库的比较。

D. 性能模型准确性

  • 实验内容:评估性能模型的准确性,使用指标为“模型排序的前k个调度中的最佳性能”(best-in-top-k)。与一个简单的“基于瓶颈分析”的模型进行比较。
  • 实验结果:如图12所示,在前10次试验中,ALCOP的模型能找到达到穷举搜索最佳性能79%的调度,而基线模型为75%。在前50次试验中(比穷举搜索节省40倍试验次数),ALCOP模型能达到92%的性能,而基线模型为88%。

    Fig. 12: 两种分析性能模型的best-in-top-k性能。“compile fail”标记表示模型提出的前10或50个调度未能成功编译。

E. 分析模型引导的调度调优

  • 实验内容:评估将分析模型与基于ML的调度调优相结合的技术。比较了四种方法:网格搜索(Grid-Search)、XGB(TVM默认)、仅分析模型(Analytical-only)和分析模型+XGB(Analytical+XGB,本文方法)。
  • 实验结果:如图13所示,在10次试验预算下,Analytical+XGB方法能找到达到穷举搜索最佳性能95%的调度,显著优于其他方法(70%-79%)。在50次试验预算下(节省超过40倍试验次数),该方法能达到99%的最佳性能,而其他方法仅为86%-92%。

    Fig. 13: 调度调优方法的搜索效率。
  • 分析结论:分析模型有助于ML(通过提供先验知识),ML也有助于分析模型(通过使用真实剖析数据进行微调)。二者结合的效果优于单独使用任何一种方法。

A5 结论

本文解决了深度学习编译器中自动流水线化的重要需求。由于缓解带宽限制所需的大尺寸分块(tiling)导致分块间并行性不足,无法实现高利用率,因此分块内流水线变得至关重要。我们提出了第一个支持多阶段、多级流水线的编译器解决方案ALCOP。通过引入自动流水线化,我们的编译器生成的GPU程序相比原生TVM【【索引编号10,TVM: An automated End-to-End optimizing compiler for deep learning+2018+OSDI 18】】平均可获得1.23倍的加速,最高可达1.73倍。此外,我们开发了一个分析性能模型,该模型显著提高了调度调优过程的搜索效率。