论文来源:Li M , Liu Y , Liu X ,et al.The Deep Learning Compiler: A Comprehensive Survey[J]. 2020.DOI:10.1109/TPDS.2020.3030548.
这是一篇关于深度学习编译器的综述类文章。
什么是深度学习编译器
深度学习(Deep Learning)编译器将深度学习框架描述的模型在各种硬件平台上生成有效的代码实现,其完成的模型定义到特定代码实现的转换将针对模型规范和硬件体系结构高度优化。具体来说,它们结合了面向深度学习的优化,例如层融合和操作符融合,实现高效的代码生成。
此外,现有的编译器还采用了来自通用编译器(例如LLVM)的成熟工具链,对各种硬件体系结构提供了更好的可移植性。
简而言之,深度学习编译器通过对深度学习模型和计算图的优化,以及对硬件资源的充分利用,提高了深度学习模型的性能和效率。这使得深度学习模型在边缘设备和嵌入式系统上能够更好地运行。
深度学习编译器的基本架构
与传统编译器类似,深度学习编译器也采用分层设计,其通用设计架构主要包括编译器前端和编译器后端,其中IR(中间语言)分布在前端和后端,它是程序的抽象,用于程序优化。
深度学习模型在深度学习编译器中被转换为多级IR,其中高阶IR驻留在前端,低阶IR驻留在后端。基于高阶IR,编译器前端负责与硬件无关的转换和优化。基于低阶IR,编译器后端负责特定于硬件的优化、代码生成和编译。
高阶IR,也称为graph IR,表示计算和控制流,与硬件无关。高阶IR的设计挑战是计算和控制流的抽象能力,能够捕获和表达不同的DL模型。高阶IR的目标是建立控制流和op(算子)与数据之间的依赖关系,以及为图级优化提供接口。它还包含用于编译的丰富语义信息,并为定制算子提供可扩展性。
低阶IR是为特定于硬件的优化和针对不同硬件目标的代码生成而设计的。因此,低阶IR应该足够细粒度,以反映硬件特征并表示特定于硬件的优化。
前端将现有的Deep Learning框架中的模型作为输入,然后将该模型转换为计算图表示(例如,graph IR)。前端需要实现多种格式转换,以支持不同框架下的多种格式。计算图优化结合了通用编译器的优化技术和Deep Learning特定优化技术,减少了graph IR的冗余,提高了图graph IR的效率。这样的优化可以分为:
-
node节点级(例如,nop消除和零维张量消除)
-
block块级(代数简化、算子融合、op下降)
-
数据流层(例如CSE、DCE、静态内存规划和布局转换)。
前端完成后,生成优化的计算图并传递给后端,后端将高阶IR转换为低阶IR并执行特定于硬件的优化。一方面,它可以直接将高阶IR转换为第三方工具链,如LLVM IR,以利用现有的基础设施进行通用优化和代码生成。另一方面,它可以利用DL模型和硬件特性的先验知识,通过定制的编译通道,实现更高效的代码生成。通常应用的特定于硬件的优化包括硬件内在映射、内存分配和获取、内存延迟隐藏、并行化以及面向循环的优化。
前端优化
在构建计算图之后,前端应用图级优化。许多优化更容易在图级识别和执行,因为图提供了计算的全局视图。这些优化只应用于计算图,而不是后端实现。因此,它们是独立于硬件的,可以应用于各种后端目标。
前端优化通常由passes定义,可以通过遍历计算图的节点并执行图转换来应用。前端提供了如下方法:
-
从计算图中捕获特定的特征;
-
重写图以进行优化。
除了预定义的passes,开发人员还可以在前端定义定制的passes。一旦一个Deep Learning模型被导入并转换为一个计算图,大多数DL编译器可以确定每个操作的输入张量和输出张量的形状。
节点级优化
计算图的节点足够粗粒度,可以在单个节点内进行优化。节点级优化包括:
-
节点消除(消除不必要的节点)
-
节点替换(用其他低成本节点替换节点)
在通用编译器中,Nop消除删除占用少量空间但不指定操作的no-op指令;在DL编译器中,Nop消除负责消除缺少足够输入的操作。例如,只有一个输入张量的和节点可以被消除,填充宽度为零的填充节点可以被消除。
块级优化
代数简化
代数简化优化包括:
-
代数识别
-
强度约简,我们可以用更便宜的算子取代更昂贵的算子
-
常数折叠,用它可以用它们的值替换常量表达式。
算子融合
算子融合是DL编译器不可缺少的优化方法。它能够更好地共享计算,消除中间分配,通过结合loop nests进一步优化,并减少launch和同步开销。
算子Sinking
算子Sinking这种优化将转置等操作置于批处理归一化、ReLU、sigmoid和channel shuffle等操作之下。
通过算子Sinking这种优化,许多相似的操作彼此之间移动得更近,为代数简化创造了更多的机会。
数据流级优化
通用子表达式消除(CSE)
如果表达式E的值之前已经计算过,并且E的值在之前的计算之后没有改变,那么它就是通用子表达式。
在这种情况下,E的值只计算一次,并且可以使用已经计算出的E的值来避免在其他地方重新计算。深度学习编译器在整个计算图中搜索公共子表达式,并用之前的计算结果替换下面的公共子表达式。
死代码消除(DCE)
如果一组代码的计算结果或副作用没有被使用,那么它就是死代码。
DCE优化会删除死代码。死代码通常不是由程序员引起的,而是由其他图优化引起的。因此DCE和CSE是在其他图优化之后应用的。
静态内存规划
使用静态内存规划优化以尽可能重用内存缓冲区,通常有两种方法:
-
in-place内存共享
-
标准内存共享
in-place内存共享使用相同的内存作为操作的输入和输出,只在计算之前分配一个内存副本。标准内存共享重用以前操作的内存,而不重叠。静态内存规划是离线完成的,这允许应用更复杂的规划算法。
布局转变
布局转换试图找到最佳的数据布局来存储张量到计算图中,然后将布局转换节点插入到图中。
注意,这里不执行实际的转换,而是在编译器后端计算计算图时执行转换。事实上,相同操作在不同的数据布局中的性能是不同的,最佳布局在不同的硬件上也是不同的。
后端优化
特定于硬件的优化
特定于硬件的优化也称为目标相关优化,用于获得针对特定硬件的高性能代码。应用后端优化的一种方法是将低层IR转换为LLVM IR,利用LLVM基础设施生成优化的CPU/GPU代码。另一种方法是使用DL领域知识设计定制优化,更有效地利用目标硬件。
以下有五种在现有DL编译器中广泛采用的方法:
硬件固有的映射
硬件内在映射可以将一定的一组低层IR指令转换为已经在硬件上高度优化的内核。在TVM中,硬件本征映射采用可扩展张量化方法实现,可声明硬件本征的行为和本征映射的降低规则。这种方法使编译器后端能够将硬件实现和高度优化的手工微内核应用到特定的操作模式中,从而获得显著的性能提升。
内存分配和获取
内存分配是代码生成中的另一个挑战,特别是对于gpu和定制加速器。例如,GPU主要包含共享内存空间(内存大小有限,访问时延较低)和本地内存空间(容量大,访问时延较高)。这样的内存层次结构需要有效的内存分配和获取技术来改善数据局部性。为了实现这一优化,TVM引入了内存作用域的调度概念。内存作用域调度原语可以将计算阶段标记为共享的或线程本地的。对于标记为共享的计算阶段,TVM生成具有共享内存分配和协同数据获取的代码,并在适当的代码位置插入内存屏障以保证正确性。
内存延迟隐藏
通过对执行管道重新排序,内存延迟隐藏也是在后端使用的一项重要技术。由于大多数DL编译器支持CPU和GPU上的并行化,内存延迟隐藏可以自然地通过硬件实现(例如,GPU上的扭曲上下文切换)。但是对于具有解耦访问库(DAE)体系结构的TPU类加速器,后端需要执行调度和细粒度同步以获得正确和高效的代码。为了获得更好的性能和减轻编程负担,TVM引入了虚拟线程调度原语,允许用户在虚拟多线程架构上指定数据并行度。然后,TVM通过插入必要的内存屏障,并将来自这些线程的操作穿插到单个指令流中,从而降低了这些实际上并行的线程,这形成了每个线程更好的执行管道,以隐藏内存访问延迟。
面向循环优化
面向循环的优化也应用于后端,为目标硬件生成高效的代码。由于Halide和LLVM已经集成了这样的优化技术,一些DL编译器在后端利用了Halide和LLVM。面向循环优化中应用的关键技术包括循环融合、滑动窗口、平铺、循环重排序和循环展开。
并行化
由于现代处理器通常支持多线程和SIMD并行性,编译器后端需要利用并行性来最大化硬件利用率,以实现高性能。
Halide使用一个叫做parallel的调度原语来指定线程级别并行化的循环的并行化维度,并通过将标记为parallel的循环维度与块和线程的注释进行映射来支持GPU并行化。它将一个大小为n的循环替换为一个宽为n的向量语句,该向量语句可以通过硬件内在映射映射到特定于硬件的SIMD操作码。
Stripe发展了一种多面体模型的变体,称为嵌套多面体模型,该模型引入了并行多面体块作为迭代的基本执行元素。在此扩展之后,一个嵌套的多面体模型可以检测平铺和跨步级别之间的层次并行性。
自动调优
由于在特定于硬件的优化中参数调优有巨大的搜索空间,因此有必要利用自动调优来确定最佳参数配置。在A Comprehensive Survey[4]这篇论文中研究的DL编译器中,TVM、TC和XLA支持自动调优。通常,自动调优实现包括Parameterization(参数化)、代价模型、搜索技术和加速四个关键部分。
Parameterization(参数化)
数据参数描述数据的规格,如输入的形状。目标参数描述在优化调度和代码生成期间要考虑的特定于硬件的特征和约束。例如,对于GPU目标,需要指定共享内存和寄存器大小等硬件参数。
优化选项包括优化调度和相应的参数,如面向循环的优化和瓷砖大小。在TVM中,既考虑了预定义调度,也考虑了自定义调度,还考虑了参数。
代价模型
自动调优中不同代价模型的比较如下:
-
黑箱模型:该模型只考虑最终的执行时间,而不考虑编译任务的特征。建立黑盒模型很容易,但是如果没有任务特性的指导,很容易导致更高的开销和更少的最优解。TC采用了这种模式。
-
基于ML的代价模型:基于ml的成本模型是一种使用机器学习方法预测性能的统计方法。它使模型能够随着新配置的探索而更新,这有助于实现更高的预测精度。TVM和XLA采用这种模型,分别是梯度树提升模型(GBDT)和前馈神经网络。
-
预定义代价模型:基于预定义代价模型的方法,期望根据编译任务的特点建立一个完美的模型,能够评估任务的整体性能。与基于ML的模型相比,预定义模型在应用时产生的计算开销更少,但需要在每个新的DL模型和硬件上重新构建模型,需要大量的工程工作。