Welder: Scheduling Deep Learning Memory Access via Tile-graph

作者/机构: Yining Shi, Peking University & Microsoft Research; Zhi Yang, Peking University; Jilong Xue, Lingxiao Ma, Yuqing Xia, Ziming Miao, Yuxiao Guo, Fan Yang, and Lidong Zhou, Microsoft Research

A1 主要贡献

核心问题:
传统的深度神经网络（DNNs）被视为计算密集型工作负载，但随着现代DNN模型处理更高保真度的数据（如大图像、长句子）以及硬件加速器中计算核心（如TensorCore）速度的提升，DNNs正变得越来越受内存限制。对一系列顶尖DNN模型的性能分析显示，GPU内存带宽利用率高达96.7%，而计算核心的平均利用率仅为51.6%。这种效率低下的原因是，现有DNN系统缺乏对模型内存访问的整体优化，尤其是在处理算子间（inter-operator）数据复用方面存在挑战，因为不同算子最优的内部数据分块（tile）形状常常不一致，导致难以在高速缓存（如共享内存）中直接复用数据。

研究目标:
本文旨在设计一个深度学习编译器WELDER，从整体内存访问的角度优化端到端DNN模型的执行效率。目标是统一和自动化内存优化，解决算子内（intra-operator）和算子间数据复用之间的冲突，发现传统基于规则的方法无法覆盖的新优化模式，并支持因输入过大而无法完全放入加速器内存的DNN模型。

创新点:
1. Tile-graph抽象: 提出了一种名为tile-graph的抽象，它将DNN计算建模为基于数据块（tile）的细粒度数据流图。每个节点处理一个数据块，使得对数据块大小和内存位置进行精细化管理成为可能。
2. 跨内存层优化的解耦: 基于“跨层独立性”（inter-layer independence）的关键观察，WELDER将复杂的组合优化空间分解为多个独立的子空间。该观察指出，针对某一内存层级的流量优化仅取决于该层级的数据块配置，从而可以独立优化每个内存层级的数据访问。
3. 基于流量的成本模型和自动化调度: WELDER利用一个基于数据块流量的成本模型来权衡算子内和算子间的数据复用。通过SetConnect和Propagate两个调度接口，WELDER实现了一个两阶段调度策略：首先探索不同算子在不同内存层级的连接方案，然后为每个连接方案下的子图搜索最优的数据块配置，从而自动生成高效的执行计划。
4. 统一的内存优化框架: WELDER将先前零散的内存优化（如基于寄存器的逐元素融合、基于共享内存的融合）统一到一个框架中。这种通用性使其能够自动发现89种未被现有方法探索过的算子融合模式。
5. 支持超大输入模型: 通过将主机内存（host memory）整合到其分层内存模型中，WELDER能够为输入任意大的DNN模型（即使单个算子也无法装入GPU内存）生成跨设备和主机内存的优化执行计划。

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

现代DNN是内存密集型的。对一系列代表性DNN基准模型在ONNXRuntime【索引8，ONNX Runtime，URL:https://github.com/microsoft/onnxruntime】上运行的GPU利用率分析显示，平均计算利用率仅为51.6%，而内存利用率高达96.7%。如图1所示，虽然ResNet和BERT等经典模型能达到较高的计算利用率（>80%），但近年来提出的流行模型由于引入了更多内存密集型模式，计算效率较低。这些新模型处理高保真度数据，产生大量中间激活值，导致内存存储流量与加载流量的比率更高。这表明模型频繁地通过全局内存交换大量中间数据，凸显了跨算子优化内存访问效率的必要性 。


图1：不同模型在NVIDIA V100 GPU上的计算FLOPS和内存带宽利用率。
图2：未融合、融合以及Matmul和Softmax各自内核的延迟数据。

算子内和算子间数据复用模式的冲突。同时优化算子内和算子间的数据复用是具有挑战性的。算子通常实现为对所有张量维度的多层嵌套循环，其内部跨多层内存的数据复用通过复杂的循环分块（loop tiling）技术【索引5，NVIDIA cutlass，URL:https://github.com/NVIDIA/cutlass；索引50 ，Ansor: Generating high-performance tensor programs for deep learning，OSDI 2020；索引52，ROLLER: Fast and efficient tensor compilation for deep learning，OSDI 2022】进行隐式优化。以Matmul和Softmax这两个连续算子为例，当它们独立优化时，在共享内存中的最优块大小不同，例如Matmul为[32×64]，Softmax为[4×128]。这导致Softmax无法在共享内存中复用Matmul的中间数据，总延迟为0.36ms，如图2所示。然而，如果强制它们同时考虑算子内和算子间的数据复用，融合后的算子延迟可以减少到0.29ms，实现了1.26倍的加速。通过检查它们对齐后的块大小（[16×128]），我们发现两个算子都牺牲了自身的效率（由于次优的数据块导致单独运行时性能分别下降15%和4%）以换取整体效率的提升。这表明需要一个能够整体优化内存访问的、高效的跨算子内和算子间的数据复用解决方案。

关键观察。通过对图2中例子的进一步分析，我们发现了三个关键观察。首先，跨算子对齐的块配置可以通过从输出块形状开始的一系列形状推断得出。例如，如果我们想计算Softmax的一个[4×128]的输出块，根据其计算逻辑（如张量表达式），可以推断出其依赖的输入块形状也是[4×128]。然后，将[4×128]作为Matmul的输出块，我们可以进一步推断出Matmul的输入块形状将是[4×k]和[k×128]，其中k是一个规约维度的大小，可以设置为不超过Matmul规约维度大小的任何数字。这样，这两个算子就可以通过在共享内存中复用中间数据块（[4×128]）而融合。

其次，给定对齐的块配置和原始张量形状，总内存流量可以很容易地进行分析性推导。在这个例子中，Matmul的输入张量A的形状为[98304×64]，B为[64×128]，输出张量C为[98304×128]。Softmax接着以C为输入，产生一个相同形状的输出张量D。输入张量A、B和输出张量D都在全局内存中。给定这些形状，我们可以首先计算计算张量D的单个输出块（即[4×128]）时的内存流量。为此，它将首先从张量A加载一个[4×k]形状的块，从张量B加载一个[k×128]的块用于Matmul，然后中间块[4×128]将在共享内存中被Softmax消耗，并向张量D写入一个[4×128]形状的块，其中k可以根据输入张量[98304×64]的形状替换为64。因此，单个输出块在全局内存中产生的总流量是35KB ((464+64128+4128)4Bytes(FP32))，其中中间块[4×128]的流量由于在共享内存中的数据复用而被节省了。为了计算完整的输出张量D，总共需要进行24,576次这样的计算（即(98304128)/(4128)），导致总的全局内存流量为840MB（即24,576*35KB）。有趣的是，遵循同样的计算方法，将输出块改为[16×128]将使总流量减少到仅264MB。

最后，我们的流量成本计算仅由目标内存层的块配置决定，例如，一旦张量形状被指定，共享内存中的输出块形状[4×128]或[16×128]。这使我们能够独立地为每一层选择块大小，以便从较低内存层优化流量成本。

总结。这些观察共同为我们提供了一种有效的方式来整体优化内存访问，即通过一个输出块形状对齐一组相邻的算子，根据内存流量决定最佳块形状，并为每个内存层独立进行优化。通过这种方式，WELDER能够将原始的粗粒度算子间依赖关系转变为更细粒度的块级依赖关系，这实质上消除了算子之间的一些虚假障碍，并实现了更多的并发性。

A2 方法细节

3. WELDER 设计

WELDER是一个深度学习编译器，旨在通过一个整体的内存访问调度空间来提高现代DNN的性能，其设计动机源于第2节的观察。图3展示了其系统概览。WELDER接收一个完整的DNN模型作为输入，并将其转换为一个基于数据块（tile）的计算任务数据流图（即算子块），称之为tile-graph（§3.1）。tile-graph提供了对数据块配置和内存布局的细粒度控制。对于一个tile-graph，WELDER通过一种“先连接后调度”的方法来解决算子内和算子间数据复用的冲突：它首先假设两个相邻的算子可以在某个内存层级复用数据块（即连接），然后推导出最佳的公共块形状，以判断总内存流量是否可以减少。为实现这一目标，WELDER提供了两个tile-graph调度接口：SetConnect和Propagate（用于形状推断链）。基于此，我们提出了一个两步调度算法，即图连接和子图调度，为多个内存层级递归地决定一个高效的tile-graph执行计划，这被称为分层tile-graph（§3.2）。最后，这个计划通过硬件层定义的四个抽象计算接口，即Allocate、LoadTiles、ComputeTile和StoreTiles，映射为特定硬件加速器的可执行代码（§3.3）。抽象加速器的内存规格被tile-graph调度层用来指导优化过程。

图3：WELDER的系统概览。

3.1 算子块和Tile-graph

算子块定义。WELDER以一种名为算子块（operator-tile）的细粒度任务粒度来定义DNN计算。一个DNN算子，如卷积，可以被实现为多个同构的算子块，这些算子块以流式或并行方式执行，以计算输出张量中的所有数据块【索引31，Rammer: Enabling holistic deep learning compiler optimizations with rtasks，OSDI 2020】。每个算子块从输入张量中切片一个数据块作为输入，并计算输出张量中的一个数据块，其计算逻辑由基于索引的张量表达式【索引15，TVM: An automated end-to-end optimizing compiler for deep learning，OSDI 2018】描述。图4(a)和(b)展示了Conv和MaxPool的算子块示例，其中Conv算子通过接收一个[3 × 3 ×C]的数据块作为输入来计算一个[1 × 1 ×C]的数据块，而MaxPool算子接收一个[2 × 2 × F]的输入块并计算一个[1 × 1 × F]的输出块。

Tile-graph的连接。为了提高分层内存资源（如共享内存）的利用率，WELDER允许两个相邻的算子块通过一个共同的中间数据块“连接”起来，这个中间数据块也称为复用块（reuse-tile）。这使得第二个算子块可以直接消费第一个算子块产生的数据，而无需将其物化为一个完整的中间张量。图4(c)展示了Conv和MaxPool两个算子块之间通过一个[2 × 2 × F]的复用块进行连接的例子。多个算子块可以沿着每个相邻的边连接起来，形成一个算子块的数据流图，即tile-graph。

图4：两个算子块的图示：(a) Conv和(b) MaxPool；以及(c) 将它们连接成一个tile-graph（为简化，省略了Conv的权重张量）。

Tile传播。一旦连接，tile-graph中的大多数数据块都是相关的，可以通过将一个输出块的形状传播到整个图来自动推断。这是通过从输出节点到输入节点的一系列形状推断来实现的。对于每个算子块，通过分析其张量表达式和输出块大小，可以准确确定输入张量的依赖区域。在输入区域可能包含不规则模式（如稀疏或非连续访问，例如Gather或带步长的Convolution）的情况下，我们的表达式分析提供了一个保守的上限作为输入块形状。如果tile-graph有多个输出节点，它们的输出形状也可能相关，因为它们可能共享图中的一个共同祖先节点。在这种情况下，在传播第一个输出块之后，我们为剩余的输出节点传播单独的形状，并使它们与第一个对齐。如果两次传播之间存在不一致的块形状，我们不将后一个输出节点连接到当前图中。

内存流量与占用。在块传播之后，可以确定一个tile-graph的内存流量和占用。首先，单个tile-graph的内存流量可以通过其输入和输出块大小的总和来计算。总流量是通过将这个值乘以计算完整输出张量所需的tile-graph数量（例如，通过将张量大小除以输出块大小）得到的。其次，tile-graph的最小内存占用可以使用内存分配算法（例如，best-fit【索引19，Worst-case analysis of memory allocation algorithms，STOC 1972】）通过按拓扑顺序分配所有数据块来计算。作为一种占用优化，包含规约轴的输入块可以进一步划分为更小的块，这些小块可以被顺序加载和消费，并将其结果累加到输出块中。具体来说，一个特定的策略可以在块传播期间自动尝试沿规约轴的不同分块大小。

图5：将三个连续的算子映射到一个三层内存层次结构（省略了Conv的权重）。

3.2 Tile-graph 调度

分层映射。为了将一个由初始数据流图表示的DNN模型映射到加速器上，我们可以递归地将每个算子划分为多个算子块以适应每个内存层，并在更高内存层连接算子块以利用算子间的数据复用。因此，整个DNN计算可以被建模为一个在二维空间上的数据流管道，数据块在内存层次结构中垂直上下移动，并在不同层级水平传递给后续算子。

映射示例。图5展示了一个将三个连续算子（Conv、ReLU和MaxPool）映射到三层内存层次结构（例如，从L2到L0）的例子。Conv算子的输入块被重复地从L2加载到L1，然后再到L0进行计算。通过在L0层连接Conv和ReLU算子，Conv算子的输出可以被复用为ReLU算子的输入，这两个算子在L0层形成一个tile-graph。同时，它们在L1层被整合成一个虚拟节点（即Conv+ReLU）。ReLU的输出随后持续地溢出到L1层的数据块中，并通过在L1层进一步连接，被复用为MaxPool的输入。这使得所有三个算子在L1层形成一个单一的tile-graph，从而在L2层产生虚拟节点Conv+ReLU+MaxPool。经过这个递归过程，所有算子都在最低层作为单一的tile-graph连接起来。

解耦优化空间。鉴于DNN计算大多是内存受限的，我们数据流管道的主要优化目标可以转化为最小化内存流量。这使我们能够利用优化跨内存层流量的内在独立性，将整个优化空间分解为几个子空间。具体来说，对于一个给定的tile-graph，从较低内存层加载和存储到该层的总数据流量仅通过其输出块形状就可以估算出来，该形状用于推断所有输入和输出块的形状。基于此属性，来自相同或不同内存层的不同tile-graph可以通过搜索最优的块形状来独立优化它们的内存流量。例如，在图5中，L0层的Conv和ReLU的tile-graph可以独立于L1层的tile-graph（例如，由Conv+ReLU和MaxPool算子形成）进行优化，这被称为跨层独立性（inter-layer independence）。这进一步意味着，L0层的子图Conv-Relu和MaxPool的最优块配置也是独立的，因为它们与L1层的tile-graph无关，我们称之为层内独立性（intra-layer independence）。在实践中，唯一的约束是较低内存级别的块大小必须大于较高内存级别的块大小。这通常是成立的，因为较低内存级别通常比上层内存级别具有更大的容量。利用这些特性，我们可以在给定图连接计划的情况下，独立地调度每个tile-graph。

调度接口。WELDER提供了两个调度接口来控制图的连接和子图的分块，如图6所示。首先，图的连接使用SetConnect接口实现，该接口为一个tile-graph中的边分配一个内存级别（默认为最低级别）。连接后，图中的块形状通过Propagate接口推断，通过指定输出块的维度大小和输入块中可选的规约轴。例如，在图5中，我们可以使用SetConnect接口在L0连接Conv和Relu，在L1连接Relu和MaxPool。连接后，对于子图Conv+Relu，我们可以使用Propagate通过指定[1, 1]的输出块形状来推断中间的复用块形状（即[1, 1]）。类似地，我们也可以通过指定[1, 1, F]的输出块形状来推断子图Conv+Relu+MaxPool的中间复用块形状（即[2,2,F]）。这两个调度原语本质上是更新tile-graph的边和顶点的两个接口。具体来说，SetConnect用于在两个节点之间添加连接，Propagate用于为一个节点设置块配置。它们共同构成了更新tile-graph的完整接口。请注意，这些原语仅由WELDER的调度策略使用，对最终用户是透明的。WELDER还提供了两个成本接口，MemFootprint和MemTraffic，用于计算tile-graph的内存占用和总流量，它们作为我们指导调度的成本模型。

图6：WELDER中的调度接口

1 Func GraphConnecting(g:Graph, d:Device):   
2 for node : TopologySort(g.nodes()) do   
3 for edge : node.out_edges() do   
4 for level : d.MemLevels() do   
5 SetConnect(edge, level);   
6 s = ExtractSubgraph(node, 0);   
7 configs = SubGraphTiling(s, 0, tensor_shapes);   
8 if t = Min(d.Profile(configs)) < best_latency   
9 best_latency = t;   
10 best_level = level;   
11 SetConnect(edge, best_level);   
12 Func SubGraphTiling(g:Graph, level:Memory, c: Config)   
13 configs = PriorityQueue();   
14 for subtile : EnumerateSubtiles(g, c) do   
15 config = Propogate(g, subtile);   
16 if MemFootprint(g) > level.capacity   
17 continue;   
18 configs.push(config, priority=MemTraffic(g));   
19 results = Dict();   
20 for config : TopK(configs, k) do   
21 // return empty sub-graph at top level to exit recursion   
subgraphs = unique([ExtractSubgraph(node, level+1)   
for node in g.nodes()]);   
22 for subgraph : subgraphs do   
23 subgraph_configs = SubGraphTiling(subgraph,   
level+1, config);   
24 results[config].append(subgraph_configs);   
25 Return results;   
26 Func ExtractSubgraph(node:Node, level:Memory)   
27 nodes = Set();   
28 for edge : node.InOutEdges() do   
29 if edge.connect_level > level   
30 nodes.insert(ExtractSubgraph(edge.node, level));   
31 return SubGraph(nodes);

调度策略。WELDER采用一种两步调度算法来有效地优化数据流计算。具体来说，图连接调度器首先通过为每条边设置不同的内存复用级别来枚举不同的图连接方案，然后子图调度器为图连接调度器解耦出的每个子图快速搜索高效的块配置。图7展示了WELDER中的两步调度算法。首先，给定一个DNN数据流图g和一个加速器设备d，图连接调度器以拓扑顺序枚举所有图节点及其输出边（第1-3行）。对于每条边，WELDER尝试不同的连接级别（例如，使用SetConnect接口）（第5行）。然后它提取连接的子图，其中所有边的连接级别都高于0。这里，我们用数字0表示最低的内存级别，更大的数字表示更高的级别。ExtractSubgraph函数在第26-31行实现。对于提取的子图，WELDER调用SubGraphTiling函数来获取几个高效的块配置，并通过在硬件上进行性能分析来选择最优的一个（第7-10行）。在与所有其他连接级别比较后，WELDER为当前边设置最佳的连接级别。

子图调度。接着，子图调度器（即SubGraphTiling函数）接收一个子图和上一级的块配置作为输入，并为当前级别搜索高效的块配置。首先，WELDER使用类似于Roller【索引52，ROLLER: Fast and efficient tensor compilation for deep learning，OSDI 2022】中的块形状扩展方法来枚举输出维度的块大小（即第14行的EnumerateSubtiles），该方法将初始块形状（例如，大小为1）扩展到能够减少总流量并与硬件特性对齐的形状。在得到输出块形状后，我们可以使用Propagate接口推断出完整的块配置，并使用MemFootprint接口检查它是否超过内存容量，或者使用MemTraffic接口将其以内存流量为键追加到一个排序的结果列表中（第15-18行）。最后，我们为当前级别选择内存流量最小的前K个配置，然后通过调用ExtractSubgraph和SubGraphTiling递归地提取上层子图并决定它们的最佳块配置（第20-24行）。

调度结果。值得注意的是，WELDER对不同内存层次结构上的内存大小没有假设，因为我们的调度策略总能尽力确定最佳的层级和块大小来放置中间数据，以最小化整体延迟。虽然WELDER总是偏好具有大容量高层级快速内存（例如，共享内存）的硬件，因为这可以容纳足够大的中间数据块，而过小的块大小可能导致更差的算子内数据复用。WELDER中一个数据流图的调度结果是一个分层tile-graph，它在最低内存级别始于一个完整的图，并递归地在更高层级分裂为几个子图，一直到顶层。

3.3 映射到硬件加速器

抽象执行计划。WELDER生成的分层tile-graph是一个抽象的执行计划，可以被映射为特定硬件加速器的可执行代码。为了便于这种映射，WELDER提供了一个具有分层内存的抽象加速器设备。内存配置，如层数、每层内存容量和事务宽度，可以通过MemLevels接口获得（例如，在图7中使用）。有了这个抽象的内存层，可以很容易地将现有加速器扩展为带有额外内存层（例如，主机内存或SSD）的新设备，使其能够处理可能无法装入单个设备内存的非常大的张量（详见§5.4）。WELDER的性能增益主要来自于内存层之间的带宽差距。因此，只要较低级别的内存成为瓶颈，且较高级别的内存可以容纳中间数据块，WELDER就可以自动地在更快的、高级别的内存上流水线化算子间的数据传输。

表1：抽象硬件加速器中的设备接口。

执行接口与流程。为了在硬件加速器上执行一个分层tile-graph，WELDER提供了四个计算接口：Allocate、LoadTiles、ComputeTile和StoreTiles（见表1）。使用这些接口执行分层tile-graph的例程如图8所示。该过程从执行最底层的tile-graph（即完整的DNN图）开始。对于每个tile-graph，它首先在相应的内存层中分配必要的工作空间（使用Allocate接口），并将输入块加载到这个空间中（LoadTiles）。然后，它以拓扑顺序执行子图中的所有节点。如果当前内存层是顶层，节点直接在计算核心中执行（ComputeTile）。否则，递归调用上层tile-graph的执行。最后，当前空间中的结果块被存储到较低的内存层中（StoreTiles）。这个执行例程既可以作为一个代码生成过程，也可以作为一个运行时过程，取决于特定的加速器是将这些计算接口实现为代码发射器还是可执行函数调用。在WELDER中，它们目前被实现为代码发射器，以生成特定于加速器的计算逻辑。通过执行这个递归例程，整个分层tile-graph被展开，并自动生成一个包含所有必要优化的全模型计算程序。

图8：分层tile-graph的编译例程。

4. 实现

基础架构。WELDER是基于开源DNN编译器TVM【索引15，TVM: An automated end-to-end optimizing compiler for deep learning，OSDI 2018】、Roller【索引52，ROLLER: Fast and efficient tensor compilation for deep learning，OSDI 2022】和Rammer【索引31，Rammer: Enabling holistic deep learning compiler optimizations with rtasks，OSDI 2020】实现的。它利用TVM编写内核调度，利用Roller枚举高效的块配置，并利用Rammer进行端到端的图优化。WELDER的核心机制，包括tile-graph、块传播、调度算法、代码生成等，用5.2k行代码实现。WELDER以ONNX图作为输入，并执行常见的图优化，如常量折叠和简单的逐元素融合。然后，它将优化后的图转换为tile-graph，以进行整体内存调度优化。WELDER在CUDA和ROCm GPU以及GraphCore IPU上通过统一的设备接口（表1）实现。对于CUDA和ROCm GPU，WELDER在三个内存层上调度数据块：全局内存（DRAM）、共享内存和寄存器。为了在CUDA GPU和GraphCore IPU上处理大图像，我们还通过添加主机内存层来扩展它们的设备内存。

4.1 硬件对齐的Tile搜索

枚举高效的数据块大小。WELDER通过在流量成本模型中引入一个惩罚因子来考虑几个可能影响数据访问效率的硬件相关因素。首先，如果存在非合并的内存访问，总内存流量将包括这些访问所需的额外事务。例如，在CUDA GPU中，对于连续128字节的数据（一个事务），总是优先使用合并内存访问。其次，当由于块大小过大导致并行度不足时，内存流量会根据计算核心的利用率百分比成比例增加。最后，如果给定块配置的总内存占用超过内存容量，我们添加一个无限大的惩罚。为了避免枚举低效的候选配置，WELDER通过只枚举那些根据成本模型能最大程度减少流量的维度来搜索输出块，并仅检索流量最小的前k个候选配置。

确定对齐的计算并行度。在GPU中，在同一线程块中执行的顶层算子块必须统一块大小（例如，线程数）。为确保这种对齐，WELDER首先在寄存器级别强制足够的并行块以与硬件并行性对齐（通过枚举硬件对齐的块）。例如，在NVIDIA V100 GPU中，块数量应大于128，因为每个SM有4个warp调度器，每个warp有32个线程。然后，我们确定所有算子的块数量的最大公约数作为公共线程块大小，前提是它大于硬件并行度（例如128）且小于最大限制（例如1024）。否则，我们将块大小设置为等于硬件并行度的数值。一旦块大小确定，我们将寄存器级别的所有算子块绑定到这些线程。如果单个线程需要运行多个块，我们使用TVM的虚拟线程来绑定它们，从而允许在所有内存bank上并发数据访问并避免bank冲突。

支持TensorCore。WELDER利用TensorCore来加速CUDA GPU上的某些算子，如GEMM、BatchMatmul和Convolution（使用implicit GEMM【索引28，Performance analysis of gpu-based convolutional neural networks，ICPP 2016】）。我们为这些算子添加注解，指明哪些轴将绑定到CUDA的Warp级矩阵操作。对于顶层算子块，我们将它们绑定到warp（而不是线程）来执行MMA操作。此外，在枚举块大小时，我们引入了一些额外的约束，例如确保线程数是warp大小的整数倍，以及每个块中的轴（M、N和K）是MMA操作的片段大小的整数倍。

4.2 代码生成与编译

基于TVM的内核生成。WELDER的内核生成基于TVM。特别地，寄存器级别的块连接是使用TVM的compute_inline调度原语实现的。对于共享内存级别的连接，我们仅使用TVM为共享内存之上的每个连接部分生成独立的内核，然后应用几个额外的pass将这些独立的内核组合成一个单一的融合内核。

加载/存储重写。TVM生成的独立内核从全局内存加载和存储数据。我们通过向TVM的lowering过程中添加一个额外的TIR【索引11，TensorIR，URL:https://discuss.tvm.apache.org/t/rfc-tensorir-a-schedulable-ir-for-tvm/7872】pass，将这些全局内存访问重写为共享内存访问。此外，我们添加内存栅栏以防止竞争条件，并应用填充来处理缓冲区中的bank冲突。结果是，原始的全局内核可以被转换为一个设备函数，并包含在最终的融合内核中 。

块/线程索引重映射。一些算子不能直接与其他算子连接，需要重映射它们的blockIdx和threadIdx值。BlockIdx重映射用于像Transpose这样的算子。重映射关系从它们的张量表达式中推导出来。ThreadIdx重映射用于将二维线程块连接到一维线程块。当线程间规约或TensorCore原语需要使用二维线程块（threadIdx.x和threadIdx.y）时，这是必需的，而其他算子可能使用一维线程块（仅threadIdx.x）。只要它们的总线程数相等，二维线程块就可以映射到一维线程块。

内存管理。我们以统一的方式管理所有共享内存，包括在每个独立内核中分配的内存和算子间的复用缓冲区。首先，我们根据拓扑执行顺序分析每个缓冲区的生命周期，并将其转换为一系列的分配和释放操作。然后，我们使用best-fit算法计算每个共享内存分配的偏移量，同时考虑到数据类型和TensorCore操作的任何对齐要求（例如，对齐到32字节以避免未对齐的地址访问）。

编译加速。WELDER通过并行编译和子图缓存来优化编译速度。首先，利用配置之间的独立性，WELDER可以使用多进程并行构建和评估每个配置。其次，在大多数DNN模型中，一些子图模式经常多次重复出现。为了避免冗余优化，WELDER利用子图签名来缓存每个唯一的图模式。例如，在一个12层的BERT模型中，我们可以缓存第一层的优化结果（内核代码和分析的延迟），并将其复用于所有剩余的11层。

A4 实验环境

硬件配置:
* NVIDIA GPU服务器:
1. Azure NC24s_v3 VM: Intel Xeon E5-2690v4 CPU, NVIDIA Tesla V100 (16GB) GPU。
2. 本地工作站: Intel Xeon E5-2678 v3 CPU, NVIDIA GeForce RTX 3090 GPU。
* AMD GPU服务器: Intel Xeon E5-2640 v4 CPU, AMD Radeon Instinct MI50 (16GB) GPU。
* Graphcore IPU服务器: Azure ND40s_v3 VM: Intel Xeon Platinum 8168 CPU, 16个IPU。

软件配置:
* 操作系统: Ubuntu 16.04 / 18.04。
* GPU驱动/库: CUDA 11.0 / 11.3, ROCm 5.2.3。
* IPU软件: Poplar-sdk 3.0。
* 依赖框架/库: PyTorch (v1.12), ONNXRuntime (v1.12), Ansor (v0.9), Rammer, TensorRT (v8.4), FasterTransformer (v5.2), BladeDISC (v0.3.0), Nimble。

数据集/模型:
* 模型: 评估了10个具有不同结构类型的DNN模型，包括CNN、Transformer、CNN-Transformer和MLP，涵盖图像分类、NLP、语音识别、3D场景生成、图像修复/超分等任务。具体模型见表2。
* 来源: 所有模型均使用其官方的PyTorch实现，未经修改。

表2：WELDER中评估的DNN模型。

实验方法:
* 模型首先在PyTorch中被追踪并导出为ONNX格式。
* 该ONNX模型作为WELDER、Ansor、ONNXRuntime和TensorRT等框架的输入。
* 为保证公平性，输入和输出张量均放置在GPU设备内存中，避免数据移动开销。
* 在评估前进行预热迭代，然后每个工作负载重复运行至少5秒，报告平均速度。
* 跨模型的平均性能（如加速比）使用几何平均值计算。

A4 实验结果

5.2 NVIDIA GPU评估

端到端性能 (V100, SIMT Core, FP32):
* 实验内容: 在V100 GPU上（仅使用SIMT Core），对比了WELDER与多个基线在10个模型、批大小为1和64时的端到端推理性能。
* 实验结果 (图9):
* 批大小=1: WELDER相较于PyTorch、ONNXRuntime、Rammer、BladeDISC、Nimble、Ansor和TensorRT，几何平均加速比分别为4.29×、2.07×、1.96×、2.70×、1.79×、1.44×和1.47×。尤其在较新的模型（如NAFNet）上对TensorRT的加速比高达3.09×。
* 批大小=64: WELDER继续保持领先，几何平均加速比分别为1.83×、1.90×、2.1×、1.57×、1.49×、1.47×和1.21×。
* 结论: WELDER的通用内存优化方法在各种模型和批大小下均显著优于现有框架、编译器及专业库。

图9：NVIDIA V100 GPU（仅SIMT Core）上的端到端模型推理性能。（左：批大小为1，右：批大小为64）。

端到端性能 (V100, TensorCore, FP16):
* 实验内容: 在V100 GPU上启用TensorCore，使用FP16精度评估模型性能。
* 实验结果 (图10):
* 批大小=1: WELDER的性能优势进一步扩大，相较于PyTorch、ONNXRuntime、BladeDISC、Nimble、Rammer和TensorRT，平均加速比分别为7.18×、3.08×、5.29×、2.72×、2.76×和1.53×。
* 批大小=64: WELDER同样表现优异，平均加速比分别为1.98×、2.13×、1.97×、3.84×、3.45×和1.16×。
* 结论: TensorCore加速了计算，使得内存访问成为更显著的瓶颈，凸显了WELDER内存优化的重要性。

图10：NVIDIA V100 GPU（启用TensorCore）上的端到端模型推理性能。（左：批大小为1，右：批大小为64）。

自动发现的优化模式:
* 实验内容: 统计并分析WELDER在所有测试用例中自动发现的融合子图模式。
* 实验结果 (表4): WELDER自动发现了约300种不同的融合子图，其中89种包含至少两个基于规约的算子，这些是传统基于规则的融合难以覆盖的。融合的算子数量从2到48不等，与Ansor等基础融合方法相比，平均可带来1.87×的加速。
* 结论: WELDER的通用性使其能够自动发现并利用大量新颖且高效的算子融合机会，性能甚至超过了为Transformer模型（FasterTransformer，表3）和NeRF模型（fully-fused MLP）手工优化的专家级实现。
表3：WELDER和FasterTransformer的性能对比

表4：WELDER发现的融合模式示例。

消融与敏感性研究:
* 实验内容:
1. 创建WELDER的两个变体：WELDER-none（无算子间融合）和WELDER-base（仅寄存器层融合），与完整版WELDER和Ansor对比。
2. 改变三个模型的输入大小，观察融合带来的增益变化。
* 实验结果:
1. 消融研究 (图13): 与WELDER-none相比，WELDER-base（寄存器融合）平均延迟降低52%；完整版WELDER（共享内存融合）在WELDER-base基础上进一步降低29%的延迟，并显著减少了内核启动次数、全局内存事务和中间结果大小（IRS）。
2. 敏感性研究 (图14): 对于NAFNet，输入图像越大，融合增益越明显；对于BERT和Conformer，输入序列越长，模型越趋向于计算密集型，融合增益减小。
* 结论: 共享内存层的数据复用是WELDER性能提升的关键来源。融合的效益与模型的内存密集程度密切相关。

图13：3个选定模型的延迟、GPU内核数量、执行的全局内存事务和中间结果大小（IRS）（FP32，批大小64）。

图14：改变输入大小，与Welder-base比较。

编译时间:
* 实验内容: 对比WELDER和Ansor在两个模型上的端到端编译时间。
* 实验结果 (表5): WELDER的编译速度比Ansor快一个数量级以上，因为它使用分析性成本模型而非机器学习模型指导搜索，需要的调优试验次数显著减少。
* 结论: WELDER的调度策略在保证优化效果的同时，具有很高的搜索效率。
表5：Ansor和WELDER的编译时间

计算密集型模型性能:
* 实验内容: 在ResNet、UNet等传统计算密集型模型上对比WELDER、Ansor和TensorRT。
* 实验结果 (表6): WELDER在ResNet上与TensorRT性能相当。但在UNet和VGG16上，由于TensorRT使用了WELDER目前不支持的Winograd等优化数值算法，且这些模型融合机会较少，因此TensorRT表现更优。
* 结论: WELDER的核心优势在于内存访问优化，对于纯计算密集且无融合空间的算子，其性能依赖于基础的算子生成能力。对特定数值算法的支持是其未来可改进的方向。
表6：计算密集型模型的性能

5.3 AMD ROCm GPU评估

  • 实验内容: 在AMD MI50 GPU上对比WELDER与PyTorch、ONNXRuntime、Rammer和Ansor的性能。
  • 实验结果 (图15): WELDER全面胜出。批大小为1时，平均加速比分别为2.62×、1.71×、2.14×和1.53×。批大小为64时，加速比分别为1.69×、1.23×、1.86×和1.47×。
  • 结论: WELDER的优化方法具有良好的跨硬件平台适应性。在MI50上的加速比略低于V100，原因是MI50的计算能力相对较弱而内存带宽更高，使得工作负载更偏向于计算密集。

    图15：AMD ROCm MI50 GPU上的端到端模型推理性能（左：批大小为1，右：批大小为64）。

5.4 使用主机内存进行扩展

  • 实验内容: 评估WELDER在处理超出GPU内存容量的超大输入（如高分辨率图像）时的性能。通过将主机内存加入内存层次结构，对比在设备内存层启用/禁用数据复用的效果。
  • 实验结果 (表7):
    • GPU: 对于UNet和VGG16处理8k*8k图像，启用设备内存层的数据复用后，WELDER分别获得2.63×和1.89×的加速，并减少了3.11×和2.90×的主机内存传输。
    • GraphCore IPU: 对于处理2k*2k图像，WELDER分别获得3.63×和3.09×的加速。
  • 结论: WELDER的抽象设备层和整体流量优化使其能有效利用主机内存，将大型DNN任务扩展到单个加速器内存容量之外，并通过最大化设备内存中的数据复用来显著提升性能。

    表7：将大型DNN模型扩展到主机内存

A7 补充细节

6. 讨论

动态模型支持。WELDER的设计和实现主要集中在静态模型上。对于动态模型的执行，有两种实用的方法来解决。首先，可以通过即时编译（JIT）将动态图转换为静态子图，例如PyTorch JIT编译，这已成为PyTorch 2.0中的标准实践。然后，WELDER可以专注于优化这些通常是计算主导部分的静态子图。其次，即使张量形状是动态的，每个算子内部的数据块（tile）可以静态确定。这为WELDER提供了一个机会，即生成一个静态的块级融合计划，但将并行任务的数量留给输入张量形状在运行时确定。

7. 相关工作

算子融合。算子融合是DNN计算中一种广泛使用的技术，用于减少内核启动开销和提高在快速内存中的数据局部性。像TVM【索引15，TVM: An automated end-to-end optimizing compiler for deep learning，OSDI 2018】、Ansor【索引50，Ansor: Generating high-performance tensor programs for deep learning，OSDI 2020】、XLA【索引12，XLA，URL:https://www.tensorflow.org/xla】、DNNfusion【索引36 ，Dnnfusion: accelerating deep neural networks execution with advanced operator fusion，PLDI 2021】都支持在寄存器级别进行算子融合。其他编译器尝试在共享内存中进一步融合算子，依赖于针对一组已知算子类型的融合规则（如AStitch【索引51，Astitch: Enabling a new multi-dimensional optimization space for memory-intensive ml training and inference on modern simt architectures，ASPLOS 2022】、Apollo【索引49，Apollo: Automatic partition-based operator fusion through layer by layer optimization，MLSys 2022】、DeepCuts【索引24，Deepcuts: a deep learning optimization framework for versatile GPU workloads，PLDI 2021】）或针对少数算子组合的特定模板（如Bolt【索引47，Bolt: Bridging the gap between auto-tuners and hardware-native performance，MLSys 2022】）。专门的DNN运行时，如TensorRT【索引7，NVIDIA TensorRT，URL:https://developer.nvidia.com/tensorrt】和ONNXRuntime【索引8 ，ONNX Runtime，URL:https://github.com/microsoft/onnxruntime】，已经为流行模型（如基于Transformer的模型）中的一些常见模式整合了专家设计的融合规则。相比之下，WELDER适用于以张量表达式实现的通用算子，不假设算子类型，并自动决定最佳的融合内存层级。这是因为算子的资源使用行为（内存密集型或计算密集型）通常取决于其形状，并因此影响融合决策 。

并行性利用。像Rammer【索引31，Rammer: Enabling holistic deep learning compiler optimizations with rtasks，OSDI 2020】、HFuse【索引27，Automatic horizontal fusion for gpu kernels，CGO 2022】、Nimble【索引25，Nimble: Lightweight and parallel gpu task scheduling for deep learning，NeurIPS 2020】等系统，通过水平融合或调度并行任务（通过多流和CUDA图）来更好地利用硬件并行性并减少内核启动。WELDER在Rammer的基础上，进一步探索了对这些系统的补充性优化，即通过垂直融合进行整体内存优化，从而为内存密集型模型带来进一步的加速。

算子内优化。Ansor【索引50，Ansor: Generating high-performance tensor programs for deep learning，OSDI 2020】和Roller【索引52，ROLLER: Fast and efficient tensor compilation for deep learning，OSDI 2022】是代表性的张量编译器，它们通过循环优化或分块优化专注于算子内优化。特别是，Roller【索引52】和Triton【索引44，Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations，MAPL 2019】也利用了数据块（tile）的概念来优化内核性能（例如，算子内数据复用）。相比之下，WELDER通过整体优化算子内和算子间的内存访问来对它们进行补充。WELDER将Roller中的数据块概念推广为tile-graph抽象，公开了一个整体的块级调度空间，并提出了一个在该整体空间和显式内存层次结构上的高效调度机制。

特定模式优化。一些工作针对特定类型的模型，通过更激进的算子融合来优化特定模式，例如用于NeRF模型的全融合MLP【索引35，Real-time neural radiance caching for path tracing，arXiv 2021】，用于CNN模型的手动融合内核【索引46，Accelerating deep learning inference with cross-layer data reuse on gpus，Euro-Par 2020】，以及用于Transformer模型的注意力融合【索引2，FasterTransformer，URL:https://github.com/NVIDIA/FasterTransformer；索引18 ，Turbotransformers: an efficient gpu serving system for transformer models，PPoPP 2021】。我们的评估表明，WELDER可以自动实现这些融合中的大部分，甚至产生新的融合模式以帮助进一步优化。

其他领域。此外，内核融合技术已在传统图像处理【索引38，Automatic kernel fusion for image processing dsls，SCOPES 2018；索引39，From loop fusion to kernel fusion: A domain-specific approach to locality optimization，CGO 2019】或高性能计算（HPC）【索引45，Scalable kernel fusion for memory-bound GPU applications，SC 2014】领域中使用。这些工作通常利用其工作负载的领域特定融合规则。WELDER专注于DNN工作负载，但它适用于由张量表达式表示的通用算子。它也可能对其他领域中可以用张量表达式实现的工作负载有所帮助。

A5 结论

本文观察到现代DNN模型正变得越来越受内存限制，并基于此引入了WELDER，一个基于新型tile-graph抽象来优化执行效率的DNN编译器。WELDER能够跨多级内存层次结构，整体性地优化算子内和算子间的高效数据复用。WELDER首次将所有常见的算子融合方式统一到一个框架中，从而能够发现89种不常见的融合模式，其中最大的一次融合将48个算子合并为一个内核。这种通用性使得WELDER能够显著超越当前最先进的基线系统。更重要的是，WELDER为未来AI加速器利用新兴的内存层次结构趋势（如更大、连接更紧密的片上内存）提供了一种系统性的方法。

A6 附录

A Artifact Appendix

摘要。WELDER通过其新的tile-graph抽象提供端到端的DNN模型编译。本附件旨在复现论文在NVIDIA V100 GPU上的主要评估结果。

范围。本附件将验证以下声明：
* 端到端模型性能。通过复现图9、图10、图11、表3和表6的实验。
* 图1和图2中的动机实验。
* 图13中的消融研究。
* 表5中的编译时间。
* 表7中的GPU扩展实验。

内容。本附件包含实现WELDER的所有源代码。我们提供了一个docker文件来设置环境。对于上面提到的每个图和表，我们提供了一个脚本来复现其结果。由于完全复现结果需要编译超过50个模型测试用例，这将花费很长时间（特别是对于Ansor基线），我们也为NVIDIA V100 GPU提供了预编译的日志和模型。更多细节请参考仓库中的README.md文件。

托管。附件托管在GitHub仓库中。请使用git克隆该仓库并切换到osdi2023welder分支。

要求。本附件需要一台配备NVIDIA V100 GPU的设备，其CUDA驱动程序需支持高于11.0的CUDA运行时。