Optimizing Deep Learning Inference via Global Analysis and Tensor Expressions

文章标题：通过全局分析和张量表达式优化深度学习推理
作者/机构：Chunwei Xia, Jiacheng Zhao, Qianqi Sun, Zheng Wang, Yuan Wen, Teng Yu, Xiaobing Feng, Huimin Cui
(SKLP, ICT, CAS; UCAS, China; University of Leeds, U.K.; University of Aberdeen, U.K.; Thewake Systems, China; Zhongguancun Laboratory, China)

A1 主要贡献

深度神经网络（DNN）的执行优化至关重要，但随着模型复杂性的增加而变得愈发困难。现有的DNN编译器无法有效地利用跨算子边界的优化机会，留下了改进空间。为了应对这一挑战，本文提出了Souffle，一个开源的编译器，旨在跨算子边界优化DNN推理。

核心问题与研究目标：
现有的DNN编译器大多采用自底向上的方法进行算子融合，这种方法将算子融合和代码生成阶段分开，可能导致算子被错误地分配到不同的内核中，从而产生额外的内存访问开销并阻碍了潜在的优化。本文的目标是提出一种新颖的自顶向下方法，通过全局分析来确定内核边界，从而克服现有方法的局限性。

核心创新点：
Souffle通过一种新颖的自顶向下方法来识别和利用跨DNN算子边界的优化机会。其主要贡献如下：
1. 提出了一种新的自顶向下方法：Souffle首先将整个计算图视为一个单一的、合并的内核，然后通过顶层的全局分析，并综合考虑共享内存/寄存器中的数据重用以及生成的指令，将图划分为多个分区（即子程序）。这种方法通过全局性地考虑内核的生成代码来确定内核边界，优于现有的自底向上策略。
2. 基于张量表达式（TE）的全局分析：Souffle利用张量表达式（TEs）【12，TVM: An Automated End-to-End Optimizing Compiler for Deep Learning，2018，OSDI 18】来编码算子和张量的数据流信息，并在此基础上对整个DNN模型生成的张量依赖图进行全局依赖分析。TEs提供了简洁的语义，使得分析和优化复杂算子的问题转化为分析和优化更简单的数学表达式，从而能够优化具有复杂数据依赖（如多对多）的算子。
3. 语义保持的变换：在图分区之后，Souffle对每个子程序应用仿射变换（affine transformations）来组合多个TEs，形成一个单一的TE，从而消除冗余的内存访问。这种变换是全自动且灵活的，因为它依赖于TE对算子数学计算的精确描述。
4. 跨算子边界的底层优化：Souffle利用TE级别的精确依赖信息，通过重用张量缓冲区来优化内存访问延迟，并通过重叠内存加载和算术指令来提高指令级并行性。由于优化是基于融合了多个算子的子程序进行的，因此消除了单个算子的优化边界。

下图展示了TensorRT、Apollo和Souffle如何将一个BERT计算图映射到内核。Souffle的优化策略（c）相比于TensorRT（a）和Apollo（b）产生了更少的GPU内存访问和更快的执行时间。

表1显示了为图1中生成的内核的性能。Souffle的方法将从全局内存加载的数据量减少到8.87MB，执行时间为57.7μs，分别比TensorRT和Apollo快1.1倍和3.1倍。

Table 1. 图1中生成内核的性能。

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

2.1 工作示例

BERT模型优化案例。本文以一个在NVIDIA A100 GPU上优化标准BERT模型【14，Bert: Pre-training of deep bidirectional transformers for language understanding，2018，arXiv preprint arXiv:1810.04805】的例子作为动机。该模型基于Transformer架构【10, 49】，并使用FP16进行推理。图1展示了TensorRT和Apollo如何将BERT中一个简化的子计算图的算子映射到内核。这个子图包含了代表性的DNN算子，如通用矩阵乘法（GEMM）、reshape、permutation、元素级算术算子（如add或exp）以及归约算子（如reduce_sum）。编译器将这些算子映射到各自的内核，这对性能有显著影响。

2.2 性能评估

现有方法的性能瓶颈。我们使用NVIDIA Nsight Compute【39】测量了生成的内核。表1显示，无论是TensorRT还是Apollo都未能为所评估的DNN提供最优的映射。图1中TensorRT和Apollo创建的子图分别从全局内存加载16.52MB和27.78MB的数据，执行时间分别为62.34μs和179.1μs。而本文提出的更好策略是将该子图优化并映射为单个内核。这一策略将从全局内存加载的字节数减少到8.87M，执行时间为57.7μs，分别比TensorRT和Apollo快1.1倍和3.1倍。需要强调的是，TensorRT专门为基于Transformer的模型进行了调优，并使用了闭源、手动优化的底层算子实现（如GEMM）。因此，考虑到Souffle无法访问TensorRT使用的一些NVIDIA优化算子实现，其在BERT上对TensorRT的改进是显著的。

2.3 错失的优化机会

现有方法的优化局限。通过仔细检查分析数据和内核融合结果，我们确定了TensorRT和Apollo错失的几个优化机会：
* 未能在内存密集型和计算密集型内核之间进行优化。如图1所示，BERT的部分计算需要执行元素级的内存操作，例如reshape和permutation。TensorRT和Apollo利用手动制定的规则来融合相邻的元素级内存算子，但它们都未能进一步在融合后的算子与其前驱的计算算子（如图1中的GEMM算子）之间进行优化。Souffle则在内存密集型和计算密集型内核之间进行优化，并最终消除了所有元素级内存算子。总之，手动制定的规则无法覆盖多样的计算模式，在这种情况下错失了优化机会。
* 归约算子的次优融合策略。图1(a)和(b)显示了TensorRT和Apollo对归约算子采用的次优内核融合策略。两种策略都选择将GEMM和归约算子映射到独立的内核中，这需要在归约发生前将归约算子所依赖的整个张量数据存储到昂贵的全局内存中。Souffle积极地将归约算子与相邻的计算密集型算子融合，例如将R0-2与GEMM0和GEMM1融合，如图1(c)所示。这是通过一种两阶段归约方法实现的：以每个块（per-block）的方式执行部分归约，并使用atomicAdd进行全局归约。结果，整个张量数据可以保留在芯片上，只有部分结果存储在全局内存中。同时插入一个全局同步（例如CUDA中的grid synchronization【40】）来同步运行的块。此优化适用于图1中的所有归约算子。此外，Souffle可以缓存算术算子1的输出在芯片上，以供算术算子2重用。
* 计算密集型内核间的优化不佳。与许多其他DNN框架一样，TensorRT和Apollo试图融合多个相同类型的计算密集型算子，但未能利用不同类型算子之间的机会。考虑图1(d)，它展示了两个相互依赖的GEMM算子在两种不同策略下被分组到内核中时，如何执行异步内存复制和张量核心计算。第一种策略是将GEMM算子映射到两个独立的内核中，因为它们不考虑融合计算密集型算子。第二种策略是将它们映射到单个内核中。TensorRT和Apollo使用前者，而Souffle使用后者。通过将两个GEMM算子放入一个内核（图1(d)的右侧部分），Souffle允许在计算GEMM2的同时流水线执行加载GEMM3的W3。Souffle旨在利用这种跨算子的流水线优化。

2.4 我们的洞见

核心设计思想。基于上述观察，有必要对DNN模型进行分析，以融合算子、对张量操作执行自动变换，并在一个融合的内核内进行优化。一个更有效的内核融合策略使得提取关键的张量信息（如生命周期和张量数据重用）成为可能。这些信息随后可用于在元素级分析细粒度的依赖关系，从而实现更好的内核级优化。

3 张量表达式（TE）预备知识

Souffle的中间表示。Souffle利用TVM的张量表达式（TE）【12，TVM: An Automated End-to-End Optimizing Compiler for Deep Learning，2018，OSDI 18】作为分析和优化的中间表示。TE指定了如何从输入张量计算输出元素。在图2所示的TE程序中，TE0是GEMM的一个示例TE，其中rk参数定义了归约轴（即张量的哪个维度将被遍历以执行归约），归约索引范围从0到64。输出张量O0使用compute操作计算，该操作指定了对每个数据元素执行的计算以及输出形状。迭代变量i和j对应于输出形状的维度，它们的迭代域可以自然推断出来。本质上，TE使用一种纯函数式语言【47】来描述张量计算，允许对每个输出张量元素进行独立计算。需要注意的是，我们的优化也适用于其他具有类似简洁语义和表达能力的DSL，如antares【34】和tensor comprehension【48】。我们选择TVM是因为它的流行度和成熟的工具链。

A2 方法细节

4 Souffle 概览

Souffle框架介绍。Souffle是我们用于DNN代码优化的开源框架。它旨在克服第2节中确定的现有DNN推理框架的三个局限性，通过增强数据重用、优化归约操作以及实现跨算子边界的优化。目前，它支持TensorFlow模型，并在单个NVIDIA GPU上优化DNN推理，但我们的许多分析和优化方法也可以应用于AMD GPU和其他加速器。

Souffle工作流。图2展示了Souffle的概览工作流，它以一个模型作为输入，并使用TVM将模型降级为TE，我们在此之上进行分析和优化。
* TE降级：对于一个DNN模型，Souffle首先将每个算子降级为其对应的TE，形成一个TE程序。图2显示，五个算子被降级为标记为TE0到TE4的五个TE。
* 全局计算图分析：降级后的TE程序被传递给Souffle分析模块。Souffle对TE程序进行两级分析。在张量级别，Souffle提取重要的张量信息，如TE的形状、生命周期和计算强度。在元素级别，Souffle分析每个TE的输出和输入张量之间的细粒度依赖关系，如第5节所述。图2展示了包括元素级数据依赖和五个TE的计算复杂度的分析结果。在张量级别，它发现O0被TE1和TE3访问，这揭示了跨多个TE的数据重用机会。
* 资源感知程序分区：Souffle分析张量依赖图，并使用Ansor【57，Ansor: Generating High-Performance Tensor Programs for Deep Learning，2020，OSDI 20】来调度计算密集型TE。它根据资源使用情况将输入的TE程序划分为多个子程序，并将每个子程序转换为一个计算内核。例如，在图2中，如果TE4的块数超过了每波最大块数限制，Souffle会将TE程序划分为两个子程序。第一个子程序包括TE0、TE1、TE2和TE3，而第二个子程序包括TE4。
* TE变换：子程序连同数据流分析和张量信息被发送到TE变换模块，以生成语义保持但经过优化的TE。在图2中，TE1和TE3的计算分别被调度到TE0和TE2的内层循环中。TE的调度和变换在第6节中解释。
* 联合优化和代码生成：变换后的TE子程序被送入一个调度器优化器（在我们的案例中是Ansor【57，Ansor: Generating High-Performance Tensor Programs for Deep Learning，2020，OSDI 20】），为TE子程序生成一个调度。接下来，Souffle将一个子程序内的调度组合成一个由TensorIR【15，TensorIR: An Abstraction for Automatic Tensorized Program Optimization，2023，ASPLOS 2023】表示的单一函数，以进行子程序内的指令和数据重用的联合优化。最后，优化后的子程序被传递给后端代码生成器以产生CUDA内核。在图2中，ldg2s代表从全局内存加载到共享内存，wmma_16x16代表warp矩阵乘法累加，sts2g代表将共享内存存储到全局内存。Souffle将TE对应的代码包装在if语句中以匹配启动维度，并插入全局同步原语（在此例中为grid.sync()）以同步线程块之间的数据。SO0被缓存在共享内存中，并在算子边界（此工作示例中为TE1和TE3）之间重用。我们在第6.5节中描述这些过程。

# 1. TE Lowering → TE Program
rk = te.reduce_axis((0, 64),)
TE0: O0 = te.compute((64,64), lambda i, j: te.sum(I0[i,rk]*W0[rk,j]),axis=[rk])
TE1: O1 = te.compute((64,64), lambda i, j: te.sigmoid(O0[i, j]))
TE2: O2 = te.compute((64,64), lambda i, j: te.sum(O1[i,rk]*W2[rk,j]),axis=[rk])
TE3: O3 = te.compute((64,64), lambda i, j: O0[i,j] + O2[i,j])
TE4: O4 = te.compute((64,256),lambda i, j: te.sum(O4[i,rk]*W4[rk,j]),axis=[rk])

# 4. TE transformation for sub-program 0 ( [TE0, TE1, TE2, TE3])
TE0: s.reorder(io, jo, ko, ii, jj, ki)
TE1: s = te.create_schedule(O1.op)
TE1: io, ii, jo, jj = s.split(i, j, 16, 16) # Inherit tile shape from TE0's schedule
TE1: s[O1.op].compute_at(jo) # Move computation of TE1 into TE0's loop
# All memory-intensive TEs are fused into compute-intensive TEs

5 全局计算图分析

分析层次。Souffle在TE的张量依赖图上应用两级分析。第一级识别数据重用机会，第二级提取元素级的数据依赖。

5.1 识别数据重用机会

时空数据重用。张量通常在时间和空间维度上被重用，这为利用数据重用以消除昂贵的内存操作提供了机会。如第2.3节所讨论，在图1(a)和图1(b)的工作示例中有两种类型的张量重用。首先，三个GEMM0算子共享相同的输入张量，可以在空间上重用。将三个GEMM0算子融合成一个内核将允许我们加载一次输入并在GEMM0算子之间重用三次。这种重用机会在DNN中很常见，包括循环神经网络【44】，卷积神经网络【29, 45, 55】和Transformer模型【31, 49】。空间数据重用优化适用于由没有数据依赖关系的算子消耗的张量，这些算子将如第6.1节所述进行水平变换。第二类重用机会可以体现在时间维度上。时间数据重用机会适用于被具有数据依赖关系的算子多次使用的张量，并指导张量重用优化，这在第6.5节中描述。再次考虑图1中的工作示例，元素级算术算子1（称为A1）的结果被两个依赖算子R1和A2使用。如果我们将会A1的计算输出缓存在寄存器/共享内存中，同样可以消除对全局内存的访问。

重用机会识别方法。Souffle通过首先遍历计算图来收集所有被多个TE访问的张量，从而在张量级别从TE张量依赖图中识别这些数据重用机会。它记录了与张量 $T_i$ 共享的算子集合 $O(T_i) = \{TE_{i_1}, ..., TE_{i_n}\}$，以实现代码优化，如第6节所述。

5.2 TE内部元素级依赖分析

依赖分析方法。Souffle通过将迭代空间定义为输出形状，并将数据空间定义为每个输入张量的迭代和归约变量的域，来捕获TE内从输出到输入张量的元素级数据依赖。这简化了从输入到输出张量的元素级数据流，因为我们只需要记录输出张量的一个元素所对应的相关输入张量元素。该信息还使得在TE变换阶段能够进行归约算子融合，这是TensorRT和Apollo等其他优化工具不支持的。

依赖类型分类。我们的关键观察是，TE内部的数据依赖分为两类。首先，对于没有归约轴的TE（参见第3节），每个输出张量元素仅依赖于一个输入张量元素（称为“一对一依赖”）。其次，对于有归约轴的TE，每个输出元素依赖于输入张量所有归约维度的所有元素（称为“一对多依赖”）。基于这一观察，与依赖于源代码或算子级别分析的其他内核融合方法相比，我们可以大大简化依赖分析过程。

依赖关系的形式化表示。我们使用多面体模型表示法【46】来表示从输出张量元素到输入张量元素的元素级依赖关系。每个张量都有一个关联的集合 $S = [i_0, ..., i_n : C_0 \wedge ... \wedge C_k]$，代表其数据空间，其中 $i_j$ 是迭代变量，$C_j$ 是来自TEs的循环边界。关系表示输出元素依赖于输入张量元素。一对输出和输入张量与一个关系 $R = \{[o_0, ..., o_m] \mapsto [i_0, ..., i_n] : C_0, ..., C_k\}$ 相关联。对于图2中的TEs，TE1给出 $R_1 = \{O_1[i, j] \mapsto O_0[i, j], 0 \le i < 64, 0 \le j < 64\}$，TE0产生 $R_0 = \{O_0[i, j] \mapsto I_0[i, rk], 0 \le i < 64, 0 \le j < 64, 0 \le rk < 64\}$。TE1属于一对一依赖类型，TE0属于一对多依赖类型。

一对一依赖的TEs。Souffle采用准仿射映射（quasi-affine maps）【7, 36】来表示一对一依赖TE的元素级依赖。从输出到输入的映射可以表示为 $\vec{v} \mapsto M\vec{v} + \vec{c}$ 的形式，其中 $\vec{v}$ 是输出张量的索引，M是一个来自 $Z^{n \times m}$ 的常数矩阵，$\vec{c}$ 是一个来自 $Z^n$ 的常数向量。这里，m是输出张量的维度，n是相应输入张量的维度。注意，输出张量的多个索引可能依赖于输入张量的相同索引。例如，关系 $R_1$ 可以表示为：
$[1 \ 1] \begin{bmatrix} i \\ j \end{bmatrix} + [0 \ 0], 0 \le i < 64, 0 \le j < 64$

一对多依赖的TEs。对于一对多依赖的TE，Souffle通过结合迭代空间和输入张量的索引函数来提取被访问的输入张量区域。由于归约轴的迭代域是一个常量值，映射可以表示为 $R = \{[o_0, ..., o_m] \mapsto \{[i_0, ..., i_n], [r_0, ..., r_p]\} : C_0, ..., C_k\}$ 的形式，其中 $[r_0, ..., r_p]$ 是一组归约变量及其范围。例如，关系 $R_0$ 可以表示如下：$R_0 = \{O_0[i, j] \mapsto \{I_0[i, rk], [0 \le rk < 64]\}, 0 \le i < 64, 0 \le j < 64\}$，其中 $\{I_0[i, rk], [0 \le rk < 64]\}$ 代表一组rk从0到64的元素。我们强调，像GEMM和卷积这样的计算密集型算子的元素级依赖可以很容易地分析，因为张量表达式明确给出了归约轴。

后续处理。具有一对一依赖的TE将在第6.2节中进行变换，而具有一对多依赖的TE将在第6.3节和第6.4节中进行调度。

5.3 TE 特征化

计算与内存密集型分类。Souffle通过估算一个TE的计算-内存比率，将其分类为内存密集型（例如，reduce_sum）或计算密集型（例如，GEMM）。该比率通过将算术指令的数量除以内存访问的次数来计算。因此，该比率单位表示每个被读取和写入的张量元素的算术操作数。在这项工作中，分类阈值经验性地设置为3。小于该阈值的比率表明该TE是内存密集型的。

5.4 TE 程序分区

分区目标与约束。Souffle试图生成大型内核以最大化数据重用并消除额外的内核启动。然而，使用全局同步施加了一个约束，即线程块计数不能超过每波的最大块数。如果无法满足此约束，Souffle会将TE程序划分为多个TE子程序。在Souffle中，一个TE子程序是高级TE变换、中端调度优化和后端代码生成的基本单元。它可以包括映射到一个GPU内核的多个算子。

分区点的选择。我们只考虑计算密集型算子作为候选分区点。计算密集型TE通常比内存密集型TE使用更多的共享内存和寄存器。过度使用共享内存和寄存器会推高占用率，从而限制每波的最大块数，使其不适用于全局同步。Souffle变换内存密集型TE，并使用其计算密集型生产者TE的调度来实现更好的数据重用（第6节）。

获取所需资源。Souffle从调度优化器（本工作中为Ansor）生成的TE调度中获取内核启动维度和寄存器/共享内存占用率。

分区算法。Souffle使用分析模型确保在TE程序分区中满足资源约束。给定一个总共有 $R$ 个寄存器/共享内存的GPU，Souffle提取当前TE子程序中所有计算密集型TE的最大启动维度 $D_{max}$ 和最大寄存器/共享内存占用率 $O_{max}$。然后它检查是否能满足约束 $D_{max} \times O_{max} < R$ 对于子程序内所有选定的TE。Souffle使用贪心算法对TE程序进行分区，从一个空的 $SP_i$ 开始，并使用广度优先搜索（BFS）将TE $TE_j$ 添加到 $SP_i$。如果将 $TE_j$ 添加到 $SP_i$ 违反了约束，Souffle会通过将此TE添加到新的子程序 $SP_{i+1}$ 来创建一个新的子程序，并重复此过程，直到所有TE都已分配到子程序中。

6 语义保持的TE变换

变换目标。在Souffle收集了第5节所述的重用和依赖信息后，它会寻找机会自动变换TE，以提高每个TE子程序内的性能。Souffle支持水平和垂直TE变换，并在同一子程序中变换TE。水平融合将算子的分支融合成单个内核【33, 52】。Souffle中的水平变换类似于水平融合，并应用于独立TE的分支。垂直变换类似于垂直融合【37, 58】，并应用于具有一对一数据依赖的多个连续TE。我们强调，我们的水平和垂直变换比IREE和Rammer【33】使用的融合策略更灵活，后者不会将具有一对多依赖的算子融合成一个内核。此外，语义保持变换确保在满足数据依赖要求的同时保留算术运算（例如add, exp）。相比之下，其他DNN优化方法使用的一些变换可能不保留语义。例如，TASO【20】通过子图替换来优化DNN图，并可能用concat+convolution替换add。

6.1 独立TEs的水平变换

变换方法。Souffle尝试将多个独立的TEs变换为单个TE以增加并行性。Souffle首先比较每个独立TEs的输出张量形状，并尝试将它们连接成单个TE。Souffle将多个独立TEs的输出张量连接成一个，因为每个TE只能产生一个输出张量。Souffle根据输出区域添加谓词并重写TE。随后，Souffle评估这些TEs是否消耗相同的输入张量。值得注意的是，会检查第5.1节中讨论的张量重用机会。因此，共享张量在生成的内核中只需要加载一次。因此，内核数量和全局内存事务都可以减少。图3给出了一个水平TE变换的例子，两个TEs共享相同的归约变量，rk。第一个和第二个TEs的输出分别为(4, 16)和(2, 16)，可以在第二个轴上连接成一个形状为(6, 16)的单个张量。如果独立TEs的输出不能连接，它会添加一个if_else语句来为连接的TEs选择相应的输入张量，类似于Rammer【33】。

6.2 一对一依赖TEs的垂直变换

变换方法。Souffle将具有一对一数据依赖的TEs垂直变换为一个TE，以减少生成的内核并重用寄存器上的数据。这得益于准仿射映射表示（第5.2节）。为此，Souffle首先通过将子TEs的索引映射函数应用于其父TEs来变换所有一对一依赖的TEs。然后，它应用索引映射函数并创建一个单一的语义保持的TE。假设我们有n个TEs，$TE_0, TE_1, \dots, TE_i, \dots, TE_{n-1}$，其中 $TE_i$ 的输出是 $TE_{i+1}$ 的输入。映射函数可以表示为 $f_i(\vec{v_i})$。从 $TE_{i+1}$ 到 $TE_i$ 的变换后TE的映射函数可以通过以下函数计算：
$$ f_{i+1,i}(\vec{v_i}) = f_{i+1}(f_i(\vec{v_i})) = M_{i+1} \times (M_i + \vec{c_i}) + \vec{c_{i+1}} $$

变换示例。对于图4中的例子，三个TEs，A, B和C的索引映射函数可以转换为一个单一的数学等效函数——有效地将TE的数量减少了3倍——如下：
$$ \begin{bmatrix} 0 & 1 \\ 1 & 0 \end{bmatrix} \begin{bmatrix} 2 & 0 \\ 0 & 1 \end{bmatrix} \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix} \begin{bmatrix} i \\ j \end{bmatrix} \Rightarrow \begin{bmatrix} 0 & 1 \\ 2 & 0 \end{bmatrix} \begin{bmatrix} i \\ j \end{bmatrix} $$

迭代优化。使用上述方法，Souffle从一组连续的TEs中迭代地优化多个一对一依赖的TEs，直到找不到进一步的优化可能。然后，它应用其计算密集型父TE的调度，将内存密集型的一对一依赖TEs附加到计算密集型TE上，如下一小节所述。与TesorRT、Apollo和Ansor使用的手工变换规则相比，我们的语义保持变换具有更好的泛化能力。

6.3 TEs调度

调度策略。Souffle使用Ansor为计算密集型和内存密集型TEs生成优化的调度。注意，在第5.4节的TE程序划分期间，我们已经为计算密集型TEs生成了调度。它为内存密集型TEs传播计算密集型生产者的调度，以最大化数据重用。对于一对一依赖的TEs，Souffle首先根据其计算密集型TEs的tile大小进行调度，然后安全地将它们内联到其生产者的计算语句中。对于一对多依赖的TEs，Souffle在本地进行归约以重用生产者在共享内存/寄存器上的数据，然后使用原子原语在线程块之间进行归约。

6.4 合并TEs调度

调度合并与同步。在调度TEs之后，Souffle将一个子程序内的TEs的调度合并，使用TensorIR【15，TensorIR: An Abstraction for Automatic Tensorized Program Optimization，2023，ASPLOS 2023】创建一个整体函数。如果TEs的启动维度不同，它会添加谓词，并在具有一对多依赖的TEs之间插入全局同步原语。最后，它执行下一节描述的几种优化。

6.5 子程序内的优化

优化类型。Souffle在一个TE子程序内支持两种类型的优化。第一种是指令级优化，旨在将异步GPU内存加载与算术指令重叠。请注意，这种流水线执行是跨TEs调度的，没有全局数据依赖分析，优化无法完成。第二种是跨TEs（并可能跨多个原始算子）重用张量缓冲区。Souffle在底层调度器（本工作中为Ansor）为一个子程序内的TE调度生成之后，执行子程序级优化。它利用全局计算依赖分析结果（第5节）来应用这两种优化。

指令级优化。Souffle在一个包含多个原始算子的融合子程序内重组指令，以并行执行内存和算术指令。这是通过跨算子边界调度加载/存储和计算指令以实现流水线执行来完成的。例如，在第2.1节和图1(d)中讨论的BERT模型中，Souffle生成的调度并行发出NVIDIA指令LDGSTS.E.BYPASS.128和HMMA.16818.F16，其中前者将128位从GPU全局内存复制到共享内存，后者在NVIDIA张量核上计算GEMM。

张量重用优化。Souffle通过一个简单的软件管理的缓存最大化跨TEs的张量缓冲区重用，使用最近最少使用（LRU）策略在运行时替换共享内存中的张量缓冲区（例如，矩阵和向量）。它线性扫描指令直到共享内存耗尽，如果共享内存耗尽，则将共享内存溢出到全局内存并添加内存屏障。

6.6 整体流程

算法1：语义保持的TE变换。算法1概述了TE变换。它接收TE程序和相应的分析结果：OR（一对一依赖TEs），MR（一对多依赖TEs），MI（内存密集型TEs），CI（计算密集型TEs），TR（时间重用张量和TE元组），以及SR（空间重用张量元组）。它根据计算密集型TEs的调度与资源限制来划分TEs（第2-9行）。然后，它在每个分区内通过垂直变换进行水平变换和优化空间重用（第11-12行）。它将计算密集型TEs的调度传播到内存密集型TEs，并合并每个子程序内的调度（第13-18行）。最后，它通过时间数据重用和跨原始算子边界优化张量重用（第19-21行）。

6.7 实现细节

代码实现。我们用10K行C++和1K行Python代码实现了Souffle。我们将Souffle与TVM【12】集成，并使用Ansor【57】作为其调度优化器。然而，Souffle可以与其他兼容TEs的调度器一起工作。Souffle支持元素级算子、广播、归约（包括reduce_sum、GEMM和Conv）、重组算子如reshape和shuffle算子如transpose。Souffle不支持非线性代数算子如TopK或Conditional。我们使用直接卷积，这是Ansor的默认实现。

A4 实验环境

• 硬件配置:
  • CPU: 双路20核 Xeon Gold 6248 @ 2.50 GHz
  • 内存: 768GB DDR4 RAM
  • GPU: 1块 40GB NVIDIA A100 GPU
• 软件配置:
  • 操作系统: Ubuntu 18.04.5 (Linux kernel 5.4.55)
  • CUDA: 11.7
  • 编译器选项: -O3
• 模型与数据集:
  • 实验评估了6个具有代表性和多样性的DNN模型，涵盖自然语言处理（BERT）、计算机视觉（Swin-Transformer）和知识发现（MMoE）等领域，也包括经典的卷积网络（ResNeXt）和循环网络（LSTM）。
  • 除GEMM使用半精度浮点（FP16）以利用张量核外，所有算子均使用单精度浮点（FP32）。
  • 所有实验均针对模型推理场景，批处理大小（batch size）设为1。

• 对比基线 (Baselines):
  • XLA (Tensorflow v2.10): TensorFlow的优化编译器，支持算子融合。
  • TensorRT (v8.2): NVIDIA官方的DNN推理优化框架。
  • Ansor (TVM v0.8): TVM中的先进DNN优化器，使用自动调优技术。
  • Rammer (v0.4): 也称为NNFusion，通过时空调度进行优化的DNN编译器。
  • Apollo: 先进的基于分区的算子融合框架。
  • IREE (2022年12月30日发布): 基于LLVM MLIR项目的中间表示执行环境，用于优化模型推理。

A4 实验结果

8.1 总体性能

Souffle在所有被评估的DNN模型上均优于所有六个竞争基线。如表3所示，Souffle在TVM+Ansor的基础上实现了平均3.9倍（最高8.5倍）的加速。与NVIDIA官方调优的TensorRT相比，Souffle取得了平均3.7倍（最高7.9倍）的加速，与Rammer、Apollo和IREE相比也取得了类似的性能提升。这表明Souffle在不同DNN负载下都能提供稳定且强大的性能改进。

XLA、Rammer和Apollo等融合技术虽然在某些情况下优于Ansor，但它们能融合的算子集合有限，并且缺乏跨算子的有效指令级优化。例如，Rammer依赖手工规则，XLA将GEMM等算子映射到库调用而无法融合，Apollo的循环融合规则受限，IREE无法融合某些计算密集型算子。相比之下，Souffle操作于语义简单明确的TEs，不依赖僵化的融合规则，能够识别更多数据重用机会，并进行更灵活的融合和指令级优化。

8.2 性能分解

为了量化各项优化的贡献，实验从基线TVM+Ansor（V0）开始，逐步添加Souffle的优化：TE水平变换（V1）、TE垂直变换（V2）、全局同步API（V3）和子程序级优化（V4）。
如表4所示，水平和垂直TE变换对所有DNN模型都有益，它们增加了SIMD并行度并减少了内存访问。对于基于Transformer的BERT和Swin-Trans.模型，全局同步和子程序级优化（重叠加载和计算、张量缓冲区重用）也带来了显著的性能提升。

8.3 性能优势分析

Souffle的性能优势主要来自两个方面：减少GPU内核调用和减少GPU内存数据传输。
* 微基准测试 (EfficientNet):
* 使用EfficientNet的一个重复子模块进行测试，比较了未融合（Unfused）、Ansor融合（Fused）、Souffle全局同步（Global-sync）和Souffle数据重用（Data-reuse）四个版本（图5）。
* 结果（图6）显示，仅使用全局同步，Souffle就比Ansor融合版本平均快1.31倍。启用数据重用后，加速比进一步提升至1.84倍。
* 内核调用和内存传输分析 (端到端模型):
* 减少GPU内核调用：Souffle通过资源感知的TE程序划分创建了更大的子程序，从而生成了更少的内核，显著降低了内核启动开销。如表5所示，对于BERT模型，Souffle将内核数从TensorRT的120个和Apollo的240个减少到24个。
* 减少GPU内存数据传输：Souffle通过TE程序划分和TE变换最大化了张量缓冲区的重用。如表5所示，对于BERT模型，Souffle将内存传输量从TensorRT的361.8MB和Apollo的880.5MB减少到226.8MB。
* BERT案例分析:
* 相比TensorRT将一个BERT层映射为10个内核，Souffle能将其划分为2个内核，并进行指令级优化，将内存密集型内核的延迟从31.0μs降低到25.5μs。
* 相比IREE（启动180个内核，耗时2.22ms），Souffle因能融合GEMM和softmax等算子，只启动24个内核，耗时1.22ms。

8.4 LSTM 案例研究

在LSTM模型上，Souffle比TensorRT快4.3倍，比表现最好的基线Rammer快2.2倍。如图7所示，Rammer利用了沿对角线算子间的“波前并行”，但每个波前都需要重新加载权重张量。而Souffle通过基于TE的全局分析，发现了权重张量在所有时间步中的时间重用性，利用全局同步将整个模型生成为一个内核。如表6所示，Souffle优化后的代码相比Rammer，内存加载量从1911MB锐减至21.11MB，同时显著提高了加载存储单元（LSU）和融合乘加单元（FMA）的流水线利用率。

8.5 编译开销

Souffle的额外编译开销主要来自依赖分析、模型拆分、调度调整和全局优化。在六个评估模型上的测量表明，Souffle在Ansor的基础上增加了最多63秒的开销，与Ansor本身需要数小时的调度搜索时间相比，这部分开销可以忽略不计。

A7 补充细节

9 讨论

• TE的表达能力。Souffle依赖于张量表达式的表达能力，目前TE尚不支持所有DNN算子（例如resize）。对于这些TE不支持的算子，Souffle会将其映射到一个计算内核并使用后端算子库的实现，但不会将它们与其他算子融合。我们期望随着TVM社区的活跃发展，这一限制将在未来的TVM版本中得到解决。
• TE程序分区的成本模型。Souffle通过编译原始TE程序来提取张量信息。这可以通过构建一个成本模型【53】来改进，从而直接从TE程序估算占用率。
• 重用动态形状的张量。某些DNN算子在编译时具有未知的张量形状。我们当前的实现不支持重用动态形状的张量。为了解决这个问题，我们可以生成一个内核的多个版本，并在执行时根据可用的形状信息选择合适的版本。
• 深度学习训练中的融合。像TensorFlow XLA这样的DL编译器也在训练（前向推理和后向参数更新）中启用了算子融合。我们基于TE的变换可以集成到DL编译器中，以加速训练中的前向和后向传播。然而，在DL训练中，为了进行像Adam【24】这样的基于梯度的后向优化，中间张量必须保留在全局内存中，这限制了算子融合的机会。我们的主要重点是优化DNN训练后的模型推理。在DL训练中支持TE变换是未来的工作。
• 性能下降。当Souffle将调度从计算密集型TE扩展到内存密集型归约TE时（在第6.3节中讨论），可能会发生性能下降。这在块之间引入了同步，可能妨碍归约TE的并行性。一个潜在的补救措施是创建一个成本模型来决定融合这些TE是否有利。

10 相关工作

• 循环和内核融合。循环融合通常用于提高CPU程序的性能【3, 8, 9, 23】。最近的研究也利用内核融合来通过减少与片外内存的数据传输来优化GPU程序【42, 50】。已经为数据中心应用【54】、网格计算【6】、机器学习工作负载【4】和图像处理【35】等工作负载提出了各种领域特定的内核融合策略。Souffle在这些先前的研究工作基础上，利用循环融合通过编译器变换来优化DNN推理。
• DNN编译器中的算子融合。算子融合可以通过改善数据重用和减少片上数据流量来提高性能。为了寻找融合机会，DNNFusion对算子进行分类，并根据其分类定义规则【37】。Astitch【58】和Rammer【33】融合独立的算子以利用算子间的并行性，而Deepcuts【22】使用规则根据GPU硬件参数融合内核。Apollo【56】采用基于分区的方法在子图中搜索融合机会。然而，这些方法依赖于手工制定的规则，需要大量的工程努力，并可能错失优化机会，如第2节所讨论。Jeong等人【19】提出了一种动态规划算法，来决定在具有全局缓冲区的DNN加速器上是固定还是融合一个激活图。DNNFusion【37】根据从输入到输出的元素级映射对算子进行分类，但它不能融合多对多与多对多的算子（如GEMM和Softmax），而Souffle可以进一步减少内核启动的开销。此外，DNNFusion缺乏对张量重用机会的全局分析，可能错失时间和空间数据重用机会，而这对于提高性能至关重要，如第8.2节所示。Souffle通过利用张量依赖图上的控制流和数据流分析将TEs划分为子程序，从而改进了以前的算子融合技术。TEs具有清晰简单的关系，它们可以组合起来表示众多的DNN算子。Souffle利用TEs中明确定义的语义来进行精确的数据依赖分析，以进行指令调度和数据重用优化。此外，Souffle应用语义保持变换来优化TEs。由于TEs可以组合以表示更复杂的算子，其优化能力具有更好的泛化性。
• 全局分析和融合优化。TensorFlow XLA【27，XLA: TensorFlow, compiled，2017，TesorFlow Dev Summit】和MLIR【43】也对输入程序图进行全局分析。XLA在其高级操作中间表示上利用盈利能力分析，然后决定分块和融合。然而，XLA依赖于手工制作的启发式方法来融合算子，这对于高级算子可能具有挑战性。例如，XLA的融合启发式不能融合BERT模型中两个连续的归约算子。此外，由于XLA在算子级别操作，并且一些算子被映射到低级库调用，因此它无法跨库进行优化。相比之下，Souffle采取了不同的方法，将高级算子降级为具有简洁语义的低级张量表达式（TEs）。在TEs上操作而不是假设高级算子，使Souffle能够灵活地跨算子边界进行优化。与XLA不同，Souffle可以合并GEMM和Softmax算子，并跨归约算子进行优化。MLIR的-affine-loop-fusion pass利用基于切片的方法来识别生产者-消费者和兄弟融合机会。Souffle在TEs上实现了一个轻量级的、专门的全局分析，可以轻松集成到DNN推理引擎中。此外，Souffle提供了比仅仅融合更多的优化机会。例如，它支持在TE子程序中跨多个计算密集型TEs进行联合优化，并促进了兄弟融合的水平和垂直变换。
• 优化单个算子。存在许多基于编译器的方法来优化单个算子，包括TVM【12, 51】、XLA【27，XLA: TensorFlow, compiled，2017，TesorFlow Dev Summit】、Tiramisu【5】和TACO【25】。这些编译器通常以高级形式（如TEs或线性代数）表示算子，通过领域特定知识实现激进的优化，而无需复杂的分析。Souffle与这些技术是正交的。

A5 结论

本文提出了Souffle，一个用于改进DNN推理的自顶向下、基于编译器的方法。Souffle通过对由张量表达式构建的整个张量依赖图进行数据流分析，来识别跨DNN算子的优化机会。它将张量表达式分组为子程序，并通过语义保持的变换、指令调度和张量缓冲区重用进行局部优化。我们在NVIDIA A100 GPU上对六个DNN模型评估了Souffle，并将其与六个最先进的DNN优化框架进行了比较。实验结果表明，Souffle的性能优于它们，相较于TensorRT最高可实现7.9倍的加速。

未来的工作将致力于解决当前的一些局限性，包括扩展Tensor Expression（TE）以支持更多DNN算子，构建更精确的成本模型以指导程序分区，支持对动态形状张量的重用，以及将基于TE的变换技术应用于深度学习训练过程的优化。