👾深度神经网络编译器原理简介
Last updated
Last updated
source from https://microsoft.github.io/AI-System/,thanks Microsoft for their contribution to AI education.
编译器（compiler）在计算机语言编译中往往指一种计算机程序，它会将某种编程语言写成的源代码（原始语言）转换成另一种编程语言（目标语言），在转换的过程中进行的程序优化就是编译优化过程。下图展示了经典计算机程序语言中较成功的开源编译框架项目LLVM的编译过程示意图。
随着深度学习的应用场景的不断泛化，深度学习计算任务也需要部署在不同的计算设备（如服务器，个人电脑，手机，手表，机器人等）和不同的硬件架构上（如X86, ARM, RISC-V等）；同时，实际部署或训练场景对性能往往也有着更为激进的要求，例如一个大规模的上线部署的模型在计算性能上的优化可以直接转换为计算成本的节省，同时，对深度学习任务来说，性能的优化能够为算法提供更大的空间，从而能让算法开发人员在合理的时间内尝试更大或更复杂的模型。这些需求在通用的深度学习计算框架中已经难已得到满足。由于深度学习计算任务在现有的计算框架中往往以DSL（Domain Specific Language）的方式进行编程和表达，这本身使得深度学习计算任务的优化和执行天然符合传统计算机语言的编译和优化过程。因此，与传统程序语言编译器类似，深度神经网络编译器的提出主要是解决多种设备适配性和性能优化的问题。具体来说，深度神经网络编译就是将当前的深度学习计算任务通过一层或多层中间表达进行翻译和优化，最终转化成目标硬件上的可执行代码的过程。
上图展示了一个典型的深度神经网络编译器架构图。与传统编译器类似，深度神经网络编译器也分为前端（Frontend）、后端（Backend）、中间表达（Intermediate Representation, IR）和优化过程（Optimizaiton Pass）等。
深度神经网络编译器的前端一般共享深度学习框架的前端表达，如TensorFlow和PyTorch，即一般为基本Python的DSL（Domain Specific Language）。这样的编程语言中的基本数据结构一般为张量（Tensor），用来描述由基本数据类型（如int, float, string等）元素构成的高维的数组，其元数据可以由一个元素类型和张量形状（如[128, 512]）来表示。在张量上进行的基本计算操作称作算子（Operator），通常为一些基本的线性代数计算组成，如矩阵乘法、向量加减乘除等。下图列举了一些深度学习计算中常用的算子。
Add | Log | While |
Sub | MatMul | Merge |
Mul | Conv | BroadCast |
Div | BatchNorm | Reduce |
Relu | Loss | Map |
Tanh | Transpose | Reshape |
Exp | Concatenate | Select |
Floor | Sigmoid | ..... |
基于张量和基本算子，当前深度学习框架一般利用Python作为胶水语言，把一个深度学习的计算模型描述成一系列算子的操作。下面展示了一个简单的MNIST训练模型在PyTorch中的实现。
深度神经网络编译器的后端指最终变化后的代码要执行的设备或神经网络加速器，目前常见的支持深度学习的计算设备有CPU、GPU、FPGA、TPU等其它专用加速器。不同的类型的计算设备往往采用完全不同的芯片架构，从而对应的编程模型和优化也完全不同。如CUDA GPU采用的是多个并行的流式处理器（Streaming Multiprocessor）和共享内存的架构，在GPU上执行的代码需要符合SIMT（Single Instruction Multiple Threads）的计算模型；而CPU一般采用的是多核架构，以及多线程模型（如线程池）来实现高性能的计算任务。更进一步，尽管是相同类型的设备，不同的型号都会有不同的硬件参数，如内存大小、算力、带宽等，这些参数都会极大的影响编译优化过程。
从前端语言到后端代码的编译过程，和传统编译器类似，需要经过若干中间表达（Intermediate Representation, IR）。目前在神经网络编译器中较为常用的中间表达主要包括计算图（DAG）和算子表达式等。计算图作为连接深度学习框架和前端语言的主要格式，也是标准化深度学习计算模型的常用格式，如ONNX格式即为一种深度学习模型的标准可交换格式，目前主流框架如TensorFlow和PyTorch的大部分程序都可以被转换或导出成ONNX格式。除此之个，每个DNN编译器也会定义自己的计算图格式，一般这些计算图都可以进行等价互相转换。计算图的节点是算子，边表示张量，所有的节点和边构成一张有向无坏图，节点之间的依赖关系表示每个算子的执行顺序。下图为一个简单的计算图示例。
除计算图之外，在将算子继续向下转换到下层并生成设备代码时，算子表达式（Tensor Expression）作为另一类中间表达被广泛使用在不同的神经网络编译器中，如TVM、Ansor、Tensor Comprehension等。算子表达式的主要作用是描述算子的计算逻辑，从而可以被下层编译器进一步翻译和生成目标设备的可执行代码。下面代码展示了TVM中的算子表达式形式，其表达了一个矩阵乘法算子的计算逻辑。
最后，深度神经网络编译器的优化过程（Optimization Pass）是构成一个编程器的最核心部分，是定义在每一种中间表达上的函数，其输入是某一种中间表达，经过一系统优化和变化过程，输出一个新的被优化过后的中间表达。如在计算图上就有非常多的经典优化过程，如常数传播、公共子表达式消除等，这些过程对输入的计算图经过一系列等价变化从而输出新的计算图。也有一些优化过程是将一个高层的中间表达翻译到低层的中间表达，甚至最终的可执行代码，这些优化过程又可分为设备相关和设备无关的优化。
计算图作为连接深度学习框架和前端语言的主要中间表达，被目前主流框架如TensorFlow和PyTorch所使用或者作为标准文件格式来导出模型。 计算图是一个有向无环图（DAG），节点表示算子，边表示张量或者控制边（control flow），节点之间的依赖关系表示每个算子的执行顺序。
计算图的优化被定义为作为在计算图上的函数，通过一系列等价或者近似的优化操作将输入的计算图变换为一个新的计算图。其目标是通过这样的图变换来化简计算图，从而降低计算复杂度或内存开销。在深度神经网络编译器中，有大量优化方法可以被表示为计算图的优化，包括一些在传统程序语言编译器中常用的优化方法。下图列举了一些常见的图优化方法：
一类最常见的计算图优化就是算术表达式化简，在计算图中的一些子图所对应的算术表达式，在数学上有等价的化简方法来简化表达式，这反应在计算图上就是将子图转化成一个更简单的子图（如更少的节点），从而降低计算量。下图展示了一个利用算术表达式化简计算图的例子，左边的子图包含了两个算法：Const算子（返回元素值为0的常量张量）和Mul算子（计算两个相同形状的算子的元素乘积），通过表达式化简，这个子图可以直接被化简成右边的只包括Const算子的子图。下表i列举了一些常见的算术表达式化简规则，其中X和Y表示张量，0和1表示常量张量，其它操作符均对应张量上的算子。
变化前 | 变换后 |
X * 0 | 0 |
X * Broadcast(0) | Broadcast(0) |
X * 1 | X |
X * Broadcast(1) | X |
X + 0 | X |
X + Broadcast(0) | X |
Log(Exp(X)/Y) | X - Log(Y) |
公共子表达式消除（Common Subexpression Elimination, CSE）也是经典编译优化中常用的优化。其目的是通过找到程序中等价的计算表达式，然后通过复用结果的方式消除其它冗余表达式的计算。同理，在计算图中，公共子表达式消除就等同于寻找并消除冗余的计算子图。一个简单的实现算法是按照图的拓扑序（保证一个访问一个节点时，其前继节点均已经访问）遍历图中节点，每个节点按照输入张量和节点类型组合作为键值进行缓存，后续如果有节点有相同的键值则可以被消除，并且将其输入边连接到缓存的节点的输入节点上。下图为一个公共子表达式消除的示例，图左边蓝色椭圆中的节点为不在缓存中的节点，也就是必须要执行的节点，而红色椭圆中的节点的计算和前面蓝色节点的计算重复，也就是其键值可以被缓存到，因此可以安全的消除这些节点，于是最终左边的计算图可以被优化成右边的计算图。
考虑到下列代码：
可以观察到 b * c 是两项表达式中的公共子表达式。如果计算这个子表达式并将其计算结果存储起来的开销，低于重复计算这个子表达式的开销，则能够将以上代码转换成以下代码：
常数传播（constant propagation）就叫常数折叠（constant folding），也是经典编译优化中的常用优化，其主要方法是通过在编译期计算出也是常数表达式的值，用计算出的值来替换原来的表达式，从而节省运行时的开销。在计算图中，如果一个节点的所有输入张量都是常数张量的话，那么这个节点就可以在编译期计算出输入张量，并替换为一个新的常数张量。如下图所示为一个常数传播的示例，其中红色方框内的两个节点都可以被提前计算出来，因此可以在编译期优化掉。值得注意的是，常数传播需要编译器具有计算的能力，甚至对于一些较大的算子还需要能够在加速硬件上（如GPU）上计算，否则优化的过程就会非常的慢。常数传播的优化在深度学习尤其是模型推理的时候非常有用，因为在推理时，模型中的参数张量全部固定为常数张量，大量计算可以在编译期计算好，极大的化简了推理运算时的计算开销。但是，在深度学习的场景中，常数传播有时候也会带来否优化，如增加计算内存甚至计算时间，一个典型的例子就是一个标量常数张量后面跟一个Broadcast的算子时，如果做了常数传播就会增加内存占用，如果后面是访存密集型的算子的话，也会增加内存压力，从而增加计算时间。
矩阵乘在深度学习计算图中被广泛应用，如常见的神经网络的线性层、循环神经网络的单元层、注意力机制层等都有大量的矩阵乘法。在同一个网络里，经常会出现形状相同的矩阵乘法，根据一些矩阵的等价规则，如果把些矩阵乘算子融合成一个大的矩阵乘算子，可以更好的利用到GPU的算力，从而加速模型计算。下图为其中一种常见的矩阵乘自动融合的示例，其中，如果有两个矩阵乘法共享同一个输入张量（图中方框内左侧），我们就可以自动把其中的两个输入张量拼接成一个大的矩阵乘算子（图中方框内右侧），其计算的结果刚好是原算子计算结果的拼接。利用这种规则，图中最右侧的GRU网络中的两组矩阵乘算子可以分别融合成两个大的矩阵乘算子。类似的融合规则还有BatchMatMul，可以把两个相同形状的矩阵拼接成一个新的BatchMatMul算子。
上小节中介绍的算子融合方法是针对矩阵乘算子特有的，在深度学习模型中，针对大量的小算子的融合都可以提高GPU的利用率，减少内核启动开销、减少访存开销等好处。例如，Element-wise的算子（如Add，Mul，Sigmoid，Relu等）其计算量非常小，主要计算瓶颈都在内存的读取和写出上，如果前后的算子能够融合起来，前面算子的计算结果就可以直接被后面算子在寄存器中使用，避免数据在内存的读写，从而提交整体计算效率。下图展示了一个Mul算子和一个Add算子融合的示例，还有其对应的融合前后的CUDA代码示例，在没有融合前，执行两个算子需要启动两个GPU内核，前一个计算的结果需要写出到主存中，下一个内核计算的时候需要再次读取到计算核上。然后，融合后的代码只需要启动一个内核，并且可以有效复用中间计算结果。
鉴于算子融合在深度学习计算中能够带来较好的性能优化，然而在实际的计算图中有太多算子无法做到自动的算子融合，主要原因包括算子的内核实现逻辑不透明、算子之前无法在特有加速器上融合等等。为了在这些的情况下还能进行优化，用户经常会实现一些手工融合的算子来提升性能。那么，编译器在计算图中识别出一个子图并替换成一个等价的新的算子或子图的过程就是子图替换优化。下图展示的是基于规则的子图替换示例，需要在系统中注册系列替换规则，如Conv和Relu的子图可以替换为Conv+Relu融合后的算子。
深度学习的计算任务大多执行在像GPU这样的加速器上，一般这样的加速器上的内存资源都比较宝贵，如几GB到几十GB的空间。随着深度学习模型的规模越来越大，从近来的BERT，到各种基于Transformer网络的模型，再到GPT-3等超大模型的出现，加速器上的内存资源变得越来越稀缺。因此，除了计算性能之外，神经网络编译器对深度学习计算任务的内存占用优化也是一个非常重要的目标。
一个深度学习计算任务中的内存占用主要包括输入数据、中间计算结果和模型参数，在模型推理的场景中，一般前面算子计算完的中间结果所占用的内存，后面的算子都可以复用，但是在训练场景中，由于反向求导计算需要使用到前向输出的中间结果，因此，前面计算出的算子需要一直保留到对应的反向计算结束后才能释放，对整个计算任务的内存占用挑战比较大。所幸的是，在计算图中，这些所有的数据都被统一建模成计算图中的张量，都可以表示成一些算子的输出。计算图可以精确的描述出所有张量之前的依赖关系以及每个张量的生命周期，因此，根据计算图对张量进行合理的分配，可以尽可能的优化计算内存的占用。
上图展示了一个根据计算图优化内存分配的例子，在上图中，默认的执行会为每一个算子的输出张量都分配一块内存空间，假设每个张量的内存大小为N，则执行该图需要4N的内存。但是通过分析计算图可知，其中的张量a可以复用张量x，张量c可以复用a，因此，总的内存分配可以降低到2N。 基于计算图进行内存优化的方法有很多，本章中主要以三类不同的方法为例具体介绍如果介绍深度学习计算中的内存优化。
计算图中的张量内存分配可以分成两个部分：张量生命期的分析和内存分配。首先，给定计算图之后，唯一决定张量生命期的就是节点（算子）的执行顺序。在计算框架中，由于执行顺序是运行时决定的，所以内存也都是运行时分配的。但在编译器中，我们可以通过生成固定顺序的代码来保证最终的节点以确定的顺序执行，因此在编译期就可以为所有张量决定内存分配的方案。一般只要以某种拓扑序要遍历计算图就可以生成一个依赖正确的节点的执行顺序，如BFS、Reverse DFS等，进而决定出每个张量的生命期，即分配和释放的时间点。
接下来，就是根据每个张量的分配和释放顺序分配对应的内存空间，使得总内存占用最小。一种常用的内存分配方法是建立一个内存池，由一个块内存分配管理器（如BFC内存分配器）管理起来，然后按照每个张量的分配和释放顺序依次向内存池申请和释放对应大小的内存空间，并记录每个张量分配的地址偏移。当一个张量被释放回内存池时，后续的张量分配就可以自动复用前面的空间。当所有张量分配完时，内存池使用到的最大内存空间即为执行该计算图所需要的最小内存。在真实的运行时，我们只需要在内存中申请一块该大小的内存空间，并按照之前的记录的地址偏移为每个张量分配内存即可。这样即可以优化总内存的占用量，也可以避免运行时的内存分配维护开销。 值得注意的是，不同拓扑序的选择会同时影响模型的计算时间和最大内存占用，同时也强制了运行时算子的执行顺序，可难会带来一定的性能损失。
上面的方法中只考虑了张量放置在加速器（如GPU）的内存中，而实际上如果内存不够的话，我们还可以将一部分张量放置到外存中（如CPU的内存中），等需要的时候再移动回GPU的内存中即可。虽然从CPU的内存到GPU的内存的拷贝延时和带宽都比较受限，但是因为计算图中有些张量的产生到消费中间会经过较长的时间，我们可以合理安排内存的搬运时机使得其和其它算子的计算重叠起来。
给定上述假设以及必要的数据（如每个内核的执行时间、算子的执行顺序等），关于每个张量在什么时间放在什么地方的问题就可以被形式化的描述成一个最优化问题。AutoTM就提出了一种把计算图中张量在异构内存环境中的问题建模成一个整数线性规划的问题并进行求解。上图展示了一个利用整数线性规划优化计算图内存分配的优化空间示例，图中每一行表示一个张量，每一列表示算子的执行顺序。每一行中，黄色Source表示张量的生成时间，紫色的SINK表示张量被消费的时间，每个张量都可以选择是在内存中（DRAM）还是外存（PMM）中。那么问题优化目标为就是给定任意的计算图最小化其执行时间，约束为主存的占用空间，优化变量就是决定放在哪个存储中，在有限的节点规模下，这个问题可以通过整数线性规划模型求解。同时，该文章中还扩展该方法并考虑了更复杂的换入换出的情形。
深度学习计算图的大多算子都是确定性的，即给定相同的输入其计算结果也是相同的。因此，我们可以进一步利用这个特点来优化内存的使用。当我们对连续的多个张量决定换入换出的方案时，如果产生这些张量的算子都具有计算确定性的话，我们可以选择只换出其中一个或一少部分张量，并把剩下的张量直接释放，当到了这些张量使用的时机，我们可以再换入这些少量的张量，并利用确定性的特点重新计算之前被释放的张量，这样就可以一定程序上缓解CPU和GPU之前的带宽压力，也为内存优化提供了更大的空间。如果考虑上换入换出，内存优化方案需要更加仔细的考虑每个算子的执行时间，从而保证重计算出的张量在需要的时候能及时的计算完成。
前面的编译优化基本都是在计算图的上进行的，当一个计算图被优化过后，就需要继续向下编译。其中一个最主要的问题就是如何对计算图中的每一个算子生成相应的代码。在计算框架中，每个算子都是预先实现并注册到框架中的，这样计算图在执行时只需要调用相应的代码即可。然而，计算框架的缺点是无法快速适配到一个新的硬件上，其需要为每一种硬件都实现一套算子代码，这不仅需要大量人力和时间成本，并且算子实现的性能也无法得到保证，因为，在对每个后端平台针对每个算子实现内核代码的时候都需要考虑不同的编程模型、数据排布、线程模型、缓存大小等等因素。
为了解决这个问题，就有了张量编译（或算子编译）的研究工作以及张量编译器。算子编译的核心思想是首先为通用算子找到一种能够描述算子与硬件无关的计算逻辑的表示，然后由编译器根据这种逻辑描述再结合具体的硬件生成相应的内核代码。近年来，有较多的研究工作都在围绕这个问题出现，例如TVM, Halide, TACO, Tensor Comprehension, FlexTensor等。
对深度学习中的大多数算子，其计算逻辑都可以描述成针对输出张量中的每一个元素的独立同构计算。以矩阵乘算子为例，矩阵C中的每一个元素（如坐标为[i,j])的值都可以通过对应的一行（第i行）和一列（第j列）的内积来计算得出。也就是说，大多数的算子的计算逻辑都要以通过描述其中的元素的计算逻辑来表示，这就是算子表达式的作用。
一个算子表达式主要包括以下几个部分：1）所有输入和输出张量，2）输出张量的计算形状，3）输出张量中每一个元素的计算表达式，其中包括元素的在张量中的位置参数，一般以lambda表达式的形式描述为坐标参数的匿名函数。如下面表中每一行为上述矩阵乘算子在TVM中的算子表达式。
算子 | 算子表达式 |
矩阵乘 |
|
仿射变换 |
|
卷积 |
|
ReLU |
|
有了算子表达式之后，我们就得到了一个算子的计算逻辑。为了生成硬件上的最终代码，我们需要把算子表达式的逻辑计算变化成符合硬件编程模型的代码，并考虑硬件特性进行代码优化，这个过程就叫作表达式的调度（Schedule）。 通常来说，一个最简单的调度方案就是通过生成多重循环来遍历一个算子表达式中输出张量中的每一个元素，然后调用其提供的lambda函数，即可完成一个简单的内核代码的生成。图5-4-2展示了一个简单的张量加算子的表达式，以及为其在TVM中创建一个默认调度的示例（上半部分），同时调度后产生出的内核代码（下半部分）。
可以看到，上面生成的内核代码只是一个简单的循环，实际中这样的代码往往性能不好。我们希望对上述循环进行一系列的变化，如把一个循环拆分成两重循环、或者把两个循环合并一个循环、或者把两个循环的顺序颠倒等等。为了方便这些优化，算子编译器也提供了一些相应的调度操作接口，如下图中的split操作即可以上述循环按照32为因子进行拆分成一个两重循环。
除了优化，我们还希望一个算子表达式能生成特定硬件上符合其编程模型的代码。这就需要我们能针对这些硬件提供一些调度操作。例如，当我们想让上述代码能在CUDA GPU上执行，我们就需要把一些循环绑定到CUDA编程模型中的threadIdx或blockIdx上，同样，我们可以使用算子编译器中的bind接口来完成，如下代码所示，最终我们就可以得到一个简单的可以GPU执行的内核代码。
有了算子表达式和对表达式的调度机制，我们就可以较容易的在一个新的硬件设备上生成一个算子的内核代码了。然而，我们可以看到，在调度的时候，有非常多种决定需要抉择，而且这些决定都会根据硬件的不同而产生不一样的性能影响，这些都需要经验非常丰富的专家才能知道一个较好的调度方案。为了进一步克复这个问题，一类利用机器学习进行自动调度搜索的方法被广泛应用。
给定一个算子表达式，我们首先需要针对该表达式自动生成出一个调度的代码模板，模板中可以预留出大量的可配置的参数。生成的模板需要能够尽可能包括各种代码的可能性，也就是保证足够大的搜索空间。给定了代码模板后，剩下的事情就是决定哪一个配置可以生成最优的代码，实际中，一个代码模板可能有成千上万种可选配置，因此，一般的编译器会采用机器学习的方法通过不断尝试，生成代码、测量性能、反馈给机器学习模型、再生成下一个（一批）代码的方式不断迭代搜索，直到搜索到一定的步数后找到一个较优的代码配置，并生成最终代码。通过机器学习的方法的好处是可以针对特别的问题输入和硬件，利用黑盒的方式找到一个较好的专用代码，但其缺点也很明显，在编译的过程中需要大量的编译和尝试，需要花费较长的编译时间和较多的算力。
前面的优化和算子生成分别在计算图和算子表达式两个层次完成。这种分层的优化给编译器的设计和实现带来更清楚的模块化和可维护性，但是同时也由于上下层的分离损失了一些更进一步的优化机会，例如硬件利用率低，无法完全发挥硬件的计算性能。造成这些低效的主要原因包括：
单个Op的调度时间与计算时间相比不可忽略，造成较大的调度开销；
OP的并行度不足以占满GPU的计算核心。
为了解决上述问题，一个很自然的想法就是：能不能对任意的算子进行融合，从而提高硬件利用率，降低算子的调度开销。一种最简单的方法就是实现更加激进的自动算子融合。如下图所示的为一个简单的计算图，与前面的算子融合不同的是，我们为了对任意算子做融合，引入了非element-wise算子，如矩阵乘法，这会给之前的融合方法增加难度。
为了实现任意算子融合，我们需要每一个算子的内核函数，如该例子中需要的Sigmoid、Relu和MatMul算子，下面展示的为这三个算子的内核函数：
为了按照前面讲到的的方法进行这三个算子的融合，我们需要将上述三个函数生成到同一个全局核函数内，如下图示例。值的注意的是，为了保证任意算子之前有正确的数据依赖，我们有时候需要在两个算子之间插入一个全局的数据同步。
上面的实现可以看出有很大的局限性和问题，例如这种方法打破了现有的模块化设计，内核融合过程需要对每个算子的内核函数有一定要求，并需进行二次修改，还需要获取其一些额外的隐式参数：如threadBlock的个数，大小等。更进一步，这种方法也引入了一些“非标准”的GPU用法，如在kernel内部做全局同步可能会引入死锁的问题 。 尽管有学术界可以使用持久化线程（Persistent threads）的方法来实现同步，但这种方法和GPU有较强的绑定，无法把优化过程通用化到其它硬件上，有大量GPU相关的实现细节混在其中。
为了更好的解决这种问题，我们就是需要一种能根据计算流图中的并行度以及算子内部的并行度的整体信息来进行一个全局的任务调度方法。本章中以Rammer的技术为例，介绍一种全局算子调度的优化方法。
首先，在计算表达层，为了打开现有算子的黑盒实现，Rammer引入rOperator来代替原有算子抽象，暴露出每一个算子中的所有并行任务（rTask）。 在硬件层，引入虚拟设备（vDevice）的抽象，并提供计算单元（vEU）级别的的粒度调度接口，中以允许将一个rTask调度到任意指定的vEU上。然而，rTask粒度的调度可能带来更严重的调度开销，Rammer利用DNN计算性能有较强的确定性，即算子的计算时间在数据流图运行前就可以通过测量得到。因此，在编译时可以将整个数据流图的所有rTask静态的编排成一个确定性执行方案，通过vDevice映射到物理Device的执行单元进行执行。
这种全局调度的抽象解耦了调度机制与优化策略，通过暴露出rTask粒度的调度接口，从而可以基于该接口设计任意的编排方案来优化整体性能。
