POWERFUSION: A Tensor Compiler with Explicit Data Movement Description and Instruction-level Graph IR
POWERFUSION: A Tensor Compiler with Explicit Data Movement Description and Instruction-level Graph IR
作者/机构: Zixuan Ma, Haojie Wang, Jingze Xing, Liyan Zheng, Chen Zhang, Huanqi Cao Kezhao Huang, Shizhi Tang, Penghan Wang and Jidong Zhai (清华大学)
A1 主要贡献
深度神经网络（DNNs）在各领域应用广泛，为加速其计算，张量编译器应运而生。现有张量编译器主要关注计算效率优化，但随着加速器计算性能的增长远超内存性能，内存访问正成为关键瓶颈。当前张量编译器的中间表示（IR）缺乏对内存访问和数据依赖的直接描述，这给生成内存高效的代码带来了巨大挑战。
核心问题与研究目标：
近年来，加速器的计算性能与内存带宽之间的差距急剧扩大。如图1所示，从2016年到2023年，计算与内存性能的比率增加了11.3倍，使内存性能成为DNN模型的主要瓶颈。同时，不同加速器在架构和内存层次上的差异也给设计可移植的高性能DNN系统带来了优化挑战。因此，优化内存效率对于充分发挥硬件性能至关重要。
图 1: 近年来计算性能和内存性能的性能趋势。
现有张量编译器（如TVM）在处理内存密集型程序时面临三大挑战：
1. 隐式数据移动表示：数据移动模式通过调度（schedule）隐式表示，难以在优化过程中评估和有效优化内存性能。
2. 调度搜索顺序：TVM采用先融合循环再应用调度的顺序，这妨碍了对不同计算操作应用不同的内存访问模式，从而错失了潜在的优化机会。
3. 粗粒度依赖分析：依赖分析仅限于循环级别，与内存层次中的数据粒度不匹配，导致错失通过数据复用减少内存访问的机会。
本文贡献：
为解决上述问题，本文提出了POWERFUSION，一个通过同时考虑计算和数据移动优化来为内存密集型算子生成高性能代码的张量编译器。其核心贡献如下：
* 提出GIR（Graph IR）：一种通过指令级图描述显式表示数据移动模式和细粒度数据依赖的中间表示，旨在优化内存性能。GIR使用并行、计算和数据移动三种原语来描述算子。
* 设计了一套基于GIR的优化策略：包括在GIR上的搜索和变换，以减少内存访问并提高内存效率。通过显式表示指令级依赖，POWERFUSION可以分析细粒度依赖并应用图优化来减少冗余内存访问。
* 提出硬件抽象并开发端到端编译器：设计了一个通用抽象，能够适应具有不同并行结构和内存层次的各种架构。通过配置GIR-Graph的生成和搜索过程，POWERFUSION可以为特定硬件平台生成优化代码，开发者只需实现计算和数据移动的指令到指令映射即可适配新硬件。
POWERFUSION已在NVIDIA GPU、AMD GPU和Cambricon MLU上实现，并与各硬件上最高效的DNN框架相比，分别取得了高达1.97倍、2.93倍和16.91倍（平均1.28倍、1.23倍和2.31倍）的加速比。
A3 背景知识
TVM及其局限性。TVM和Ansor采用计算/调度分离的思想，并通过自动调优技术搜索不同的调度方案来生成代码。如图2所示，在优化多个计算操作时，TVM首先合并不同操作的循环，然后为合并后的循环搜索执行计划。这种方法导致两个局限性：首先，在依赖分析时，无法分析对应于高层内存层次上每次数据移动操作的数据块的重用关系，导致许多操作无法融合。其次，先合并后调度的策略使得不同算子无法采用不同的执行计划，导致次优的内存性能。
多面体方法。以PPCG【索引24，Scheduling for PPCG+2017+CW Reports】和PLUTO【索引5，A practical automatic polyhedral program optimization system+2008+ACM SIGPLAN Conference on Programming Language Design and Implementation (PLDI)】为代表的多面体方法，能够通过将程序数学化地表示为高维整数空间中的点来分析循环程序中的元素级依赖，并通过变换该空间来优化程序。这类工作通常采用整数线性规划来寻找数学上最优的变换。然而，由于缺乏对特定硬件约束的考虑，优化后的程序在实际硬件上运行时可能表现不佳。同时，考虑到整数线性规划的巨大开销，也很难采用自动调优技术。
图 2: TVM与POWERFUSION方法的图示。
图级融合方法。另一种方法，以TensorFlow-XLA【索引3，Xla: Optimizing compiler for tensorflow+2017+https://www.tensorflow.org/xla】和DNNFusion【索引15 ，Dnnfusion: accelerating deep neural networks execution with advanced operator fusion+2021+Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation】为代表，在计算图级别使用核融合来减少内存访问。它在计算图上生成融合计划，然后使用代码生成工具生成设备代码。TensorRT【索引22，NVIDIA TensorRT: Programmable inference accelerator+2017+https://developer.nvidia.com/tensorrt】是另一类工作，它使用一系列规则将DNN模型中的复杂计算操作映射到手动优化的核。这种方法以操作为分析粒度，其可扩展性受限于后端算子库或代码生成的能力 。
POWERFUSION的分析粒度。为了解决这些局限性，POWERFUSION采用指令块和内存切片作为分析粒度。这个粒度与硬件上实际执行的计算和数据移动操作相匹配，为数据重用提供了所需的信息。通过将程序组织成一个包含这种粒度的计算和数据移动操作的图，依赖分析的开销显著降低。如图2所示，与TVM相比，POWERFUSION首先搜索不同操作的执行计划，然后通过细粒度的依赖分析组合不同的执行计划，以实现更高的内存性能。
异构加速器的挑战与POWERFUSION的抽象。各种领域特定的加速器，如NVIDIA GPU和Cambricon MLU，通常采用基于暂存存储器（scratch-pad memory）的分层内存结构，包括寄存器文件、共享内存和全局内存等多层。这种复杂多样的内存层次结构给加速器的内存性能优化带来了两个主要挑战：1）内存层次与并行结构的深度耦合，难以确定数据使用的最佳内存层次。2) 不同内存层次之间的数据移动复杂，难以优化。为了应对这些挑战，我们将不同硬件的内存层次抽象为一个统一的抽象，并在其属性中表示不同硬件的特性。这种方法使我们能够提供统一的优化，并简化不同架构的内存性能优化。此外，通过在GIR中加入同步操作，我们可以有效地分析数据可以使用的最高内存级别，使POWERFUSION能够为不同架构提供最佳的内存性能优化。
3. 概览
图3展示了POWERFUSION的概览，并通过一个真实DNN模型中的运行示例说明了它如何优化内存性能。给定的模型片段来自ShuffleNet，它将两个张量合并为一个，对其进行洗牌（shuffle），然后将其分割成两个新的张量。为了优化这个模型，POWERFUSION将此模型片段作为输入，并将其转换为称为GIR的指令级图IR。然后，POWERFUSION通过基于GIR的优化来优化该模型，并最终生成高性能代码。GIR抽象将在第4节中介绍。
GIR优化流程。基于GIR的优化将计算图作为输入，并对每个算子分别应用图生成，以获得每个算子的GIR-Graph。接下来，它搜索并合并这些GIR-Graph，以获得对应于计算图的GIR-Graph。利用合并后的GIR-Graph，POWERFUSION接着应用图重写规则，使用三种图重写规则来优化GIR-Graph，从而减少内存访问量，显著提高内存性能。图重写和基于GIR的优化过程将分别在第5节和第6节中介绍。
代码生成。最后，POWERFUSION重新排列优化后的GIR-Graph，并顺序地为特定硬件生成代码。代码生成技术将在第7节中讨论。
图 3: POWERFUSION的概览和动机示例。
图 4: GIR抽象的图示。
A2 方法细节
4. GIR 抽象
本节介绍GIR的设计，旨在解决现有张量编译器中缺乏对内存层次间数据移动的全面描述，从而阻碍内存访问性能优化的问题。GIR的核心思想是将一个张量程序表示为一个由计算和数据移动原语（称为gOperators）组成的图，其粒度是在特定并行级别下的指令块。图4展示了一个RELU操作的GIR-Graph及其设备代码。
GIR的结构。GIR被表示为一个由并行原语包裹的数据流图。GIR-Graph的每个节点表示一个指令级的计算或数据移动gOperator，每条边代表特定内存层次上的一个内存切片。并行原语指明了该指令级数据流图的并行策略。这种方法有两个优点：首先，所有的计算和数据移动操作都以适当的粒度表示，使其更易于优化；其次，计算和数据移动之间的依赖关系清晰且易于分析。中间表示中的图结构也使得可以使用图算法来优化GIR-Graph并实现端到端的优化效果。
内存切片（Memory slice）。在GIR中，数据被表示为内存切片，即张量上一组连续的元素。每个内存切片有四个属性：内存层次（memory hierarchy）、数量（num）、宽度（width）和步长（stride）。内存层次指明了该内存切片存储在哪一级内存中，例如DRAM或SRAM。num属性指定了连续段的数量，每个段的长度为width个元素，段与段之间有固定的stride步长。这种格式有两个优点。首先，它符合内存性能的最佳实践，因为大多数内存层次（如DRAM和SRAM）都为连续访问进行了优化。其次，它与计算指令块的粒度和形状相匹配，这些指令块通常也设计成这种格式。这种二维格式能够在不显著损失性能的情况下对数据移动操作进行微调。
计算gOperators。GIR在指令块级别使用计算抽象，即计算gOperator，其中每个计算gOperator对应于代码中优化计算核的一组指令。GIR中的计算由三类gOperator定义：逐元素（element-wise）、归约（reduce）和广播（broadcast）。逐元素gOperator执行元素到元素的计算，如算术和激活函数。归约和广播gOperator处理维度扩展和收缩。需要注意的是，大多数封装的计算也可以表示为上述三种gOperator，例如，我们可以用逐元素和归约操作来表示矩阵乘法。POWERFUSION将自动选择表示方法以获得更好的性能（一个例子见图8）。使用这种细粒度的计算原语，我们可以描述程序计算并优化其调度以获得更高性能，而无需考虑与硬件相关的信息，如指令排序。
数据移动gOperators。GIR中的数据移动gOperators描述了数据在不同内存层次内部和之间的传输。这些操作要求输入和输出具有相同的模式，包括元素数量、每个元素的宽度以及连续元素间的步长。通过显式地表示数据移动操作，POWERFUSION能够高效地处理数据传输并优化内存性能。
同步gOperators。GIR还显式引入了一种特殊的数据移动gOperator，称为同步gOperator，以确保数据移动gOperator之间的正确性。例如，在CUDA中，当不同的warp访问同一个共享内存对象时，线程块内的同步操作是必需的。GIR中同步gOperator的scope属性指定了其作用范围，这决定了操作的实现方式。在CUDA中，warp内同步使用warp shuffle，而块同步使用共享内存。在GIR中，同步操作不仅意味着等待，还表明内存切片模式的改变，反映了不同的数据移动操作使用不同的内存切片模式来读写相同的内存数据。通过引入同步操作，GIR能更好地分析数据间的依赖关系并确保程序正确性。
并行（Parallel）。在POWERFUSION中，我们将并行粒度定义为最小的内存性能粒度。例如，在GPU上，它对应于一个流式多处理器（SM），在CUDA编程模型中等同于一个warp。而在CPU和Cambricon MLU上，它对应于一个核心或IPU，等同于一个线程。我们将这个最小单元称为并行单元，并为每个单元分配一个索引。然后，我们将硬件的嵌套并行结构映射到索引空间上，每个结构连续分配不同的段。例如，在CUDA中，同一块中所有warp的索引是连续的。这种方法使我们能够分析每个数据移动操作的并行范围，并消除跨架构的不同并行嵌套结构的影响。因此，我们可以简化代码优化过程以实现更高效的执行。
5. GIR-Graph 重写
如第4节所讨论，GIR可以表达计算和数据移动gOperators之间的数据依赖，从而带来了额外的分析和优化机会。为了优化一个GIR-Graph，我们提出使用一组称为重写规则的规则进行图重写技术，这些规则通常用于图变换。在POWERFUSION中，我们提出了三种类型的重写规则和一个在GIR上部署这些规则的优化策略。
5.1. 重写规则
同步插入。同步插入的目标是提升GIR中内存切片的内存层次，并添加同步操作以确保程序的正确性。当多个数据移动gOperator访问同一个内存切片时，参与该gOperator的并行单元之间必须进行同步。同步的需求源于同步前后不同的并行单元访问同一个内存切片。本质上，读-同步-写的过程是在不同并行单元之间交换数据。不同的内存层次对应不同的并行范围。例如，在NVIDIA GPU上，同一warp内的数据交换存储在寄存器中，并使用warp shuffle进行同步。当数据在线程块内的不同warp之间交换时，则使用共享内存，并在线程块级别进行同步。因此，通过分析同一内存块的读写操作的数据交换并行单元，可以确定数据存储的最高内存层次。
图 5: POWERFUSION如何确定同步范围的演示。
同步范围的确定。内存层次和并行单元之间的映射是特定于每个硬件架构的，并通过一组条件来表示。虽然这些条件与核函数执行配置相关，但它们的格式是独立于硬件的。例如，在一个每个线程块有四个warp的CUDA核中，同步范围条件如图5所示。图中的每个方块代表一个内存切片，其图案代表其内存切片模式。不同的颜色代表访问相应内存中数据的不同并行单元。同步范围由参与读写操作的单元决定。如果写入的单元与读取的单元相同，同步范围被限制在warp级别。在读写模式相同的情况下，同步gOperator的范围被限制在lane级别，该操作没有效果。如果读取或写入单元的ID除以4的结果相同，则同步范围被限制在块级别。否则，范围被限制在设备级别。
同步插入操作。使用上述条件，我们定义了一个同步插入操作。当遇到两个连续的数据移动gOperator时（允许它们之间有同步gOperator，本节中所有连续gOperator都是如此），我们分析两个内存切片之间的同步关系，并将它们的内存层次提升到最高级别。然后，一个新的同步gOperator被插入到重写后的图中，唯一的例外是lane同步，它只插入一个新的内存切片而没有同步。这个过程提高了GIR的内存性能并确保了程序的正确性。
图 6: GIR图重写的示例。
gOperator合并。虽然我们可以通过应用同步插入来提高内存切片的内存层次，但GIR中仍存在大量冗余的数据移动gOperator，并引入了额外的数据移动和同步开销，同时生成了冗余的内存切片，从而降低了执行性能。这种冗余主要分为两类：写后读（read-after-write）和读后读（read-after-read）。写后读意味着两个连续的数据移动gOperator顺序地操作同一个内存切片，先写后读。这种数据移动可能是冗余的。读后读意味着两个数据移动gOperator读取同一个内存切片。当两个gOperator之间没有同步关系时，这种组合会导致冗余的数据移动。为了解决这个冗余问题，我们提出了gOperator合并规则，该规则由两条规则组成：
* 用一个新的数据移动gOperator替换两个连续的数据移动gOperator，新gOperator的输入是第一个gOperator的输入，输出是最后一个gOperator的输出。如果读写的内存切片模式一致，则它将被一个空的gOperator替换。
* 将两个读取相同内存切片且无依赖关系的数据移动gOperator替换为一个gOperator。
通过将多个gOperator合并为一个，该规则可以减少GIR-Graph的内存访问，从而优化内存性能。
gOperator交换。应用上述两条规则并不能保证程序的内存性能达到最优。事实上，如图6所示的一个GIR，无法被进一步优化。这是因为程序受到两个全局同步gOperator的约束，导致生成了四次DRAM读写操作。这是由于数据移动和同步gOperator之间的脱节，使得优化规则无法应用。为了解决这个问题，我们可以移动计算gOperator。例如，由算术计算、激活函数等表示的逐元素gOperator可以与数据移动、同步、广播以及部分归约gOperator互换。这种可互换性使我们能够移动逐元素gOperator，从而提供更多的优化机会。图6展示了通过gOperator交换优化GIR-Graph的过程。通过应用gOperator交换，GIR可以进一步应用gOperator合并规则，从而减少程序的总执行时间。
5.2. 优化策略
为了优化程序的内存使用，我们可以在GIR上应用三种图重写规则。这些规则确保了应用规则后程序的理论内存性能不会下降。因此，我们可以在优化过程中尽可能多地应用重写规则。然而，如果所有规则都是不可逆的，我们可以使用贪心策略而不是基于搜索的方法来快速找到最优解。在这三种规则中，同步插入和gOperator合并是不可逆的，只能单向应用。另一方面，gOperator交换规则是可逆的，因此我们的优化方法只允许逐元素gOperator向前交换，以使其不可逆。
贪心优化算法。基于修正后的规则，我们提出了一种贪心优化策略，其伪代码如算法1所示。该策略尽可能多地应用同步插入和gOperator交换，然后进行gOperator合并，直到gOperator合并无法再应用为止。由于规则是不可逆的，优化是定向的并且保证终止。该算法使我们能够快速获得一个优化的GIR-Graph，并且在实践中，求解时间可以忽略不计。因此，我们可以搜索复杂的图结构并尝试不同的规则组合来优化各种应用中的GIR-Graph。这种快速的优化方法是使用GIR优化张量程序的宝贵工具。
1: 输入: 一个输入GIR G
2: 输出: 优化后的GIR G
3:
4: function INSERTSYNC(G, n0, n1)
5: if n0 和 n1 是连续的 then
6: 在 n0 和 n1 之间添加同步gOperator
7:
8: function MERGEOP(G, n0, n1)
9: if n0, n1 可合并 then
10: 合并 n0, n1 并将合并后的节点插入G
11: return true
12: else
13: return false
14:
15: function SWAPOP(G, n0, n1)
16: if n0, n1 可交换 then
17: 交换 n0 和 n1
18:
19: while true do
20: for n0, n1 in G.nodes do
21: INSERTSYNC(G, n0, n1)
22: for n0, n1 in G.nodes do
23: SWAPOP(G, n0, n1)
24: flag = false
25: for n0, n1 in G.nodes do
26: if MERGEOP(G, n0, n1) then
27: flag = true
28: if flag == false then
29: break
30:
31: return G
6. GIR-Graph 生成
POWERFUSION的主要目标是从端到端优化张量程序的内存性能。为实现这一目标，我们为输入模型生成GIR-Graph，这些图可以通过图重写技术进一步优化并生成高效代码。本节将解释如何从模型生成GIR-Graph。
计算图的预处理。POWERFUSION的输入是一个表示为计算图的模型。由于POWERFUSION的主要优化目标是内存性能，对于受计算性能限制的算子，我们直接使用DNN库，如NVIDIA GPU上的cuDNN或Cambricon MLU上的CNNL。我们根据算子的计算量和内存访问量来决定是否调用DNN库，而不是其类型。如图7所示的一个典型例子，这是EfficientNet模型中的一个子图，深度可分离卷积算子是内存密集型的。因此，它将使用GIR表示，并与其前面的RELU算子联合优化。另一方面，两个逐点卷积是计算密集型的，所以POWERFUSION选择直接调用DNN库。POWERFUSION在调用库函数时还会调整算子数据布局以获得更好性能，通过引入额外的转置算子来实现。这些转置算子将与其他内存密集型算子联合优化，以进一步提高性能。经过这些优化后，模型被分割成只包含内存密集型算子的独立子图。
图 7: EfficientNet中模型片段的示例。
图 8: GIR-Graph生成和合并的图示。
GIR-Graph生成。POWERFUSION的下一步是从由内存密集型算子组成的子图中生成GIR-Graph。如第2节所述，搜索不同算子的执行计划然后进行组合可以有效地扩大程序的搜索空间，从而带来更高的性能。为了实现这一点，需要对复杂算子进行拆分。例如，SILU算子是一个由SIGMOID和MUL算子组成的复合算子。POWERFUSION将所有这类复合算子拆分为四个基本算子：逐元素（element-wise）、广播（broadcast）、归约（reduce）和转置（transpose）。请注意，这些基本算子是计算图上的高级算子，与GIR的gOperator不同。这些算子可以应用预定义的代码模板来生成GIR-Graph。对于每个基本算子，我们定义了几个相应的GIR-Graph结构。以REDUCE算子为例，使用一个并行单元和多个并行单元的归约需要不同的计算和数据移动gOperator，对应于不同的图结构。此外，对于一个固定的图结构，每个算子有几种实现方法，包括执行时的并行度以及特定实现的分块大小等参数。POWERFUSION遍历图结构上这些参数的可能值。图结构和参数的组合被定义为一个GIR模板，可以为一个特定的基本算子生成一组GIR-Graph。如图8所示，SIGMOID和ADD的图生成都产生了多个GIR-Graph。
GIR-Graph合并。为每个基本算子生成GIR-Graph后，POWERFUSION尝试将多个GIR-Graph合并成一个更大的图，从而联合优化多个算子。我们的设计允许GIR-Graph合并，条件是两个GIR-Graph的并行结构匹配，并且相应计算图上的算子之间没有外部依赖。例如，两个算子可以是独立的或直接相连的。根据这个规则，当我们为每个基本算子选择一个GIR-Graph时，可以高效地生成一个合并计划。在合并两个GIR-Graph时，我们还在读写相同张量的数据移动gOperator之间插入全局同步操作。对于顺序的读写，需要在两个数据移动gOperator之间插入一个全局同步gOperator。对于并行的读取，需要在两个gOperator之前插入一个全局同步操作。通过应用图重写技术，得到的合并后的GIR-Graph可以被快速优化以实现高内存性能。图8展示了如何将同步操作插入到合并的GIR-Graph中的一个例子。
搜索空间与性能模型。为了尽可能地探索搜索空间，POWERFUSION遍历所有基本算子的所有GIR-Graph，并生成所有可能的合并计划。然而，直接搜索这一步的计算成本很高。为了解决这个问题，我们引入了一个性能模型。由于GIR-Graph中的所有gOperator都是规则且性能独立的，我们可以使用性能模型准确地预测GIR-Graph的运行时间。然后，我们可以修剪搜索空间，并高效地获得一个能保证优化结果的优化计划。
7. GIR 代码生成
POWERFUSION的设计原则包括将不同的硬件抽象成一个统一的结构，并将硬件特性提取为参数。这种抽象使我们能够优化程序并为特定硬件生成并行计算和数据移动gOperator的代码。通过解耦硬件和优化，系统变得更具扩展性。
代码生成流程。在生成设备代码之前，POWERFUSION对GIR-Graph进行重排序，以确定计算和数据移动gOperator的执行顺序。这决定了物理内存层次中内存切片的重用。重排序过程类似于传统的编译器指令重排序，并可以利用现有方法进行优化。确定执行序列后，POWERFUSION将GIR-Graph生成为本地编程语言的代码，并使用本地编译器将其编译成设备可执行文件。在代码生成阶段，POWERFUSION首先生成与硬件信息相关的并行结构代码，然后按顺序为计算和数据移动gOperator生成设备代码。生成的代码使用本地编译器进行编译和评估，以确保最佳性能。这种设计使得POWERFUSION移植到新硬件的成本很低，因为移植只需要映射并行抽象以及计算和数据移动的原语。例如，移植到NVIDIA GPU和AMD GPU各需要不到1000行新代码。我们将以NVIDIA GPU和Cambricon为例，解释POWERFUSION如何适应不同平台并生成高效代码。
图 9: Cambricon MLU的架构。
NVIDIA GPU适配。GPU是AI领域最常用的加速器，NVIDIA和AMD等厂商具有相似的硬件架构和编程模型。为了说明POWERFUSION如何移植到GPU，我们展示它如何映射到NVIDIA GPU上的CUDA。CUDA的并行结构包括三层：线程块（thread block）、线程束（warp）和线程（thread）。为了将GIR的并行单元映射到CUDA上，我们将其映射到warp，如第4节所述。warp是CUDA中具有独立性能的最小调度单元。相比之下，CUDA中的线程具有依赖的内存访问和调度，并对应于SIMD架构中的SIMD通道，这使它们不适合作为GIR中的并行单元。在搜索过程中，CUDA中线程块的数量和大小可以被配置并用作搜索参数。对于计算和内存访问操作的指令，CUDA程序表示在特定线程上执行的程序。因此，在GIR中，一个表示“读取包含n × 32个元素的连续内存”的实现将生成“以32为步长读取n个元素”。
Cambricon MLU适配。Cambricon MLU是一种领域特定架构，其架构与GPU有显著不同。它有两层暂存存储器（scratch-pad memory），NRAM和SRAM，分别位于一个IPU或一个MTL Cluster的作用域内。然而，MLU平台上的数据移动操作需要直接内存访问（DMA），这会引入开销并限制在每个IPU上同步执行计算和数据移动操作所能达到的内存性能。为了提高内存效率，POWERFUSION应用了流水线并行，这是MLU平台上一种常见的优化技术。在这种技术中，多个流水线在一个IPU上同时执行，使数据移动与计算重叠。在代码生成期间，POWERFUSION选择IPU上的一个流水线作为并行单元，并交错多个流水线来生成执行代码。然后插入同步操作以确保正确执行。这种流水线并行优化在MTP Cluster和IPU下引入了一个新的并行维度，并且不影响GIR上任何现有的优化。POWERFUSION成功移植到Cambricon MLU平台证明了其强大的可扩展性。
A4 实验环境
我们在三个不同的硬件平台上评估了POWERFUSION：NVIDIA GPU、AMD GPU和Cambricon。
* NVIDIA GPU: 使用一块Tesla A100 40GB PCIe GPU，支持Tensor Core加速矩阵计算。Tesla A100在TF32数据类型上的峰值性能为156 TFLOPS，峰值DRAM带宽为1555 GB/s。我们使用CUDA 11.7.0版本进行评估，并将内存和应用时钟设置为最大值。
* AMD GPU: 使用INSTINCT MI100，它为高性能矩阵计算进行了优化。MI100在矩阵核心上的峰值性能为46.1 TFLOPS，峰值DRAM带宽为1200 GB/s。我们使用ROCm 4.3版本进行评估。
* Cambricon: 使用MLU-370x4，其峰值FP32性能为24 TFLOPS，峰值DRAM带宽为307.2 GB/s。我们使用BANGC 1.0版本进行评估。
选择这些硬件平台是为了代表多样化的架构，并对POWERFUSION的性能进行全面评估。
A4 实验结果
8.2. 端到端性能
我们在POWERFUSION上评估了七个模型，以与不同架构上的各种基线进行比较。
DNN模型：我们评估了七个模型，包括三个Transformer模型和四个CNN模型，以确保模型多样性。
* BERT【索引7，BERT: pre-training of deep bidirectional transformers for language understanding+2018+CoRR】和GPT-2【索引18，Language models are unsupervised multitask learners+2019+】是流行的语言模型。
* Vision Transformer (ViT)【索引9，An image is worth 16x16 words: Transformers for image recognition at scale+2021+】将Transformer模型应用于计算机视觉任务。
* SAR-DRN【索引27，Learning a dilated residual network for sar image despeckling+2018+】和EfficientNet【索引20，Efficientnet: Rethinking model scaling for convolutional neural networks+2020+】是流行的CNN模型。
* ShuffleNet【索引28，Shufflenet: An extremely efficient convolutional neural network for mobile devices+2018+2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018】和RedNet-50【索引13，Involution: Inverting the inherence of convolution for visual recognition+2021+Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition】为CNN模型引入了更复杂的内存密集型算子。
除GPT-2的批处理大小为128（作为自回归模型）外，所有模型的批处理大小均设为1。
基线：
* 传统框架: TensorFlow【索引4，Tensorflow: A system for large-scale machine learning+2016+12th USENIX symposium on operating systems design and implementation (OSDI 16)】和PyTorch【索引17，Tensors and Dynamic neural networks in Python with strong GPU acceleration+2017+https://pytorch.org】 。
* 优化库/编译器: TorchScript【索引2，TorchScript+https://pytorch.org/docs/stable/jit.html】，TensorFlow-XLA【索引3 ，Xla: Optimizing compiler for tensorflow+2017+https://www.tensorflow.org/xla】，TensorRT【索引22 ，NVIDIA TensorRT: Programmable inference accelerator+2017+https://developer.nvidia.com/tensorrt】（NVIDIA），MagicMind【索引1 ，MagicMind+https://www.cambricon.com/index. php?m=content&c=index&a=lists&catid=378】（Cambricon）。
* 张量编译器: TVM【索引6，TVM: end-to-end optimization stack for deep learning+2018+CoRR】/Ansor【索引29，Ansor: generating high-performance tensor programs for deep learning+2020+14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)】。
结果：
* 在NVIDIA A100上 (图10(a))：POWERFUSION的性能优于所有基线。与PyTorch、TorchScript、TensorFlow和TensorFlow-XLA相比，平均加速比分别为9.7倍、7.3倍、8.2倍和4.1倍。与TensorRT相比，POWERFUSION平均加速28%，在GPT-2模型上达到1.98倍的加速比。与TVM相比，平均加速比为1.98倍。对于BERT和ViT，由于GEMM占主导，性能与TensorRT相似。对于ShuffleNet和EfficientNet，由于卷积操作耗时较多，性能与TVM相似。
* 在AMD MI100上 (图10(b))：POWERFUSION的性能优于所有基线。与PyTorch、TorchScript和TensorFlow相比，平均加速比分别为3.1倍、2.7倍和18.0倍。与TVM相比，平均加速比为1.24倍。加速比较低的原因是AMD GPU的计算性能与内存带宽之比较低，内存访问优化效果不那么显著。
* 在Cambricon MLU-370上 (图10(c))：POWERFUSION相比PyTorch平均加速3.1倍，相比MagicMind平均加速2.3倍。这表明POWERFUSION支持不同架构，并具有良好的跨平台优化能力。
图 10: 端到端性能。
8.3. 案例研究：GPT-2
我们以GPT-2为例，展示POWERFUSION如何超越现有框架。图11展示了GPT-2在A100 GPU上生成应用中的性能。
图 11: 不同索引下GPT-2的性能。
无KVCache优化：在不使用KVCache的情况下，POWERFUSION的性能与TensorRT相当。这是因为TensorRT对该模型计算图中的三个固定的内存密集型算子进行了手动优化，而POWERFUSION提出了类似的优化策略。
KVCache优化：KVCache是一种在生成模型中广泛使用的技术【索引16，Efficiently scaling transformer inference+2022+CoRR】，它通过存储和重用中间计算结果来减少重复计算。如图11所示，应用KVCache后，TensorRT的执行时间显著下降，特别是在输入索引较大时。
POWERFUSION的进一步优化：应用KVCache后，每次迭代只需计算一个token，序列长度始终为1。这使得注意力层中的两个矩阵乘法操作退化为内存性能受限的矩阵-向量乘法。POWERFUSION能够联合优化注意力层的所有操作，并生成一个内存高效的核，从而显著提高程序性能。如图12所示，POWERFUSION优化注意力层时，两个形状和计算模式相同的矩阵乘法，在优化后性能最佳的程序中却使用了不同的GIR-Graph。这些GIR-Graph在数据切片模式和计算操作轴向上有所不同，这种优化对于基于循环的方法来说很难实现。因此，在应用KVCache后，POWERFUSION展现出比TensorRT更显著的性能优势，性能提升高达3.16倍。
图 12: GPT2的案例研究。
8.4. 性能分解
为了说明POWERFUSION性能提升的来源，我们分解了TensorRT、TVM和POWERFUSION启动的核（kernel）数量。如图13所示，我们将核库中的矩阵乘法和2D卷积等计算操作定义为计算核，其他操作归类为内存核。
图 13: TensorRT、TVM和POWERFUSION的性能分解
- BERT, ViT, RedNet50: POWERFUSION和TensorRT使用相同数量的计算核，但POWERFUSION生成的内存核更少，尤其是在RedNet50中，内存核数量比TensorRT减少了59%。TVM的总核数则高于两者。这表明POWERFUSION具有更强的核融合能力。
- SAR-DRN: TVM的核数与POWERFUSION和TensorRT的计算核数相同，说明TVM能将此模型中的所有内存核融合到计算核中。而POWERFUSION受限于厂商提供的核库，无法进行此类融合。
- GPT-2, ShuffleNet, EfficientNet: POWERFUSION通过将计算核与内存核优化相结合，减少了计算核的数量，从而降低了总核数。例如，在GPT-2中，POWERFUSION生成的核数比TensorRT和TVM分别减少了73%和71%。
这些发现证明了POWERFUSION强大的优化和代码生成能力。
8.5. 搜索时间
我们比较了POWERFUSION和TVM在优化各种模型时的搜索时间。如图14所示，POWERFUSION的搜索时间比TVM短了两个数量级。
图 14: TVM和POWERFUSION的搜索时间。
原因分析：
1. 利用原生库：我们的系统直接使用原生库中计算密集型操作的实现，无需对其进行搜索，从而减少了搜索时间。
2. 性能模型剪枝：我们在优化过程中加入了内存性能模型剪枝，这避免了对涉及大量内存访问的解决方案的探索，从而降低了搜索开销。
评估结果表明，POWERFUSION能用更少的时间为模型生成代码，为需要快速部署模型的用户提供了显著优势。
A7 补充细节：相关工作
算子级张量编译器。许多张量编译器能够为独立的深度学习算子生成高性能代码，包括TVM【索引6，TVM: end-to-end optimization stack for deep learning+2018+CoRR】、FlexTensor【索引31，Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system+2020+Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems】、Ansor【索引29，Ansor: generating high-performance tensor programs for deep learning+2020+14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)】、AMOS【索引30，Amos: enabling automatic mapping for tensor computations on spatial accelerators with hardware abstraction+2022+ISCA】、Roller【索引33，{ROLLER}: Fast and efficient tensor compilation for deep learning+2022+16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】、TensorIR【索引10，Tensorir: An abstraction for automatic tensorized program optimization+2023+】和Hidet【索引8，Hidet: Task-mapping programming paradigm for deep learning tensor programs+2023+Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2】。然而，它们的方法没有对算子间的数据依赖进行建模，无法探索相邻算子的数据局部性。多面体方法，包括PPCG【索引24，Scheduling for PPCG+2017+CW Reports】和PLUTO【索引5，A practical automatic polyhedral program optimization system+2008+ACM SIGPLAN Conference on Programming Language Design and Implementation (PLDI)】，在元素级别对数据依赖进行建模，这导致搜索空间过大，难以找到高效的解决方案。
图级DNN加速。其他工作在图级别加速DNN模型。TensorRT【索引22，NVIDIA TensorRT: Programmable inference accelerator+2017+https://developer.nvidia.com/tensorrt】和AITemplate【索引26，AITemplate+2022+】手动优化常用模式。TASO【索引12 ，Taso: optimizing deep learning computation with automatic generation of graph substitutions+2019+Proceedings of the 27th ACM Symposium on Operating Systems Principles】、Rammer【索引14，Rammer: Enabling holistic deep learning compiler optimizations with rtasks+2020+Proceedings of the 14th USENIX Conference on Operating Systems Design and Implementation】和PET【索引25，Pet: Optimizing tensor programs with partially equivalent transformations and automated corrections+2021+15th USENIX Symposium on Operating Systems Design and Implementation (OSDI 21)】可以组合现有算子以生成更高效的代码。这些预定义的模式和算子集无法覆盖现有深度学习模型中多样化的内存访问模式。Astitch【索引32，Astitch: Enabling a new multidimensional optimization space for memory-intensive ml training and inference on modern simt architectures+2022+Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, ASPLOS ’22】和DNNFusion【索引15，Dnnfusion: accelerating deep neural networks execution with advanced operator fusion+2021+Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation】可以为未见过的模式生成内存优化代码，但它们的融合是基于规则的，忽略了许多更高效的调度方案。
领域特定语言。一些领域特定语言，如Triton【索引23，Triton: an intermediate language and compiler for tiled neural network computations+2019+Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages】、FreeTensor【索引21，Freetensor: a free-form dsl with holistic optimizations for irregular tensor programs+2022+Proceedings of the 43rd ACM SIGPLAN International Conference on Programming Language Design and Implementation】和Graphene【索引11，Graphene: An ir for optimized tensor computations on gpus+2023+Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】，允许用户以比基于算子的计算图更细的粒度来表达张量程序，但这需要更多的专家知识和人力来明确定义程序应如何执行。
A5 结论
本文提出了POWERFUSION，一个通过GIR来优化内存性能的编译器，GIR通过指令级图描述显式地表示数据移动模式和细粒度的数据依赖。通过设计基于GIR的优化策略和硬件抽象，我们能够在不同硬件上实现显著的加速效果。我们的评估结果表明，与当前性能最佳的框架相比，POWERFUSION分别取得了高达1.97倍、2.93倍和16.91倍的加速（平均分别为1.28倍、1.23倍和2.31倍）。
💬 评论讨论
欢迎在这里分享您的想法和见解！