Effectively Scheduling Computational Graphs of Deep Neural Networks toward Their Domain-Specific Accelerators

有效调度面向领域特定加速器的深度神经网络计算图

作者/机构: Jie Zhao, Information Engineering University; Siyuan Feng, Shanghai Jiao Tong University; Xiaoqiang Dan, Fei Liu, Chengke Wang, Sheng Yuan, Wenyuan Lv, and Qikai Xie, Stream Computing Inc.

A1 主要贡献

本文提出了一种在领域特定架构（DSA）平台上有效调度深度神经网络（DNN）计算图的系统性方法——GraphTurbo。现有方法在划分DNN计算图时忽略了硬件架构，导致了冗余的核外数据移动和硬件资源利用不足。本文的核心思想是在图划分阶段就充分考虑硬件架构，从而实现网络与硬件架构的协同。

核心问题与研究目标:
现有DNN调度器将计算图划分为多个子图，但这种划分抽象了硬件架构，导致了以下问题：
1. 冗余的核外数据移动：生成了过多细粒度的内核（kernels），导致大量的中间数据需要在本地缓存（Local Buffer, LB）和全局缓存（Global Buffer, GB）或DDR之间移动。
2. 硬件资源利用不足：
* 本地内存利用不足：未能利用DNN网络架构中内存使用不均衡的分布特性，导致在处理张量尺寸较小的网络阶段时，高速本地内存被闲置。
* 并行性利用不足：错失了跨层指令调度的机会，无法充分利用不同专用计算单元（Compute Units, CUs）之间的并行性。

创新点与主要贡献:
为了解决上述问题，本文提出的GraphTurbo方法实现了网络与硬件架构的协同，其主要贡献如下：
1. 考虑硬件架构的图划分：在图划分阶段就整合硬件信息，生成更大但数量更少的内核。这一协同作用将大量原先的核外数据移动（off-core data movements）转化为核内数据交换（on-core data exchanges），显著减少了对内存带宽的压力。
2. 利用内存使用不均衡性：通过识别和利用DNN网络架构中内存使用分布不均衡的现象，GraphTurbo能够更好地饱和DSA的内存层次结构，尤其是在处理不同网络阶段时，避免了高速本地内存的浪费。
3. 实现跨层指令调度：通过构建更粗粒度的子图，实现了以往研究中未曾探索过的跨层指令调度，从而能更好地隐藏内存延迟，并进一步挖掘不同专用计算单元之间的并行性。
4. 系统设计与实现：设计并实现了一种新颖的调度方法GraphTurbo，为在DSA芯片上部署DNN提供了解决方案，并为其他平台提供了洞见。实验结果表明，GraphTurbo的性能优于两种最先进的工具（TVM和AStitch），并能达到与厂商手工优化的代码相媲美的性能。

下图展示了通用的DNN DSA抽象及其在不同商业加速器上的定制配置。

图 1: DNN DSA及其厂商定制。左侧为通用抽象，右侧为Habana Goya、华为Ascend 910和Graphcore IPU的具体配置。

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

DNN领域专用架构（DSA）的通用抽象。由于摩尔定律放缓，转向DSA被认为是满足DNN对计算能力巨大需求的有前景的方法【24，Norman P. Jouppi et al. A domain-specific architecture for deep neural networks. Commun. ACM 2018】。经过多年对DNN专用加速器的研究【7, 8, 15, 22, 25, 29, 32, 55, 58】，已形成一种普遍使用的DSA抽象（如图1左侧所示），大多数现有DNN加速器均基于此制造。

DSA抽象实例解析。以Habana Goya加速器【32，Eitan Medina. Habana labs presentation. In 2019 IEEE Hot Chips 31 Symposium (HCS)】为例（其定制配置如图1右侧所示），该抽象包含d个集群（cluster），每个集群有c个核心（core），每个核心包含u个用于不同DNN任务的计算单元（CU）。CU可以是带有暂存本地缓冲（LB）的张量处理核心（TPC），也可以是无LB的通用矩阵乘法（GEMM）引擎。核心通过集群内互连机制连接，并配备一个暂存全局缓冲（GB）。若d > 1，集群堆叠并通过DMA与DDR通信。图1中也展示了华为Ascend 910【29，Heng Liao et al. Ascend: a scalable and unified architecture for ubiquitous deep neural network computing. HPCA 2021】和Graphcore IPU【22，Zhe Jia et al. Dissecting the graphcore ipu architecture via microbenchmarking. arXiv preprint 2019】的定制配置。Graphcore IPU使用“tile”来表示一个核心及其独有的LB。其他加速器【7, 8, 15】的硬件架构也可根据图1的抽象推导。

现有DSA调度方法的局限性。因此，有效地将DNN调度到这种抽象上对于DSA编译器至关重要。这些为机器学习（ML）应用专门设计的加速器，具有基于暂存器的内存层次结构以及跨多核和多个CU的并行性。然而，先前为在这些DSA加速器上调度DNN而设计的工作【5, 13, 31, 54】在划分DNN计算图时未考虑硬件架构，从而引入了冗余的核外数据移动（即LB与GB/DDR之间），并且未能充分利用更快的本地内存和跨CU的并行性。

计算图相关概念与符号。为了解释先前工作的问题，我们首先介绍计算图，它被现有ML框架【1, 38】用来表示DNN模型，如图2所示。由于计算图可能包含大量节点，每个节点在多个张量上执行计算任务，其内存占用通常超过目标平台的本地内存容量，因此无法作为一个整体进行调度。现有调度器首先将其划分为子图，然后为每个子图分配资源。一个子图，也称为融合算子（fused op），首先由一个图节点初始化，然后根据其与其他节点的生产者-消费者关系进行分组，最终在目标平台上实现为一个内核函数或内核可执行文件。


图 2: ResNet-50的计算图【16，Kaiming He et al. Deep residual learning for image recognition. CVPR 2016】。节点（圆圈）代表算子，边（实线箭头）表示生产者-消费者关系。底部是一个由三个节点组成的3×3卷积（conv）层。

层、块和阶段的定义。我们使用“层（layer）”一词表示一组以直线方式连接的节点，其中至多一个节点包含应使用损失梯度学习的参数。虽然传统上图节点代表的算子被称为神经网络层，但某些神经网络层（如ReLU）在训练过程中不需要学习参数，或无需使用损失梯度学习（如批归一化），因此可被视为需要参数的算子（如卷积）的辅助算子。我们将层定义如此，是因为这一定义总结了现有编译器【5, 53】广泛使用的算子融合规则。我们还使用“块（block）”来表示单个层或由多个层组成的组件，在DNN计算图中被递归使用。例如，图2中由五层组成的conv块在每个阶段中使用一次，而identity块在每个阶段内使用多次。“阶段（stage）”是ResNet-50模型架构中使用的一种逻辑、高级抽象，它接收前一阶段的结果作为输入并生成输出张量，通常不被优化编译器考虑。

先前工作的挑战：细粒度子图与核外数据移动。先前工作【5, 19, 54】因模糊了硬件架构，将子图分组限制在层（layer）级别内（图2的底层），从而产生细粒度的子图。由于每个子图由一个内核实现，这导致了更多的内核和更多的核外数据移动。在图2中，conv和identity块分别包含五层和四层，整个网络可能需要数百到数千个内核和同等数量级的核外数据移动【21，Andrei Ivanov et al. Data movement is all you need: A case study on optimizing transformers. MLSys 2021】，给DSA平台有限的内存带宽带来巨大压力。此外，管理如此大量的核外数据移动对于DSA来说并非易事，因为与通用架构成熟的硬件机制不同，DSA的层次化暂存内存仍需手动或软件控制【37，Young H. Oh et al. Layerweaver: Maximizing resource utilization of neural processing units via layer-wise scheduling. HPCA 2021】。

先前工作的挑战：错失跨层指令调度机会。即使可以管理DSA内存层次中的数据移动，生成细粒度的子图仍然会错失跨层指令调度的机会。每个子图形成后，会被降级为一个循环嵌套流水线，并应用内存优化和循环变换（如分块和融合）来更好地利用硬件资源。尽管这些不同组合构成了现有自动调谐器【2, 6, 26, 57】的搜索空间，但跨层的指令调度机会，例如将在图2中一个3×3卷积层的权重内存提升语句和计算任务与其之前的1×1卷积层重叠，并未被这些搜索空间覆盖。

先前工作的挑战：未利用内存使用不均衡性。最后，由于内存使用不均衡分布（指内存使用在网络架构中变化【30，Ji Lin et al. Memory-efficient patch-based inference for tiny deep learning. NeurIPS 2021】的现象）未被暴露或利用，上述调度范式也未能充分利用DSA的快速本地内存。如图2顶部所示，ResNet-50被划分为四个阶段。conv块执行下采样操作（当S > 1时），将其输入形状从[N,C,H,W]转换为[N, 2C, H/S, W/S]，而identity块不改变张量形状。这导致stage1的内存使用量是stage2的$S^2$倍。如果目标DSA的快速内存在执行stage1时饱和，那么在执行其余阶段时将被严重低利用。

A2 方法细节

2 核心思想与概览

2.1 核心思想示例

核心思想概述。本文以图2中对一批输入图像进行分类的任务为例阐述核心思想。数据并行性在图1的d个集群间被利用，这通过分割张量的一个或多个可并行化维度来实现。假设大小为N=32的批处理维度被分割到这些集群中，每个集群处理n=8张已加载到GB的图像。本文研究一个DNN模型如何在一个集群内被调度。核心思想是最大程度地将输入张量保留在LB中，以便将尽可能多的核外数据移动转换为核内数据交换。为清晰起见，我们在图3a中重现了图2的各个阶段，并假设每个阶段执行S=2的下采样操作。


图 3: 不同调度方法下LB的利用率。(b)中的时间戳定义了GraphTurbo的调度。

现有方法的问题。现有方法【5, 13】虽然可以通过分组小算子来构建比层更大的子图，但由于缺乏硬件架构信息，它们不知道何时应停止分组。即使能为图3a的每个阶段构建大子图，这些方法也只能根据粗粒度的依赖关系（如图3a中箭头所示）来调度这些子图，这种策略将每个阶段的子图分布到图1的c个核心上。假设批处理维度被分割，每个核心处理一张图像。如果一张图像在执行stage1时饱和了LB，那么当其核心执行后续三个阶段时，这个本地内存将被低利用，因为前面的阶段通过下采样操作已将张量大小分别减小了2倍、4倍和8倍。

GraphTurbo的解决方案。GraphTurbo通过综合网络和硬件架构，可以轻松地为每个阶段构建大子图。它将这些大子图分别分割成八个、四个、两个和一个实例，从而将大子图之间的粗粒度依赖关系转换为子图实例之间的细粒度依赖关系。通过消除冗余的细粒度依赖，GraphTurbo在时间戳1和2执行stage1的两个实例（如图3b所示），在饱和LB的同时，通过将其他可并行化维度分布到核心间来利用核心间的并行性。接着，GraphTurbo在时间戳3执行stage2的一个实例，该实例处理两张图像，这两张图像都在LB中，因为图像尺寸已减小2倍，因此LB未被低利用。通过这种方式，LB始终得到充分利用，同时核心间的并行性也得到充分挖掘。这种调度方法在我们的实验平台上比TVM实现了7.97倍的加速。

2.2 GraphTurbo概览

GraphTurbo的调度流程。为了获得类似于图3b的调度策略，GraphTurbo接收一个经过标准图优化【13, 23】简化的计算图，并在图级别进行调度（§3）。为了构建更粗粒度的子图（例如图3a中的阶段），GraphTurbo首先收集硬件信息（§3.1），以指导其通过综合网络和硬件架构进行分组（§3.2）。这些子图随后被分割成实例，并进行排序（§3.3），以实现如图3b所示的调度顺序。然后解释了如何在多个核心间利用并行性并保证负载均衡（§3.4），同时自动推断核心绑定和内存作用域。接着，一个连接步骤被用来收集生产者子图实例的张量（§3.5），并进行了一些泛化讨论（§3.6）。

内核生成流程。图调度器生成有序的子图实例，这些实例被传递给内核生成器（§4），通过结合循环融合（§4.1）和缓冲拼接（§4.2）来生成内核，并自动管理内存分配和复用（§4.3）。对于§2.1中的例子，图调度器专注于输入图像。模型的卷积权重虽然越来越大，但只在层内使用，不会导致阶段间的通信。GraphTurbo只在LB中分配小的、固定大小的缓冲区，以允许在处理每个层时将这类张量提升到本地内存，而这种内存提升语句的开销则被计算所隐藏（§4.4）。

3 调度子图实例

本节构建更粗粒度的子图并调度它们的实例。为实现此目标，我们需解决五个问题，因此本节分为五个步骤，并在每一步开始时解释其难点。

3.1 收集切分信息

计算SplitInfo。GraphTurbo依赖生产者-消费者关系来将子图分组为更大的子图，但若无硬件架构信息，此策略不知何时停止。因此，本节首先为子图收集硬件知识。由于这些知识也用于切分子图（§3.3），我们称之为切分信息SplitInfo，它是一组4D元组(split_d, n_d, f_d, d)的集合。算法1总结了其计算方法。

算法 1: 计算SplitInfo

SplitInfo计算逻辑。在分组前，一个子图SG是一个代表算子的节点。当其被降级（lowered）时，可能会产生多个循环嵌套，因为它所代表的算子可能很复杂，需要多个中间张量来实现其功能【53，Jie Zhao et al. Apollo: Automatic partition-based operator fusion through layer by layer optimization. MLSys 2022】。除了定义输出张量的最后一个循环嵌套，其余都写入中间张量。算法1首先切分输出张量（第3-6行），然后将其切分信息传播到每个中间张量（第8-11行），原因如下：首先，一个子图只有一个输出但可能有多个输入张量，仅考虑输出张量简化了算法设计。其次，输出决定了该子图的循环嵌套应如何切分或分块【41, 52】。一旦输出和中间张量的信息确定，输入张量的切分方式也随之确定。这种计算SplitInfo的方式可能会引入中间或输入张量的重计算，在融合多个卷积或矩阵乘法算子时开销较大。因此，我们使用一个简单的成本模型来防止过度的重计算抵消融合带来的好处。

SplitInfo计算细节。depth代表输出张量的循环嵌套深度。通过从最外层到内层迭代循环维度d（第2行），算法1尽早地利用核心间的并行性。接着，算法1定义了三个度量（第3行）：split_d、f_d、n_d，分别代表当前维度d是否可切分、该维度的切分因子以及被其切分的算子数量。v遍历第4行的值来实例化f_d。我们以当前维度d的循环范围size(d)output为上界，因为v > size(d)output不会切分当前维度。前三个值将维度d分解到八、四和两个核心上，同时保证负载均衡，如图3b中前三个阶段的切分方式。8 ≤ v ≤ size(d)output之间的值不利用当前维度d在核心间的并行性，而是并行化其他维度，并充分考虑负载均衡，如图3b中stage1的切分。如果一个子图实例所需的内存占用峰值⌈peak/v⌉不超过LB的内存容量（第5行），则该值被用于实例化f_d（第6行）。n_d和split_d也相应更新。由于较小的v将peak划分为较大的块，这里的v循环是一种贪心策略。t遍历每个中间张量（第7行）。它接收维度d并首先检查该维度是否能切分输出张量（第8行）。若满足条件，n_d增加t中的算子数量（第9行）。此外，如果维度d与t的一个循环维度匹配，且匹配维度的循环范围size(d)t能被f_d整除（第10行），则该4D元组被添加到SplitInfo中（第11行）。

选择最佳切分维度。通过精确计算子图的SplitInfo，GraphTurbo为其生成的内核确定了合适的大小。由于SplitInfo通常有多个元素，每个元素编码一个可切分的循环维度，GraphTurbo需要为子图选择最佳维度。我们考虑以下标准：首先，能切分更多算子的循环维度更好。其次，切分因子较小的循环维度更优，因为它倾向于更好地饱和LB。最后，循环深度较小的维度更优，因为它表现出外部并行性且通信较少。这些标准被建模为求解一个优化问题的字典序最大值：
$\operatorname{lexmax}_{\forall d \in \operatorname{SplitInfo}}\left(n_d,-f_d,-d\right)$
其中三个度量的顺序定义了优先级。

3.2 子图分组

分组算法概述。现在我们可以对子图进行分组。GraphTurbo仍然根据生产者-消费者关系执行此步骤，但它通过利用SplitInfo来确定分组的终止条件，从而构建更大的子图，这种分组不受限于层内【5, 19, 54】。算法2概述了此过程。

算法2首先通过拓扑排序对计算图G的所有g个节点进行排序（第1行），每个节点被视为一个子图SG，并传递给算法1以计算其SplitInfo（第2-3行）。接着，算法2考虑三种不同的合并模式（第5-15行）来对这些子图进行分组，这将子图数量分别从g减少到s、d和b。每次分组后，子图索引和BestSplit都会更新。

算法 2: 分组子图

1 SG[1, · · · , g] ← topological_order(G); b ← g;
2 foreach i in [1, · · · , g] do
3     SplitInfo[i] ← Algo.1 (SG[i]);
4     BestSplit[i] ← Eq. (1) (SG[i], SplitInfo[i]);
5 repeat
6     {G, s} ← straight_merge(SG[1, ··· , b], SplitInfo[1, ·· · , b]);
7     foreach i in [1, · · · , s] do
8         BestSplit[i] ← Eq. (1) (SG[i], SplitInfo[i]);
9     {G, d} ← diamond_merge(SG[1, ··· , s], SplitInfo[1, ·· · , s]);
10     foreach i in [1, · · · , d] do
11         BestSplit[i] ← Eq. (1) (SG[i], SplitInfo[i]);
12     {G,b} ← branch_merge(SG[1,··· ,d], SplitInfo[1,··· ,d]);
13     foreach i in [1, · · · , b] do
14         BestSplit[i] ← Eq. (1) (SG[i], SplitInfo[i]);
15 until s, d and b all do not decrease;

合并模式定义。算法2考虑的合并模式定义如下。首先，两个子图SG1和SG2构成一个直线模式（图4a），当且仅当SG1是SG2的唯一生产者，且SG2是SG1的唯一消费者。其次，四个子图构成一个钻石模式（图4b），当且仅当只有一个入口、一个出口，并且从入口到出口至少有两条路径。入口（SG1）/出口（SG4）可以是多个外部子图的消费者/生产者。第三，三个子图构成一个分支模式（图4c），当且仅当只有一个出口和多条通往它的路径，且SG1可以是但不必是SG2的生产者。


图 4: GraphTurbo考虑的合并模式。实线箭头是生产者-消费者关系，虚线箭头表示可能与另一个子图连接。

分组启发式规则。算法2并非合并一个模式中的所有组件，而是将其中一部分融合成一个更大的子图SG，如图4中红色虚线框所示。算法2使用两个启发式规则来决定是否允许分组：
(i) 图4a中的SG1和SG2（或图4c中的SG2和SG3）可以合并，需满足两个前提条件。第一，图4a中SG2（或图4c中SG3）的切分因子不小于SG1（或SG2）的切分因子。由于较大的切分因子将内存峰值peak划分为更小的块，此条件防止了在下采样操作中，一个子图的大切分因子传播给其后续节点。第二，合并后子图SG的切分因子不大于SG2的切分因子，确保分组结果不会恶化SG2所利用的LB。SG的SplitInfo可使用算法1计算。
(ii) 图4b中的SG2、SG3和SG4可以合并成SG，如果SG的切分因子不大于SG2、SG3和SG4切分因子中的最大值，这也是为了保证良好的LB利用率。

分组示例。我们现在解释图2中的conv和identity块如何被分组。首先，算法2识别每层中的直线模式，得到图5a。具体来说，图2中conv块左侧的三层被识别为直线模式，融合成一个标记为1的子图。conv块中的conv和ReLU层均未被识别为直线模式，而是各自形成独立的子图，分别用标签2和3表示。类似地，identity块的三个conv层构成一个直线模式，用标签4的子图表示；其ReLU层用标签5表示。

图 5: 合并图2中的conv和identity块。

接下来，算法2找到两个分支模式（即图5a中的子图3和5）并生成图5b，其中的直线模式首先被合并成图5c，最终形成一个单一子图{1, 2, 3, 4, 5}。在合并这些模式时，上述两个规则的前提条件都得到满足。在实践中，算法2还可以将多个identity块和一个conv块分组成一个单一子图，从而为图3a中的阶段生成四个大子图。这些大子图不再被分组，因为规则(i)的第二个前提条件不再满足。

3.3 子图实例排序

生成子图实例图。通过§3.1和§3.2实现的软硬件协同，将一个计算图划分为更大的子图。此外，GraphTurbo还使用SplitInfo将每个子图分割成实例，本节确定这些实例的顺序。例如，图3a中的每个阶段被转换为一个或多个带有相同颜色标签的子图实例，形成图6中的新图。同一水平层上的所有实例都是同构的，因此可以按任意顺序执行。实例间的边继承自子图，但灰色的边是冗余的，可以轻易消除。我们通过检查一个生产者实例的输出是否被一个消费者实例使用来判断边是否冗余。结合SplitInfo和输出张量的形状可以完成此检查。若检查结果为真，则该边是冗余的并被消除。


图 6: 图3a中四个阶段的子图实例。

排序算法。GraphTurbo目前使用两种方法来调度子图实例的计算图G。第一，它以广度优先搜索（BFS）方式访问子图实例，产生如图7a所示的调度顺序。第二，GraphTurbo以深度优先搜索（DFS）方式访问它们，如算法3所述。我们目前仅使用BFS和DFS搜索，因为它们简化了GraphTurbo的算法设计。如§5将报告的，这一选择取得了有希望的结果。我们目前正在研究一种整数线性规划方法，以找到比这两种搜索启发式更好的解决方案，该方法将很快发布。

算法3首先以任意顺序访问G来实例化列表visit（第1行），然后将所有入度为零的子图实例从visit移动到另一个列表ready中（第3-4行）。SGI的入度和出度仅由图6中的黑边决定。visit中最后一个入度为零的SGI（第6-8行）被提取并添加到有序列表order中（第9行），其消费者被更新（第12-13行）并从visit移动到ready中（第14行）。

算法 3: 调度子图实例

排序示例。作为一个例子，图7展示了应用算法3时三个列表的变化情况。特别地，由于G可以按任何顺序访问，我们假设为便于说明，它按BFS顺序访问。order从左到右检查。GraphTurbo最终在两个结果中选择内存利用率更好的一个（图3b）。如果G按其他顺序访问，图7f中每个圆圈的标签会改变，但颜色不变，这不重要，因为同一子图的所有实例都是同构的。


图 7: 图6的不同调度顺序。为清晰起见，我们仅展示了算法3的前四个步骤。

3.4 推断核心绑定和缓冲作用域

核心绑定策略。下一步是将每个排好序的子图实例绑定到多个核心。一个子图实例的绑定策略可以通过检查其输出张量的循环维度来确定。然而，单独确定绑定策略可能会因策略不匹配而引入更多通信。我们使用算法4来推断绑定策略。由于算法4可以检测一对生产者-消费者子图实例之间的不匹配，它还利用这些信息来推断一个子图实例的输出张量应在哪个缓冲作用域声明。它首先以逆序访问像图7f这样的调度顺序O，并将所有相应的子图实例记录在列表visit中（第1行）。默认定义为空元组[]和LB的bind[i]和scope[i]分别用于记录子图实例的绑定策略和其输出张量声明的缓冲级别（第2行）。由于一个输出张量被另一个子图实例作为输入，我们不关心子图实例的输入张量应在哪里声明。一个子图实例也产生算法4未考虑的中间张量，GraphTurbo使用其内核生成器来管理（见§4.1）。

算法4：绑定和作用域推断逻辑。算法4为visit[i]表示的每个子图实例推断绑定策略和缓冲作用域（第3-16行）。如果bind[i]为空（即没有信息可用于推断）或scope[i]不是LB（即已知信息不能用于推断）（第4行），算法4使用一个朴素策略（plain strategy）实例化bind[i]（第5行），并推断当前子图实例输入张量的绑定策略。朴素绑定策略通过从最外层到内层循环维度贪心地分配更多核心来获得。沿多个维度的绑定因子形成一个多维元组。算法4通过迭代整数8、4、2和1来尝试实例化绑定因子，这保证了核心间的负载均衡。如果一个循环范围不能被当前值整除，则迭代进入下一个值。注意，某些张量可能在其循环维度上附加了特定要求，这些维度不能绑定到核心。

绑定策略传播与冲突处理。接下来，算法4使用infer_binding来推断visit[i]输入张量的绑定策略，这些输入张量是visit[i]生产者的输出张量（第8行）。推断是通过首先像算法1的第10行那样匹配输入张量的循环维度，然后从输出张量传播匹配维度的绑定因子来实现的。如果一个绑定策略不满足附加要求，则可能无效。因此，当推断的绑定策略既不为空也非无效时，算法4会进入以下三种情况之一：(1) 如果bind[j]尚未定义，则用推断的绑定策略实例化它（第9-10行）；(2) 如果已定义的bind[j]与推断的绑定策略不匹配，则scope[j]被GB覆盖（第11-12行），因为这里需要通信（见§3.5）；(3) 已定义的bind[j]与推断成功匹配，算法4进入下一次迭代（第13-14行）。在else情况下（第15行），bind[i]已被推断且不为空，此时算法4会尝试通过重新考虑visit[i]输出张量可能附加的要求来更新bind[i]，并为其计算一个朴素绑定策略。最终，在这两种绑定策略中使用更多核心的那个被用来重写bind[i]（第16行）。

3.5 连接实例输出

引入连接算子。由于GraphTurbo将一个子图分割成多个实例，子图的输出张量也被分割成多个部分。因此，本节引入连接这些子图实例输出的步骤。为了实现这一步，GraphTurbo检测子图实例之间的细粒度依赖关系，并在多个生产者的每个消费者之前引入连接（concatenation）算子，为运行示例得到图8。连接算子是轻量级的，因为它的输入和输出都停留在LB或GB中。如果一个连接算子接收的张量的绑定策略和内存作用域彼此不同和/或与其输出不同，GraphTurbo需要插入额外的算子来跨内存层次结构移动数据。


图 8: 连接实例输出。

插入数据移动算子。图9描绘了两种需要插入此类（灰色）算子的场景。一个子图实例或辅助算子用一个椭圆表示，显示其输出张量的形状、作用域和绑定策略。在左侧，一旦捕获到内存作用域之间的不匹配，就引入一个复制（copying）算子，将数据从其输入（GB）提升到其输出（LB）。在右侧，由于绑定策略的差异而插入一个重分布（redistribution）算子，触发核心间的通信。这可以通过使用每个维度上较小的绑定因子重置元组来实现。


图 9: 插入数据移动算子。

3.6 方法的泛化

泛化讨论。现在GraphTurbo可以有效地调度一个DNN计算图，但我们在其设计中做了一些假设。本节讨论如何将它们泛化。首先，§3.1假设一个子图只有一个输出张量，这在实践中通常是成立的。对于有多个输出张量的子图，可以对每个输出张量重复算法1。其次，§3.2中定义的规则考虑了DNN模型中的下采样操作。然而，我们的调度器也可以通过简单修改规则(i)来泛化到上采样操作。可能需要一个进一步的分割算子（可以看作是§3.5中引入的连接算子的反操作）来将一个子图实例的输出分发给其多个消费者。我们在这项工作中没有实验上采样操作。第三，§3.3使用两种方法对子图实例进行排序，这些方法简单但有效，如§5将要展示的。我们正在研究另一种启发式方法，即为内存占用较重的子图实例分配更高的优先级，并计划在未来发布。最后，GraphTurbo贪婪地使用LB，但作为一个建议，它可以用GB替换来在图1的d个集群间调度计算图。它也可以被其他平台的更快的内存替代，例如GPU的共享内存。有趣的是，更大的LB尺寸会简化GraphTurbo的算法流程。即使在不需要分割子图实例的情况下，GraphTurbo也能获得与TVM类似的调度，但充分考虑了跨层内存优化（§4.4）。此外，利用更高级别的缓冲，例如图1中CU内的缓冲（如果有的话），是有益的，因为在这种情况下，这类缓冲和LB之间的数据交换可能主导通信。

4 生成子图实例的内核

内核生成面临的挑战。一旦获得排好序的子图实例，内核生成器就可以将它们作为输入，并降级为循环嵌套流水线。我们的内核生成器的任务是通过实现循环变换和在DSA平台的快速本地内存中拼接中间张量来生成更大的内核。为最小化工程成本，我们在TVM【5，Tianqi Chen et al. TVM: An automated Endto-End optimizing compiler for deep learning. OSDI 2018】中原型化了我们的内核生成器，但它可能无法为一个子图实例生成单个内核。图10通过逐步展开图8中的一个子图实例b3来例证这个问题。


图 10: 展开b3以为其生成单个内核。b3是图2中stage2的子图实例。它被展开为一个conv块和三个identity块，每个块又被展开为多个层。

TVM的局限性。最右边的部分是TVM内核生成器的输入，但图10显示b3由38个算子组成（conv块11个，每个identity块9个），其中13个是conv算子。跨多个conv算子执行循环融合超出了TVM内核生成器的能力。尽管最近将CUTLASS【27，Andrew Kerr et al. Cutlass: Fast linear algebra in cuda c++. NVIDIA Developer Blog 2017】集成到TVM中以缓解GPU上的这个问题【50，Jiarong Xing et al. Bolt: Bridging the gap between auto-tuners and hardware-native performance. MLSys 2022】，但可接受的conv算子数量仍然有限。此外，即使在DSA平台上可以提供类似的厂商定制库，内核生成器仍会将子图实例的输出放在DDR中，这反过来会削弱GraphTurbo创造的优势。

4.1 层内循环融合

自动循环融合机制。内核生成器可以轻松地融合一个层内的算子。我们以图10中的#2层为例，将其降级为图11中间所示的循环嵌套流水线。由于TVM要求用户为这些循环嵌套编写调度模板，简单采用其工作流程无法完全自动化GraphTurbo。为了解决这个问题，GraphTurbo从每层中选择一个锚点算子（anchor op），并自动对该算子执行循环分块。如果当前层中出现conv/矩阵乘法，则应将其设为锚点算子；否则，应使用该层的最后一个算子（通常是元素级算子）。在两种情况下，都选择最外层的两个循环维度进行分块，因为它们是可并行的。然后，锚点算子两个维度上的块大小的选择方式与§3.4中确定朴素绑定策略的方式类似，即贪婪地最大化LB的内存利用率。换言之，块大小被实例化为一个既能整除当前循环范围又允许分块后的张量保留在LB中的最大整数。一旦锚点算子的块形状和大小确定，其他算子的循环边界就可以被推断出来，并与锚点算子融合，就像现有技术【41, 52】所做的那样，产生图11红色区域所示的融合和分块后的循环嵌套。


图 11: 层内循环融合。左：获取子图实例b3的作用域信息。中：第2层的循环嵌套流水线。右：分块和融合后的循环嵌套，并自动插入权重内存提升语句。

自动内存管理。通过将conv和batchnorm算子写入的张量转换为中间张量，该方法会自动将它们分配在更快的内存中，如§3.4所述。在这样做之前，内核生成器获取当前子图实例（即图11中的b3）的作用域信息，在定义的内存级别分配中间张量。权重张量的内存提升语句也以同样的方式自动注入，如图11蓝色区域所示。注意，像批归一化这样的一些算子可以被折叠，但我们在这里为了说明目的而保留它。

4.2 跨层/块的缓冲拼接

缓冲拼接机制。在一个层内部融合后，我们不将其输出放回DDR，而是让它保留在LB中，例如图11中的res_l1。因此，conv块的所有层都可以通过LB交换数据，我们称之为缓冲拼接（buffer stitching），实现这些层的内核可以被包装成一个。conv块的输入张量也放在LB中，如§3.4中的scope所声明。identity块可以类似地处理。由于每个块最后一层的输出张量也保留在LB中，这四个块都可以被拼接在一起。AStitch【60，Zhen Zheng et al. Astitch: Enabling a new multi-dimensional optimization space for memory-intensive ml training and inference on modern simt architectures. ASPLOS 2022】通过针对内存密集型算子也实现了类似功能。然而，我们的工作也考虑了计算密集型算子。此外，我们还试图在子图实例之间最大化快速本地内存。通过结合循环融合和缓冲拼接，我们的实现为子图实例b3生成了单个内核。相比之下，TVM为每层生成一个内核，增加了生成内核的数量（总共65个），因此需要更多的核外数据移动。

4.3 内存分配与复用

内存复用策略。剩下的任务是在scope定义的内存级别为每个张量分配空间，这很容易实现。然而，通过将许多张量放在更快的本地内存中，GraphTurbo需要一个仔细的机制来复用有限的LB容量。我们总是在一个层/块/子图实例的输出张量不再被使用时立即释放其消耗的空间。这样空间就可以被其他张量复用。LB只需要同时容纳有限数量的张量。如果这些张量的总大小超过LB的容量，具有最长生命周期的张量将被首先溢出到GB，然后是DDR。图12描绘了张量在计算任务间的生命周期。我们的启发式方法与先前工作【40, 46】不同，后者总是将内存大小最大的张量溢出到较低的内存层次级别。选择生命周期最长的张量有更高的概率溢出一个较小的张量，这很可能减少数据移动的开销。

(a) 执行ct1. (b) 执行ct4. (c) 执行ct5. (d) 执行ct7.


图 12: 张量在计算任务间的生命周期。一个（椭圆形）计算任务（ct）可以是一个子图实例、一个块或一个层。一个（矩形）张量在绿色时是存活的，灰色时是释放的。一个ct在黄色时正在执行，灰色时已完成。

4.4 跨层指令调度

跨层优化的实现。结合循环融合和缓冲拼接不仅为子图实例生成了单个内核，还允许不同层计算任务的重叠。在图11的右侧，提升一个权重是通过首先使用DMA将其张量从DDR复制到GB，然后将张量从GB提升到LB来实现的，这可能还会进一步分派到图1的各个CU。这些提升语句的延迟可以被一个更早执行的层隐藏起来，并且多个CU可以同时执行计算任务。图13展示了这种优化是如何执行的。一个圆角矩形代表一个层的分块计算任务。跨层内存延迟隐藏发生在两条垂直线之间，CU可以执行两个块的不同计算任务。我们的方法通过将优化粒度提高到超越先前工作【5, 13, 31, 54】所研究的层级，显著增强了这种内存延迟隐藏和并行性的机会。


图 13: 跨层指令调度。

A4 实验环境

  • 硬件配置:

    • DSA平台: STCP920 [51]，一个SoC DSA平台。
      • 集群 (d): 4个
      • 核心 (c): 每个集群8个
      • 计算单元 (u): 每个核心3个 (CU1: 向量核心, CU2: VME, CU3: MME)
      • 本地缓存 (LB): 64 KB L1
      • 全局缓存 (GB): 8MB Last Local Buffer (LLB)
      • 连接: 8个核心通过LLB双向连接。
    • GPU平台 (案例研究): NVIDIA A100。

  • 软件配置:

    • 编译器: LLVM v12.0.0。
    • 框架/库:
      • GraphTurbo原型基于TVM v0.8 [5]实现。
      • 模型实现：TensorFlow v1.13 [1] (大部分模型), Pytorch v1.8.1 [38] (DLRM)。
      • GPU案例研究：CUDA toolkit v11.4, CUTLASS v2.9 [27]。
    • 代码实现: GraphTurbo由19k行Python、44.2k行C/C++和2k行其他代码构成，其中图调度算法约7k行。

  • 模型与数据集:

    • 模型来源: 主要来自MLPerf v2.0 [43]基准测试套件，以及其他流行模型。
    • 具体模型列表:
      1. ResNet-50 v1.5 [16]
      2. BERT [9] (序列长度为128, 256, 384, 512)
      3. DLRM [33]
      4. MobileNet v2 [17]
      5. Vision Transformer (ViT) [11]
      6. DenseNet [18]
      7. Conformer [14]
    • 说明: 实验主要针对推理任务，未考虑动态形状张量的模型。

A4 实验结果

5.1 跨集群任务分解与最优批大小选择
- 实验内容: 使用BERT-128模型，在STCP920平台上比较TVM和GraphTurbo在不同批大小下的吞吐量和延迟。
- 实验结果 (Table 1):
- 当批大小较小（8, 16）时，TVM和GraphTurbo性能相当，因为集群的LLB足以容纳所有批次。
- 当批大小增加到32时，TVM只能为每个集群分配4个批次并通过循环执行，导致延迟增加；而GraphTurbo能分配8个批次，通过创建和切分大子图，吞吐量达到4052 sentences/s，延迟仅为7.67ms，远优于TVM（512 sentences/s, 18.96ms延迟）。
- GraphTurbo的性能与厂商手工优化的代码非常接近（4048 sentences/s, 7.62ms延迟）。
- 分析结论: GraphTurbo通过其调度策略，能够有效地处理更大的批大小，从而在不显著增加延迟的情况下大幅提升吞吐量，这证明了其方法在利用硬件资源上的优越性。

5.2 性能比较
- 实验内容: 在7个DNN模型（10种配置）上，比较GraphTurbo与TVM、AStitch以及厂商手工优化代码的吞吐量。
- 实验结果 (Fig. 14):
- vs. TVM: GraphTurbo平均性能超出TVM 11.15倍。这是由于GraphTurbo生成了更少的内核，减少了核外数据移动，更好地饱和了L1缓存，并利用了跨层指令调度。
- vs. AStitch: GraphTurbo平均性能超出AStitch 6.16倍。AStitch未能通过切分子图实例来利用内存使用不均衡的特性。
- vs. 厂商代码: GraphTurbo平均性能超出厂商手工优化代码1.04倍。虽然手工优化在某些模型（如BERT-512）上表现更佳，但GraphTurbo的全自动方法在整体上更具优势，且能处理厂商尚未完成优化的复杂模型。
- 分析结论: GraphTurbo的性能显著优于现有的自动化工具，并且达到了与专家级手动优化相媲美甚至超越的水平。

5.3 性能分解
- 实验内容: 通过逐步开启GraphTurbo的各项优化（数据保存在LLB、保存在L1、实例排序、跨层指令调度），分析各优化对总体性能提升的贡献。
- 实验结果 (Fig. 15):
- 数据保留在L1: 将数据从DDR移至L1（而不是LLB）带来了显著的性能提升，平均加速比为3.67倍，证明了将核外数据移动转换为核内数据交换的有效性。
- 实例排序: 在L1优化的基础上，通过子图实例排序来利用内存不均衡分布，带来了额外的2.20倍平均加速。
- 跨层指令调度: 最后，开启跨层指令调度带来了额外的1.72倍平均加速。
- 分析结论: GraphTurbo的每一项核心优化都对最终性能有重要且独立的贡献，其中最大化利用最快的本地内存（L1）是基础，而实例排序和跨层调度则进一步放大了性能优势。

5.4 硬件利用率
- 实验内容: 评估GraphTurbo对DSA硬件资源的利用效率，包括内存层次结构和计算单元（VME和MME）。
- 实验结果:
- 内存层次 (Table 3): TVM几乎总是将数据写入DDR，而GraphTurbo和厂商代码则大量使用LLB和L1，极大地减少了对DDR的访问。
- 计算单元 (Fig. 16, Fig. 17): 随着批大小的增加，GraphTurbo能持续提升VME和MME的利用率。在ResNet-50的四个阶段中，GraphTurbo在后三个阶段的硬件利用率远高于TVM，这得益于其对内存不均衡分布的利用。
- 分析结论: GraphTurbo通过其调度策略，显著提高了DSA平台内存层次和计算单元的利用率，将硬件潜力更充分地发挥出来。

5.5 编译开销比较
- 实验内容: 比较TVM、AStitch和GraphTurbo在不同模型上的编译时间。
- 实验结果 (Table 4): GraphTurbo的编译时间与基线方法相当，没有因为引入复杂的调度算法而显著增加编译开销。
- 分析结论: GraphTurbo在实现巨大性能提升的同时，保持了可接受的编译效率。

5.6 GPU案例研究
- 实验内容: 在NVIDIA A100 GPU上，使用ResNet18-Tailor模型，将GraphTurbo的调度思想与TVM+CUTLASS进行比较。
- 实验结果 (Table 5): GraphTurbo的调度方法平均比TVM（使用CUTLASS）快1.23倍，比CUTLASS自身的卷积融合快1.06倍。
- 分析结论: 尽管GPU的高内存带宽削弱了减少核外数据移动带来的部分优势，但GraphTurbo通过调度子图实例来利用内存不均衡分布的思想仍然有效，证明了其方法具有跨平台的泛化潜力。

A5 结论

本文提出了一种名为GraphTurbo的DNN模型调度器，其核心创新在于实现了网络与硬件架构的协同。与先前工作在图划分时忽略硬件信息的做法不同，GraphTurbo在调度早期就整合硬件特性，从而带来了三大优势：
1. 减少核外数据移动：通过生成更粗粒度、数量更少的内核，将大量数据交换限制在片上缓存中。
2. 饱和本地内存：通过切分和重排子图实例，有效利用了DNN模型中普遍存在的内存使用不均衡现象，提高了高速本地内存的利用率。
3. 实现跨层指令调度：创建了更大范围的优化机会，允许在不同层的计算和内存访问之间进行重叠。

在七个DNN模型上的实验结果表明，GraphTurbo的性能显著优于现有的先进工具，并能达到甚至超越厂商手工优化的水平。在GPU上的案例研究也证明了其核心思想的泛化能力。

局限性与未来工作:
- 最优批大小选择: 目前集群的最优批大小仍通过简单的自动调优选择，未来计划开发更智能的技术来解决此问题。
- 动态形状张量: GraphTurbo当前不支持动态形状的张量，未来考虑与处理动态张量的最新方法【12, 56】相结合以缓解此问题。

A7 补充细节

6 相关工作

与图调度器的比较。调度计算图是部署DNN模型的第一步。我们的调度器与先前工作【13, 19, 44, 49】的区别在于，我们在分组子图时考虑了硬件架构，实现了网络架构和DSA的协同，而现有方法没有。通过生成更粗粒度的子图并将其分割成实例，GraphTurbo暴露了内存使用不均衡的分布，这是MCUNetV2【30，Ji Lin et al. Memory-efficient patch-based inference for tiny deep learning. NeurIPS 2021】首次研究的网络属性。然而，MCUNetV2仅讨论了微控制器上的微型DNN模型，而本文考虑了大规模DNN模型并面向云端DSA芯片。MCUNetV2未考虑的一些只能在子图降级为循环嵌套流水线时才能实现的优化，在§4中进行了研究。

与循环级融合方法的比较。当子图降级为循环嵌套流水线时，像TVM【5，Tianqi Chen et al. TVM: An automated Endto-End optimizing compiler for deep learning. OSDI 2018】这样的现有方法借助手动编写的调度模板来融合循环嵌套。由于TVM在子图内算子数量增加时扩展性不佳，我们仅使用TVM的循环融合来分组层内的循环嵌套。我们的实现还避免了手动编写调度模板和注入权重张量内存提升语句的需要，后者通过与图调度器交互实现自动化。

与XLA及类似方法的比较。XLA【13，Google. Xla: Optimizing compiler for machine learning. 2017】通过将高级子图展开为独立的低级算子，不将融合限制在层内。然而，通过低级算子检索高级信息对XLA融合低级算子至关重要，并且在该编译器中，手动形成有益的高级子图被认为比通过自动模式匹配更可靠【45，Bjarke Roune. Compiling ml with xla (slides). 2019】。GraphTurbo完全自动化了此过程，并取得了比AStitch更好的性能，而AStitch已被证明优于XLA【60，Zhen Zheng et al. Astitch: Enabling a new multi-dimensional optimization space for memory-intensive ml training and inference on modern simt architectures. ASPLOS 2022】。

与其他相关工作的比较。IREE【20，IREE. Iree. 2021】是另一项在低级并行流水线硬件/API之间通信数据时利用其图调度逻辑的工作。我们的工作与IREE的不同之处在于，我们专注于调度更大子图的实例，这倾向于产生更少的内核。Zero-Infinity【42，Samyam Rajbhandari et al. Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning. SC 2021】通过动态维护一个算子序列的内部映射，在执行当前算子时预取未来算子所需的参数，从而利用细粒度的重叠。我们工作中所考虑的层通常包含多个算子，其执行更有可能完全隐藏数据传输开销。

与自动调优框架的比较。一些自动调优框架【6, 57, 59】也被设计用来增强TVM的能力，减少或无需手写调度模板。这些自动调优器使用其搜索启发式来调整内存优化，以进一步提高生成代码的性能。不幸的是，它们的搜索空间都限制在层内【39，Phitchaya Mangpo Phothilimthana et al. A flexible approach to autotuning multi-pass machine learning compilers. PACT 2021】，而我们的工作实现了跨层指令调度（§4.4）。特别是，Ansor【57，Lianmin Zheng et al. Ansor: Generating high-performance tensor programs for deep learning. OSDI 2020】代表了这类工作的最前沿，它与我们的工作是正交的，既没有考虑GraphTurbo的调度，也没有利用多个conv/矩阵乘法算子的融合可能性。多面体框架【3, 47, 54】也研究了循环融合，但它们没有考虑§4.2中讨论的缓冲拼接。与GraphTurbo类似，DNNFusion【34，Wei Niu et al. Dnnfusion: Accelerating deep neural networks execution with advanced operator fusion. PLDI 2021】也研究了移动设备上的跨层融合。我们未能获取其代码库进行实验比较。

与水平融合的比较。另一类工作【10, 28, 35】研究了没有生产者-消费者关系的算子之间的水平融合，以更好地利用其目标的硬件资源。GraphTurbo用不同的思想解决了同样的问题。我们的调度器在一个子图实例内利用并行性，该实例总是与同一子图的其他实例同构。它总是分解一个或多个张量维度来利用并行性。相反，通过水平融合分组的算子是异构的，这需要一个更复杂的并行化机制。

与DSA平台编译器的比较。最近，DSA平台的调度器和代码生成器被广泛研究。Rammer【31，Lingxiao Ma et al. Rammer: Enabling holistic deep learning compiler optimizations with rtasks. OSDI 2020】和Roller【62，Hongyu Zhu et al. ROLLER: Fast and efficient tensor compilation for deep learning. OSDI 2022】为Graphcore IPU【22，Zhe Jia et al. Dissecting the graphcore ipu architecture via microbenchmarking. arXiv preprint 2019】生成代码。它们通过组合不能饱和硬件资源的算子来最大化利用更快的内存。AKG【54，Jie Zhao et al. Akg: Automatic kernel generation for neural processing units using polyhedral transformations. PLDI 2021】使用多面体模型【4, 48】为Ascend 910【29，Heng Liao et al. Ascend: a scalable and unified architecture for ubiquitous deep neural network computing. HPCA 2021】进行代码生成和循环融合。XLA【13，Google. Xla: Optimizing compiler for machine learning. 2017】和NaaS【61，Yanqi Zhou et al. Towards the co-design of neural networks and accelerators. MLSys 2022】利用子图的调度为TPU【25，Norman P. Jouppi et al. In-datacenter performance analysis of a tensor processing unit. ISCA 2017】生成代码。我们的工作与这些方法的区别在于，GraphTurbo沿一个输出张量的维度划分一个子图，而这些方法沿多个维度按块划分张量。GraphTurbo这样做的主要原因是图1的DSA抽象中的核心是以一维形式组织的。例如，这个级别对应于GPU共享内存的32个硬件核心。GraphTurbo的划分方法也可扩展到处理多维核心网格组织，通过逐步划分和映射多个循环维度到这些硬件维度。此外，由于它们的目标共享图1中的DSA抽象，本文提出的思想也可以在它们的目标上使用。