ACROBAT: OPTIMIZING AUTO-BATCHING OF DYNAMIC DEEP LEARNING AT COMPILE TIME

• 文章标题：ACROBAT：在编译时优化动态深度学习的自动批处理
• 作者/机构：Pratik Fegade 1, Tianqi Chen 1 2, Phillip B. Gibbons 1, Todd C. Mowry 1

A1 主要贡献

本文提出了 ACROBAT，一个为动态深度学习计算设计的自动批处理（auto-batching）框架，它通过新颖的混合静态+动态优化和端到端的张量核代码生成来解决现有方法的局限性。

• 核心问题：对于包含动态控制流的深度学习（DL）计算，批处理（Batching）是一项重要的性能优化，但手动实现既困难又容易出错。现有的自动批处理框架通常面临两难：要么严重依赖动态分析，导致高昂的运行时开销（如 DyNet）；要么针对特定模式进行静态优化，缺乏通用性（如 Cortex）。此外，这些框架过度依赖供应商库（如 cuDNN），限制了诸如核融合等优化。

• 研究目标：设计一个通用的自动批处理框架，既能减少动态分析的运行时开销，又能生成高效的、针对特定计算优化的张量核，从而在支持广泛的动态控制流模式的同时实现高性能。

• 主要创新点与贡献：

  1. 调研与表征：对不同深度学习计算中存在的动态控制流结构进行了调研和特征分析。
  2. ACROBAT 框架设计：设计了 ACROBAT，一个利用新颖的混合静态+动态优化和自动化端到端核代码生成的自动批处理框架。
     • 混合分析：核心思想是，即使在编译时缺乏完整的执行信息，编译器仍可通过静态分析（如上下文敏感性分析、污点分析）来辅助动态分析，从而减少运行时开销并有效利用并行性。
     • 端到端代码生成：ACROBAT 不依赖供应商库，而是通过端到端的张量核生成，利用静态分析识别并利用数据重用机会，生成针对整体计算进行优化的专用核，例如融合内存 gather 操作。
  3. 原型实现与评估：实现了 ACROBAT 的原型，并与最先进的深度学习框架（如 DyNet, Cortex, PyTorch）在Nvidia GPU上进行了评估，取得了显著的性能提升，最高比 DyNet 快 8.5 倍。

下表是 ACROBAT 与相关工作的定性比较：
Table 1. ACROBAT 与其他动态 DL 计算自动批处理方案的比较。纯静态或纯动态方法可能分别过于保守或开销过高，而 ACROBAT 的混合分析则不同。

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

2.1 深度学习计算中的动态控制流

DL计算中的动态控制流种类。本节从自动批处理问题的角度，探讨了各种深度学习计算中存在的不同类型的控制流动态性。这将为我们如何设计一个系统来利用动态深度学习计算批处理执行中张量算子间的并行性提供信息。
Table 2. 深度学习计算中发现的控制流属性。图例：ITE：迭代控制流，REC：递归控制流，TDC：模型表现出张量依赖的控制流（控制流决策取决于中间张量的值），IP：计算表现出高实例并行性，ICF：模型推理表现出控制流，TCF：模型训练表现出控制流。

实现的自然性。需要注意的是，对于一个涉及控制流的计算，通常有多种实现方式。我们考虑最自然的实现方式。例如，一个自顶向下的树遍历可以实现为广度优先遍历（BFS）或深度优先遍历（DFS）。虽然 BFS 遍历可能更高效，但基于 DFS 的遍历实现起来更自然。下面的讨论也总结在表2中。

静态子图周围的控制流。我们观察到，对于大多数表现出控制流动态性的深度学习计算，动态控制流都围绕着张量计算。以Listing 1中@rnn函数实现的简单序贯RNN模型为例。在这里，我们看到序贯控制流围绕着第5行和第6行的RNN单元，这是一个没有中间控制流的静态张量计算子图。

张量依赖的控制流。在深度学习计算中，控制流决策通常依赖于中间张量的值。这类模型和计算的例子包括机器翻译中的束搜索（beam search）、StackLSTMs【23, Dyer et al. Transition-based dependency parsing with stack long short-term memory. 2015. CoRR】、Tree-to-Tree神经网络（T2TNN）【15, Chen et al. Tree-to-tree neural networks for program translation. 2018b. CoRR】、带有提前退出的模型【42, Kaya & Dumitras. How to stop off-the-shelf deep neural networks from overthinking. 2018. CoRR; 81, Teerapittayanon et al. Branchynet: Fast inference via early exiting from deep neural networks. 2017. CoRR; 86, Xin et al. Deebert: Dynamic early exiting for accelerating BERT inference. 2020. CoRR; 24, Elbayad et al. Depth-adaptive transformer. 2019. CoRR】和混合专家模型（Mixture-of-Experts）【71, Shazeer et al. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. 2017. CoRR; 52, Ma et al. Modeling task relationships in multi-task learning with multi-gate mixture-of-experts. 2018. KDD; 25, Fedus et al. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. 2021. CoRR】。与此同时，在TreeLSTM【76, Socher et al. Recursive deep models for semantic compositionality over a sentiment treebank. 2013a. EMNLP】、DAG-RNN、序贯RNN及其变体等模型中，控制流仅依赖于输入，而不依赖于中间张量。

重复性控制流。如果一个模型可以表示为迭代或递归计算，我们就说它表现出重复性控制流。这包括迭代模型，如RNN及其变体（例如LSTM和GRU【17, Cho et al. Learning phrase representations using RNN encoder-decoder for statistical machine translation. 2014. CoRR】）和StackLSTMs，以及递归模型，如TreeLSTM、Tree-to-Tree神经网络和DAG-RNNs【73, Shuai et al. Dagrecurrent neural networks for scene labeling. 2015. CoRR】。另一方面，混合专家模型和提前退出模型不表现出重复性控制流。这类模型在其他方面是静态的前馈网络中包含条件执行。重复性控制流通常也可以是嵌套的。例如，GraphRNN模型执行两个RNN，一个嵌套在另一个内部。类似地，用于自顶向下递归树生成的DRNN模型，涉及到为给定树节点迭代生成子节点。

递归与迭代控制流的差异。递归控制流的存在，相较于迭代控制流，常常会使静态分析变得复杂，因为后者更容易利用并行性。我们在§4.2中看到，例如，在运行时利用递归调用间的并行性，可能需要多个并发执行上下文，类似于fork-join并行范式【53, McCool et al. Chapter 8 - fork–join. In Structured Parallel Programming, 2012】。

训练和推理中的控制流。我们在表2中看到，许多模型的计算在训练和推理期间都涉及动态控制流。然而，对于带有提前退出的模型来说，情况并非如此。在训练期间，我们通常希望训练所有的退出分支，而不是像推理时那样只评估一个。此外，像束搜索这样的搜索过程通常只在推理时使用，因此底层模型在训练期间可能不表现出动态性（除非模型计算本身涉及动态性，例如在RNN模型的情况下）。

控制流并行性。动态控制流可以在深度学习计算中引入并行性。这样的计算可能表现出（1）批处理并行性（Batch Parallelism），存在于小批量处理中不同输入实例之间，和/或（2）实例并行性（Instance Parallelism），指由于动态控制流依赖而产生的并行性，例如递归并行性。这种并行性的数量在不同计算中差异很大。递归模型通常（虽然不总是）在不同的递归调用之间具有显著的并行性。相应地，迭代计算可能包含可以并发执行的循环。一个例子是Listing 1中RNN实现中对@map函数的调用。

2.2 动态批处理

动态批处理技术。ACROBAT 建立在动态批处理（dynamic batching）【51, Looks et al. Deep learning with dynamic computation graphs. 2017. CoRR; 56, Neubig et al. On-the-fly operation batching in dynamic computation graphs. 2017b.】的基础上，这是一种在存在动态控制流的情况下执行自动批处理的先前技术。给定一小批输入实例，动态批处理涉及为每个输入实例懒惰地执行模型计算，同时在后台为每个实例构建张量算子的数据流图（DFG）。当请求特定张量的值时（例如，当模型包含张量依赖的控制流时），会触发这些DFG的执行。在此执行期间，运行时可以识别DFG内的批处理机会，并适当地启动批处理后的核函数。

ACROBAT：概览与API

控制流的动态性使得自动批处理必须依赖可能开销高昂的运行时分析。在ACROBAT中，我们观察到，积极的静态分析通常能提供足够的信息来降低这类分析的开销。这样的分析进一步使我们能够以端到端的方式生成专用且更高效的张量核。

我们现在来看一下ACROBAT利用上述见解的编译和执行工作流程（如图1所示）。ACROBAT被设计为接受一个以简单的图灵完备函数式语言表示的非批处理DL计算作为输入。这使得ACROBAT用户可以轻松表达具有动态控制流的模型，例如§2.1中讨论的那些。例如，Listing 1展示了一个简单的RNN模型，ACROBAT可以将其作为输入。

给定一个输入计算1，ACROBAT的编译从批处理核生成2开始。在这里，ACROBAT执行新颖的静态分析（§5.1）以识别数据重用机会，并相应地生成批处理核3，实现输入程序中使用的张量算子。此外，收集算子融合（gather operator fusion）（§5.2）使我们能够生成专门的核，从而最小化数据移动。这些未优化的核随后由一个自动调度器4进行优化。优化后，可以为批处理核生成目标代码10，如CUDA C++。同时，输入程序被进一步优化和编译5，以预先（AOT）的方式生成C++代码7。作为此编译的一部分，ACROBAT生成的代码能够（1）通过我们的内联深度计算方法实现低开销调度，以及（2）在存在张量依赖控制流的情况下自动启用并发执行（§4.2）。

图 1. ACROBAT 工作流程概览。附录中的图 7 展示了 DyNet（一种先前的完全动态方法）的相应概览。请注意 ACROBAT 如何在编译时执行大量新颖的分析和代码生成，以减少运行时开销。

在运行时，ACROBAT 在一小批输入6上懒惰地执行AOT编译的输入程序7，并构建数据流图（DFG）8。然后，ACROBAT运行时库将调度这些DFG（使用上面提到的内联深度计算）9，同时寻找批处理机会。接着，它将为每个识别出的DFG节点批次调用优化后的批处理核10。如果输入程序表现出张量依赖的控制流，执行将循环回到AOT编译的程序，该程序将进一步执行并创建更多的DFG。

我们现在将在§4中看一下ACROBAT的混合优化，在§5中看一下它的张量核生成。

A2 方法细节

4 混合静态+动态优化

动态控制流通常会阻碍静态程序转换。因此，ACROBAT采用了一种混合方法，通过以下方式利用静态程序知识：（1）为动态分析提供提示（§4.1），或（2）生成代码，使动态分析在利用并行性方面有更大的自由度（§4.2）。此外，静态分析还使我们能够执行诸如核融合之类的优化，这对于高性能至关重要（§7.4）。下面，我们提供有关我们混合分析的更多细节。

4.1 内联深度计算

降低调度开销。正如先前工作【26, Fegade et al. Cortex: A compiler for recursive deep learning models. 2021. MLSys】所指出的，以往的完全动态方法会产生显著的调度开销。例如，我们将在表5中展示，对于TreeLSTM模型，DyNet的调度开销主导了张量计算的时间。相反，如下所述，ACROBAT设计了一种在构建DFG时执行调度的方案，从而极大地降低了调度开销（§7）。

调度算法目标。一个DFG调度算法有两个目标：G.1 正确性：调度任务时尊重任务间的依赖关系。G.2 性能：识别并利用并行性。给定一个或多个DFG，我们可以通过按拓扑深度的递增顺序执行DFG节点（每个节点代表一个张量算子）来满足这两个目标，使得相同深度的节点可以并发执行【55, Neubig et al. Dynet: The dynamic neural network toolkit. 2017a; 51, Looks et al. Deep learning with dynamic computation graphs. 2017. CoRR】。为了在DFG构建期间计算这些深度，我们做出以下两个观察：

深度计算观察。O.1 未批处理程序调用张量算子的顺序，即节点被添加到DFG的顺序，是一个有效的依赖顺序。O.2 关于实例并行性的信息（例如，如表2所示的TreeLSTM模型中的递归并行性）通常在编译期间是可用的。

深度计算实现。基于这些观察，我们将一个算子的深度设置为其在未批处理程序执行所引发的依赖顺序中的位置，从而满足目标G.1。然后，我们依赖上述观察O.2，通过使用以下技术来发现和利用并行机会：

/* 清单 1：RNN 模型实现。为简洁起见，省略了类型注解。 */
def @rnn(inputs, state, iweight, hweight, bias) {
  match(inputs) {
  Nil -> Nil,
  Cons(input, tail) -> {
    let input_transformed = nn.dense(input, iweight);
    let new_state = sigmoid(input_transformed + nn.dense(state, hweight));
    Cons(new_state, @rnn(tail, new_state, iweight, hweight, bias))
  }}
}
def @main(%data: Tensor[(?, 128), "float32"],
           %rnn_iweight: Tensor[(128, 128), "float32"],
           %rnn_hweight: Tensor[(128, 128), "float32"],
           %rnn_bias: Tensor[(128), "float32"],
           %rnn_init: Tensor[(128), "float32"],
           %out_weight: Tensor[(128, 1), "float32"],
           %out_bias: Tensor[(1), "float32"]) {
  let %encoded_states = @rnn(%data, %rnn_init, %rnn_iweight,
                            %rnn_hweight, %rnn_bias);
  let %out_fn = fn (%state) {
    nn.bias_add(nn.dense(%state, %out_weight), %out_bias)
  };
  @map(%out_fn, %encoded_states)
}

实例并行性。我们注意到实例并行性通常源于递归或在独立项列表上使用函数式@map函数（观察O.2）。我们确保这类并发算子在未批处理程序执行期间被赋予相同的深度。我们依赖简单的用户注解来获取关于递归并行性的信息。图2展示了一个例子，其中对funA的两个递归调用被注解为并发。还需注意，以往的自动并行化工作【35, Hogen et al. Automatic parallelization of lazy functional programs. 1992. ESOP; 4, Aleen & Clark. Commutativity analysis for software parallelization: Letting program transformations see the big picture. 2009. SIGARCH Comput. Archit. News】可能可以替代此类注解。清单2展示了为清单1中的RNN模型生成的AOT编译代码。我们在第23行看到，@map函数内所有对relu_bias_dense核的调用都被赋予了相同的深度。

/* 清单 2：为清单 1 中 RNN 模型生成的 AOT 编译输出，内联深度计算代码已高亮显示。*/
AcrobatVal rnn(AcrobatVal inputs, AcrobatVal state, ...) {
  if (inputs.is_nil()) return nil;
  ...
  // 静态确定的深度
  acrobat_rt.invoke_kernel("bias_dense", 0, ...);
  ...
  return rnn(tail, new_state, ...);
}

AcrobatVal main(...) {
  ...
  AcrobatVal encoded_states = rnn(data, ...);
  ...
  auto out_fn = [&](AcrobatVal state) {
    acrobat_rt.invoke_kernel("relu_bias_dense",
      acrobat_rt.curr_depth, // 动态确定的深度
      ...);
  };
  // @map 函数的 AOT 编译输出
  // 所有 out_fn 的调用将被赋予相同的深度
  acrobat_rt.curr_depth++;
  for(auto state: encoded_states) {
    out_fn(state);
  }
  acrobat_rt.curr_depth++;
  ...
}

图 2. 并发调用注解。

应对深度调度的急切性。正如先前工作【56, Neubig et al. On-the-fly operation batching in dynamic computation graphs. 2017b.】所指出的，像ACROBAT使用的这种基于深度的调度方案，在执行张量算子时常常会过于急切，导致利用的并行性次优。先前的工作依赖于基于议程的调度（agenda-based scheduling）【56, Neubig et al. On-the-fly operation batching in dynamic computation graphs. 2017b.】，这是一种更昂贵的调度方案，作为深度方案的替代来缓解此问题。ACROBAT则依赖于编译时分析，如下所述。

幽灵操作。在存在条件if语句的情况下，急切的批处理会导致次优的批处理，如图3的上半部分所示。我们看到急切的批处理导致了一个次优的批处理调度，因为输入Inp1和Inp2的操作B实例被急切地批处理，并且更重要的是，与输入Inp3和Inp4的操作B实例分开了。在这种情况下，ACROBAT可以静态地插入幽灵操作，以实质上延迟某些算子的调度和执行，如图的下半部分所示。注意，幽灵操作仅影响ACROBAT的调度行为，在张量核执行期间被忽略。

图 3. 幽灵算子可以实现更好的批处理。

程序阶段。另一方面，当存在重复性（递归或迭代）控制流时，我们依赖于程序阶段（program phases）【72, Sherwood et al. Phase tracking and prediction. 2003. SIGARCH Comput. Archit. News】来解决上述调度的次优性问题。在知道这些程序阶段的情况下，ACROBAT会等待调度和执行一个阶段中的算子，直到所有先前阶段的算子都已被调度和执行。我们发现，将输入DL计算的各个语义阶段视为独立的阶段是将计算划分为阶段的一个良好启发式方法。ACROBAT也提供了一种方式让用户通过手动注解程序阶段来覆盖这个启发式方法，尽管在我们的评估中，我们并不需要这样的注解。我们在附录§A.3中提供了关于程序阶段和幽灵操作的更多细节和解释。

静态算子提升。此外，ACROBAT也能够静态地提升算子，我们在附录§A.1中有更详细的描述。例如，在清单2中，第5行对核bias_dense的调用被赋予了一个静态计算的深度0，这在运行时有效地将核调用提升到了递归之外。

4.2 张量依赖控制流

张量依赖控制流的挑战。ACROBAT懒惰地执行未批处理的程序，为批次中的每个输入实例创建DFG。在没有张量依赖控制流的情况下，我们可以首先为每个实例顺序执行未批处理的程序，然后一次性触发所有DFG的批处理和执行。然而，在存在张量依赖控制流的情况下，这种顺序执行将无法利用任何批处理并行性，因为我们必须在依赖于中间张量值的控制流决策点触发执行。虽然先前的工作将重构输入计算以缓解此问题的负担放在用户身上，ACROBOT却能自动生成代码，通过使用协程（fibers）为每个输入实例并发地执行未批处理的程序。

ACROBAT的并发执行方案。通过这种方式，可以为每个实例执行未批处理的程序，直到没有任何一个可以继续进行而不触发DFG的求值。此时，可以执行求值，然后在之后恢复并发执行，如图4所示。相应地，为了在存在张量依赖控制流的情况下利用实例并行性，ACROBAT会启动并发的协程，类似于fork-join并行模型【53, McCool et al. Chapter 8 - fork–join. In Structured Parallel Programming, 2012】。因此，ACROBAT将其混合分析的一部分，结合了静态的并行性知识和动态的并发执行，以在存在张量依赖控制流的情况下有效利用并行性。

图 4. 在存在张量依赖控制流的情况下，并发执行未批处理的程序。

5 端到端张量核生成

正如我们上面提到的，ACROBAT通过避免使用供应商库，实现了端到端、统一和自动的张量核代码生成。这使得ACROBAT能够在不需要额外编译器开发工作的情况下支持更大范围的算子。下面提供了有关ACROBAT张量核生成的更多细节。

5.1 利用参数重用

参数重用挑战。给定输入的未批处理计算，ACROBAT需要生成实现计算中使用的张量算子的批处理核。生成这些核并非易事，因为一些输入张量（通常是模型参数）可能在算子调用之间共享。例如，在图1的输入计算1中使用的元素加法算子add3的多次调用中，bias参数将被共享（因为它是一个模型参数），因此应该在参数input和state的所有值之间重用。这可以在图1中相应的批处理核（3和10）中看到。

ACROBAT的静态分析方法。一个完全动态的自动批处理方法，如DyNet中使用的方法，无法准确识别这种参数重用，而是依赖于启发式方法，这可能很脆弱，导致次优性能（§7.3）。另一方面，ACROBAT使用1-上下文敏感的污点分析来识别张量算子的这类共享参数。这里使用静态分析使得ACROBAT能够获得关于参数重用模式的准确知识。

进一步的数据重用机会。除了上述分析，ACROBAT还通过采用代码复制和水平融合来进一步探索数据重用的机会，详见§B.1。

5.2 融合内存Gather操作

融合内存Gather操作。由于ACROBAT动态地识别跨DFG的批处理机会，一个批次中所有DFG节点的输入张量在加速器内存中可能不是连续布局的。在这种情况下，先前的工作在操作张量之前执行内存收集（gather）（通过调用供应商库核），导致大量数据移动（§7.4）。相反，ACROBAT生成专门的批处理核，直接在内存中分散的张量上操作，实际上是将昂贵的收集操作与批处理核融合在一起。图1中生成的批处理核10说明了这一点。这种融合可以带来显著的性能提升，如§7所示。

A4 实验环境

• 模型与数据集：评估中使用了7个模型，如Table 3所示，包括TreeLSTM、MV-RNN、BiRNN、NestedRNN、DRNN、Berxit和StackRNN。数据集包括斯坦福情感树库、XNLI等，部分模型使用随机生成的数据。
  
• 模型参数：每个模型评估了两种尺寸（小尺寸和 大尺寸）。对于TreeLSTM、BiRNN、NestedRNN、StackRNN，小尺寸和大尺寸的隐藏层大小分别为256和512。MV-RNN的小、大尺寸隐藏层大小分别为64和128。Berxit的小尺寸模型与BERT-BASE超参数相同，大尺寸模型与BERT-LARGE（18层）超参数相同。
• 硬件配置：实验在一台Linux工作站上进行，配备一个AMD Ryzen Threadripper 3970X CPU（64个逻辑核心，支持2路超线程）和一个Nvidia RTX 3070 GPU。
• 软件配置：
  • 操作系统：Ubuntu 20.04
  • 依赖库：CUDA 11.1, cuDNN 8.0.5, Eigen library (v3.3.90)。
  • 比较框架：DyNet (commit 3e1b48c7), PyTorch (v1.9.0a0+gitf096245)。
  • ACROBAT实现：基于TVM v0.9.dev0构建。

A4 实验结果

AOT编译的益处 (7.2节)
- 实验内容：比较了ACROBAT在使用Relay VM和其AOT编译器时的性能。
- 实验结果：如Table 4所示，对于TreeLSTM、MV-RNN和BiRNN模型，AOT编译相比VM执行带来了高达13.45倍的性能提升。这证明了VM在处理动态控制流时会引入巨大的开销。后续所有实验均开启AOT编译。

与PyTorch的性能比较 (7.3节)
- 实验内容：将ACROBAT与不进行自动批处理的PyTorch在TreeLSTM, MV-RNN, BiRNN模型上进行比较。
- 实验结果：如Fig. 5所示，ACROBAT取得了显著的加速。原因在于ACROBAT能够利用实例和批处理并行性，而PyTorch不能。此外，ACROBAT的核融合等静态优化也贡献了性能提升。对于小模型，加速比更高，因为并行性利用的相对重要性更大。

与DyNet的性能比较 (7.3节)
- 实验内容：将ACROBAT与完全动态的自动批处理框架DyNet在所有7个模型上进行比较。
- 实验结果：如Table 5所示，ACROBAT在大多数情况下都优于DyNet，在所有模型配置上平均性能高出2.3倍。Table 6的细分数据显示，ACROBAT在DFG构建、调度、内存拷贝和CUDA API调用上的开销都显著低于DyNet，这得益于静态核融合、粒度粗化和内联深度计算等优化。
- 分析结论：
- 准确的参数重用推断：ACROBAT的静态分析能准确识别参数重用，而DyNet的启发式方法可能失效（如MV-RNN），导致性能低下。
- 更广的算子覆盖：ACROBAT的端到端核生成支持更多算子的批处理（如argmax），而DyNet依赖的库不支持，导致串行执行（如StackRNN, DRNN）。
- 自动处理张量依赖控制流：ACROBAT能自动利用协程（fibers）发掘实例并行性（如DRNN），而DyNet无法做到。Table 7显示，在手动修复DyNet的这些问题后，其性能有显著提升，但仍不及ACROBAT。

与Cortex的性能比较 (7.3节)
- 实验内容：将ACROBAT与专门用于递归计算的Cortex框架在TreeLSTM, MV-RNN, BiRNN模型上进行比较。
- 实验结果：如Table 8所示，这不是一个完全公平的比较。Cortex高度特化，不支持通用控制流，且需要大量手动优化。对于TreeLSTM和BiRNN，Cortex性能更高（最高1.87倍），因为它能进行更激进的融合。但在MV-RNN上，ACROBAT因其更灵活的接口而远超Cortex。
- 分析结论：ACROBAT在提供与特化框架相当性能的同时，支持更广泛的模型且开发成本低得多。

各项优化的益处 (7.4节)
- 实验内容：逐步开启ACROBAT的各项优化，评估它们对性能的贡献。
- 实验结果：如Fig. 6所示：
- 标准核融合：对所有模型都有显著好处。
- 粒度粗化和内联深度计算：对控制流密集的模型（如TreeLSTM, MV-RNN）最有效，因为它们减少了调度开销。
- 程序阶段/幽灵算子：对特定控制流模式（如BiRNN中的输出变换，StackRNN中的条件语句）有明显效果。
- Gather算子融合：效果不一。对于迭代执行且实例并行性较少的模型，可能因引入间接内存访问而导致性能下降。
- 分析结论：对于控制流较少或张量计算量大的模型（如Berxit），减少调度开销的优化带来的好处较小。

A5 结论

本文介绍了ACROBAT，一个用于自动批处理动态深度学习计算的编译器和运行时框架。ACROBAT的核心优势在于其创新的混合静态+动态分析方法和端到端的代码生成能力。通过在编译时利用静态分析辅助动态分析，ACROBAT有效地降低了运行时开销，实现了高效的批处理。同时，其端到端的代码生成摆脱了对供应商库的依赖，能够生成高度优化的张量核以实现高效执行。

尽管本文中的评估主要集中在批处理推理上，但作者认为这些技术同样适用于深度学习训练。随着动态性在深度学习计算中变得越来越重要，ACROBAT为实现深度学习框架各组件（如张量编译器、高级语言编译器和运行时）之间更紧密的协作迈出了重要一步。

A6 附录

A 关于混合静态+动态优化的更多细节

A.1 算子提升

静态识别可提升的算子。对于一个递归计算，例如清单1中的@rnn函数，通常某些张量算子并不属于递归所引发的顺序依赖。例如，清单1中第5行对输入的线性变换可以被提升到递归之外。与先前工作中依赖运行时调度算法来识别这一点不同，ACROBAT能够静态地发现这些可以被提升的算子。我们通过依赖1-上下文敏感的污点分析来静态计算这些算子的深度，从而实现这一点。我们在清单2中看到，第5行对核bias_dense的调用被赋予了一个静态计算的深度0。在运行时，这些算子因此被有效地提升到递归之外。对于RNN的例子，这使得我们可以将所有输入词嵌入的线性变换一起进行批处理，而不是一次执行一个。

A.2 粒度粗化

以静态子图为调度单元。通常，调度是在单个张量算子的粒度上进行的，即DFG中的每个节点对应一个张量核调用。我们在§2.1中看到，深度学习计算经常包含嵌入在动态控制流中的较大静态子图。因此，ACROBAT在静态子图的更粗粒度上执行调度，从而减少调度开销。由于这些块不包含任何控制流，以这种方式粗化粒度不会导致已利用并行性的损失。这项优化在先前的工作中也得到了探索【90, Zha et al. Just-in-time dynamic-batching. 2019. CoRR; 88, Xu et al. Cavs: An efficient runtime system for dynamic neural networks. 2018. USENIX ATC; 26, Fegade et al. Cortex: A compiler for recursive deep learning models. 2021. MLSys; 29, Gao et al. Low latency rnn inference with cellular batching. 2018. EuroSys; 74, Silfa et al. E-BATCH: energy- efficient and high-throughput RNN batching. 2020. CoRR】，并如图8所示。

图 8. 对清单 1 中 @rnn 函数的粒度粗化。

A.3 应对深度调度的急切性

我们在§4.1中看到ACROBAT如何依赖幽灵操作和程序阶段来应对深度调度的急切性。下面，我们对此进行更详细的解释。

幽灵算子。在图3的上半部分，我们看到在存在条件语句的情况下，急切的批处理会导致次优的批处理调度。具体来说，输入Inp1和Inp2的算子B实例被急切地批处理，并且更重要的是，与输入Inp3和Inp4的算子B实例分开了。在下半部分，我们插入了一个对幽灵算子的调用，从而得到了一个最优的调度。ACROBAT会静态识别这类情况并根据需要插入幽灵算子。请注意，幽灵算子仅影响调度，在核执行期间会被忽略。

图 7. DyNet 的运行时流水线概览。请注意，它缺乏任何静态或编译时分析，并且 DyNet 依赖显式的内存收集操作，导致高数据移动成本，如我们在 §7 中所示。

程序阶段。对于清单1中的RNN示例，为了最大限度地利用第19行输出算子的并行性，应该等到@rnn函数中调用的所有算子都已对所有输入实例执行完毕。这样，所有输入实例中所有单词对应的所有输出算子都可以作为一个批处理核调用来执行。这要求所有这些输出算子都被赋予相同的深度。然而，由于每个输入句子的长度可能不同，情况可能并非如此。从语义上讲，我们可以将RNN计算分为两个语义阶段——初始的递归计算和随后的输出转换。给定这样的程序阶段，ACROBAT会先调度和执行一个阶段中的算子，然后再进入下一个阶段。通过这种方式，ACROBAT确保在处理输出算子之前，所有的@rnn函数都已对所有输入实例执行完毕。

B 关于ACROBAT张量核生成的更多细节

B.1 利用数据重用

通过代码复制实现更好的数据重用。输入程序中的代码重用常常会阻碍上面提到的参数重用。考虑下面的代码清单，我们实现了一个双向RNN（BiRNN）【70, Schuster & Paliwal. Bidirectional recurrent neural networks. 1997. IEEE transactions on Signal Processing】计算，类似于清单1中实现的RNN模型。在这里，我们用不同的模型参数调用同一个@rnn函数来实现前向和后向RNN。在这种情况下，@rnn函数调用的张量算子将不会被静态地确定为在多次调用中有任何常量参数，从而排除了模型参数的数据重用。为了解决这个问题，在生成批处理核之前，ACROBOT识别出这类数据重用的情况（同样使用上下文敏感的污点分析），并传递性地复制必要的函数，以便在稍后生成批处理核时启用数据重用。例如，在BiRNN的例子中，ACROBAT会传递性地复制@rnn函数（包括它调用的张量算子），并为下面清单中的两个前向和后向调用使用@rnn函数的不同副本。

/* 为简洁起见，清单中省略了类型注解 */
def @main(f_rnn_bias, f_rnn_i_wt, f_rnn_h_wt, f_rnn_init,
          b_rnn_bias, b_rnn_i_wt, b_rnn_h_wt, b_rnn_init,
          inps_list) {
  let rinps_list = reverse_list(inps_list);
  let forward_res = @rnn(inps_list, f_rnn_init,
                        f_rnn_bias, f_rnn_i_wt, f_rnn_h_wt);
  let backward_res = @rnn(rinps_list, b_rnn_init,
                         b_rnn_bias, b_rnn_i_wt, b_rnn_h_wt);
  ...
}

静态块内的重用。对于一个给定的张量算子，上面讨论的分析考虑了在小批量处理中由不同输入实例进行的调用所共享的参数。这通常适用于模型参数，因为它们在多个输入实例之间共享。然而，通常情况下，在同一个静态块内对同一个张量算子的多次调用会共享一个参数。例如，在常用的LSTM单元中就是这种情况，其中四个门的计算都涉及到对同一个输入向量的并发线性变换。在这种情况下，ACROBAT会水平融合这些调用，以利用参数数据的重用。这在图9中有所说明。

图 9. 水平融合促进参数重用。

C 更多实现细节

C.1 张量核优化

下面，我们讨论ACROBAT如何依赖TVM的自动调度器【93, Zheng et al. Ansor: Generating high-performance tensor programs for deep learning. 2020. OSDI】来自动生成输入程序中使用的（可能融合的）张量算子的批处理版本的优化实现。

自动调度器算子优先级。给定一个由多个张量算子组成的DL计算，自动调度器会根据它们的相对估计执行成本来优先优化张量算子。在其他因素中，这个估计成本与算子在输入程序执行期间被调用的次数成正比。为了在存在控制流（如重复或条件控制流）的情况下准确估计给定算子的执行频率，ACROBAT依赖于基于剖面的优化（PGO）。当无法进行PGO时，ACROBAT也提供了一个简单的静态分析，根据算子调用在递归中的嵌套深度来启发式地进行此估计。

处理可变循环范围。由于ACROBAT调度的动态性，生成的未优化批处理核中对应于批处理维度的循环具有可变的范围（例如，图1中的核3）。为了优化这些核，ACROBAT自动调度一个对应的具有静态批处理维度循环范围的核，并自动将生成的调度应用于具有可变范围的原始核。此外，在为具有可变范围的循环生成代码时，我们通常必须插入条件检查以避免越界访问。我们依赖于DietCode【91, Zheng et al. Dietcode: Automatic optimization for dynamic tensor programs. 2022. MLSys】中提出的局部填充和局部分区技术来在适当时消除这些条件检查，因为它们可能严重损害性能。

C.2 AOT（Ahead-of-time）编译

我们在§6中看到，ACROBAT以AOT方式将输入的Relay计算编译成C++。作为此编译的一部分，ACROBAT将所有动态控制流以及不规则数据结构降低为原生的C++控制流和类。Relay通过将标量建模为零维张量来处理它们。ACROBAT的AOT编译器也将这类零维张量及其上的常见算术运算符降低为原生的C++标量。我们在§7.2中看到，这种AOT编译显著降低了动态控制流的执行开销。

C.3 其他细节

如正文§6所述，我们通过扩展TVM来构建ACROBAT的原型。我们发现TVM的算子融合过程有限，并且常常无法融合内存复制算子，如张量重塑、连接和转置。因此，在我们的DL计算实现中，我们手动向编译器提供融合提示，以强制将这类算子与其消费者融合。此外，我们当前的原型只支持Relay的函数式子集。具体来说，目前不支持通过可变引用产生的副作用。ACROBAT的运行时系统经过了大量优化以减少运行时开销。我们使用区域分配（arena allocation）（在CPU和GPU上都使用）和GPU上的异步执行。我们还在可能的情况下批处理CPU和GPU之间的内存传输操作，以减少CUDA API的开销。

D 补充评估和附加细节

D.1 PGO在张量核自动调度中的益处

我们在§C.1中提到，ACROBAT使用调用频率（通过PGO获得）来在自动调度期间优先进行张量算子优化。为了评估这项优化的好处，我们观察了NestedRNN在有和没有这项优化时的性能。该基准计算在外部GRU循环的每次迭代中平均执行30次内部RNN循环的迭代。因此，RNN循环中调用的算子对基准性能的影响远大于GRU循环中调用的算子。表9显示了在自动调度器的不同迭代次数下，有和没有PGO时基准的执行时间，这表明ACROBAT在开启PGO时可以更好地优先为RNN算子进行自动调度。
Table 9. NestedRNN（小模型，批大小为8）的执行时间（无/有PGO），展示了在自动调度期间使用PGO调用频率的好处。