SoD2: Statically Optimizing Dynamic Deep Neural Network Execution

文章标题：SoD2: 静态优化动态深度神经网络执行
作者/机构：Wei Niu (University of Georgia), Gagan Agrawal (University of Georgia), Bin Ren (William & Mary)

A1 主要贡献

本文提出了一个名为 SoD2 的综合框架，用于优化动态深度神经网络（DNN）。

核心问题：现有的DNN编译和运行时系统主要关注静态DNN，其每个层的输入输出形状和大小都是先验已知的，且执行路径固定。然而，动态DNN正变得越来越普遍，其张量形状、大小甚至使用的算子集都依赖于输入或执行过程。现有系统处理动态模型时，通常采用保守假设或昂贵的运行时分析，导致高昂的开销。例如，TFLite 和 MNN 在输入形状改变时会进行重新初始化（相当于重新编译），其开销甚至可能超过推理本身（如表1所示）；或者采用保守的最大内存分配，造成严重浪费。

研究目标：提出一种精细化的方法来优化具有动态特性的DNN推理，重点在于降低推理延迟和内存需求，后者对于移动设备尤为重要。

创新点：
1. DNN算子分类：本文对现代DNN中使用的算子（特别是ONNX中的150个算子）进行了深入研究，并根据输出形状与输入形状和值的关系将其分为四类：
* 输入形状决定输出 (Input Shape Determined Output)
* 输入形状决定输出形状 (Input Shape Determined Output Shape)
* 输入形状和值决定输出形状 (Input Shape & Value Determined Output Shape)
* 执行决定输出 (Execution Determined Output)
这种分类为静态分析动态DNN奠定了基础。

1. 秩和维度传播的数据流分析（RDP）：基于算子分类，本文开发了一个名为秩和维度传播（Rank and Dimension Propagation, RDP）的静态分析框架。该框架通过前向和后向分析，在编译时推断中间张量的形状和维度，这些形状可以是已知常量、符号常量或涉及这些常量的表达式。

2. 基于RDP的全面静态和动态优化：利用RDP分析的结果，SoD2实现了一系列优化，包括：

   • 算子融合与融合代码生成：即使在形状动态的情况下也能进行算子融合，并在必要时生成多版本代码。
   • 静态执行规划：利用RDP结果对计算图进行分区，并基于RDP输出的启发式规则选择执行顺序，以优化内存使用。
   • 运行时内存分配规划：为运行时生成内存分配计划，减少内存开销。
   • 多版本代码生成：为关键算子生成针对不同形状优化的多个代码版本。

通过在10个新兴的动态DNN模型（包括StableDiffusion和SegmentAnything等）上进行评估，SoD2相较于四个现有系统（ONNX Runtime, MNN, TVM with Nimble, TensorFlow Lite），实现了高达3.9倍的执行速度提升和高达88%的峰值内存消耗降低。

Table 1: 动态形状下执行重新初始化的推理开销。SL：形状传播和布局选择。ST：调度和调优。Alloc：内存分配。Infer：推理时间。实验在三星Galaxy S21上使用MNN [26]进行。

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

2 现有框架及其局限性

静态解决方案。许多现有的移动平台DNN推理引擎（特别是TFLite 【1, TensorFlow: A system for large-scale machine learning. 2016. OSDI】和MNN 【26, MNN: A Universal and Efficient Inference Engine. 2020. MLSys】）通过扩展其静态模型执行来支持动态特性。为了处理动态输入形状，这涉及到当输入形状改变时执行重新初始化，或者在输入形状未知时进行保守的（最大）内存分配。为了处理动态控制流，通常需要执行所有可能的路径，并剥离无效结果。毫不奇怪，这种对动态特性的简单处理会产生显著的执行和/或内存开销。为了进一步说明，表1展示了三个可以接受动态形状输入的模型（YOLO-V6 【36, YOLOv6: A single-stage object detection framework for industrial applications. 2022. ArXiv】, Conformer 【20, Conformer: Convolution-augmented Transformer for Speech Recognition. 2020. Interspeech】, 和CodeBERT 【16, CodeBERT: A Pre-Trained Model for Programming and Natural Languages. 2020. EMNLP】）的性能研究。MNN 【26, MNN: A Universal and Efficient Inference Engine. 2020. MLSys】在三星Galaxy S21上运行这些模型，通过执行重新初始化来处理变化的输入形状。这些结果表明，重新初始化通常比推理本身花费的时间还要长得多。这种方法对于重新初始化的开销可以被多次推理任务摊销的情况可能是可以接受的（例如，某些视频处理场景）。然而，许多应用场景（跨越图像、音频和语言处理）涉及不断变化的输入。另一种方法，如上所述，是保守地分配大的内存空间。然而，这会造成显著的内存浪费，可能限制执行大型模型的能力或有效执行的能力，尤其是在内存有限的移动（或边缘）设备上。

运行时解决方案。TVM（带有Nimble扩展） 【5, TVM: An automated end-to-end optimizing compiler for deep learning. 2018. OSDI】【57, Nimble: Efficiently compiling dynamic neural networks for model inference. 2021. MLSys】通过在虚拟机内提供一组优化来改进静态解决方案的局限性。此功能的一个例子是使用形状函数来推断输出张量形状，并使用此信息进行动态内存分配。然而，这样的函数和随后的动态内存分配会引入显著的执行开销。

3 基于动态性的算子分类

我们的观察是，DNN算子具有不同的动态程度，这导致了优化它们时面临的不同层次的挑战和机遇。更具体地说，这项工作将DNN算子分为四种类型：输入形状决定输出、输入形状决定输出形状、输入形状和值决定输出形状，以及执行决定输出。本节给出了正式定义。

背景和符号。通常将DNN表示为计算图，它是一个有向无环图（DAG）。每个张量（可以是输入和/或输出）可以通过形状（包括维度）和内容或值来分类。每个算子表示为$L_l$，其中$l$是算子索引。假设$L_l$有$m$个输入张量（其中，$1, \dots, k$是常量张量，而$k, \dots, m$是来自先前算子的输出张量）和$n$个输出张量。第$l$个算子的输入张量$i$的形状表示为$IS_i^l$，相应的张量值可以表示为$IV_i^l$。类似地，每个输出张量的形状和值分别表示为$OS_i^l$和$OV_i^l$。现在，直观上，有一类函数将输出的形状和值与输入的形状和值联系起来——$F_s$用于形状，$F_v$用于值。

• 输入形状决定输出 (Input Shape Determined Output)：输出（张量），由其形状和值共同表征，对输入具有以下依赖性。输出张量的形状依赖于输入张量的形状，而输出张量的值由输入张量的形状和可能的一些常量张量决定——输入值不影响输出。例子包括Shape和EyeLike。形式上，存在一对函数$F_s^f$, $F_v$，使得：
  $$OS_i^l \xleftarrow{F^{fs}} IS_1^l, \dots, IS_m^l$$
  $$OV_i^l \xleftarrow{F^v} IS_1^l, \dots, IS_m^l, IV_1^l, \dots, IV_{k-1}^l$$
  其中 $1 \le k \le m$。

• 输入形状决定输出形状 (Input Shape Determined Output Shape)：与前一类别类似，输出形状依赖于输入形状。然而，不同之处在于输出值依赖于所有的输入值（包括中间输入值和常量输入值）。例子包括Conv、Add和Pooling。与后续类别相比，此类别的重要性在于，如果该算子的输入形状已知，编译器优化（例如，算子融合、执行/内存优化）便可启用。形式上，存在一对函数$F_s^f$, $F_v$，使得：
  $$OS_i^l \xleftarrow{F^{fs}} IS_1^l, \dots, IS_m^l$$
  $$OV_i^l \xleftarrow{F^v} IS_1^l, \dots, IS_m^l, IV_1^l, \dots, IV_m^l$$

• 输入形状和值决定输出形状 (Input Shape & Value Determined Output Shape)：与前一类别类似，输出值依赖于输入形状和所有输入值。不同之处在于输出形状还依赖于部分输入值。例子包括Extend和Range。形式上，存在一对函数$F_s^f$, $F_v$和一个输入张量的子集($p, \dots, q$)，其值指定了输出形状，使得：
  $$OS_i^l \xleftarrow{F^{fs}} IS_1^l, \dots, IS_m^l, IV_p^l, \dots, IV_q^l$$
  $$OV_i^l \xleftarrow{F^v} IS_1^l, \dots, IS_m^l, IV_1^l, \dots, IV_m^l$$
  其中$1 \le p \le q \le m$。如果$p \le q \le k$，这与“输入形状决定输出形状”类别相同，并且所有依赖的输入张量都是常量。在这种情况下，可以在不知道其他中间输入张量的情况下计算输入形状。如果只知道该算子的输入形状，只能对其应用带有保守分析的部分编译器优化，而完全优化需要动态执行结果。

• 执行决定输出 (Execution Determined Output)：与前两个类别类似，输出值依赖于输入形状和所有输入值。例子包括Nonzero和If。形式上，存在一个函数$F_v$，使得：
  $OV_i^l \xleftarrow{F^v} IS_1^l, \ldots, IS_m^l, IV_1^l, \ldots, IV_m^l$
  并且第i个输出张量的形状只能在其实际值生成后才能测量：
  $OS_i^l \leftarrow SHAPE\_OF OV_i^l$
  这意味着在生成输出张量（即执行该层之后）之前无法知道输出形状。只能对该算子应用带有保守分析的部分优化，而完全优化需要动态执行结果。

反向传递。尽管这些算子类型是根据前向传递定义的，即输出张量的形状和值与输入张量的形状和/或值相关。在实践中，也使用反向传递，即我们可以（并且需要）将已知的输出形状（秩、维度或两者）反向传播到未知的输入形状。例如，如果我们知道Add算子的输出形状，由于广播规则【11, Numpy developers. 2023. Tensor Broadcasting】，其输入维度可能是1或与相应的输出维度相同。我们将反向传递函数定义为：
$IS_i^l \xleftarrow{F^{bs}} OS_1^l, \dots, OS_n^l.$

算子分类示例。表2显示了根据上述分类对ONNX 【48, ONNX. 2017. Open Neural Network Exchange.】中典型算子的分类。作为进一步的说明，图1显示了四个子图，它们代表具有不同动态程度的算子（用红色边界标记）及其连接。图1(a)显示了一个“输入形状决定输出”算子Shape。一旦其输入形状已知，其值结果就可以直接推断出来（实际上，这个值可以从Shape传播到BiasAdd，因为所有后续算子都属于“输入形状决定输出形状”组）。类似地，图1(b)意味着如果Conv的输入形状已知，这个形状信息可以传播到整个子图，因为该子图中的所有算子都属于“输入形状决定输出形状”组。对于(a)和(b)中表示的情况，即使确切的形状未知，我们仍然可以执行编译器优化，如算子融合和融合代码生成、执行顺序优化和内存优化，这将在下一节详细阐述。在图1(c)中，TopK的输出形状取决于其输入值（在示例中是左侧前驱分支），即TopK（及其后继者）的输出形状在其左侧前驱分支被执行之前是未知的。图1(d)表示一个涉及动态控制流的子图。Switch的结果决定了路径①、②和/或③是否会被执行，而Combine则合并已执行路径的结果。(c)和(d)都需要动态执行，因此更难进行静态优化。

Table 2: 基于动态程度的DNN算子分类。算子来自ONNX（开放神经网络交换）[48]。

讨论。尽管上面提到的表2和图1中的例子简单地将每个算子归入一个类别，但还有其他考虑因素。例如，一个Upsample算子可能属于“输入形状决定输出形状”或“输入形状和值决定输出形状”，这取决于其某些输入张量是否是常量。因此，通过常量传播，一个算子可能会从一个更动态的分类转变为一个较不动态的分类，为我们提供更激进的优化机会。这激发了SoD2的某些方面。此外，因为我们的算子分类本质上是通过研究其计算逻辑以及输入/输出张量的形状和值来模拟算子的动态程度，所以可以基于现有的中间表示（如张量表达式，例如TVM表达式）创建一个自动和通用的工具来将算子分类。

A2 方法细节

4 SoD2的设计

基于上文介绍的DNN算子分类，SoD2引入了一个新的静态数据流分析框架，以推断中间结果张量的形状。这种分析是多项优化的基础，包括动态DNN算子融合、执行路径规划、内存规划和多版本代码生成。所有这些优化都确保在给定特定输入的情况下，具有确定的运行序列和一致的输出。在较高层面上，我们的方法不需要保守的静态假设或运行时开销，从而比现有技术有显著改进。

4.1 部署前数据流分析

RDP框架。为了促进动态DNN的静态优化，一个关键要求是（可能以符号形式）知道中间结果张量的形状（即秩和维度）。我们的关键观察是，对于许多算子和算子组合（例如，一个“输入形状决定输出”算子和一个“输入形状决定输出形状”算子），即使不知道输入张量的形状，仍然有可能在一定程度上推断出中间结果张量的形状。我们的框架基于这一观察，并被称为算子秩和维度传播（operator Rank and Dimension Propagation, RDP）。虽然RDP与经典的（符号）常量传播框架 【4, Interprocedural constant propagation. 1986. ACM SIGPLAN Notices】有某些相似之处，但它需要处理DNN操作和计算图的细微差别。RDP还将其格中的可能性考虑为对多个（符号）常量的操作，并需要迭代的前向和后向分析。

RDP的正式定义。整个RDP算法表示为一个四元组 < G, D, L', F >。
- G 是一个扩展的计算图（DAG），带有控制流算子<Switch, Combine>。如果这个扩展的计算图涉及多个需要全部执行的分支，我们假设执行顺序总是从左到右。很容易证明G等价于一个算子上的控制流图，这作为该数据流分析的基础。
- D 是数据流的方向，可以是FORWARD（前向）和BACKWARD（后向）。与大多数经典的数据流公式（例如，常量传播或到达定义）不同，RDP在前向和后向方向上迭代处理G，直到结果收敛。这是因为一个张量的形状可以从其生产者算子和/或消费者算子推断出来，并且它们的推断结果应该相同以保证此DNN执行的正确性。

• L' 本身是一个三元组 < V', ∧, m >。V' 是值的域（如图2所示），包括形成一个格（lattice）的已知常量、符号常量和操作推断常量。该格还包括作为格的顶部（⊤）的undef（未定义）和作为底部（⊥）的nac（非常量）。∧ 是一个meet操作符，它遵循积格（product lattice）的通用定义。m 是一个映射函数，将格中的值映射到两个变量，Shape (S) 和 Value (V)，代表中间张量的这两个属性。更具体地说，RDP是一种数据流分析，其中L'描述了其分析范围，即如何将每个形状和值属性映射到一种构成格结构（如图2）的常量。这个格保证了RDP分析的属性遵循格理论，因此RDP分析将收敛到一个唯一的解。
• F : V' → V' 是传递函数的域。F是为每个算子（类型）设计的，并根据算子类型将Shape (S)和Value (V)从输入张量传递到输出张量。与常量传播的数据流分析类似，RDP有两种传递函数，Update和Merge。Update为单个算子从输入张量传递到输出张量；而Merge操作于分支控制流，并合并来自多个可能执行路径的（输出）张量。因为RDP有FORWARD和BACKWARD两个方向，所以F也包含两个方向的传递函数。

传递函数示例。SoD2包含16种Update传递函数。这些函数基于算子动态程度的分类，如表3所示。该表包括四种动态程度，涵盖了两个方向（前向和后向）和两种传播类型（形状和值）。在前向传递期间，每个算子根据其相关的动态程度使用其形状和值传递函数来分别推断其输出张量的形状（例如，$F_{ISDO}^{fs}$）和值（例如，$F_{ISDO}^{fv}$）。类似地，$F_{ISDO}^{bs}$和$F_{ISDO}^{bv}$是后向传递过程中形状和值的两个后向传递函数示例。图3说明了几个常见的例子。左侧（图3(a)）显示了一个包含四个前向传递的例子，它使用了三种类型的Update传递函数。类似地，右侧（图3(b)）显示了一个包含两个属于同一类型的后向传递函数的例子。值得注意的一点是，应用于算子的适当传递函数不仅取决于计算图，还取决于在RDP分析过程中推断出的常量，这些常量决定了算子的动态分类。Merge传递函数很简单——它根据图2中的格合并来自多个控制流分支的S-map和V-map。表4总结了RDP的关键组成部分。

Table 3: 不同算子类型的前向和后向传递函数说明。

Table 4: 秩和维度传播（RDP）的定义。

RDP求解过程。该方法如算法1所示，涉及将传递函数（F）沿两个方向迭代地应用于扩展计算图（G）。详细来说，算法1首先以深度优先顺序对计算图G中的节点（即算子）进行排序，并将每个节点的输出形状和值映射初始化为undef（第1至2行）。接下来，它通过对每个节点（n）的前驱节点的输出形状和值映射（即n的输入形状和值映射）应用前向传递函数来处理每个节点（第13行）。此外，如果任何前驱节点的分析结果为undef，它会通过后向传递函数将n的输出形状和值映射传播到n的前驱节点的输出形状和值映射（第14至15行）。这些前向和后向传递函数是根据DNN算子的动态分类定义的（如第20至32行所示）。算法1需要处理两种特定类型的节点（算子）：i) 控制流节点（如Combine或Switch），对此，它需要调用Merge函数来合并来自多个控制流路径的分析结果（第9至10行），以及 ii) “输入形状决定输出”节点，对此，它为值映射分配一个符号常量以方便后续分析（第16至18行）。算法1持续处理G中的节点，直到任何节点的形状/值映射不再更新。与其他数据流分析类似，RDP遵循格理论【28, A unified approach to global program optimization. 1973. POPL】，因此一个优化的混沌实现（基于工作列表）保证会收敛。

4.2 基于RDP的动态DNN算子融合

利用RDP进行算子融合。尽管融合已成为DNN上的一项成功优化【46, DNNFusion: accelerating deep neural networks execution with advanced operator fusion. 2021. PLDI】，但众所周知，在动态DNN上实现它非常困难【57, Nimble: Efficiently compiling dynamic neural networks for model inference. 2021. MLSys】。一个常见的问题是，在不知道两个算子的张量形状的情况下，DNN编译器要么根本无法融合它们，要么必须生成大量的代码版本，每个版本对应两个算子形状的一种可能组合。事实上，由于通常会合并两个以上的算子，需要单独生成代码的可能组合会迅速增加。我们提出的RDP分析可以通过使用（可能是符号的）形状信息来解决这个问题。诸如两个算子具有相同形状的张量之类的信息可以启用和/或简化融合，即使确切的维度在运行时才知道。

融合示例。图4展示了一个简化的例子，包含两个常见的DNN算子（Sigmoid和Add），它们操作的张量形状在运行时才可知。Sigmoid接受一个输入张量A，其动态形状为[I', J', K']。Add对Sigmoid的输出和另一个输入张量B执行逐元素加法，B的形状恰好是[I, J, K]。现在，如果A和B的形状不同，需要在逐元素加法之前立即对Sigmoid的输出张量进行形状广播操作。没有我们的RDP分析，A和B的动态形状（以及可能的形状广播操作）会阻止DNN编译器以高效的方式融合这两个算子，即编译器要么生成没有融合的代码（如图4蓝色框所示），要么生成多个代码版本（此例中为8个版本）并在运行时选择一个版本。假设我们的RDP分析结果是I' = I, J' = 1, K' = 1，即符号常量（I）和已知常量的混合，DNN编译器可以进一步生成一个唯一的融合代码版本（如图4绿色框所示）。SoD2将RDP和上述基于RDP的算子融合整合到一个先进的静态DNN算子融合框架（DNNFusion 【46, DNNFusion: accelerating deep neural networks execution with advanced operator fusion. 2021. PLDI】）中，以生成融合计划和优化的融合代码。

4.3 基于RDP的静态执行规划

执行规划的挑战。计算图（DAG）通常允许多种不同的算子执行顺序。顺序的选择会影响峰值内存使用（用于中间结果），进而对缓存性能和执行延迟产生影响。之前已有关于这个问题的工作，事实上已经表明生成一个最优的执行计划（按内存消耗等指标衡量）是一个NP完全问题【2, Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices. 2020. MLSys】。因此，对于具有成百上千个算子的现代大型DNN来说，选择一个最优计划可能很困难。

SoD2的启发式方法。动态属性（例如，动态形状和控制流）进一步使这个问题复杂化。在SoD2中，我们开发了一系列由提出的RDP分析驱动的启发式方法。总体思想是，由于全局最优解几乎不可行，基于图划分的方法是合理的。事实证明，RDP的结果能够指导图的划分和每个子图内解的选择。特别是，我们观察到，已知常量、符号常量、操作推断常量和⊥或nac会逐渐增加生成最优执行计划的障碍。更具体地说，对于一个包含有限数量算子的子图sg：

RDP驱动的启发式规则。
1. 首先，如果sg中所有张量的形状都是已知常量，那么sg的最优执行计划可以通过穷举搜索静态获得——sg的有限大小可以使这种搜索变得可行。
2. 其次，如果sg中张量的形状是已知常量、符号常量和操作推断常量的混合，仍然可以比较内存需求，从而生成一个（接近）最优的执行计划。如果这些形状是从同一组符号常量派生出来的，这一点尤其适用。
3. 第三，如果一个算子有一个nac（非常量）输出张量形状，它会禁用进一步的分析和执行规划。事实证明，这样的算子为将原始图划分为可以独立分析的子图提供了机会。

4.4 其他优化

4.4.1 内存分配计划

动态DNN内存规划。除了执行（顺序）规划，DNN的内存规划也是一个关键步骤【2, Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices. 2020. MLSys】【51, Efficient memory management for deep neural net inference. 2020. ArXiv】。一个内存分配计划决定了每个中间张量在线性内存空间中的分配位置以及何时被释放，可以限制峰值内存使用并提高执行性能——后者通过减少内存碎片、避免内存移动和限制内存分配/释放来实现。与执行规划（即使对于动态DNN也可以在编译时进行）相反，动态DNN的内存规划只能在执行时当所有张量大小都已知时才能进行。静态DNN执行的内存规划也已被证明是NP完全的【2, Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices. 2020. MLSys】，而DNN模型的动态性进一步使内存规划复杂化。

RDP驱动的内存规划策略。现有的动态DNN执行内存规划方法（例如，Nimble 【57, Nimble: Efficiently compiling dynamic neural networks for model inference. 2021. MLSys】）已经解决了这个问题。在不知道确切张量形状的情况下，这些方法通常依赖于贪心策略【51, Efficient memory management for deep neural net inference. 2020. ArXiv】，（例如，找到当前可用的能容纳新张量的最小内存槽）。相比之下，我们使用RDP结果和以下两个关键见解。首先，我们的方法基于静态执行规划方法生成的子图。事实证明，对于具有已知常量形状的子图，以及那些形状仅由子图的输入张量定义的符号/操作推断常量形状的子图，峰值内存需求可以从静态RDP分析结果和随后的执行计划生成中推断出来。其次，我们观察到，对于大多数子图，内存需求从图中使用峰值内存的位置向前向和后向方向单调递减。因此，从峰值内存消耗位置开始进行内存规划，并向前向和后向遍历，并选择可用的内存槽进行重用，是一种很好的策略，并且不会导致需要额外的内存空间。

轻量级贪心法。基于这些见解，一个从峰值内存需求位置开始的轻量级贪心方法可以帮助为许多/大多数子图找到最优的内存使用方案。我们的评估（由于空间限制细节省略）在ConvNet-AIG 【62, Convolutional Networks with Adaptive Inference Graphs. 2018. ECCV】上表明，我们基于RDP的内存分配计划需要最优峰值内存消耗的1.05倍（该结果来自穷举搜索）；而基于上述贪心策略的方法（MNN）需要最优峰值内存消耗的1.16倍。

4.4.2 基于RDP的多版本代码生成

减少代码版本数量。正如我们在4.2节中讨论的，RDP分析通过揭示（可能是符号的）张量形状来启用和/或简化算子融合。在单一（融合）版本不可行的情况下，通过RDP获得的信息的一个优点是可以显著减少生成的融合代码的不同版本的数量。

优化热点算子。SoD2通过生成多版本代码来优化主导DNN执行的热点算子（例如，CONV和GEMM），从而进一步受益于RDP的这一特性。先前的努力【1, TensorFlow: A system for large-scale machine learning. 2016. OSDI】【26, MNN: A Universal and Efficient Inference Engine. 2020. MLSys】已经表明，这些算子的优化机会取决于输入/输出张量的形状和大小。因此，对于静态DNN执行，现有框架（如TensorFlow Lite 【1, TensorFlow: A system for large-scale machine learning. 2016. OSDI】和MNN 【26, MNN: A Universal and Efficient Inference Engine. 2020. MLSys】）通常采用涉及不同优化的多版本代码（例如，分块、展开、选择线程块数量等）。然而，这种优化对于动态DNN来说是具有挑战性的，因为未知的张量形状和/或张量大小意味着需要太多的版本。RDP提供的张量形状（或形状关系）有助于仅为更具体的张量形状生成代码，从而导致更少的代码版本。

自动调优器。更具体地说，SoD2依赖于一个基于遗传算法的自动调优器来为内核代码生成生成探索空间（例如，分块形状、循环置换和展开设置），如同DNNFusion 【46, DNNFusion: accelerating deep neural networks execution with advanced operator fusion. 2021. PLDI】。这个自动调优器的一个特点是更有效地利用硬件中可用的并行性。为了应对动态形状的挑战，SoD2采用了一种多版本方法，其中版本的选择基于与不同形状对性能影响相关的经验证据。例如，我们的自动调优器为GEMM和CONV内核都考虑了胖、常规和瘦矩阵。

A4 实验环境

动态模型与数据集：
实验在10个动态模型上进行，分为三类：
1. 形状动态：StableDiffusion Encoder (SDE), SegmentAnything, Conformer, CodeBERT, YOLO-V6 (YL-V6)。
2. 控制流动态：DGNet。
3. 形状与控制流均动态：SkipNet (SNet), ConvNet-AIG (CNet), RaNet, BlockDrop (BDrop)。
数据集包括：ImageNet, MS COCO, SA-1B, Librispeech等。模型具体信息见表5。

硬件配置：
- 主要平台：三星Galaxy S21智能手机，搭载高通骁龙888处理器（八核Kryo 680 CPU，Adreno 660 GPU）。
- 可移植性测试平台：搭载高通骁龙835处理器的设备（八核Kryo 280 CPU，Adreno 540 GPU）。

软件配置：
- SoD2实现：基于一个仅支持静态DNN的执行框架DNNFusion [46]进行扩展。
- 对比框架：ONNX Runtime (ORT) V1.14.1, MNN Vdcb080c, TVM w/ Nimble extension (TVM-N) V7831a79, TFLite V2.11.1。
- 执行设置：移动CPU上使用8个线程，GPU上使用流水线执行。GPU执行使用16位浮点数，CPU执行使用32位浮点数。每个实验重复50次，报告平均值。

A4 实验结果

总体内存消耗对比 (表5)
在移动CPU上，SoD2的内存消耗显著低于其他框架。与SoD2相比，ORT、MNN和TVM-N的归一化几何平均内存消耗分别为3.64倍、1.37倍和8.62倍。SoD2在图像模型上内存节省更多，因为这些模型内存占用更大，优化空间也更大。值得注意的是，由于优化限制或算子缺失，其他框架无法支持所有模型，例如SegmentAnything。

Table 5: ONNX Runtime、MNN、TVM with Nimble extension (TVM-N) 和 SoD2 在移动 CPU 上的内存消耗（为中间结果分配的内存）对比。“-”表示该模型尚不被该框架支持。“S”代表形状动态，“C”代表控制流动态。

总体延迟对比 (表6)
在移动CPU上，SoD2相比ORT、MNN和TVM-N分别取得了2.5倍、1.7倍和2.7倍的平均加速。在移动GPU上，SoD2相比ORT和MNN分别取得了3.9倍和2.3倍的加速（TVM-N不支持移动GPU上的动态模型）。SoD2的优化能够有效处理不同执行路径变化带来的影响。表7显示，即使在不同输入分布下（以YOLO-V6为例），SoD2的加速比依然稳定。

Table 6: ONNX Runtime、MNN、TVM-N 和 SoD2 在移动 CPU 和移动 GPU 上的端到端执行延迟对比。“-”表示该模型尚不被该框架支持。

Table 7: 输入分布对YOLO-V6延迟的影响。每个单元格显示SoD2相对于ORT、MMN或TVM-N的相应基线的延迟加速比。

优化分解分析
- 内存减少 (图5)：在基线（No opt.）之上，RDP启用的算子融合（Fusion）带来18%-30%的内存减少，静态执行规划（SEP）额外带来22%-37%的减少，动态内存规划（DMP）再额外带来3%-7%的减少。
- 延迟降低 (图6)：在移动CPU上，Fusion、SEP、DMP和多版本代码生成（MVC）分别带来1.3-1.9倍、1.1-1.3倍、1.04-1.1倍和1.3-1.6倍的额外加速。在对内存和数据移动更敏感的移动GPU上，加速效果更明显。
- RDP融合效果 (图7)：与仅使用静态融合（SFusion）相比，RDP融合（RDP Fusion）能进一步将层数减少16%-46%，并将中间结果（IR）大小额外减少13%-40%。
- 子图分析 (图8)：超过90%的子图属于可以被SoD2优化的“全已知常量”或“混合常量”类别，证明了执行和内存规划优化的广泛适用性。

进一步性能分析
- 相同执行路径对比 (图9)：在禁用动态分支选择，与MNN进行同路径比较时，SoD2依然能在CPU上实现1.5-2.0倍的加速和1.2-1.5倍的内存减少，验证了其核心优化的有效性。
- 不同输入尺寸对比 (图10)：对于YOLO-V6，随着输入尺寸增大，SoD2的延迟增长更平缓，表现出比MNN更好的性能和稳定性。
- 固定内存预算对比 (图11)：当TFLite的内存预算被限制为与SoD2相同时（通过重物质化处理内存不足），SoD2的性能优势更大。
- 静态模型对比 (图12)：在处理静态模型时，与专为静态模型优化的基线DNNFusion相比，SoD2有3%-7%的轻微性能开销，这归因于动态内存规划的开销和相对保守的融合策略。

可移植性 (图13)
在更早一代、资源更受限的骁龙835平台上，SoD2取得了相似甚至更高的加速比，证明了其内存优化在资源受限设备上的重要性，以及框架良好的可移植性。

A7 补充细节

6 相关工作

动态神经网络优化。类型分析和类型推断【8, Flexible type analysis. 1999. ICFP】【22, Compiling polymorphism using intensional type analysis. 1995. POPL】【33, Mlir: Scaling compiler infrastructure for domain specific computation. 2021. CGO】【44, A theory of type polymorphism in programming. 1978. JCSS】【58, Gradual typing for objects. 2007. ECOOP】被广泛用于分析张量形状，从而辅助动态神经网络优化。Nimble 【57, Nimble: Efficiently compiling dynamic neural networks for model inference. 2021. MLSys】是一个已集成到TVM中的基于编译的动态DNN框架。该框架依赖昂贵的动态函数在运行时解释动态形状。这种实现（我们已对其进行了广泛比较）限制了优化代码生成的机会，例如执行算子融合。DISC 【74, DISC: A dynamic shape compiler for machine learning workloads. 2021. MLSys Workshop】扩展了MLIR-HLO 【33, Mlir: Scaling compiler infrastructure for domain specific computation. 2021. CGO】并为具有某些约束（例如，相同维度（激活函数的情况）和相同大小（转置的情况））的算子传播形状信息。SoD2基于动态程度提供了更全面的算子分类，带来了显著增强的优化机会。Axon 【6, Axon: A Language for Dynamic Shapes in Deep Learning Graphs. 2022. ArXiv】是一种允许为计算图的输入和输出张量指定符号形状的编程语言。它使用约束求解器来找到形状，而SoD2使用前向和后向数据流分析（RDP），这也减轻了程序员的额外参与。此外，SoD2还包括一系列由RDP启用的优化。与SoD2关系不那么密切的，DietCode 【73, DietCode: Automatic optimization for dynamic tensor programs. 2022. MLSys】提出了一个基于TVM的自动调度器框架，用于动态形状。该框架建立了一个成本模型来预测运行时性能，并减少搜索空间以找到最优的运行时参数（例如，循环分块）。Cortex 【15, Cortex: A compiler for recursive deep learning models. 2021. MLSys】、Cavs 【67, Cavs: An efficient runtime system for dynamic neural networks. 2018. USENIX ATC】和另一项工作【25, Improving the expressiveness of deep learning frameworks with recursion. 2018. EuroSys】主要旨在解决神经网络的递归动态性，这与SoD2的重点不同。其他工作则侧重于推理的动态批处理【14, TurboTransformers: an efficient GPU serving system for transformer models. 2021. PPoPP】【17, Low latency RNN inference with cellular batching. 2018. EuroSys】【41, Deep Learning with Dynamic Computation Graphs. 2017. ICLR】【72, ByteTransformer: A High-Performance Transformer Boosted for Variable-Length Inputs. 2022. ArXiv】或专为动态DNN训练设计【45, Dynet: The dynamic neural network toolkit. 2017. ArXiv】。

DNN执行与内存优化。存在一些关于算子执行顺序调度的研究，如【2, Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices. 2020. MLSys】【38, Neural networks on microcontrollers: saving memory at inference via operator reordering. 2019. ArXiv】【39, MCUNet: Tiny Deep Learning on IoT Devices. 2020. NeurIPS】。其中，【38, Neural networks on microcontrollers: saving memory at inference via operator reordering. 2019. ArXiv】【39, MCUNet: Tiny Deep Learning on IoT Devices. 2020. NeurIPS】侧重于通过重新排序算子来最小化资源受限设备（如MCU）的峰值内存消耗，而工作【2, Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices. 2020. MLSys】为复杂模型（不规则连接的神经网络）提出了一个优化的调度框架。这些方法仅依赖于静态形状。最近也有一些关于优化DNN内存分配规划和内存管理的工作。诸如【35, Memory Planning for Deep Neural Networks. 2022. ArXiv】【51, Efficient memory management for deep neural net inference. 2020. ArXiv】等工作仅为静态DNN设计了各种启发式内存规划算法。TelaMalloc 【43, TelaMalloc: Efficient On-Chip Memory Allocation for Production Machine Learning Accelerators. 2023. ASPLOS】为具有已知中间张量形状和大小的静态控制流图动态执行内存管理。它没有完全考虑DNN的控制流动态性和动态形状。未来的一项可能工作是将我们的RDP分析与TelaMalloc的启发式方法和基于求解器的方法相结合，以进一步改进我们的内存规划。当可用内存有限时，重物质化（rematerialization）【24, Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization. 2020. MLSys】【30, Dynamic Tensor Rematerialization. 2021. ICLR】和重计算（recomputation）【3, In-Place Activated BatchNorm for Memory-Optimized Training of DNNs. 2018. CVPR】方法在内存消耗和执行延迟之间实现了权衡。这些方面未来可以考虑用于动态DNN。

移动端DNN推理引擎。近年来，支持移动设备上的DNN推理已成为一个活跃的研究领域。诸如MCDNN 【21, Mcdnn: An approximation-based execution framework for deep stream processing under resource constraints. 2016. MobiSys】、DeepX 【32, DeepX: A Software Accelerator for Low-Power Deep Learning Inference on Mobile Devices. 2016. IPSN】、DeepMon 【23, Deepmon: Mobile gpu-based deep learning framework for continuous vision applications. 2017. MobiSys】、DeepSense 【69, DeepSense: A Unified Deep Learning Framework for Time-Series Mobile Sensing Data Processing. 2017. WWW】和DeepCache 【66, DeepCache: Principled Cache for Mobile Deep Vision. 2018. MobiCom】等工作主要集中在优化具有静态形状和控制流的静态DNN的执行。TensorFlow Lite (TFLite) 【1, TensorFlow: A system for large-scale machine learning. 2016. OSDI】、Pytorch-Mobile 【50, PyTorch: An Imperative Style, High-Performance Deep Learning Library. 2019. NeurIPS】、TVM 【5, TVM: An automated end-to-end optimizing compiler for deep learning. 2018. OSDI】和MNN 【26, MNN: A Universal and Efficient Inference Engine. 2020. MLSys】通过重新初始化或保守（最大）内存分配来支持动态形状。它们要么不支持动态控制流，要么需要执行所有路径并剥离无效结果。如我们的评估所示，这些方法引入了很高的运行时开销。之前的一个静态DNN系统DNNFusion 【46, DNNFusion: accelerating deep neural networks execution with advanced operator fusion. 2021. PLDI】也涉及DNN算子的分类，然而，这里介绍的分类是正交的。

7 讨论与未来工作

推广到其他平台。所提出的技术，如RDP分析、基于RDP的融合以及执行和内存规划，对各种平台（包括数据中心GPU）具有广泛的适用性。这对于单输入推理场景尤其如此。可能出现的一个细微差别是数据中心GPU和移动GPU在执行批量推理能力上的区别。与移动GPU不同，数据中心GPU能够同时处理多个输入，从而最大化其计算能力。然而，批次中的不同输入样本可能需要使用不同的执行路径。因此，将动态批处理与动态神经网络相结合是未来研究的一个潜在方向。

处理LLM的可扩展性。SoD2中的优化也可以应用于大规模语言模型（LLM）。我们采用的主要程序之一是图划分，如4.3节所述。该过程涉及将整个计算图划分为一组子图，每个子图包含有限数量的层。每个子图的最优解是离线确定的。然而，语言模型（LLM）的特点是参数数量极其庞大，达到数十亿【7, Free Dolly: Introducing the World’s First Truly Open InstructionTuned LLM. 2023. Databricks Blog】【60, Llama 2: Open foundation and fine-tuned chat models. 2023. ArXiv】【70, Glm-130b: An open bilingual pre-trained model. 2022. ArXiv】。这在计算和资源需求方面对移动设备构成了重大挑战。我们未来的工作将通过将SoD2与模型剪枝和量化方面的进展【27, F8net: Fixed-point 8-bit only multiplication for network quantization. 2022. ArXiv】【47, Patdnn: Achieving real-time dnn execution on mobile devices with pattern-based weight pruning. 2020. ASPLOS】【64, SparCL: Sparse Continual Learning on the Edge. 2022. NeurIPS】相结合来增强SoD2，以实现更好的性能。

扩展到ONNX之外。算子分类和相关的优化设计也不限于ONNX或其他推理格式（例如，TFLite，Caffe2）。这是因为我们提出的分析是基于算子的计算逻辑及其输入和输出之间的关系所定义的动态程度，而不是依赖于算子的特定表示或格式。一些格式尚未完全支持动态计算图。例如，PyTorch支持将具有动态形状（如输入形状决定输出、输入形状决定输出形状、输入形状和值决定输出形状）的模型导出到ONNX。然而，它无法将具有动态控制流的模型转换为ONNX。为了解决这个限制，我们添加了一对自定义的ONNX算子<Switch, Combine>（如图1d所示），并在PyTorch上为具有动态控制流的模型专门注册了一个自定义的导出例程。SoD2在处理可以在PyTorch中很好表示的非常复杂（或用户定义）的动态模型（如图神经网络或涉及递归执行的DNN）方面确实存在局限性。我们将这一进一步的优化作为未来的工作。

A5 结论

本文提出了一个用于优化DNN的综合框架SoD2。SoD2将动态DNN的常见算子分为四种类型，并包含一种新颖的静态数据流分析（RDP）。在此基础上，SoD2实现了一系列由RDP赋能的动态DNN优化，包括算子融合、静态执行（顺序）规划、动态内存分配规划和多版本代码生成。SoD2在一个移动系统上对10个新兴的动态DNN进行了广泛评估，评估结果表明，与四个最先进的DNN执行框架相比，它节省了高达88%的内存消耗，并带来了高达3.9倍的执行加速。由于底层技术是通用的，也适用于其他设备，我们未来的工作将在其他设备（如边缘GPU和树莓派）上评估SoD2的功效。