SIRIUS: HARVESTING WHOLE-PROGRAM OPTIMIZATION OPPORTUNITIESFOR DNNS

作者/机构: Yijin Li, Jiacheng Zhao, Qianqi Sun, Haohui Mai, Lei Chen, Wanlu Cao, Yanfan Chen, Zhicheng Li, Ying Liu, Xinyuan Zhang, Xiyu Shi, Jie Zhao, Jingling Xue, Huimin Cui, Xiaobing Feng

A1 主要贡献

本文提出了一种名为SIRIUS的全程序分析与优化编译器框架，旨在解决现有AI生态系统中（如TensorFlow、PyTorch、TVM等）从计算图和算子表示出发进行优化的局限性，即未能充分利用经典的编译器优化技术。

核心问题: 现有AI系统的优化方法（如新编程语言、新内核融合启发式、自动调优等）虽然取得了成果，但忽视了从DNN模型的全部源代码（包括主机端和设备端）出发，利用经典的全程序分析来发掘最大化性能潜力的机会。以BERT的一个子图为例（图1），TensorFlow不进行融合，执行时间为54.00µs；Apollo和TensorRT进行部分融合，执行时间分别为46.20µs和29.91µs。而SIRIUS通过全程序分析，能够将9个算子全部融合成一个内核，指令数从470万减少到77万，执行时间降至15.82µs，这得益于更广阔的优化视野带来的冗余加载消除和死代码消除等经典优化。

研究目标: 本文的目标是设计并实现一个编译器框架SIRIUS，该框架能够：
1. 统一建模主机（host）和设备（kernel）的计算，形成一个统一的多面体表示（polyhedral representation）。
2. 从全局视角寻找最大的内核融合机会。
3. 利用经典的编译器优化技术来提升融合后内核的性能。

主要创新点与贡献:
1. 性能瓶颈识别: 通过对代表性加速器应用的性能剖析，识别出扩大代码优化范围对加速器内核性能有显著影响。
2. 统一抽象表示: 提出了一个统一的多面体抽象表示（UAR），它统一了来自应用程序主机端和设备端的数据依赖和同步，从而实现了跨主机驱动和多个设备内核的过程序优化（inter-procedural optimizations）。
3. 全局融合启发式: 提出了两种启发式方法来全局地、贪婪地寻找融合机会，并对融合后的代码片段进行有效安排。论文在6个代表性的DNN模型上定量地证明了这些方法的有效性。
4. 显著性能提升: 在6个DNN模型上，SIRIUS相比TensorRT平均取得了4.32倍（最高7.12倍）的加速比。特别地，对于BERT模型，SIRIUS相比TensorRT实现了1.46倍的加速。


图1. 动机示例（来自BERT的子图）。(a)展示了QKV计算的子图；(b)展示了不同融合策略下的性能特征，SIRIUS通过将9个算子融合成1个，显著减少了指令数和执行时间；(c)展示了融合与优化后的伪代码，指令数从27条减少到13条。

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

SIRIUS的整体架构。如图2所示，SIRIUS首先将输入的CUDA程序的主机端和设备端代码都编译成LLVM IR。接着，SIRIUS利用全程序分析，在一个统一的抽象表示（Unified Abstract Representations, UAR）上捕捉整个程序的LLVM IR的数据和控制依赖（详见第3节）。UAR在多面体域【索引8，A practical automatic polyhedral program optimization system+2008+PLDI】【索引55，Optimizing the memory hierarchy by compositing automatic transformations on computations and data+2020+MICRO】中表示依赖关系，以描述细粒度的内存布局和访问模式。


图2. SIRIUS的整体架构。

UAR中的依赖宽度标注。为了精确表示跨内核的数据依赖，UAR用宽度来标注这种依赖，以指明承载依赖的作用域。SIRIUS定义了三种级别的标注：线程级（thread-level）、块级（block-level）和全局级（global）。具体来说，线程/块级数据依赖意味着依赖关系是从源内核的线程/块到目标内核的线程/块的一对一映射。线程级数据依赖意味着两个内核可以被融合而无需引入任何同步语句，而块级数据依赖则意味着融合这两个内核时需要一个syncthreads。全局依赖意味着两个内核之间存在跨块的依赖关系，例如，后续内核某个块的输入数据是前一个内核多个块的输出数据。图3(a)展示了图1中BERT子图的伪代码，图3(b)展示了其中的块级、线程级和全局依赖。


图3. SIRIUS处理BERT子图的工作流程。

SIRIUS的融合三大原则。SIRIUS中融合的关键洞察基于三个原则：
* 原则1 (P1)：如果两个内核只有线程级/块级依赖，它们可以通过按顺序合并内核语句来进行局部融合。
* 原则2 (P2)：如果两个内核有全局依赖，则禁止融合，因为GPU对全局同步的支持很差【索引24，Optimizing parallel reduction in cuda+2007】【索引70，Inter-block gpu communication via fast barrier synchronization+2010+IPDPS】。
* 原则3 (P3)：如果两个内核没有依赖关系，SIRIUS仅在它们具有高代码相似度时（本文中限制为代码相同）才将它们融合在一起。

两阶段贪心融合调度。SIRIUS通过两个阶段计算一个贪心的融合调度：首先是局部融合，然后是全局融合。局部融合阶段通过对UAR进行广度优先搜索（BFS）遍历来实现。全局融合阶段则利用ISL求解器，通过引入一个调度约束来指定UAR中具有全局数据依赖的两个内核的执行顺序。对于图3(b)中的例子，由于存在全局数据依赖，GEMM(2)没有被融合；而三个分支因为没有依赖关系且代码相同而被融合。

内核片段的编排。SIRIUS根据输出的调度来融合内核的调用。内核可以被水平编排（即，在不同的块中执行片段以实现并行性）或垂直编排（即，在同一个线程中执行片段以实现顺序性能）。不同的编排方式在并行性和顺序性能之间呈现出不同的权衡。最优的权衡高度依赖于硬件配置。因此，SIRIUS使用一个自动调优器来选择性能最佳的调度，并将其报告给用户。

A2 方法细节

3 构造统一抽象表示（UAR）

UAR在一个统一的依赖图上描述了主机驱动和设备内核的同步与数据依赖。它在多面体域中将依赖表示为仿射关系，以捕捉依赖和访问模式。UAR的信息驱动了SIRIUS中的过程序分析和转换。

3.1 背景

多面体分析与调用图。本小节提供了多面体分析【索引6，Tiramisu: A polyhedral compiler for expressing fast and portable code+2019+CGO】【索引8，A practical automatic polyhedral program optimization system+2008+PLDI】【索引9，Effective automatic parallelization and locality optimization using the polyhedral model+2008】【索引20，Polly - performing polyhedral optimizations on a low-level intermediate representation+2012】和调用图的基本背景。一个依赖图 $G = (V, E)$ 是一个有向多重图，记录了特定过程的数据和控制依赖。顶点 $v \in V$ 代表过程中的一个语句。一个迭代向量 $\vec{i} = (i_0, i_1, . . . , i_n) \in D_v$ 描述了围绕顶点 $v$ 的所有循环的索引值。一条边 $e = (S_i, S_j) \in E$ 定义了从语句 $S_i$到语句$S_j$的边，代表从 $S_i$ 的一个动态实例到 $S_j$ 的一个实例的多面体依赖。这种多面体依赖由依赖多面体 $P_e$ 【索引8，A practical automatic polyhedral program optimization system+2008+PLDI】【索引9，Effective automatic parallelization and locality optimization using the polyhedral model+2008】来刻画，它基于迭代向量捕捉了边 $e$ 的精确依赖信息。如果 $\vec{s}$ 和 $\vec{t}$ 是相互依赖的源和目标迭代，我们可以表示为：

3.2 UAR

UAR的定义与构建。UAR被定义为一个带有附加信息的依赖图G，用于统一建模主机驱动和设备内核。在UAR中，顶点 $v \in V$ 由其所代表的主机或内核语句的内存效应来表征，而边 $e$ 则由依赖多面体 $P_e$ 来表征，它描述了主机驱动和设备内核语句动态实例之间的同步和数据依赖。SIRIUS分两个阶段构建UAR。首先，SIRIUS对每个函数进行局部分析，分析每个语句的依赖和效应。其次，SIRIUS从底层向上遍历调用图，将局部分析的结果传播到每个调用点。最终，它生成一个描述程序依赖的统一图。

3.3 局部分析

局部分析的功能。局部分析阶段执行两个功能：为每个函数执行依赖分析，并计算每个函数的内存访问效应。

分析局部依赖。SIRIUS通过两种方式增强了LLVM【索引40，Llvm: A compilation framework for lifelong program analysis & transformation+2004+CGO】中的依赖分析，以计算每个函数的局部依赖图。首先，SIRIUS在设备函数中引入循环来表示块和线程的并行性。其次，如果SIRIUS能够证明准仿射访问（如模运算表达式）仅发生在同一个块内部，它会通过用其值域替换模运算表达式，将准仿射访问重写为仿射，以计算内存效应。

调用点与同步屏障处理。在局部分析阶段，SIRIUS不解析调用点。它将调用点标记为具有未知效应，并添加控制依赖边以防止调用点重排序。SIRIUS将像syncthreads这样的同步屏障视为特殊的函数调用。它添加了控制依赖边，但标注其效应为空。

虚拟数据依赖。由于计算语句和同步屏障之间不存在数据依赖，因此在多面体模型中直接表示它们的执行顺序并不直观。然而，同步屏障给程序带来了隐式的控制依赖。因此，我们通过将隐式控制依赖转换为数据依赖，在同步屏障和执行语句之间引入一种虚拟数据依赖，我们称之为虚拟数据依赖。

主机端过程处理。SIRIUS对主机上的过程遵循类似的工作流程，不同之处在于它能识别内存操作（包括拷贝和malloc）以及同步原语（如cudaDeviceSynchronize()）等内在函数，并相应地进行标注。

计算内存效应。对于每个内核，SIRIUS扩展了LLVM中的标量演化算法，以在线程、块和内核级别计算其内存效应。具体来说，一个内核的内存效应表示为一个读集合和一个写集合，分别代表该内核在线程/块/内核级别读取和写入的内存位置。SIRIUS调用ScalarEvolutionAnalysis pass直接计算线程级别的内存效应。对于每个线程块，SIRIUS在线程ID维度上聚合块中所有线程的内存效应，得到块级别的内存效应。类似地，SIRIUS在块ID维度上聚合所有线程块的内存效应，得到内核级别的内存效应。

3.4 自底向上分析

信息传播与依赖宽度确定。自底向上分析沿着调用图传播依赖和效应，以计算整个程序的UAR，并确定每个数据依赖的多面体及其依赖宽度。SIRIUS通过克隆来传播信息。从调用图的底部开始，自底向上分析解析调用点的参数，将被调用者的局部UAR内联到调用者的UAR中，并确定数据依赖宽度。在解析内核启动的调用点时，SIRIUS只内联被调用者在全局内存堆上的效应，而不是其UAR，从而丢弃了在主机端无关的共享内存上的效应。

依赖宽度的定义。通过检查每个调用点的内存效应，SIRIUS增强了LLVM中的依赖分析算法来计算依赖多面体 $P_e$，并用线程级、块级或全局级的标签对其进行标注。对于一个数据依赖 $e_{<s,t>}$，其宽度由承载该依赖的级别决定。例如，考虑由数组A引起的依赖，s写入A，t读取A。对于s，我们用 $A_{thread(tid)s}$, $A_{block(bid)s}$ 和 $A_s$ 分别表示s的一个线程、一个线程块和整个内核的访问区域，t也类似。边e的依赖宽度被定义为块级，当且仅当s的各个块的访问区域是线性不相交的

$\forall(i, j) A_{s}^{\text {block}(i)} \cap A_{s}^{\text {block}(j)}=\emptyset$

且t的每个块的访问区域仅依赖于s的一个块

$\forall j, \exists ! i, A_s^{block(i)} \cap A_l^{block(j)} \neq \emptyset$

线程级和全局级的宽度定义与此类似。

间接/递归调用支持。目前，我们尚未在SIRIUS中实现对间接调用或递归调用的支持。我们的经验表明，大多数加速器程序依赖于简单的直接调用。支持递归调用将使SIRIUS能够分析如深度优先搜索等具有不规则模式的用例。使用数据结构分析中描述的方法【索引41，Making contextsensitive points-to analysis with heap cloning practical for the real world+2007+PLDI】来实现这种支持是可能的。

示例。图3(a)描述了BERT子图的实现，(c)展示了GEMM、add和reshape的局部UAR，(d)展示了顶层UAR。实线和虚线箭头分别代表虚拟数据依赖和真实数据依赖。SIRIUS还用相应的多面体和宽度标注来注解虚拟依赖的边。

3.5 别名分析

别名分析挑战与策略。我们的别名分析中的关键挑战是跟踪和解析主机句柄与设备内存实际区域之间的映射。SIRIUS首先将整个程序的指针分为主机指针和设备指针，分析cudaMemcpy语句以确保主机和设备指针之间不存在别名，然后分别对主机指针和设备内存执行别名分析。

具体算法。为了获得更精确的分析结果，SIRIUS扩展了Andersen算法【索引5，Program analysis and specialization for the C programming language+1994】，并实现了一个域/数组/上下文敏感、路径不敏感的基于包含的点对分析算法。SIRIUS利用了(Pearce等人，2007)【索引54，Efficient fieldsensitive pointer analysis of c+2007+TOPLAS】和(Whaley & Lam，2004)【索引69，Cloning-based context-sensitive pointer alias analysis using binary decision diagrams+2004】提出的方法在LLVM中实现域/数组/上下文敏感性。

分析精度。由于DNN模型具有规则的内存访问模式，该分析足够准确。对于DNN模型，内存效应分析和依赖分析可能会引入一些不精确性。例如，如果一个内核包含模运算来计算数组下标，SIRIUS会通过使用其值域将模运算改为仿射运算。因此，内存效应分析会报告一个放大的结果。对于我们的DNN基准测试，这种不精确性不会对进一步的依赖分析和优化产生影响。

4 推导融合调度

4.1 计算贪心融合调度

融合原则的实现。如第2节所述，我们引入了三个原则来计算融合调度，并通过一个局部融合过程和一个全局融合过程来实现它们。

局部融合过程 (P1)。对于原则(P1)，局部融合过程对顶层UAR执行一次广度优先搜索（BFS）遍历。对于每个节点，如果它与其前驱节点之间只存在线程级或块级数据依赖，SIRIUS会将它们合并成一个节点，并通过按顺序放置语句和为块级依赖插入syncthreads来保证它们的执行顺序。

全局融合过程 (P2 & P3)。对于原则(P2)和(P3)，SIRIUS利用ISL求解器来计算顶层UAR的融合调度。SIRIUS生成关系并将其输入ISL求解器，然后求解器通过将有效性关系建模为整数线性规划（ILP）约束【索引65，Scheduling for ppcg+2017】【索引78，Unified polyhedral modeling of temporal and spatial locality+2017】来确定一个调度。首先，SIRIUS利用有效性关系来表示原则(P2)，通过引入一个调度约束来指定UAR中具有全局数据依赖的两个内核的执行顺序。特别地，只有全局数据依赖会被建模为约束并输入到ISL求解器中。如果两个内核有线程级/块级数据依赖，它们在局部融合阶段就已经被合并了。

4.2 编排内核片段 - 水平和垂直

水平与垂直融合的权衡。对于一组没有依赖关系的内核片段，SIRIUS可以从两个可能的方向进行融合。如图4所示，它可以水平融合内核（即，在不同的块中执行片段），也可以垂直融合内核（即，在同一个线程中执行片段）。这两个方向代表了并行性与顺序优化机会之间的权衡。水平融合会增加并行性。然而，垂直融合会减少并行性，但可以为融合后的编译器优化开辟更多机会。因此，SIRIUS生成从完全水平到完全垂直（以及一些混合版本）的一系列代码版本，应用第5节中描述的一系列优化，并利用自动调优来选择性能最佳的版本。


图4. 垂直和水平内核编排。

代码生成细节。当为垂直融合的片段生成代码时，SIRIUS会连接共享内存缓冲区以更有效地利用共享内存。当融合具有不同迭代域（即块和线程）的片段时，SIRIUS会将它们对齐到最大或最小的域，并相应地生成if或for语句。当融合涉及一个元素级（element-wise）内核时，SIRIUS会根据其他内核调整其迭代域，从而可以将元素级语句嵌入到融合内核的迭代域中。对于水平融合的代码生成，SIRIUS会为主机重新配置gridDim和BlockDim，并为设备发出到不同内核的分支。

示例。SIRIUS检测到图3中的九个算子作为融合候选。局部融合过程在BFS遍历期间垂直地融合了GEMM、add和reshape，并在GEMM和add之间插入syncthreads以确保执行顺序。这三个UAR节点将被合并为一个，记为LF_K。然后，全局融合过程利用ISL求解器，得到将三个LF_K节点融合的调度，因为它们之间没有依赖关系且代码相同。此外，如图4所示，“Q/K/V”三个分支可以从完全水平融合到完全垂直。在图4(a)中，一个输入切片在多个块中执行，每个线程执行Q、K和V的“GEMM+add+reshape”计算序列。在图4(b)中，每个块执行Q或K或V的“GEMM+add+reshape”计算序列。SIRIUS通过自动调优选择了完全垂直融合的版本作为性能最佳的版本。第7节详细讨论了并行性与顺序优化之间的这些权衡。

5 优化融合后的内核

融合UAR的生成。通过推导出的融合调度，SIRIUS将局部UAR合并在一起，得到融合后的UAR。对于水平融合的片段，SIRIUS不引入额外的同步，因为这些片段由不同的块执行。对于垂直融合的片段，如果存在块级数据依赖，SIRIUS会在这些片段之间插入一个同步节点。在融合后的UAR中，同步被表示为一条虚拟数据依赖边，当它不承载实际数据依赖时，这条边是冗余的。在移除所有冗余同步之后，SIRIUS执行代码移动和融合后优化。

算法 1 优化同步
输入: uar
输出: uar
1 for vdep in uar.vdeps do
2   for ddep in uar.ddeps do
3     if No VDep Path(ddep.src, ddep.dst) then
4       uar.remove(cdep);
5 for vdep in uar.vdeps do
6   for ddep in uar.ddeps do
7     if Exist VDep Path BypassV(ddep.src, ddep.dst, vdep) then
8       uar.remove(cdep);
9 for pnode in uar.cdeps do
10    if pnode.out degree == 0 then
11      uar.remove(pnode);
12 return uar

5.1 移除冗余同步

同步优化的关键点。优化同步的关键点在于，对于从s到t的每个数据依赖边，只应该存在一条从s到t的虚拟数据依赖边。我们将移除冗余同步的问题抽象为屏障最小化问题的一个特例，该问题已被多位研究者深入研究【索引25，Data-parallel programming on MIMD computers, volume 90+1991】【索引52，Compile time barrier synchronization minimization+2002+IEEE Transactions on Parallel and Distributed Systems】【索引67，Compiler optimizations for eliminating barrier synchronization+1995+ACM SIGPLAN Notices】。特别是，Darte和Schreiber提出了一个线性时间算法来解决这个问题【索引15，A linear-time algorithm for optimal barrier placement+2005】。本文实现了该线性时间算法的一个适用于我们UAR的定制版本（算法1）。

算法1的冗余定义。算法1遍历融合后的UAR中的每条虚拟数据依赖边，并在其冗余时将其移除。“冗余”定义如下：
* 一条虚拟数据依赖边 $e(s → t)$ 是冗余的，如果对于任何数据依赖边 $e'(s' → t')$，不存在一条从 $s'$ 到 $t'$ 的包含 $e$ 的虚拟数据依赖路径。
* 一条虚拟数据依赖边 $e(s → t)$ 是冗余的，如果对于任何数据依赖边 $e'(s' → t)$，所有从 $s'$ 到 $t$ 的虚拟数据依赖路径都包含另一条虚拟数据依赖边。
* 一个同步节点的出度为0时是冗余的。

5.2 融合后优化

扩展的融合后优化。除了LLVM中传统的编译器优化，如动机示例中讨论的冗余加载消除和死代码消除，SIRIUS还扩展了以下融合后优化。

代码移动。移除同步后，融合后的UAR会分裂成一组连通分量。这些连通分量之间没有数据依赖，因此SIRIUS应用代码移动，使它们能够共享一组公共的同步。例如，片段 $(S1 → P0 → S2 → S3)$ 和 $(S1' → P0' → S2')$ 被改为 $(\{S1, S1'\} → P0' → \{S2, S2'\} → S3)$。

常量传播。通过在UAR中统一表示主机和内核代码，SIRIUS将常量从主机传播到内核代码中，从而使供应商编译器能够利用常量值生成优化的指令。

为映射操作重塑迭代域。当融合涉及一个被标注为成对映射或归约操作的片段时，SIRIUS会简单地重用融合中其他片段的迭代域，并调整生成的代码。

A4 实验环境

数据集与模型:
* 本文评估了6个代表性的DNN模型：ResNeXt, NASNet, LSTM, DeepSpeech2, Seq2Seq, 和 BERT。
* ResNeXt 在 Cifar10 数据集上训练。
* BERT 是一个在 Wikipedia + Book Corpus 数据集上预训练的标准 Bert-base 模型，参数为：序列长度=384，头数=层数=12，隐藏维度=64，精度为float16。
* LSTM, DeepSpeech2, Seq2Seq, NASNet 的超参数与Rammer【索引46，Rammer: Enabling holistic deep learning compiler optimizations with rtasks+2020+OSDI】中的设置一致。

硬件配置:
* 主测试平台: 一台配备两颗Intel Xeon Gold 6248 CPU、768GB DDR4内存和一块NVIDIA A100 GPU的服务器。
* 参考平台: 一台配备NVIDIA Tesla V100 GPU的服务器。

软件配置:
* 操作系统: Ubuntu 22.04
* CUDA版本: 11.8
* SIRIUS实现: 使用约7000行C++代码实现，基于LLVM release/16.x。ILP求解器使用PipLib 1.3.3。
* 输入代码: SIRIUS的输入是CUDA源代码。实验中使用的245个算子中，231个来自AutoTVM/Rammer，另外14个BERT算子是为与TensorRT使用的cuBLAS性能相匹配而手动实现的。
* 对比框架: TensorFlow (TF, v1.15.5), TensorRT (TRT, v7.2), TVM (v0.12, 试验次数设置为2000), Rammer, 以及 AStitch。

A4 实验结果

7.1 总体性能
如图5所示，在NVIDIA A100上，SIRIUS相较于TensorFlow/TF-XLA/TVM/TensorRT，平均取得了41.78倍/21.76倍/19.16倍/5.15倍的加速比，最高可达154.84倍/60.80倍/43.69倍/11.98倍。相较于Rammer，SIRIUS平均取得了1.62倍（最高2.19倍）的加速。性能提升主要来源于两个方面：(1) 全程序分析使得SIRIUS能够发掘更激进的融合机会；(2) 扩大后的内核和垂直融合为单线程的顺序优化（sequential optimization）创造了更多机会。



图5. 在NVIDIA A100 (a) / V100 (b) 系统上，批处理大小为1时，六个DNN模型的端到端运行时间和加速比。

7.2 性能分解
图6展示了SIRIUS关键技术对BERT和LSTM的性能影响。
* BERT: 简单的融合（+fusion）将执行时间从2.54ms降至2.06ms，主要减少了内核启动开销。协调水平/垂直融合（+h/v fusion）性能变化不大，因为此时GPU SMs已充分利用。但垂直融合为后续的顺序优化（+seqopt）创造了条件，最终将时间降至1.79ms。
* LSTM: 垂直融合在未优化时，因降低了并行度而导致性能下降。但经过融合后优化，性能得到显著提升，并最终优于水平融合。
如图7所示，针对BERT的融合后优化效果：基线（直接拼接内核）耗时2.06ms；移除冗余同步（OPTSYNC）降至2.04ms；代码移动（CM）降至1.86ms；最终的指令优化（INSTOPT）降至1.79ms。该图右轴也显示了指令数随着优化步骤的减少。


图6. SIRIUS的性能分解。


图7. 融合后优化的效果。

7.3 开销分析
SIRIUS的开销主要来自：全程序分析（<10分钟）、ISL求解器计算融合调度（<3分钟）、融合后优化（<10分钟）以及水平/垂直融合的自动调优。对于大型模型（如NASNet, Bert），SIRIUS会将其UAR划分为小的基本单元进行处理，避免了高昂的编译成本。调优开销因模型而异，BERT约2分钟，而LSTM因候选方案多（545个）需54分钟。考虑到TVM等框架本身编译时间就很长，SIRIUS引入的额外开销（<1小时）是可以接受的。

7.4 案例研究：BERT模型
SIRIUS在A100/V100上相比TensorRT对BERT分别取得了1.46倍/1.25倍的加速（图5）。在V100平台上，SIRIUS相比AStitch、Rammer和TensorRT分别实现了1.61倍、2.35倍、1.25倍的加速。对于图1中的子图，AStitch、Rammer和TensorRT分别只能融合部分算子（最终内核数为6、7、4），而SIRIUS能将全部9个算子融合成一个。如图8所示，在不同批处理大小（1, 4, 16）下，SIRIUS相比AStitch、Rammer和TensorRT平均分别取得了1.56倍、2.92倍和1.48倍的加速，证明了其在不同场景下的有效性。


图8. 在V100上不同批处理大小下的性能结果。

7.5 案例研究：LSTM模型
对于LSTM，SIRIUS相比TensorRT和Rammer分别取得了6.13倍和1.42倍的显著性能提升。SIRIUS和Rammer都能发现沿反向对角线（即波前）的融合机会（图9b）。但SIRIUS将每个波前的gemv和solve都融合成一个内核并采用垂直编排（图9d），而Rammer生成两个内核并采用水平编排（图9c）。TensorRT仅融合了gemv，未能执行倾斜融合。SIRIUS的融合策略和垂直编排带来了显著的指令优化机会。如图10所示，SIRIUS通过全程序分析得知8个矩阵在主机端是连续分配的，从而在融合内核中将基地址重写为公共基地址的偏移量。这使得NVCC编译器能在编译时计算偏移并将其直接嵌入加载指令中，消除了运行时的MAD地址计算指令，每条FMA指令的辅助指令数从6.1（基线）和3.2（Rammer）降至1.9（SIRIUS）。


图9. SIRIUS中LSTM的融合调度。


图10. (a) LSTM单元内核三种实现（Baseline, Rammer, SIRIUS）在NVIDIA A100上的性能特征。(b)/(c) Rammer和SIRIUS的汇编代码片段。

A5 结论

本文提出了一个名为SIRIUS的优化编译器框架。通过在一个统一的多面体表示中对主机和内核代码进行建模，SIRIUS利用多面体分析来暴露最大化的内核融合机会，然后生成融合后的内核。最终，融合后的内核可以受益于许多传统的顺序优化。评估表明，SIRIUS相较于TensorRT最高可实现11.98倍的加速。具体到BERT模型，SIRIUS相较于TensorRT可实现1.46倍的加速。

SIRIUS存在两个局限性。首先，它需要看到算子的CUDA源代码。目前SIRIUS使用TVM生成的CUDA代码。未来的工作将采用张量表达式来简化分析和优化。其次，SIRIUS作为一个在CUDA代码生成后的优化编译器，无法改变线程组织。未来的工作将把它提升到TVM中，以便能够全局优化线程组织。并且，未来将利用SIRIUS的方法来支持训练过程。

A6 附录

A.1 摘要

制品内容。该制品包含了验证本文主要结果所需的软件组件。我们提供了一个docker镜像以简化环境设置。该docker镜像包含Sirius的源代码、评估推理性能的脚本以及绘制图表的脚本。它需要一个运行在配备NVIDIA A100 Tensor Core GPU的x86_64机器上的Linux系统，该系统需装有能够运行CUDA 11.8的NVIDIA驱动，才能创建docker容器。启动docker容器后，您可以运行脚本来收集性能数据并绘制图表。

A.2 制品清单（元信息）

  • 模型: ResNext, NANet, DeepSpeech2, LSTM, Seq2Seq 和 BERT（已包含在docker镜像中）
  • 运行环境: 带有NVIDIA驱动（能运行CUDA 11.8）的Linux系统
  • 硬件: NVIDIA A100 Tensor Core GPU
  • 指标: 端到端推理时间
  • 输出: 关键图表和必要数据
  • 实验: 参见docker镜像中的read-me文件。
  • 所需磁盘空间（大约）: 60GB
  • 准备工作流所需时间（大约）: 20分钟
  • 完成实验所需时间（大约）: 下载docker镜像需要数十分钟。然后您可以运行一个脚本来一次性收集所有性能结果并绘制图表。总共大约需要2小时。
  • 是否公开可用: 是。docker镜像及其中的源代码是公开的。
  • 代码许可证（如果公开）: The GNU General Public License (GPL)
  • 使用的工作流框架: LLVM
  • 已存档（提供DOI）: https://zenodo.org/record/7885573

A.3 描述

A.3.1 如何访问

访问链接。我们在dockerhub和zenodo上都提供了docker镜像。
Docker-hub URL: https://hub.docker.com/r/sunqianqi/sirius
Zenodo URL: https://zenodo.org/record/7885573

A.3.2 硬件依赖

硬件要求。配备NVIDIA A100 Tensor Core GPU的x86_64机器。

A.3.3 软件依赖

软件要求。
* 带有能够运行CUDA 11.8的NVIDIA驱动的Linux系统
* Docker版本 20.10.14

A.3.4 模型

内置模型。所需的模型已包含在docker中，无需额外工作。模型包括ResNext, NANet, DeepSpeech2, LSTM, Seq2Seq 和 BERT。

A.4 安装

安装步骤。您只需拉取docker镜像并启动一个容器：

$ docker pull sunqianqi/sirius:mlsys_ae
$ docker run -it --name=sirius_test \
  --gpus all --privileged \
  sunqianqi/sirius:mlsys_ae /bin/bash

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

A.5 实验工作流

A.6 评估与预期结果

复现指南。我们在提供的docker镜像中提供了一个README文件（/root/mlsys_ae/http://README.md），说明了如何复现结果。您可以按照README.md中的步骤复现SIRIUS论文中的实验结果。