Cocktailer: Analyzing and Optimizing Dynamic Control Flow in Deep Learning

发表时间: 2023-07

COCKTAILER：分析和优化深度学习中的动态控制流

作者/机构：Chen Zhang†£⋄∗, Lingxiao Ma⋄, Jilong Xue⋄, Yining Shi‡⋄∗, Ziming Miao⋄, Fan Yang⋄, Jidong Zhai†, Zhi Yang‡, Mao Yang⋄
(†清华大学 ‡北京大学 ⋄微软研究院)

速读

一句话结论 本文提出了一个名为 COCKTAILER 的深度神经网络编译器，通过将动态控制流与数据流统一卸载到硬件加速器上进行协同调度，在带有复杂控制逻辑的模型上实现了最高 8.22 倍的推理加速。

要解决什么问题 现有的深度学习框架通常将计算密集型的数据流算子卸载到 GPU 等加速器上，而将循环、分支、递归等动态控制流保留在 CPU 上执行。这种分离执行模式带来了严重的“并行性不匹配”问题，具体表现为三个致命卡点。首先是高昂的边界同步开销，以条件分支为例，CPU 需等待 GPU 返回决定分支目标的数据，GPU 又要等待 CPU 检查条件并发起后续操作，频繁的跨设备通信打断了异步执行，导致 GPU 长时间空闲。其次是优化范围受限，控制流在程序中形成了硬边界，使得算子融合等全局优化只能在被切割的静态数据流子程序内部进行，无法跨越循环或分支边界。最后是隐式的串行化，原本无依赖、可并发执行的算子（如递归树中不同分支的节点）受制于 CPU 单线程控制流的顺序调度，被迫串行执行，严重限制了硬件潜在并行性的发挥。

怎么做的 核心思路是打破 CPU 控制流与 GPU 数据流的物理隔离，将控制流逻辑复制并下放到 GPU 的各个并行层级中，实现两者的统一协同调度。这一方法之所以能绕开同步卡点，是因为深度神经网络中的控制流通常作用于算子级别，算子内部的所有底层并行任务实际上共享同一个控制结果，且现代 GPU 底层语言原生支持线程级控制逻辑。关键设计由三个核心部件构成：首先是引入了基本执行单元 `uTask` 抽象。数据流算子被自然拆解为与硬件并行度匹配的独立任务集合，即 $$uOperator = \{uTask_0, uTask_1, ..., uTask_{N-1}\}$$ 而控制流则被定义为嵌套结构 $$NestedUTask = (ControlFlow, \{body\_uTask_i\})$$ 包含循环、分支和用于递归的引用三种特殊类型。其次是统一的协同调度空间与策略。编译器采用自底向上的启发式遍历，将包含控制流和数据流的程序统一映射到 GPU 的多级处理单元（如线程、线程块、内核）上。在调度过程中，系统会自动执行函数内联、循环展开和递归展开，从而彻底消除原有的控制流边界。最后是针对控制流的设备端代码生成器。它将调度好的任务转换为 GPU 内核，其函数签名可抽象为 $$(A \Rightarrow B \mid C)$$ 其中 $A$ 为输入张量，$B$ 为输出张量，$C$ 为预分配的中间结果缓冲区。为了解决 GPU 上的特殊控制流执行问题，生成器引入了静态内存管理来处理循环携带依赖，设计了分支重聚类技术将偏好相似线程占用率的算子打包以缓解寄存器压力，并利用全局或共享内存实现了一个轻量级的模拟调用栈，成功绕开了 GPU 原生栈深度极浅的限制，高效支持复杂的递归调用。

效果如何 实验在 NVIDIA V100 和 AMD MI100 硬件平台上进行，测试了批处理大小为 1 和 64 的在线与离线推理场景。对比基线包括代表通用语言运行时路线的 PyTorch 和 JAX，代表静态图路线的 TensorFlow，以及一个仅加速基本块但保留 CPU 控制流的内部强基线 COCKTAILERBASE。测试模型覆盖了包含循环的 LSTM 和 Seq2seq、包含分支的 BlockDrop 和 SkipNet，以及包含递归的 RAE。量化结果显示，在 NVIDIA GPU 上，该方法相比各模型的最优基线取得了平均 1.85 倍的端到端加速。其中，在控制流密集的递归模型 RAE 上效果最为惊人，相比 PyTorch 加速 9.35 倍，相比 JAX 加速 327.62 倍；在带有动态分支的 BlockDrop 模型上（批处理大小为 1 时）也获得了 1.84 倍的加速。控制流开销分析表明，该方法将分支模型的控制流开销从基线的 34% 压降至 11%，使得动态跳过计算层的理论收益真正转化为实际的性能提升。例如在 SkipNet 中，使用 JAX 只有在执行层比例低于 20% 时才有性能提升，而该方法将这一比例门槛放宽到了 65%。该方法也有其局限和代价。当模型中数据流计算占据绝对主导地位时（如批处理大小为 64 的 Attention 模型），其性能与最快基线基本持平，加速比仅为 1.01 倍。此外，在处理递归模型时，由于采用了模拟栈机制，其最大栈深度无法在运行时动态增加，必须由用户手动设置深度限制，否则编译器会耗尽所有空闲内存来保存调用栈的中间结果。

A1 主要贡献

本文针对深度神经网络（DNN）中日益复杂的控制流逻辑（如循环、分支、递归）在加速器上执行效率低下的问题，提出了一个新的DNN编译器COCKTAILER。

核心问题：
现有的DNN框架通常在CPU上处理控制流，而将计算密集型的数据流卸载到GPU等加速器上。这种分离执行的模式导致了以下几个严重的效率问题：

  1. 高昂的同步开销：CPU和加速器之间需要频繁同步（例如，CPU等待GPU返回条件判断所需的数据），这在关键路径上引入了显著的通信开销（如跨PCIE）。
  2. 受限的优化范围：控制流在DNN程序中形成了明确的边界，阻碍了跨控制流范围的全局优化，例如无法将循环内外的两个算子进行融合。
  3. 隐式的串行化：控制流隐式地将本可以并行执行的数据流算子串行化，限制了并行性的发掘。作者观察到，在PyTorch中，这些开销有时会占据总执行时间的72%。

研究目标与核心思想：
本文的目标是解决控制流与数据流之间的“并行性不匹配”问题，实现二者在硬件加速器上的协同优化。核心思想基于以下三个关键洞察：

  1. DNN中的数据流本质上是一个多级并行的程序，其算子可以在不同硬件并行层级（如GPU的线程、线程块、内核）上执行。
  2. DNN中的控制流操作大多应用于算子级别，即所有较低级别的并行任务共享相同的控制结果。这意味着可以通过将控制逻辑复制到不同级别的所有并行任务中，来重新调度控制流。
  3. 现代硬件加速器（如GPU）在其底层编程语言中为每个线程提供了支持控制逻辑的能力，使得上述重新调度方法在实践中可行。

主要创新点与贡献：
为实现上述目标，COCKTAILER做出了以下贡献：

  1. uTask抽象：提出了一种名为uTask的抽象，作为DNN程序（包括控制流和数据流）的基本执行单元。数据流算子可以被自然地分解为与其计算并行性相符的不同粒度的uTask。对于控制流，引入了三种特殊的uTask（循环uTask、分支uTask、uTask引用）来表示。
  2. 统一的协同调度空间：通过将整个DNN程序统一为uTask粒度，COCKTAILER创建了一个整体的调度空间，用于协同调度控制流uTask和计算uTask。这使得可以将控制流uTask分配到最高效的并行级别上执行。
  3. 启发式调度策略与自动代码生成：COCKTAILER使用一种基于遍历的自底向上启发式调度策略来寻找高效的调度方案，并能够自动将循环、分支等控制流操作移入加速器内核中，从而实现跨控制流边界的优化，如函数内联、循环展开和递归展开。
  4. 通用性与高性能：COCKTAILER构建于通用的DNN张量编译器之上，利用其为uTask生成内核的能力，可以轻松适应不同的加速器（如CUDA GPU和ROCm GPU）。评估表明，COCKTAILER相较于当前最先进的DNN框架和编译器，可将带有控制流的DNN模型加速高达8.2倍，并显著降低控制流带来的开销，使得动态跳过部分计算等优化能够获得实际的性能提升。
图 1: 带有控制流的模型
图 1: 带有控制流的模型

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

DNN中控制流的重要性。随着DNN模型从顺序前馈层【15, Kaiming He等人, Deep residual learning for image recognition, 2016, CVPR】【23, Alex Krizhevsky等人, Imagenet classification with deep convolutional neural networks, 2012, NIPS】【38, K. Simonyan等人, Very deep convolutional networks for large-scale image recognition, 2014, CoRR】发展到具有复杂控制逻辑的结构【16, Sepp Hochreiter等人, Long short-term memory, 1997, Neural Comput.】【39, Richard Socher等人, Semi-supervised recursive autoencoders for predicting sentiment distributions, 2011, EMNLP】【40, Ilya Sutskever等人, Sequence to sequence learning with neural networks, 2014, NIPS】【41, Kai Sheng Tai等人, Improved semantic representations from treestructured long short-term memory networks, 2015, arXiv】【46, Xin Wang等人, Skipnet: Learning dynamic routing in convolutional networks, 2018, ECCV】【47, Zuxuan Wu等人, Blockdrop: Dynamic inference paths in residual networks, 2018, CVPR】，控制流在实现动态计算和网络架构适应性方面变得至关重要。具体应用包括：

根据对PyTorch Hub中52个模型的统计，近27%的模型包含结构化数据（如序列、树），且动态DNN模型【13, Yizeng Han等人, Dynamic neural networks: A survey, 2021, IEEE Transactions on Pattern Analysis and Machine Intelligence】的调查显示，条件计算和高效计算是很有前景的研究方向。

DNN中控制流的分类与代表模型。与编程语言类似，DNN中的控制流结构主要分为序列、分支、循环和子程序（函数）。大多数具有复杂控制流的DNN模型可分为三类：

这些代表性模型如图1所示。

现有支持方法的局限性。为了支持这些新兴的DNN模型，主流方法有两种：第一种是以TensorFlow 1.x【4, Martín Abadi等人, TensorFlow: A System for Large-Scale Machine Learning, 2016, OSDI 16】为代表，通过引入Enter、NextIteration等一系列控制流算子【52, Yuan Yu等人, Dynamic control flow in large-scale machine learning, 2018, EuroSys】来支持复杂模型架构，这些算子在框架运行时由CPU线程执行。第二种是以PyTorch【36, Adam Paszke等人, Pytorch: An imperative style, high-performance deep learning library, 2019, Advances in neural information processing systems】和JAX【11, Roy Frostig等人, Compiling machine learning programs via high-level tracing, 2018, Systems for Machine Learning】为代表，利用编程语言本身来表示和执行控制流，例如在Python中编写控制逻辑，由Python运行时执行。

并行性不匹配导致性能问题。两种主流方法都将数据流算子调度到加速器上，而将控制流保留在CPU端执行。其根本原因是数据流和控制流之间的并行性不匹配：数据流算子（如矩阵乘法）具有内部数据并行性，而控制流操作是单线程计算。现代加速器（如GPU）拥有大规模的层级化并行处理单元，这与数据流算子的并行性相符，但难以将单线程的控制流调度到这些单元上执行。因此，现有实践【4, Martín Abadi等人, TensorFlow: A System for Large-Scale Machine Learning, 2016, OSDI 16】【11, Roy Frostig等人, Compiling machine learning programs via high-level tracing, 2018, Systems for Machine Learning】【36, Adam Paszke等人, Pytorch: An imperative style, high-performance deep learning library, 2019, Advances in neural information processing systems】【52, Yuan Yu等人, Dynamic control flow in large-scale machine learning, 2018, EuroSys】将控制流操作调度到CPU端执行，这种做法在不同基本块的DNN算子之间引入了边界，导致了性能问题。如图2所示，JAX在执行Seq2Seq、BlockDrop和RAE时，与移除了所有控制流计算、只保留数据流计算的静态图相比，分别慢了1.16倍、1.54倍和56.22倍。

图 2: JAX中的控制流开销
图 2: JAX中的控制流开销
图 3: JAX中BlockDrop的时间线。从GPU到CPU的数据拷贝（MemcpyD2H）导致了同步。
图 3: JAX中BlockDrop的时间线。从GPU到CPU的数据拷贝（MemcpyD2H）导致了同步。

性能问题的具体来源。性能问题主要源于以下三个方面：

  1. 边界开销：在加速器上执行数据流算子、在CPU上执行控制流操作，会引发CPU与加速器之间的同步。以BlockDrop模型为例，CPU需要等待GPU提供用于决定分支目标的数据，然后GPU又需要等待CPU检查分支条件并发起后续操作。这种边界开销主要包括CPU与加速器间的通信和内核启动。如图3所示，JAX执行BlockDrop模型的时间线显示，CPU-GPU同步不仅开销高，还破坏了异步执行，导致GPU长时间空闲。
  2. 边界限制优化范围：将控制流和数据流分开执行，会将DNN程序分割成多个子程序，每个子程序代表一个可在加速器上执行的静态数据流。许多DNN优化（如Rammer【28, Lingxiao Ma等人, Rammer: Enabling holistic deep learning compiler optimizations with rtasks, 2020, OSDI 20】、核函数融合【35, Wei Niu等人, DNNFusion: accelerating deep neural networks execution with advanced operator fusion, 2021, PLDI】【56, Zhen Zheng等人, Astitch: enabling a new multi-dimensional optimization space for memoryintensive ml training and inference on modern simt architectures, 2022, ASPLOS】）只能在这些子程序内部进行，导致次优性能。例如，在一个多层LSTM模型中，跨不同层的LSTM单元中的DNN算子本可以并行执行，但循环控制流将DNN优化限制在单个单元内，忽视了这种并行性。
  3. 边界阻碍DNN程序的并行性：由于CPU与加速器之间的同步，不同控制流语句中的DNN程序被迫顺序执行，这可能会阻碍潜在的算子间并行性。以RAE模型为例，递归构建了一个树形架构，其中无依赖的算子可以并发执行。然而，由于控制流只能顺序执行，无依赖节点的算子也被顺序执行了。

观察与机遇。现有方法的根本局限性促使我们期望在一个统一的空间（即加速器端）内调度包含控制流和数据流的DNN程序。尽管控制流与数据流的并行性不匹配使之具有挑战性，但我们观察到，DNN程序中的控制流通常应用于DNN算子级别，即受控制流影响的算子其内部计算共享相同的控制逻辑。另一方面，大多数硬件加速器（如GPU）具备执行控制流指令的能力。因此，如果我们能将控制流表示为一种更细的粒度，并能恰当地映射到并行处理单元上执行，就可以将数据流和控制流都调度到加速器端。

A2 方法细节

3. COCKTAILER 设计

COCKTAILER系统概览。基于第2节的观察，本文提出了DNN编译器COCKTAILER，用于在单一空间内协同优化控制流和数据流。如图4所示，COCKTAILER首先接收一个包含控制流和数据流的DNN程序作为输入，其中数据流中的每个算子是一个由独立同构的uTask组成的uOperator，每个uTask可调度至加速器的一个计算单元。其次，COCKTAILER不在CPU和加速器上分别调度控制流和数据流，而是在一个统一空间内进行调度。COCKTAILER会为程序生成uProgram表示，其中包含多个独立的uTask（如图4(b)中的Loop-uTasks），这些uTask可被调度到硬件加速器的并行计算单元上执行。每个uTask同时表示了一个计算单元的控制流和数据流逻辑。

图 4: (a) 现有DNN框架（如JAX）和 (b) COCKTAILER中的DNN计算系统概览
图 4: (a) 现有DNN框架（如JAX）和 (b) COCKTAILER中的DNN计算系统概览

硬件抽象。COCKTAILER将具有大规模并行性的加速器抽象为多个级别的并行处理单元。在每个级别中，有多个并行且同构的处理单元，它们共同构成一个更高级别的处理单元。这种硬件抽象与常见的硬件加速器天然对齐。以NVIDIA GPU为例，它包含多个同构的流式多处理器（SM），每个SM又包含多个同构的核心。因此，NVIDIA GPU可以被映射为COCKTAILER硬件抽象中的一个三级并行处理单元架构，如图4所示：L0是核心（线程），L1是SM（线程块），L2是GPU设备（内核）。

调度示例。图4中的循环结构被调度为一个映射到三级加速器上的uProgram。该uProgram包含4个分别用于4个L1单元的loop-uTask，每个loop-uTask被映射到一个L1单元执行。数据流算子和循环本身都被调度进了loop-uTask中。以第一个loop-uTask为例，它包含一个循环控制流和一系列用于数据流操作的uTask，其中包括1个MatMul uTask、1个Add uTask和1个Relu uTask。

3.1 基于uTask的DNN程序

uTask、uOperator和uProgram的定义。为了将DNN程序的控制流和数据流协同调度到具有大规模并行单元的加速器上，COCKTAILER使用单元任务（uTask）的概念来精细化定义DNN程序。具体来说，uTask被定义为可调度至硬件加速器中多级处理单元之一执行的计算逻辑。一个uTask中的计算可以是其他uTask的列表，即一个嵌套的uTask。uProgram表示映射到硬件上某一级别并行处理单元的、以uTask表示的DNN程序的执行计划。

图 5: uTask、uOperator和uProgram的定义
图 5: uTask、uOperator和uProgram的定义

数据流算子的uTask和uOperator。如图5(a)(c)所示，一个数据流算子被表示为一组独立且同构的uTask，每个uTask是要调度到一个处理单元的计算。具体而言，每个uTask通过get_input_data()获取输入张量的一个切片，并执行在compute()中定义的相应计算。一个uOperator被定义为相应数据流算子所有uTask的集合。uOperator中的uTask通过逻辑uTask_id进行索引，并通过compute(uTask_id)调用。uOperator中的uTask总数由get_uTask_num()报告。当一个算子中的所有uTask都执行完毕，该算子的执行就完成了。数据流算子（如矩阵乘法）通常被实现为多个独立同构的任务，调度到加速器的大规模并行单元上执行。因此，uTask的概念不仅能自然地表示数据流算子的细粒度计算，也与加速器中多级并行处理单元的硬件架构相契合。

控制流的uTask。与数据流算子不同，控制流本身无法被划分为并行的任务。为了在并行处理单元上调度控制流，需要解决控制流计算与加速器大规模并行性之间的不匹配问题。在DNN程序中，控制流操作应用于一系列DNN算子。当这些DNN算子可以被划分为独立同构的uTask时，控制这些算子就等同于在每个uTask上应用控制流计算。例如，一个循环体中有一个矩阵乘法算子，与其对整个算子执行循环，等价于让硬件加速器的每个单元对该算子的uTask执行循环控制流。通过在这些DNN算子的细粒度表示上应用控制流，我们可以将包含控制流的计算调度到硬件加速器的并行处理单元上。也就是说，通过将控制流计算复制到多级并行单元，让每个单元独立执行控制流并控制调度到该单元上的uTask，我们就可以在uTask粒度上表示控制流。

控制流uTask的具体实现。基于上述观察，COCKTA
AILER将控制流操作表示为图5(b)中定义的NestedUtasks，其中主体内的计算由body_uTasks表示。这些uTask具有数据依赖性，应在一个处理单元上顺序执行。与数据流算子中get_input_data()提取输入张量切片不同，控制流body_uTasks中uTask的输入数据与控制流的结果相关。例如，在LSTM模型中，循环控制流主体中的uTask在不同循环步骤中需要不同值的循环计数器。因此，控制流的get_input_data()应考虑控制流结果来准备输入数据。不同控制流操作在body_uTasks中有不同的数据访问模式。

图 6: 控制流uTasks
图 6: 控制流uTasks

控制流类型的uTask表示。根据第2节的讨论，DNN程序中有三种主要控制流：循环、分支和递归。因此，COCKTAILER定义了loop-uTask、branch-uTask和uTask引用来表示这些控制流的细粒度uTask。

  1. Loop-uTask：如图6(a)所示，COCKTAILER目前支持for循环和while循环。control_flow()接口实现循环条件，循环体计算由body_uTasks中的uTask实现。body_uTasks在循环中会以不同的输入数据执行多次。例如，在LSTM模型中，每个循环步骤中LSTM单元的计算需要相同的模型参数张量，但需要不同的循环计数器张量和状态张量。get_input_data()接口需要在每个循环步骤准备相应的张量。
  2. Branch-uTask：如图6(b)所示，control_flow()接口实现分支中的条件计算。branch-uTask有then_uTasks和else_uTasks，分别表示两个分支中以uTask表示的计算。get_input_data()接口根据条件结果返回相应分支所需的数据。
  3. Function：函数可以自然地用NestedUTasks表示，其中body_uTasks中的每个uTask代表函数体内的细粒度计算。get_input_data()接口准备输入数据张量，compute()接口顺序执行body_uTasks中的uTask。
  4. Recursion and uTask reference：虽然函数可以用uTask表示，但递归是特殊情况，即一个函数可能在函数体内调用自身，意味着一个uTask可能在其body_uTasks中包含自身。为了支持递归，COCKTAILER引入了uTask引用来引用一个uTask定义。uTask引用可被视为对一个uTask的函数调用。引用是对uTask的声明，而uTask定义了函数中的计算。执行引用时，COCKTAILER会找到其uTask定义并执行该uTask。

COCKTAILER的uTask抽象是一种通用的表示控制流的方式，通过继承NestedUTask可以表示更多类型的控制流。

uProgram。整个输入DNN程序的生成执行计划由一个uProgram表示。uProgram包含多个独立的uTask，每个uTask是调度到加速器unit_level级别的一个处理单元的计算逻辑。这些uTask可以通过compute执行，uProgram的总uTask数量由get_utask_num报告。uTask抽象使得COCKTAILER能够以细粒度表示包含数据流算子和控制流的DNN程序，为协同调度开辟了新的空间。

3.2 uProgram调度

调度空间与机制。uTask表示为DNN程序开辟了一个巨大的调度空间，用于在单一空间内协同优化控制流和数据流。COCKTAILER在编译时探索这个调度空间，而非像现有框架那样采用预定义的调度策略（数据流在加速器，控制流在CPU）。为此，COCKTAILER将调度策略与其机制分离。在机制方面，COCKTAILER提供了带约束的调度接口；在策略方面，提供了一种基于遍历的调度策略。该调度是为uTask表示的算子设计的，可以自动执行。

图 7: 调度接口
图 7: 调度接口

调度接口定义。如图7所示，COCKTAILER提供了三个接口来促进调度过程：ScheduleOperator、ScheduleControlFlow和SetResource。

约束1：控制流的unit_level不应低于其主体内数据流的unit_level。
COCKTAILER还包含一个分析器（profiler），用于测量uProgram的执行时间，其分析信息可以指导策略决定是否将uProgram调度到加速器的某个unit_level。

基于遍历的自底向上调度策略。算法1描述了一种基于遍历的调度策略，展示了如何使用接口和分析器将控制流和数据流协同调度到加速器端。该策略以一个表示为控制流操作和数据流uOperator的DNN程序g以及加速器D为输入，返回在该加速器上调度好的一系列uProgram。策略还接受一个unit_level参数，表示图g内部算子已被调度的最高级别，初始为NULL。如果输入程序有多个算子，COCKTAILER会先将它们放入一个函数算子中再进行调度。
策略首先将所有数据流uOperator调度到uProgram（第4行和第21-27行）。如果分析器表明某个调度可以减少总执行时间，策略会逐步尝试将程序的更多部分调度到同一个uProgram中（第5-27行）。具体来说，策略会递归遍历程序（第7行），直到只包含一个uOperator（第3-4行），并通过ScheduleProgram将其调度到uProgram中。在遍历过程中，如果输入程序中的所有操作都已调度到加速器的单元上（第21行），策略会尝试通过ScheduleProgram将整个程序（即控制流）调度到uProgram中（第21-27行）。ScheduleProgram负责将输入程序g调度到加速器D的某个unit_level。它通过调用SetResource配置调度，并利用该配置通过ScheduleOperator和ScheduleControlFlow进行调度，调度期间始终维持约束1。

算法 1: 基于遍历的调度策略

调度时间优化。为了减少调度时间，可以采用一些优化措施，例如对不同unit_level的尝试（第22行）可以并行执行。这些优化并未在伪代码中显式展示。

调度优化。在调度过程中，有三种优化机会：

  1. 函数内联（Function inline）：为消除函数调用开销，COCKTAILER将无递归的函数控制流转换为一系列计算，这移除了函数控制流边界，使得DNN优化可以应用于更大的程序范围。
  2. 循环展开（Loop unroll）：COCKTAILER会展开循环控制流若干步，以探索更多优化机会。例如，在多层RNN模型中展开循环可以暴露RNN单元之间的并行性。循环展开在调度期间应用，并会评估其效果以决定是否启用。
  3. 递归展开（Recursion unroll）：与循环展开类似，递归也可以被展开以显式暴露递归树结构。COCKTAILER应用此优化来展开递归结构若干次，以暴露更多优化机会。例如，展开的递归树可以自然地暴露递归调用之间的并行性，从而实现并发执行。如图8所示，通过展开递归调用，无依赖的计算（如图8中的节点2和3）可以并发执行。递归展开在调度期间应用，调度器会评估展开结果以决定是否启用，并将展开的主体调度到不同的计算单元。

这些优化都包含在ScheduleControlFlow中。调度器会尝试启用这些优化，并用样本数据评估性能，以决定是否最终启用。

图 8: 递归调用的并行执行
图 8: 递归调用的并行执行

4. 实现

总体实现。COCKTAILER基于PyTorch【36, Adam Paszke等人, Pytorch: An imperative style, high-performance deep learning library, 2019, Advances in neural information processing systems】和Rammer【28, Lingxiao Ma等人, Rammer: Enabling holistic deep learning compiler optimizations with rtasks, 2020, OSDI 20】，使用约10000行Python和C++代码实现。它首先将PyTorch程序导出为带有内置循环和分支算子以及扩展的递归调用算子的ONNX图。然后，COCKTAILER自动执行数据流和控制流的调度，并应用第3节中描述的控制流相关优化。最后，COCKTAILER将生成的代码包装成一个自定义的PyTorch算子，并用对该算子的调用替换原PyTorch程序。本文实现了针对NVIDIA GPU和AMD GPU的版本。

4.1 COCKTAILER在NVIDIA CUDA GPU上的实现

硬件抽象与内核生成。如第3节所述，NVIDIA GPU可被抽象为一个三级硬件。COCKTAILER在Rammer【28, Lingxiao Ma等人, Rammer: Enabling holistic deep learning compiler optimizations with rtasks, 2020, OSDI 20】、AutoTVM【6, Tianqi Chen等人, TVM: An automated endto-end optimizing compiler for deep learning, 2018, OSDI 18】、Ansor【53, Lianmin Zheng等人, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI 20】、Roller【57, Hongyu Zhu等人, ROLLER: Fast and efficient tensor compilation for deep learning, 2022, OSDI 22】和手动实现的内核之上实现了ScheduleOperator接口。具体来说，COCKTAILER首先为给定unit_level的每个数据流算子获取源代码，来源可以是现有简单算子（如逐元素操作）的手动实现，或通过AutoTVM、Ansor或Roller对算子进行调优。然后，COCKTAILER利用Rammer将数据流算子的内核源代码转换为包含多个uTask的uOperator。之后，COCKTAILER调度程序并为控制流主体生成内核代码。

4.1.1 Nested-uTask的代码生成

整体结构。一个函数内的一系列uOperator将被调度为一个包含多个Nested-uTask的uProgram。这会被转换成一个函数，其参数为相关张量的指针。具体地，我们用(A => B | C)表示一个函数，输入张量为A，输出张量为B，中间结果保存在缓冲区C中。该函数还接受一个uTask_id参数用于索引uProgram中的uTask。图9展示了一个Function-uProgram示例，其中有一个由4个uTask实现的matmul uOperator和一个由2个uTask实现的tanh uOperator。这个Function-uProgram包含2个uTask，每个uTask的body_uTasks中包含2个matmul uTask和1个tanh uTask，并插入了适当的屏障（第5-8行，11-14行）。屏障可以通过CUDA Cooperative Groups【1, Cooperative Groups, http://nvidia.com】或扩展的无锁GPU同步技术【48, Shucai Xiao等人, Inter-block gpu communication via fast barrier synchronization, 2010, IPDPS】实现。

图 9: uTask示例
图 9: uTask示例

函数封装与设备函数。Function-uProgram会为张量y分配存储空间（第4, 10行），并用函数名和签名包装代码（第2行）。使用__device__函数限定符，以便该函数可以被其他uTask调用。为简洁起见，后续章节将省略uTask_id，只展示uOperator内一个uTask的生成代码。

线程块对齐。将多个DNN算子调度到单个GPU内核中的一个挑战是每个GPU块内的线程数（blockDim）不同。uProgram的内核的blockDim必须设置为其uOperator的最大blockDim，这会导致具有大量GPU块（gridDim）和较小blockDim的内核在被调度到具有大blockDim的算子的同一个内核中时执行效率低下。为解决此问题，我们对可配置blockDim的uOperator（如逐元素操作、归约、转置）进行了重新实现。在调度期间，COCKTAILER收集具有预定义blockDim的uOperator（如matmul和卷积）的最快内核，并将可配置uOperator的blockDim设置为所收集uOperator的最大blockDim。如果收集到的uOperator的blockDim差异很大，COCKTAILER将利用扩展的Roller【57, Hongyu Zhu等人, ROLLER: Fast and efficient tensor compilation for deep learning, 2022, OSDI 22】重新生成具有固定blockDim的内核。

寄存器压力。生成的长时间运行的GPU内核可能会面临寄存器压力。为缓解此问题，COCKTAILER使用3.2节中的分析器来检测因寄存器过度使用导致的性能下降，并停止扩大当前内核。对于没有回边的控制流图，COCKTAROILED还可以利用4.1.3节中的分支重聚类技术，既能将控制流调度到加速器端，又能减小内核大小。

4.1.2 Loop-uTask的代码生成

整体结构。图10(a)展示了一个简化的RNN模型。它被调度为一个Loop-uProgram，包含若干个Loop-uTask。uProgram中的每个Loop-uTask包含来自三种uOperator的body_uTasks：gather、matmul和融合的add-tanh操作。这个Loop-uProgram接收三个输入张量inp, wx, h（图10(b)中的h_in），并产生一个更新后的张量h（图10(b)中的h_out）。每个Loop-uTask生成的代码包含一个循环（第4行）和一系列由屏障分隔的uTask（第5-7行），用于跨GPU块同步。

图 10: Loop-uTask示例
图 10: Loop-uTask示例

内存管理。与在运行时分配张量的现有DNN框架不同，COCKTAILER需要静态分配张量内存以便在GPU上执行控制流操作。循环体中的变量可分为四类：(1) 常量 (wx, inp)；(2) 中间结果 (inpi, xi)；(3) 迭代计数 (i)；(4) 循环携带依赖 (h)。所有这些变量都由指向相应预分配张量的指针表示，并可从get_input_data获取。具体来说，指向常量的指针是相应的函数输入，中间结果从一个tmp缓冲区分配（图10(b)第2行），指向迭代计数的指针是&i。指向循环携带依赖的指针处理较为复杂，因为该变量同时存在于Loop-uOperator的输入和输出张量中。首先，插入一些CopyUOperator将输入张量(h_in)复制到相应的输出张量(h_out)。然后，仅访问输出张量来生成body_uTasks。最后，添加额外的CopyUOperator和循环体内uOperator之间的依赖关系，以确保重叠的输入和输出张量的正确性。

4.1.3 Branch-uTask的代码生成

整体结构。图11(a)包含一个带有两个分支的DNN模型。then分支接收张量x, w1, w2作为输入，产生张量y和z；else分支接收张量x和b作为输入，产生张量z。生成的Branch-uProgram的输入是两个分支输入的并集以及条件张量cond。输出是两个分支输出的并集。如果一个输出只存在于一个分支中，则会在另一个分支中添加CopyUOperator，将相应的旧值移动到输出张量（如图11(b)第9行）。中间结果保存在从tmp缓冲区分配的张量中。由于每次运行只可能执行一个分支，两个分支的中间结果可以使用相同的内存空间。

图 11: Branch-uTask示例
图 11: Branch-uTask示例

分支重聚类（Branch reclustering）。将整个ControlFlow-uProgram调度到单个GPU内核中并非总是最佳选择，因为不同的操作偏好不同的GPU占用率（SM上并发执行的线程数）。例如，matmul使用大量共享内存和寄存器来保存瓦片，导致占用率受限，而逐元素操作则偏好高占用率以提高内存带宽。COCKTAILER也会尝试将一个Branch-uProgram调度为多个Branch-uProgram，每个Branch-uProgram包含具有相似偏好占用率的uOperator，同时保持分支条件的执行在GPU上。图12中的示例模型包含占用率受限的uOperator matmul和conv（绿色）以及占用率大的uOperator sigmoid、add和copy（蓝色）。这些uOperator被调度到三个Branch-uProgram中，分别对应于受限占用率、大占用率和受限占用率。两个分支被协同调度，因此每个GPU内核可以包含来自两个分支的uOperator。这种分支重聚类技术减小了内核大小，因此也能缓解大型GPU内核的寄存器压力。

图 12: 通过分支重聚类优化Branch-uProgram
图 12: 通过分支重聚类优化Branch-uProgram
4.1.4 uTask引用的代码生成

整体结构。uTask引用是一种特殊情况，它调用在另一个uProgram中定义的uTask，专为递归设计，即一个函数可能调用其调用者，如图13(a)所示。为了支持递归，所有被uTask引用引用的uProgram的函数声明会在代码开头生成（图13(b)第1行）。然后，所有uProgram生成它们的函数定义。我们的递归实现的最大栈深度不能在运行时增加，因此用户需要手动设置栈深度限制，否则COCKTAILER将使用所有空闲内存来保存调用栈中的中间结果。基本情况检查作为分支操作保留在函数体内。

1 def Recursion(l, r, is_leaf, inp, w, root):   
2   cond = is_leaf[root]   
3   if cond:   
4     output = inp[root]   
5   else:   
6     a = Recursion(l, r, is_leaf, inp, w, l[root])   
7     b = Recursion(l, r, is_leaf, inp, w, r[root])   
8     c = a + b   
9     output = matmul(c, w)   
10  return output
图 13: 使用uTask引用的递归示例
图 13: 使用uTask引用的递归示例

模拟GPU栈。尽管NVIDIA GPU内置支持递归，但其栈速度慢且支持的深度非常有限，原因是GPU需要在函数调用期间保存所有线程的寄存器。然而，在DNN程序中，我们只需要在执行函数调用前保存指向张量的指针和当前栈帧的程序计数器。此外，同一组张量指针由多个uTask共享，只需保存一份副本。因此，我们有机会减少保存信息的大小，从而增加栈深度并减少保存栈帧的时间。为实现这一点，COCKTAILER在全局内存中实现了一个栈来模拟函数调用行为。由于直接更新程序计数器很危险，COCKTAILER选择将所有uProgram内联到单个函数中，并使用goto和插入到内联函数中的label来模拟程序计数器的更新。标签放置在函数的开头和函数内每个函数调用的末尾。栈不维护程序计数器，而是保存每个栈帧的标签。每个栈帧只消耗几十字节的内存，因此COCKTAILER也可以在可能的情况下将栈保存在GPU共享内存中，以避免为在不同uTask间维护同步栈而产生的内存栅栏和块间屏障。

4.2 COCKTAILER在AMD ROCm GPU上的实现

实现方法。AMD ROCm GPU提供了一个与CUDA类似且兼容大多数CUDA语句的HIP编程模型【2, HIP Programming Guide, http://rocmdocs.amd.com】。此外，AMD提供了一个hipify工具，可以将CUDA内核转换为HIP内核。COCKTAILER首先生成一个CUDA内核，然后利用hipify工具将其转换为HIP版本。由于CUDA和ROCm架构的差异，一些uOperator被重新实现。

A4 实验环境与结果

实验环境

(a) 一个递归模型
表 1: 模型配置。BS指“批处理大小”。

实验结果

5.1 NVIDIA GPU端到端评估

总体性能：如图14所示，COCKTAILER在所有模型上的性能均显著优于基线。与每个模型的最优基线相比，COCKTAILER的几何平均加速比为1.85倍（最高达8.22倍）。具体而言，平均加速比分别为：对TorchScript 3.98倍（最高9.35倍），对TensorFlow 18.45倍（最高196.85倍），对JAX 3.05倍（最高327.62倍）。

图 14: 在NVIDIA V100 GPU上的端到端DNN推理性能
图 14: 在NVIDIA V100 GPU上的端到端DNN推理性能
5.2 控制流开销分析

通过将动态执行控制流的真实场景与执行已追踪的、无控制流的计算图进行比较，评估了控制流边界带来的性能下降。

图 15: 带循环模型的控制流开销。
图 15: 带循环模型的控制流开销。

图 16: 带分支模型的控制流开销。“No skip”指运行ResNet模型的所有层。
图 17: 不同执行层比例下的性能

图 18: RAE递归模型的控制流开销。
图 18: RAE递归模型的控制流开销。
5.3 优化分解

图19展示了各项优化对性能的贡献。

图 19: BS=1时模型的性能分解
图 19: BS=1时模型的性能分解
5.4 AMD GPU端到端评估

如图20所示，在AMD MI100 GPU上（BS=1），COCKTAILER同样表现出色，平均加速比分别为：对TorchScript 2.97倍（最高5.86倍），对TensorFlow 21.28倍（最高112.34倍），对JAX 3.22倍（最高272.63倍）。

图 20: 在AMD MI100 GPU上（BS=1）的端到端DNN推理性能
图 20: 在AMD MI100 GPU上（BS=1）的端到端DNN推理性能

A7 补充细节

6. 相关工作

DNN中的控制流支持。支持深度学习中控制流的方法可分为两类。第一类以TensorFlow 1.x【4, Martín Abadi等人, TensorFlow: A System for Large-Scale Machine Learning, 2016, OSDI 16】和TorchScript【36, Adam Paszke等人, Pytorch: An imperative style, high-performance deep learning library, 2019, Advances in neural information processing systems】为代表，在框架运行时于CPU上执行控制流操作。控制流被实现为特殊算子（如循环的NextIteration【52, Yuan Yu等人, Dynamic control flow in large-scale machine learning, 2018, EuroSys】、分支的Switch【52, Yuan Yu等人, Dynamic control flow in large-scale machine learning, 2018, EuroSys】、递归的InvokeOp【19, Eunji Jeong等人, Improving the expressiveness of deep learning frameworks with recursion, 2018, EuroSys】）或运行时指令。第二类以Chainer【42, Seiya Tokui等人, Chainer: a next-generation open source framework for deep learning, 2015, NIPS Workshop on Machine Learning Systems】、PyTorch【36, Adam Paszke等人, Pytorch: An imperative style, high-performance deep learning library, 2019, Advances in neural information processing systems】和JAX【11, Roy Frostig等人, Compiling machine learning programs via high-level tracing, 2018, Systems for Machine Learning】为代表，利用通用编程语言（如Python）的运行时来支持控制流操作。AutoGraph【31, Dan Moldovan等人, Autograph: Imperativestyle coding with graph-based performance, 2019, Volume 1】、Janus【18, Eunji Jeong等人, JANUS: Fast and flexible deep learning via symbolic graph execution of imperative programs, 2019, NSDI 19】和Terra【22, Taebum Kim等人, Terra: Imperative-symbolic co-execution of imperative deep learning programs, 2021, NeurIPS】等工作表明，通用语言表示的控制流有时可转换为框架运行时的控制流算子。尽管支持方式不同，这些工作中的控制流操作都只能由CPU执行。

特定控制流的深度优化。一些工作对特定形式的控制流进行了深度优化。VersaPipe【55, Zhen Zheng等人, Versapipe: a versatile programming framework for pipelined computing on gpu, 2017, MICRO】优化了通用GPU程序的流水线。Cortex【10, Pratik Fegade等人, Cortex: A compiler for recursive deep learning models, 2021, MLSys】提供了描述递归数据模式（即递归树结构）的接口，但它假设所有控制流的跳转方向仅依赖于输入的递归树结构，因此不适用于依赖动态计算数据的控制流，例如Seq2seq【40, Ilya Sutskever等人, Sequence to sequence learning with neural networks, 2014, NIPS】中迭代次数未知的while循环，以及BlockDrop【47, Zuxuan Wu等人, Blockdrop: Dynamic inference paths in residual networks, 2018, CVPR】和SkipNet【46, Xin Wang等人, Skipnet: Learning dynamic routing in convolutional networks, 2018, ECCV】中运行时决定的分支。COCKTAILER不假设此类树结构的可用性，能够处理这些模型。

批处理（Batching）相关工作。过去在批处理方面的工作（如DyNet【33, Graham Neubig等人, Onthe-fly operation batching in dynamic computation graphs, 2017, NIPS】、Cavs【49, Shizhen Xu等人, Cavs: An efficient runtime system for dynamic neural networks, 2018, USENIX ATC 18】、Tensorflow Fold【27, Moshe Looks等人, Deep learning with dynamic computation graphs, 2017, arXiv】、BatchMaker【12, Pin Gao等人, Low latency rnn inference with cellular batching, 2018, EuroSys】、Program-counter-autobatching【37, Alexey Radul等人, Automatically batching control-intensive programs for modern accelerators, 2020, MLSys】和ORCA【51, Gyeong-In Yu等人, Orca: A distributed serving system for Transformer-Based generative models, 2022, OSDI 22】）通过引入调度器来批处理准备就绪的算子，从而实现不同控制流操作的并行化。这是一种可行的方法，与COCKTAILER互补。具体来说，COCKTAILER可以编译模型的子图，然后批处理可以应用于这些子图。在更粗粒度的子图上应用批处理还可以降低批处理调度的成本。

深度学习编译器与分布式框架。许多深度学习编译器致力于优化无控制流的计算图，包括TVM【6, Tianqi Chen等人, TVM: An automated endto-end optimizing compiler for deep learning, 2018, OSDI 18】、TASO【20, Zhihao Jia等人, TASO: Optimizing deep learning computation with automatic generation of graph substitutions, 2019, SOSP】、Rammer【28, Lingxiao Ma等人, Rammer: Enabling holistic deep learning compiler optimizations with rtasks, 2020, OSDI 20】、DNNFusion【35, Wei Niu等人, DNNFusion: accelerating deep neural networks execution with advanced operator fusion, 2021, PLDI】、PET【44, Haojie Wang等人, PET: Optimizing tensor programs with partially equivalent transformations and automated corrections, 2021, OSDI 21】和AStitch【56, Zhen Zheng等人, Astitch: enabling a new multi-dimensional optimization space for memoryintensive ml training and inference on modern simt architectures, 2022, ASPLOS】。这些优化与COCKTAILER兼容，COCKTAILER甚至通过减少控制流边界扩大了它们的优化范围。通用语言编译器中的函数内联【5, Pohua P Chang等人, Inline function expansion for compiling c programs, 1989, PLDI】、循环展开【9, Jack J Dongarra等人, Unrolling loops in fortran, 1979, Software: Practice and Experience】等编译优化也被实现在COCKTAILER中。COCKTAILER进一步引入了新的uTask抽象来表示数据流和控制流操作，与硬件加速器的并行性对齐，从而能够在异构加速器（即GPU）上分析和优化数据流与控制流计算。为了在分布式架构上扩展DNN模型，Tofu【45, Minjie Wang等人, Supporting very large models using automatic dataflow graph partitioning, 2019, EuroSys】、FlexFlow【21, Zhihao Jia等人, Beyond data and model parallelism for deep neural networks, 2019, MLSys】、GSPMD【50, Yuanzhong Xu等人, GSPMD: general and scalable parallelization for ml computation graphs, 2021, arXiv】、PipeDream【32, Deepak Narayanan等人, PipeDream: Generalized pipeline parallelism for dnn training, 2019, SOSP】、Tutel【17, Changho Hwang等人, Tutel: Adaptive mixture-of-experts at scale, 2022, arXiv】、FasterMOE【14, Jiaao He等人, Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models, 2022, PPoPP】、FlexMoE【34, Xiaonan Nie等人, FlexMoE: Scaling large-scale sparse pre-trained model training via dynamic device placement, 2023, arXiv】、BaGuaLu【29, Zixuan Ma等人, Bagualu: targeting brain scale pretrained models with over 37 million cores, 2022, PPoPP】、Alpa【54, Lianmin Zheng等人, Alpa: Automating inter-and Intra-Operator parallelism for distributed deep learning, 2022, OSDI 22】和SuperScaler【25, Zhiqi Lin等人, SuperScaler: Supporting flexible dnn parallelization via a unified abstraction, 2023, arXiv】等框架和编译器在多个硬件设备间并行化DNN模型的执行，但主要关注静态架构或特定类型的动态模型（如专家混合模型【30, Saeed Masoudnia等人, Mixture of experts: a literature survey, 2014, The Artificial Intelligence Review】）。COCKTAILER暴露了控制流操作的并行性，这可以被利用来支持分布式设备上的动态模型。

A5 结论

DNN框架和编译器在支持复杂的动态DNN模型时面临性能问题。控制流和数据流之间的并行性不匹配导致它们在CPU和加速器上分离执行，这不仅带来了开销，也错失了优化机会。COCKTAILER通过协同调度控制流和数据流的执行来支持复杂的DNN模型，具体做法包括：(1) 为DNN程序中的控制流和数据流提供了细粒度的uTask抽象，从而在硬件加速器上开辟了一个整体的调度空间；(2) 设计了调度机制和启发式策略来利用这个调度空间；(3) 在调度和代码生成中提供了控制流优化。评估表明，COCKTAILER在复杂的DNN模型上显著优于现有技术。通过在单一空间内实现控制流和数据流的协同优化，COCKTAILER将自己定位为深度学习基础设施的一个新增强功能。

A6 附录

A. 工件附录

摘要。本工件旨在帮助复现OSDI'23论文《COCKTAILER: Analyzing and Optimizing Dynamic Control Flow in Deep Learning》中的实验结果。

使用方法。COCKTAILER的输入是一个PyTorch程序。COCKTAILER将PyTorch程序导出为ONNX格式，其中包含ONNX的循环和分支算子，以及一个用于递归的扩展invoke算子。然后，COCKTAILER根据论文中描述的优化方法生成代码，并将代码包装成一个PyTorch自定义算子以供执行。

范围。本工件可用于复现论文中的实验，包括端到端比较（图14和20）、控制流开销分析（图2、15、16和18）、不同执行层比例下的性能（图17）以及优化分解（图19）。

内容。本工件包含COCKTAILER的代码、实验输入数据、实验环境设置指南以及运行实验的脚本。它可以帮助复现以下图表：

托管。COCKTAILER的主要内容托管在 https://github.com/microsoft/nnfusion/tree/cocktailer_artifact/artifacts，分支 为cocktailer_artifact。

要求。本工件需要两台机器：