NeoFlow: A Flexible Framework for Enabling Efficient Compilation for High Performance DNN Training

文章标题：NeoFlow: 一个用于实现高性能DNN训练高效编译的灵活框架
作者/机构：Size Zheng, Renze Chen, Yicheng Jin, Anjiang Wei, Bingyang Wu, Xiuhong Li, Shengen Yan, and Yun Liang

A1 主要贡献

本文介绍了一个名为 NeoFlow 的灵活框架，旨在为高性能深度神经网络（DNN）训练提供高效的编译能力。现有的 DNN 训练框架（如 PyTorch 和 TensorFlow）依赖于手动优化的库，对包含新算子的创新性神经网络支持不足，导致程序员必须手动替换或实现底层算子，效率低下且容易出错。NeoFlow 通过允许程序员直接使用表达式定义新算子，并自动将其映射到图表示和底层实现，从而解决了这一核心问题，兼顾了高编程生产力和高性能。

本文的主要贡献如下：
1. 提出了 NeoFlow 框架：这是一个灵活的框架，能够自动且高效地支持包含新型算子的 DNN 模型训练。
2. 设计了基于表达式的自动微分：针对包含新颖算子的 DNN，设计了一种基于表达式的自动微分方法，并提出了自动模型转换技术来优化 DNN 模型。
3. 开发了高效的 DNN 编译和运行时系统：该系统无缝集成了增量子图编译和代码生成，以实现高性能训练。
4. 实验验证：通过在多种新型 DNN 模型（如 CapsuleNet, MI-LSTM, SC-RNN）上的实验证明，NeoFlow 性能优越。与 PyTorch、TensorFlow 和 CuDNN 相比，其几何平均加速比分别达到了 3.16倍、2.43倍 和 1.92倍。

A3 背景知识与动机

2.1 深度学习框架

深度学习框架的功能与工作流程。深度学习框架，如PyTorch 【7，PyTorch: An imperative style, high-performance deep learning library，2019，Annu. Conf. Neural Inf. Process. Syst.】，TensorFlow 【8，Tensorflow: A system for large-scale machine learning，2016，12th USENIX Symp. Oper. Syst. Des. Implementation】，以及MxNet 【9，MXNet: A flexible and efficient machine learning library for heterogeneous distributed systems，2015，arXiv:1512.01274v1】，是执行张量管理、自动微分、库调用等多种功能的复杂系统。用户通过高级语言（如Python）构建DNN计算图，其中图的节点是算子，边是张量。用户定义DNN图的前向部分后，框架会自动计算每个算子的梯度，从而构建用于训练的反向图。如图1所示，一个典型的训练图包含用户提供的前向和损失部分，以及框架自动生成的反向部分，其中func1到func5是由框架自动附加的梯度算子。

对高性能库的依赖。为了加速DNN训练，深度学习框架会将训练图映射到手动优化的高性能库上执行。这些库包括用于Nvidia GPU的CuDNN 【10，cuDNN: Efficient primitives for deep learning，2014，arXiv:1410.0759v3】 和CuBLAS 【25，CUBLAS，2021，NVIDIA】，以及用于Intel CPU的MKL 【26，MKL，2020，Intel】 和OneDNN 【27，Onednn，2020，Intel】 等。基于这些库，像PyTorch这样的框架可以支持数千种算子。

2.2 深度学习编译器

深度学习编译器的功能与优化方法。近年来，深度学习编译器领域取得了重要进展，例如Halide 【16，Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，ACM SIGPLAN Conf. Program. Lang. Des. Implementation】，TVM 【17，TVM: An automated end-to-end optimizing compiler for deep learning，2018，13th USENIX Symp. Oper. Syst. Des. Implementation】，PlaidML 【18，PlaidML，2020，Intel】 和Tensor Comprehensions 【19，Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions，2018，arXiv:1802.04730v3】。这些编译器能够接收高级表达式作为输入，并自动生成底层的C或CUDA代码。为了生成高性能代码，一些编译器如Tensor Comprehensions和PlaidML依赖成本模型指导代码生成；而Halide和TVM则通过编写调度语言来控制复杂的循环变换和硬件特定优化，如循环分块、展开和向量化。在巨大的调度空间中寻找最优调度组合是困难且耗时的，先前的工作提出使用机器学习技术来指导调度搜索，典型的框架包括Halide autoscheduler 【28，Learning to optimize halide with tree search and random programs，2019，ACM Trans. Graph.】，AutoTVM 【29，Learning to optimize tensor programs，2018，Annu. Conf. Neural Inf. Process. Syst.】，Chameleon 【30，Chameleon: Adaptive code optimization for expedited deep neural network compilation，2020，8th Int. Conf. Learn. Representations】，FlexTensor 【31，FlexTensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system，2020，Int. Conf. Architectural Support Program. Languages Oper. Syst.】，和Ansor 【32，Ansor: Generating high-performance tensor programs for deep learning，2020，arXiv:2006.06762】。然而，这些框架都专注于推理场景，并依赖于其他图级框架（如Relay 【21，Relay: A high-level IR for deep learning，2019，arXiv:1904.08368v2】）来处理DNN模型的表示和优化。

现有编译器的两级IR架构及其局限性。现有编译器通常采用图IR（Graph IR）和算子IR（Operator IR）两级表示来编译整个DNN模型。图IR由代表不同算子的节点构成，例如Relay使用特定节点表示二维卷积和矩阵乘法。在编译过程中，图IR中的节点会被降低（lower）到由循环嵌套构成的算子IR。这种两级IR架构使得在同一框架内支持训练和代码生成变得困难。首先，如果编译器采用当前深度学习框架（如PyTorch和TensorFlow）中广泛使用的自动梯度算法，很可能会生成编译器无法降低的未识别算子。例如，胶囊卷积 【11，Transforming auto-encoders，2011，21st Int. Conf. Artif. Neural Netw. Artif. Neural Netw. Mach. Learn.】,【12，Dynamic routing between capsules，2017，Annu. Conf. Neural Inf. Process. Syst.】,【13，Matrix capsules with EM routing，2018，Int. Conf. Learn. Representations】 的梯度算子本身也是一个任何框架或编译器都不支持的新算子。其次，如果编译器选择在算子级别实现自动微分，那么它们很难应用算子融合等图优化，因为现有的基于循环级IR的自动微分算法要求算子必须由完美循环嵌套表示【33，Differentiable programming for image processing and deep learning in halide，2018，ACM Trans. Graph.】,【34，Automatic differentiation for tensor algebras，2017，arXiv:1711.01348v1】，而融合操作往往会产生不完美的循环嵌套。

2.3 动机

加速新型DNN训练的挑战：以胶囊卷积为例。本文使用胶囊卷积 【11，Transforming auto-encoders，2011，21st Int. Conf. Artif. Neural Netw. Artif. Neural Netw. Mach. Learn.】,【12，Dynamic routing between capsules，2017，Annu. Conf. Neural Inf. Process. Syst.】,【13，Matrix capsules with EM routing，2018，Int. Conf. Learn. Representations】 为例，阐述加速包含新算子的新型DNN训练所面临的挑战。胶囊卷积的计算表达式如下：
$C[b, k, p, q, i, j] = A[b, c, p*2+r, q*2+s, i, k] * B[k, c, r, s, k, j]$
其中A、B、C是张量。胶囊卷积是CapsuleNet 【11，Transforming auto-encoders，2011，21st Int. Conf. Artif. Neural Netw. Artif. Neural Netw. Mach. Learn.】,【12，Dynamic routing between capsules，2017，Annu. Conf. Neural Inf. Process. Syst.】,【13，Matrix capsules with EM routing，2018，Int. Conf. Learn. Representations】 的核心算子，旨在更好地建模层次关系。然而，目前没有任何库支持胶囊卷积。一种支持该算子的方法是将其拆分为一系列小的2D卷积，然后通过拼接这些2D卷积的部分结果来生成输出张量。

现有框架实现新算子的性能瓶颈。我们在Nvidia Tesla V100 GPU上展示了PyTorch和TensorFlow中胶囊卷积的性能，如图2所示，图中列出了问题规模、内核启动开销和设备利用率。PyTorch和TensorFlow中的实现使用8个2D卷积来组合胶囊卷积的结果，这种方式由于频繁的内核启动和低资源利用率而效率低下。相比之下，NeoFlow通过在单个内核内实现计算，实现了高加速比。若要在不使用代码生成的情况下在训练任务中使用胶囊卷积，我们必须手动推导胶囊卷积的梯度算子并实现底层代码，这需要大量的专业知识且非常耗时。因此，本文提出NeoFlow，利用代码生成技术来自动支持带有新算子的高性能训练。

A2 方法细节

3 NEOFLOW 概览

NeoFlow框架的三个核心部分。图3展示了NeoFlow的概览，它包含三个主要部分：模型定义、图编译和图执行。在模型定义部分，NeoFlow使用表达式来定义DNN模型。这些表达式被视为算子，而表达式的符号化输入和输出则被视为张量。定义好的模型会经过NeoFlow的自动模型转换进行优化，包括用于改善数据访问模式的布局转换和用于简化图结构的并行算子融合。模型转换之后，NeoFlow执行自动微分（autodiff）来获取梯度图，从而形成一个新的训练图。在图编译部分，NeoFlow将整个DNN图划分为子图。为了将这些子图编译成高效的库，NeoFlow在子图内部应用了复杂的优化，包括表达式融合、块/线程分解、数据缓冲区配置、循环展开和向量化。在图执行部分，NeoFlow按拓扑顺序调用编译部分生成的库来执行图计算和权重更新。NeoFlow的最终输出是生成的高性能库和一个训练好的模型。NeoFlow兼顾了编程生产力和高性能，允许用户通过在Python中编写数学表达式来设计新算子，以鼓励DNN中新算子的创新并推动突破性DNN的发展。

4 模型表示与优化

NeoFlow的单级IR系统。现有的张量编译器如TVM 【17，TVM: An automated end-to-end optimizing compiler for deep learning，2018，13th USENIX Symp. Oper. Syst. Des. Implementation】 和PlaidML 【18，PlaidML，2020，Intel】 由于其两级IR基础设施，在训练方面存在局限性。NeoFlow采用一种单级的基于表达式的IR系统来表示图和算子，这使得通过编写简单的表达式来定义DNN模型成为可能。接下来，我们将介绍使用这种基于表达式的表示法进行的自动微分和模型优化。

4.1 基于表达式的表示法

单级IR的设计理念。如第2.2节所述，现有编译器 【17，TVM: An automated end-to-end optimizing compiler for deep learning，2018，13th USENIX Symp. Oper. Syst. Des. Implementation】,【18，PlaidML，2020，Intel】,【19，Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions，2018，arXiv:1802.04730v3】 使用两级IR进行编译，这限制了它们对训练的支持。在NeoFlow中，我们为图和算子的编译使用单级IR，我们称之为基于表达式的表示法，因为该IR通过简单的算术表达式描绘了整个DNN模型的数学计算定义。为了在单级IR中捕获DNN模型的所有信息，我们关注两个方面：每层产生多少数据元素，以及每个数据元素是如何由该层产生的。第一个方面由张量声明捕获，而第二个方面由表达式捕获。在NeoFlow中，我们使用以下形式来表示一层内的计算：

$$\begin{aligned} B[x_1, \dots, x_N] = \mathbf{F}_R=\{r_1, \dots, r_L\} \\ (A_1[f_1^1(x_1, \dots, x_N, r_1, \dots, r_L), \dots, f_{M_1}^1(x_1, \dots, x_N, r_1, \dots, r_L)], \\ A_2[f_1^2(x_1, \dots, x_N, r_1, \dots, r_L), \dots, f_{M_2}^2(x_1, \dots, x_N, r_1, \dots, r_L)], \\ \dots, \\ A_K[f_1^K(x_1, \dots, x_N, r_1, \dots, r_L), \dots, f_{M_K}^K(x_1, \dots, x_N, r_1, \dots, r_L)]) \end{aligned}$$

其中B是该层的N维输出张量，$A_i$（$1 \le i \le K$）是该层的$M_i$维输入张量，F是计算操作（如ReduceAdd用于归约求和，Select用于条件操作等）。$R = \{r_1, \dots, r_L\}$是代表归约循环的索引变量（如果没有归约则为空集），而$f_{ij}$（$1 \le i \le K, 1 \le j \le M_i$）是一个关于$x_1, x_2, \dots, x_N$和$r_1, \dots, r_L$的函数，用于生成输入张量的相应访问索引。在大多数情况下，$f_i$是索引变量的（准）仿射变换。基于表达式的表示法为定义新算子和图结构提供了极大的灵活性，但同时也给高性能训练带来了巨大挑战。首先，深度学习框架围绕算子组织数据数组、梯度计算和图优化，并定义了一系列算子，每个算子都有特殊处理。然而，这种逐案设计原则无法处理各种任意表达式。其次，像CuDNN 【10，cuDNN: Efficient primitives for deep learning，2014，arXiv:1410.0759v3】 这样的手动优化库通常遵循固定的计算模式，如循环顺序、数据访问模式等，因此无法支持任意表达式。我们在NeoFlow中通过基于表达式的优化来解决这些挑战，具体将在后续章节中阐述。

表示法示例。图4展示了一个使用基于表达式表示法的图示例。图中的张量声明显示了每层产生的元素数量。图中有四个操作，每个操作对应一层。这四层包括了计算密集型和内存密集型操作。第一层（Op1）是填充操作，其填充边界宽度为2。第二层（Op2）是空洞卷积【35，Multi-scale context aggregation by dilated convolutions，2016，4th Int. Conf. Learn. Representations】，其步幅和空洞因子均为2。对于这一层，公式1中的F是ReduceAdd，归约循环是{c, r, s}。第三层（Op3）是depth2space 【36，Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network，2016，IEEE Conf. Comput. Vis. Pattern Recognit.】。最后一层（Op4）是两个张量的拼接，并带有一个到图输入张量的快捷连接。这个示例图展示了基于表达式表示法的通用性。

4.2 基于表达式的自动微分

梯度计算的挑战与NeoFlow的混合方法。梯度计算是DNN训练不可或缺的一部分。深度学习框架 【7，PyTorch: An imperative style, high-performance deep learning library，2019，Annu. Conf. Neural Inf. Process. Syst.】,【8，Tensorflow: A system for large-scale machine learning，2016，12th USENIX Symp. Oper. Syst. Des. Implementation】,【9，MXNet: A flexible and efficient machine learning library for heterogeneous distributed systems，2015，arXiv:1512.01274v1】 通常采用自动梯度（autograd） 【37，Autograd: Effortless gradients in numpy，2015，Int. Conf. Mach. Learn. Workshop】 来计算梯度。对于每个算子，算子库中都存在一个对应的梯度算子。autograd算法遍历DNN图，并找到从损失函数到输入张量的所有反向路径，每条路径都由算子库中的梯度算子组成。这种autograd算法要求梯度算子必须预先准备好。对于无法用库中现有算子表示的新算子，用户必须通过推导梯度公式和编写底层代码来手动实现相应的核函数。尽管如何自动获取科学计算表达式的梯度公式是众所周知的，但目前仍缺乏一个专为深度学习表达式设计的良好自动算法，因为在深度学习中，操作数都是张量，计算都是深度嵌套的循环 【34，Automatic differentiation for tensor algebras，2017，arXiv:1711.01348v1】。在NeoFlow中，我们开发了一种自动微分（autodiff）算法来计算表达式和整个DNN的梯度。我们的梯度算法是符号梯度算法 【34，Automatic differentiation for tensor algebras，2017，arXiv:1711.01348v1】 和autograd算法 【37，Autograd: Effortless gradients in numpy，2015，Int. Conf. Mach. Learn. Workshop】 的混合体。对于表达式，我们使用符号微分；对于整个DNN，我们使用autograd算法。

符号微分算法详解。接下来，我们重点介绍算法的符号部分，并详细解释其计算过程。形式上，对于公式1中给定的表达式，如果我们想计算对$A_i$的梯度，自动微分算法应生成一个如下表达式：

$$\begin{aligned} dA_i[z_1^i, \ldots, z_{M_i}^i] = \mathbf{H}_{R'=\{r_1', \ldots, r_P'\}}(dB[g_1(z_1^i, \ldots, z_{M_i}^i, r_1', \ldots, r_P'), \ldots, g_N(z_1^i, \ldots, z_{M_i}^i, r_1', \ldots, r_P')], \\ A_1[h_1^1(z_1^i, \ldots, z_{M_i}^i, r_1', \ldots, r_P'), \ldots, h_{M_1}^1(z_1^i, \ldots, z_{M_i}^i, r_1', \ldots, r_P')], \\ \ldots, \\ A_K[h_1^K(z_1^i, \ldots, z_{M_i}^i, r_1', \ldots, r_P'), \ldots, h_{M_K}^K(z_1^i, \ldots, z_{M_i}^i, r_1', \ldots, r_P')]) \end{aligned}$$

其中$dA_i$是$A_i$的梯度张量，dB是B的梯度张量，H是与F对应的梯度函数，$R' = \{r'_1, \dots, r'_P\}$是归约轴，$g_p$和$h_s$（$1 \le p \le N; 1 \le s \le K; 1 \le t \le M_s$）都是关于$z_1^i, \dots, z_{M_i}^i$和$r'_1, \dots, r'_P$的函数，用于生成数据访问的索引。自动微分的核心问题是找出索引函数$g_p$和$h_s$（$1 \le p \le N; 1 \le s \le K; 1 \le t \le M_s$）并确定$r'_1, \dots, r'_P$的迭代域。我们指出，当正向表达式中使用的索引函数（公式(1)中的$f_i$）是（准）仿射时，这个问题可以通过线性代数解决。形式上，该问题表示为以下线性方程组：

其中$x_1, \dots, x_N$和$r_1, \dots, r_L$被视为未知数，而$z_1, \dots, z_{M_i}$被视为常数。如果我们能解出公式3中的未知数，那么我们就可以轻易地将公式1中的未知数替换为解，并将其重组为公式2的形式。之前的工作 【34，Automatic differentiation for tensor algebras，2017，arXiv:1711.01348v1】 解决这个问题的前提是索引变换是纯仿射的，即$f_i$总是$x_1, \dots, x_N$和$r_1, \dots, r_L$的线性组合。这严重限制了DNN算子的范围，例如图4中所示的depth2space 【36，Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network，2016，IEEE Conf. Comput. Vis. Pattern Recognit.】 等重要算子因存在非纯仿射的整数除法和模运算而无法支持。

处理非仿射变换的方法。为了解决这个问题，我们将$f_{i1}, \dots, f_{iM_i}$中所有非仿射的子表达式替换为新变量，使得剩余的表达式仍然是纯仿射的。然后我们求解这个线性系统，最后通过特殊处理准仿射操作（如除法和模运算）将替换的子表达式带回到结果中。具体来说，如果一个整数除法子表达式$x/V$（V是常数）被$s_1$替换，为了解这个替换，我们首先为模运算子表达式$s_2 = x \% V$找到相应的替换，因为它们一起可以求解$x = s_1 * V + s_2$。如果我们找不到这样一对除法和模运算，我们则引入新的自由变量作为归约轴，例如$x = s_1 * V + r$，其中r是一个归约轴。最后，我们推断归约轴的迭代域来完成我们的符号梯度过程。有了所有前向算子的梯度表达式后，我们就可以像以前的框架 【7，PyTorch: An imperative style, high-performance deep learning library，2019，Annu. Conf. Neural Inf. Process. Syst.】,【8，Tensorflow: A system for large-scale machine learning，2016，12th USENIX Symp. Oper. Syst. Des. Implementation】,【9，MXNet: A flexible and efficient machine learning library for heterogeneous distributed systems，2015，arXiv:1512.01274v1】 一样，通过autograd算法构建训练图。

4.3 自动模型转换

高层模型优化。基于上述提出的表示法，我们可以在高层级执行自动模型转换，以进一步提升模型性能。我们考虑的转换包括布局转换、张量交换和并行融合。

布局转换（Layout Transformation）。布局对DNN的性能至关重要，原因有三。首先，合适的布局可以实现合并的数据访问（coalesced data access），从而提高内存访问效率。其次，通过使用具有规则地址间隔的向量加载，可以减少内存地址计算。第三，可以利用硬件特定的指令，如CPU的AVX和GPU的Tensor Core，来获得额外的加速。以往的编译器如TVM要求手动指定布局转换，且只提供有限的布局选择（如NCHW、NHWC等），这无法充分利用布局转换带来的加速。NeoFlow不要求手动指定，而是自动应用转换步骤来改变张量布局和相应的表达式。为此，NeoFlow首先在表达式中进行向量或矩阵的模式匹配，然后拆分张量的匹配维度，以产生一个打包好的最内层向量或矩阵。此外，NeoFlow能够将数据打包成可调大小的小向量或矩阵，以便进行性能调优。

布局转换示例与权衡。图5展示了两个布局转换的例子。a)部分是一个NCHW布局的卷积表达式。从此表达式生成的代码将以标量为单位加载输入数据和计算输出数据，因此数据加载和地址计算开销可能会拖慢程序。在b)部分，卷积的通道维度被拆分为两部分，内部部分被移动到最内层维度。这种布局转换允许对张量A和张量B进行向量加载，并减少了地址计算开销。在c)部分，我们展示了对卷积的额外维度拆分，以在表达式内产生一个最内层的GEMM计算。这种转换使得可以使用Tensor Core指令，但由于张量A被复制了s次（s是核宽度），增加了数据加载开销。布局转换开销与特殊加载/计算指令的效率之间存在权衡。布局转换会带来额外的数据传输开销，并需要额外的融合步骤来实现隐式转换。而且，将数据打包成向量或矩阵时使用的拆分因子也会影响性能。为了获得良好性能，NeoFlow将在代码生成过程中对相关因子（向量/矩阵大小）进行调优。

张量交换（Tensor Swapping）。仅靠布局转换不足以优化程序，因为它无法改变表达式中张量的顺序。对于大多数DNN模型，单层内的计算是可交换的。例如，如果在图5中交换张量A和张量B，结果保持不变。通过改变张量的顺序，我们有机会利用更好的向量化加速。例如，对于DNN模型中的矩阵乘法，表达式写为 $C[i, j] = \text{ReduceSum}(\{k\}, A[i, k] * B[k, j])$，因此向量化适用于张量C和B的维度j。但是，当j维度很小，不足以使用向量指令时，我们可以改变张量A和B的顺序，将表达式重写为 $C[j, i] = \text{ReduceSum}(\{k\}, B[j, k] * A[k, i])$，从而暴露维度i进行向量化，这可能带来更好的性能。在NeoFlow中，张量交换对矩阵乘法等二元操作启用。是否应用张量交换也是代码生成过程中的一个可调优因子。

并行融合（Parallel Fusion）。并行融合 【38，Optimizing DNN computation with relaxed graph substitutions，2019，Mach. Learn. Syst】,【39，TASO: Optimizing deep learning computation with automatic generation of graph substitutions，2019，27th ACM Symp. Oper. Syst. Princ.】 是将两个没有直接数据依赖但共享至少一个共同输入张量的算子合并为一个算子。并行融合有助于提高并行性。使用基于表达式的表示法，我们可以轻松检测到这种情况。例如，两个表达式 A1[i, k] * A2[k, j1] 和 A1[i, k] * A3[k, j2] 可以合并为 A1[i, k] * A[k, j]，其中A是A2和A3沿第二维度（称为可融合维度）的拼接（为了得到原始表达式的输出，还应附加一个额外的张量切片操作）。我们的并行融合算法工作如下。首先，对于每个算子（由一个带有输出B的表达式定义，为简单起见我们用B表示一个算子），我们首先通过枚举所有与B共享输入张量的算子并检查是否存在直接数据依赖来找到可融合的算子。如果没有，我们记录下这个可融合算子。然后，对于每个可融合算子B1和B的每个维度$x_p$，我们检查除了维度$x_p$之外的B的形状是否与B1的形状相同，如果是，我们称$x_p$为可融合维度。可能存在许多可融合维度，但我们只取第一个，并沿该维度将B与B1融合。

5 编译系统

编译系统的优化技术。为了为DNN模型生成高效的代码，NeoFlow采用了包括融合和调优在内的多种优化技术。对于训练任务，由于重计算和数据传输开销之间存在非平凡的权衡，融合和调优具有挑战性。我们首先介绍子图划分和融合，然后解释NeoFlow为代码生成进行的调优。

5.1 子图划分与融合

子图划分的必要性与方法。一个DNN图可能包含数十到数百个算子。对于一些大型网络，如ShuffleNet 【40，ShuffleNet: An extremely efficient convolutional neural network for mobile devices，2018，IEEE/CVF Conf. Comput. Vis. Pattern Recognit.】，可能有多达946个算子。基于探索的编译器 【17，TVM: An automated end-to-end optimizing compiler for deep learning，2018，13th USENIX Symp. Oper. Syst. Des. Implementation】,【19，Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions，2018，arXiv:1802.04730v3】 将花费数小时甚至数天来编译和生成整个训练图的代码。因此，这种编译方式会产生高得令人望而却步的开销，并严重延迟训练过程。NeoFlow通过将原始图划分为小的子图，并将这些子图融合成更大的子图来解决这个问题。最终，NeoFlow为每个融合后的子图生成一个内核。为了划分图，NeoFlow从输出到输入遍历表达式IR，以捕获整个图的生产者-消费者关系。对于两个操作之间的每条边（张量），NeoFlow创建一个新的占位符来替代原始张量，从而使这两个操作相互分离，形成独立的子图。结果，整个图被划分为一系列子图，每个子图包含一层的操作。这一系列子图随后将被传递给融合过程进行进一步优化。

训练场景下融合的挑战。为了高效地融合这些子图，NeoFlow需要解决重计算和数据传输开销之间的权衡。这种挑战源于训练场景，并且在之前的融合框架如TVM中的Relay 【21，Relay: A high-level IR for deep learning，2019，arXiv:1904.08368v2】 和TASO 【39，TASO: Optimizing deep learning computation with automatic generation of graph substitutions，2019，27th ACM Symp. Oper. Syst. Princ.】 中并未揭示，因为它们只关注推理任务。对于训练任务，我们面临两种计算：前向和后向。优化前向图和优化后向图是相互对偶的问题。如果前向图中存在一个张量，那么后向图中就会有一个操作来计算该张量的梯度，反之亦然。此外，前向图和后向图通过大量中间张量连接。例如，为了计算卷积层权重的梯度，需要该卷积层的原始输入，而这很可能是前几层（如ReLU 【41，Deep learning using rectified linear units (ReLU)，2018，arXiv:1803.08375v2】 层或另一个卷积层）的中间结果。因此，如果我们在前向图中融合两层并消除一个中间张量（减少数据传输开销），我们将需要在后向图中为需要该中间张量作为输入的相应操作重新生成这个张量（带来重计算开销）。同时，这个被消除张量的梯度操作也应该被融合，因为不再需要计算该张量的梯度。前向图和后向图之间的这种联系使得训练任务的融合更加困难。

融合决策分析与奖励函数。NeoFlow通过分析每个融合决策的收益和成本来应对这一挑战。为简化问题，我们一次考虑融合两个具有单个输出的操作。多输出的分析是一个对偶问题，我们在此省略讨论。我们根据输入数量（扇入）将算子分为两类：只有一个输入的我们称之为One，有多个输入的称之为Many。One和Many之间有四种组合，如图6所示。我们用$T_A, T_B, \dots$表示数据传输量，用$Op_1, Op_2, \dots$表示计算量。gA是张量A的梯度，其他张量类似。#Kernels显示了生成的内核数量。通过比较这些情况下的数据传输量、计算量和内核数量，我们有以下观察：
1. One + One 融合总是有益的。我们可以在不增加任何计算的情况下减少$4T_B$的数据传输开销。这种融合可以应用于ReLU 【41，Deep learning using rectified linear units (ReLU)，2018，arXiv:1803.08375v2】，Padding，Reshape等层。
2. One + Many 和 Many + One 的融合决策很可能是有利的，因为它们以增加一个操作的计算量为代价，减少了4-5个张量的传输需求。然而，这类融合决策是否有利，应通过精确计算来判断。
3. Many + Many 融合（例如卷积+卷积）增加了一个前向操作和一个后向操作的计算量，而数据传输量是否减少取决于中间张量的大小（如果$6T_C > T_D + T_E$，则数据传输量减少）。所以我们仍然需要进一步计算来判断这种融合是否有益。

奖励函数的定义与应用。在编译期间，NeoFlow采用一个简单的奖励函数来判断一个融合决策是否有利。预测真实的运行时行为是困难的，但奖励函数的原则是过滤掉那些注定性能低下的融合决策。因此，奖励函数使用峰值性能对每个融合决策进行乐观评估，如果在乐观假设下奖励仍然为负，NeoFlow就可以安全地拒绝该融合计划。我们用$P_d$表示片外内存和片上内存之间的峰值数据传输带宽，用$P_c$表示设备的峰值计算性能。我们还用L表示平均内核启动开销。$P_d, P_c$和L是设备特定的，可以通过基准测试【42，Dissecting the NVIDIA volta GPU architecture via microbenchmarking，2018，arXiv:1804.06826v1】来测量。奖励函数如下：

其中$\Delta T$是减少的数据传输量，$\Delta C$是减少的计算量，$\Delta K$是消除的内核启动次数。如果相应的指标增加，这些值可以为负。我们在图6中展示了四种融合决策的奖励R。One + One的奖励总是正的，而对于其他情况，则取决于具体的问题规模。NeoFlow在编译期间计算这些奖励，并贪婪地选择奖励最佳的决策。

5.2 调优与库生成

优化空间探索与代码生成。为了提供高性能，必须在优化空间内选择好的调优因子并生成高性能库。NeoFlow探索了几个关键的优化，包括布局转换因子、交换因子、分块因子、展开和流水线因子等。所有这些优化调度及其参数共同构成了一个巨大的调度空间。NeoFlow使用遗传算法来探索该调度空间。它首先生成一组初始参数，然后在硬件上评估性能。根据评估的性能，NeoFlow对参数进行变异以产生新的因子。这个过程重复数千次，最终找到一组好的参数。在GPU上，NeoFlow使用TVM生成CUDA代码，并使用NVCC 【43，CUDA NVCC compiler，2020】 生成可执行的二进制代码。NeoFlow也可以生成LLVM IR 【44，LLVM: A compilation framework for lifelong program analysis & transformation，2004，2nd IEEE/ ACM Int. Symp. Code Gener. Optim.】 通过LLVM后端直接生成PTX代码。此外，由于LLVM支持多种硬件设备，因此也可以支持其他平台。为了使用GPU中的Tensor Core等特殊指令，NeoFlow使用CUDA WMMA指令来执行矩阵加载、计算和存储。还支持对Tensor Core的额外优化，如寄存器分块和流水线。具体来说，NeoFlow在一个warp内加载1到4个片段（可调），这些片段存储在寄存器中。加载和计算通过双缓冲重叠，形成一个三级流水线（如图7所示）。

6 运行时系统

6.1 系统设计

运行时系统的两种模式。运行时系统可以在静态模式或动态模式下使用编译系统生成的优化库。

静态模式（Static Mode）。在静态模式下，子图被提前编译，库在运行时被调用。然而，为整个DNN图编译和生成库通常需要数小时或数天，但好处是生成的库可以在后续执行中重用，执行期间没有编译开销。这种模式通常在DNN模型稳定后使用。

动态模式（Dynamic Mode）。在动态模式下，子图被动态编译以生成库。我们的洞见是，我们可以立即以一个中等的调度开始训练，并随着训练的进行迭代地改进调度。早期迭代的模型执行也为后续迭代的子图和优化调度的选择提供了反馈。基于这一洞见，我们开发了一个调用增量编译的运行时。当训练开始执行生成的库并更新权重时，我们同时开始下一轮的子图编译，将编译和执行重叠起来。

6.2 实现细节

静态模式与动态模式的实现。对于静态模式，我们可以轻松地将子图替换为生成的库调用。对于动态模式，我们需要减轻运行时的编译开销。编译过程进一步分为探索和代码生成两部分。我们在运行时使用三个线程：一个用于探索，一个用于代码生成，一个用于执行。为了在线程间通信，我们使用一个名为schedule-queue的消息队列。我们还分配了一个名为function-cache的中间缓冲区来缓存生成的代码，该缓冲区由代码生成线程和执行线程共享。这三个线程的工作流程如下：探索线程探索优化调度空间，并在每轮结束时将当前最优调度放入schedule-queue。代码生成线程在schedule-queue等待调度，当新调度到来时，它根据该调度生成代码。每个生成的代码对应一个子图。如果function-cache中没有该子图的记录，生成的代码会立即保存在function-cache中。否则，其性能将与之前保存的代码进行比较，只有更好的才会被保存。为了获得性能评估，代码生成线程在目标设备上运行生成的代码并收集性能数据。这些性能数据被发送回去以指导搜索更好的调度。至于执行线程，它通过使用function-cache中的函数按顺序执行整个图。

动态模式示例。图8展示了一个动态模式的例子。在这个例子中，我们有四个子图要运行，我们将每个子图的探索分为四轮（每轮包含多个探索步骤）。方案1是静态模式，它一次性完成四轮探索，并为每个子图生成一个库。在所有库都生成后，训练才能开始。方案2是动态模式的一种实现，没有function-cache，我们重叠了探索、代码生成和训练线程（t1、t2和t3）。每轮的最优调度被用于代码生成和训练，但它们不会被存储供以后使用，因此执行线程必须等待探索线程提供新的调度。方案3是我们动态模式的最终实现，我们通过使用function-cache改进了方案2，因此除了第一次迭代外，执行不再等待探索。对于训练任务，动态模式可以有效地减少等待时间。我们在运行时使用多线程来加速编译和评估。对于GPU，我们还支持使用CUDA Graph 【45，CUDA grah，2021，Nvidia】 来进一步提高性能。

A4 实验

7.1 实验设置

评估基准与硬件。我们使用多种DNN来评估NeoFlow，基准测试包括CNN和RNN。一些DNN包含新算子，给深度学习框架带来了挑战。其中，CapsuleNet 【11，Transforming auto-encoders，2011，21st Int. Conf. Artif. Neural Netw. Artif. Neural Netw. Mach. Learn.】,【12，Dynamic routing between capsules，2017，Annu. Conf. Neural Inf. Process. Syst.】,【13，Matrix capsules with EM routing，2018，Int. Conf. Learn. Representations】 包含两层胶囊卷积和它们之间的三轮动态路由 【12，Dynamic routing between capsules，2017，Annu. Conf. Neural Inf. Process. Syst.】；ShuffleNet 【40，ShuffleNet: An extremely efficient convolutional neural network for mobile devices，2018，IEEE/CVF Conf. Comput. Vis. Pattern Recognit.】 不仅包含深度可分离卷积，还有通道混洗算子；MI-LSTM 【14，On multiplicative integration with recurrent neural networks，2016，Annu. Conf. Neural Inf. Process. Syst.】，SCRNN 【15，Learning longer memory in recurrent neural networks，2015，3rd Int. Conf. Learn. Representations】，subLSTM 【47，Cortical microcircuits as gated-recurrent neural networks，2017，Annu. Conf. Neural Inf. Process. Syst.】，和LLTM 【48，PyTorch LLTM Implementation，2020，Facebook】 是LSTM 【4，Long short-term memory recurrent neural network architectures for large scale acoustic modeling，2014，15th Annu. Conf. Int. Speech Commun. Assoc.】 的变体。我们还评估了常见模型，如ResNet 【2，Deep residual learning for image recognition，2016，IEEE Conf. Comput. Vis. Pattern Recognit.】，MobileNet 【49，MobileNets: Efficient convolutional neural networks for mobile vision applications，2017，arXiv:1704.04861v1】，和Bert 【5，BERT: Pre-training of deep bidirectional transformers for language understanding，2018，arXiv:1810.04805v2】。所有性能实验均在Nvidia Tesla V100 GPU上完成。

对比基线与编译时间。我们将NeoFlow与PyTorch 【7，PyTorch: An imperative style, high-performance deep learning library，2019，Annu. Conf. Neural Inf. Process. Syst.】，TensorFlow 【8，Tensorflow: A system for large-scale machine learning，2016，12th USENIX Symp. Oper. Syst. Des. Implementation】，CuDNN 【10，cuDNN: Efficient primitives for deep learning，2014，arXiv:1410.0759v3】，和AutoTVM 【29，Learning to optimize tensor programs，2018，Annu. Conf. Neural Inf. Process. Syst.】 进行比较。我们基线的设置如表1所示。接下来，我们将评估不同批量大小下的训练和推理性能。除了Tensor Core外，所有网络都使用FP32数据类型。在Tensor Core上，我们使用FP16数据类型。每个基准测试的编译时间大约需要8-10小时（约60-260轮）。我们还在第7.3节讨论了我们运行时的优势。

7.2 性能结果

训练性能。图9比较了NeoFlow与PyTorch在七个模型上的性能，批量大小从1到64不等。我们还展示了ResNet-50的性能。对于批量大小为1的情况，NeoFlow的性能是PyTorch的1.61倍。但对于更大的批量大小，NeoFlow无法超越PyTorch，因为反向图中的转置卷积算子在PyTorch中通过隐式GEMM算法进行了特别优化，而NeoFlow目前尚未实现该算法。但对于其他由新算子组成的模型，NeoFlow的加速效果显著。对于CapsuleNet，由于频繁的内核启动和较低的设备利用率，PyTorch运行缓慢。对于ShuffleNet，NeoFlow可以探索更深层次的融合，例如将ReLU 【41，Deep learning using rectified linear units (ReLU)，2018，arXiv:1803.08375v2】 与通道混洗算子融合，以及将卷积层与批量归一化层融合。ShuffleNet的深度可分离卷积 【50，Xception: Deep learning with depthwise separable convolutions，2016，arXiv:1610.02357v3】 在库中由于优化不佳而运行缓慢。但在NeoFlow中，通过探索各种分块、绑定和展开调度，该算子得到了特别优化。对于MI-LSTM、SCRNN、subLSTM和LLTM，它们的计算是LSTM的变体，因此它们的PyTorch实现必须使用多步的GEMM、加法和激活函数来获得最终结果。但在NeoFlow中，我们通过将GEMM和逐元素操作融合在一起，生成了一个高效的库。总体而言，NeoFlow在这些DNN上实现了3.16倍的几何平均加速比。当启用CuDNN时，对PyTorch的加速比为1.92倍。

与TensorFlow的比较。图10显示了与TensorFlow的比较结果。与PyTorch的实现类似，我们通过用多个现有算子替换新算子来手动实现基线。我们展示了带和不带XLA的TensorFlow性能。相对于带XLA的TensorFlow，平均加速比为2.43倍。XLA是一个代码生成编译器，能为新模型生成高效代码。但XLA的性能不如NeoFlow，因为XLA无法调整优化参数以获得更高性能。NeoFlow可以在代码生成期间探索一个调度空间以找到更好的优化选择。

使用LLVM后端的推理性能。我们在图11的a)部分展示了NeoFlow的推理性能。特别地，我们使用LLVM 【44，LLVM: A compilation framework for lifelong program analysis & transformation，2004，2nd IEEE/ ACM Int. Symp. Code Gener. Optim.】 后端在NeoFlow中生成PTX代码。结果显示，对于批量大小为1和16的情况，NeoFlow可以超越PyTorch（带CuDNN）和TensorFlow（带XLA）。相对于带CuDNN的PyTorch，几何平均加速比为6.72倍，相对于带XLA的TensorFlow，加速比为4.96倍。

大型模型评估。我们还使用大型模型评估了性能。我们在图11的b)部分展示了Bert 【5，BERT: Pre-training of deep bidirectional transformers for language understanding，2018，arXiv:1810.04805v2】 编码器的性能。Bert的解码器目前无法支持，因为解码器有动态循环，而代码生成编译器需要静态循环结构。我们使用Bert-base配置，该配置在批量大小为1时可以占满V100的设备内存（批量大小为16将超出内存容量）。根据结果，NeoFlow的推理和训练执行时间与PyTorch相当。

与AutoTVM的比较。我们与AutoTVM进行了推理性能比较，结果如图11的c)部分所示。AutoTVM需要Relay 【21，Relay: A high-level IR for deep learning，2019，arXiv:1904.08368v2】 来处理图级别的代码生成。Relay的IR与AutoTVM使用的IR不同，且Relay的autograd模块与AutoTVM不兼容。因此，我们无法使用AutoTVM获得训练性能。在推理性能方面，NeoFlow与AutoTVM相当（几何平均加速比为1.08倍）。

使用FP16 Tensor Core的代码生成。在V100 GPU上，NeoFlow还可以通过生成CUDA WMMA指令来利用FP16 Tensor Core加速。我们在图11的d)部分展示了使用Tensor Core的NeoFlow性能。相对于PyTorch的几何平均加速比为2.31倍。

7.3 案例研究：MI-LSTM

MI-LSTM模型与NeoFlow的动态模式。为了评估NeoFlow的运行时系统，我们以MI-LSTM 【14，On multiplicative integration with recurrent neural networks，2016，Annu. Conf. Neural Inf. Process. Syst.】 基准作为案例研究。MI-LSTM包含数百个子图（35个是唯一的）。它将LSTM 【4，Long short-term memory recurrent neural network architectures for large scale acoustic modeling，2014，15th Annu. Conf. Int. Speech Commun. Assoc.】 中的原始方程 $f(Wx + Uz + b)$ 替换为 $f((Wx \odot Uz) + b)$，其中$\odot$是哈达玛积 【14，On multiplicative integration with recurrent neural networks，2016，Annu. Conf. Neural Inf. Process. Syst.】。这样一个新方程在深度学习框架中没有得到很好的支持，但NeoFlow可以用一个表达式来支持它。我们使用运行时的动态模式来训练MI-LSTM，并展示动态模式如何让训练提前开始并提高整体系统效率。

动态模式的性能优势。结果如图12所示。我们将批量大小设置为64，并运行训练任务10小时。图12的上半部分显示了每次迭代的延迟（单位为毫秒）。蓝色虚线是我们的基线（PyTorch），红色实线是NeoFlow。NeoFlow的训练因编译延迟了3.96分钟开始，初始延迟为27.791毫秒。但在19.11分钟后，延迟降至5.35毫秒，优于PyTorch的5.49毫秒。图12中的波动是由目标设备上的评估噪声引起的，但延迟在大约1.5小时后变得稳定。动态模式允许用户立即（几分钟内）开始训练，而无需等待整个DNN图的编译完成（数小时）。另一方面，我们在图12的下半部分展示了随时间完成的迭代次数（以$10^6$次迭代为单位）。PyTorch以稳定的速度进行，而NeoFlow在第一个小时内就超过了它。尽管开始时落后，NeoFlow完成的迭代次数在41分钟后超过了PyTorch，并且在10小时内完成的最终迭代次数大约是PyTorch的2.22倍。

7.4 模型转换结果

模型转换优化的效果评估。NeoFlow使用了三种重要的优化：布局转换、张量交换和并行融合。我们使用2D卷积、GEMM和胶囊卷积来评估它们的效果。对于布局转换，我们测试了一个2D卷积（特征图尺寸14x14，通道数1024，核尺寸1x1）在不同批量大小下的情况。图13显示了布局转换的效果。默认情况下，我们使用NCHW布局，当批量大小大于1时，它总是比PyTorch慢。NeoFlow自动将卷积转换为CHWN布局并利用向量化优化，比PyTorch快1.25倍。对于张量交换，我们测试了四种GEMM形状。如图13所示，x轴是形状M x N x K。通过交换操作数的顺序，性能平均提高了20.19%。对于并行融合，我们测试了一个胶囊卷积。我们特意在NeoFlow中用8个2D卷积和张量拼接来实现它。我们比较了有并行融合和没有并行融合的性能。并行融合将8个2D卷积融合在一起。通过并行融合，我们可以将性能提高到9.66倍。

A7 补充细节

8 相关工作

现有深度学习框架与编译器的局限性。深度学习框架如PyTorch 【7，PyTorch: An imperative style, high-performance deep learning library，2019，Annu. Conf. Neural Inf. Process. Syst.】，TensorFlow 【8，Tensorflow: A system for large-scale machine learning，2016，12th USENIX Symp. Oper. Syst. Des. Implementation】，和MXNet 【9，MXNet: A flexible and efficient machine learning library for heterogeneous distributed systems，2015，arXiv:1512.01274v1】 依赖于手动优化的库，包括CuDNN 【10，cuDNN: Efficient primitives for deep learning，2014，arXiv:1410.0759v3】 和MKL 【26，MKL，2020，Intel】，来加速DNN的推理和训练。这些库只为少数算子进行了特别优化，因此框架受限于一个不透明且不灵活的算子库，难以加速新颖的算子和网络架构。另一方面，深度学习编译器如Halide 【16，Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines，2013，ACM SIGPLAN Conf. Program. Lang. Des. Implementation】，TACO 【51，The tensor algebra compiler，2017，Proc. ACMon Program. Languages】，TVM 【17，TVM: An automated end-to-end optimizing compiler for deep learning，2018，13th USENIX Symp. Oper. Syst. Des. Implementation】，PlaidML 【18，PlaidML，2020，Intel】，和Tensor Comprehensions 【19，Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions，2018，arXiv:1802.04730v3】 为深度学习代码生成提供了高度的灵活性。它们通过高级表达式和编译器优化为定制算子生成底层代码。Latte 【52，Latte: A language, compiler, and runtime for elegant and efficient deep neural networks，2016，37th ACM SIGPLAN Conf. Program. Lang. Des. Implementation】 和SWIRL 【53，SWIRL: High-performance many-core CPU code generation for deep neural networks，2019，Int. J. High Perform. Comput. Appl.】 可以在CPU上为推理和训练生成高性能代码，但它们专注于卷积和池化等典型层，而我们则专注于新颖的模型。为了通过这些编译器为DNN算子获得高性能，Halide autoscheduler 【28，Learning to optimize halide with tree search and random programs，2019，ACM Trans. Graph.】 使用树搜索和随机程序；Chameleon 【30，Chameleon: Adaptive code optimization for expedited deep neural network compilation，2020，8th Int. Conf. Learn. Representations】 和FlexTensor 【31，FlexTensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system，2020，Int. Conf. Architectural Support Program. Languages Oper. Syst.】 提出使用强化学习方法，而Ansor 【32，Ansor: Generating high-performance tensor programs for deep learning，2020，arXiv:2006.06762】 使用进化搜索算法。然而，这些优化框架要么只关注单个算子，要么只关注推理场景，不适用于DNN训练加速。

代码生成技术与DNN框架的集成。将代码生成技术集成到DNN框架中以加速DNN越来越受到关注。PyTorch 【7，PyTorch: An imperative style, high-performance deep learning library，2019，Annu. Conf. Neural Inf. Process. Syst.】 和Relay 【21，Relay: A high-level IR for deep learning，2019，arXiv:1904.08368v2】 使用TVM为DNN图生成代码，但仅限于推理任务。Relay提供的训练支持有限，尚不可用。nGraph 【22，Intel ngraph: An intermediate representation, compiler, and executor for deep learning，2018，arXiv:1801.08058v2】 接收用深度学习框架（如TensorFlow）编写的模型定义，并通过PlaidML生成底层代码。由于其固定的接口和有限的算子支持，它无法直接支持新颖的DNN。TensorFlow 【8，Tensorflow: A system for large-scale machine learning，2016，12th USENIX Symp. Oper. Syst. Des. Implementation】 和JAX 【23，Compiling machine learning programs via high-level tracing，2018，Syst. Mach. Learn.】 使用XLA 【24，TensorFlow XLA，2020，Google】 为训练生成代码，但用户只能用受支持的（或可追踪的）原语定义DNN。为了支持表达式的自动微分，【33，Differentiable programming for image processing and deep learning in halide，2018，ACM Trans. Graph.】,【54，DLVM: A modern compiler infrastructure for deep learning systems，2017，arXiv:1711.03016】 提供了基于表达式的自动微分，但它们在图级别的优化是有限的。

A5 结论

随着深度学习的快速发展，加速新型DNN的训练需求巨大。本文介绍了一个名为NeoFlow的灵活框架，旨在为高性能DNN训练提供高效的编译。NeoFlow支持直接通过表达式定义模型，并提供基于表达式的自动微分来支持训练任务。在GPU上对新型DNN进行的训练实验表明，NeoFlow可以取得更好的性能。与PyTorch、TensorFlow和CuDNN相比，其几何平均加速比分别达到了3.16倍、2.43倍和1.92倍。