一、introduction:
现有流行的DL框架:TensorFlow,PyTorch,MXNet和CNTK,互操作性可以减少重复的模型训练工作。ONNX(The Open Neural Network Exchange)就是这个目的,它定义了表示深度学习模型的统一格式,以促进不同框架之间的模型转换。
DL硬件分为:1)具有软硬件协同设计的通用硬件; 2)完全为深度学习模型定制的专用硬件;3)受生物脑科学启发的类神经硬件。 举个例子,在通用硬件(CPU/GPU)添加特殊硬件组件以加速DL模型,如AVX512矢量单元和张量核。 Google TPU之类的专用硬件,设计专用集成电路(例如矩阵乘法引擎和高带宽内存)将性能和功耗效率提升到了极致。 在可预见的将来,深度学习硬件的设计将变得更加多样化。 对于某些独特的计算特性(矩阵乘法),芯片设计师会为其开发独特的深度学习硬件平台,提高DL模型的性能
由于硬件多样性,重要的工作就是如何将计算有效地映射。通用硬件对高度优化的线性代数库,例如BLAS库(MKL和cuBLAS)依赖比较多。以卷积运算为例,深度学习框架将卷积转换为矩阵乘法,然后在BLAS库中调用GEMM函数。此外,硬件供应商还发布了特别的优化库(例如MKL-DNN和cuDNN),包括正向和反向卷积、池化、规范化和激活等操作符。他们还开发了更高级的工具来进一步加快深度学习操作。例如,TensorRT支持图形优化(层融合)和基于大量优化GPU内核的低比特量化。专用硬件,则会提供类似的开发库。但是,过于依赖库无法有效利用深度学习芯片。为解决这些依赖库和工具的缺点,减轻手动优化每个硬件运行模型的负担,深度学习社区诉诸于专门的编译器。已经有几种流行的编译器出现,例如TVM,Tensor Comprehension,Glow,nGraph和Tensorflow XLA(Accelerated Linear Algebra)。 编译器将深度学习框架描述的模型在各种硬件平台上生成有效的代码实现,其完成的模型定义到特定代码实现的转换将针对模型规范和硬件体系结构高度优化。具体来说,它们结合了面向深度学习的优化,例如层融合和操作符融合,实现高效的代码生成。此外,现有的编译器还采用了来自通用编译器(例如LLVM)的成熟工具链,对各种硬件体系结构提供了更好的可移植性。与传统编译器类似,深度学习编译器也采用分层设计,包括前端、中间表示(IR)和后端。但是,这种编译器的独特之处在于多级IR和特定深度学习模型实现优化的设计。
二、主要内容:
1. 常用编译器设计思想;
2. 面向深度学习的多级中间表示;
3. 前端优化;
4. 后端优化。
三、DL编译器的通用设计体系结构
DL编译器的通用设计架构主要包括两个部分:编译器前端和编译器后端,中间表示(IR)分布在前端和后端。通常,IR是程序的抽象,用于程序优化。具体地说,DL模型在DL编译器中被转换为多级IR,其中高阶IR驻留在前端,低阶IR驻留在后端。基于高阶IR,编译器前端负责与硬件无关的转换和优化。基于低阶IR,编译器后端负责特定于硬件的优化、代码生成和编译。
高阶IR,也称为graph IR,表示计算和控制流,与硬件无关。高阶IR的设计挑战是计算和控制流的抽象能力,能够捕获和表达不同的DL模型。高阶IR的目标是建立控制流和op(算子)与数据之间的依赖关系,以及为图级优化提供接口。它还包含用于编译的丰富语义信息,并为定制算子提供可扩展性。
低阶IR是为特定于硬件的优化和针对不同硬件目标的代码生成而设计的。因此,低阶IR应该足够细粒度,以反映硬件特征并表示特定于硬件的优化。它还应该允许在编译器后端使用成熟的第三方工具链,如Halide[33]、多面体模型[34]和LLVM[30]。
前端将现有的DL框架中的DL模型作为输入,然后将该模型转换为计算图表示(例如,graph IR)。前端需要实现多种格式转换,以支持不同框架下的多种格式。计算图优化结合了通用编译器的优化技术和DL特定优化技术,减少了图IR的冗余,提高了图IR的效率。这样的优化可以分为node级(例如,nop消除和零维张量消除),block级(代数简化、算子融合、op下降),数据流层(例如CSE、DCE、静态内存规划和布局转换)。前端完成后,生成优化的计算图并传递给后端。
后端将高阶IR转换为低阶IR并执行特定于硬件的优化。一方面,它可以直接将高阶IR转换为第三方工具链,如LLVM IR,以利用现有的基础设施进行通用优化和代码生成。另一方面,它可以利用DL模型和硬件特性的先验知识,通过定制的编译通道,实现更高效的代码生成。通常应用的特定于硬件的优化包括硬件内在映射、内存分配和获取、内存延迟隐藏、并行化以及面向循环的优化。为了在大的优化空间中确定最优参数设置,现有的DL编译器普遍采用两种方法:自动调度(如多面体模型)和自动调优(如AutoTVM)。优化后的底层IR采用JIT或AOT编译,为不同的硬件目标生成代码。
四、 面向深度学习的多级中间表示;
4.1 graph IR表示
graph IR的表示方式影响着graph IR的表达能力,也决定着DL编译器对graph IR的分析方式。
4.1.1 基于DAG的IR
基于DAG的IR是编译器构建计算图的最传统方法之一,它将节点和边组织成有向无环图(DAG)。在DL编译器中,DAG的节点代表原子DL算子(卷积,池化等),边代表张量。借助于DAG计算图,DL编译器可以分析各种运算符之间的关系和依赖关系,并利用它们来指导优化。DDG上已经有很多优化,比如公共子表达式消除(CSE)和死代码消除(DCE)。通过将DL的领域知识与这些算法相结合,进一步的优化可以应用到DAG计算图中,基于dag的IR算法简单,便于编程和编译,但存在计算范围定义不明确导致语义歧义等不足。
4.1.2 Let-Binding-Based IR
let绑定是一种解决语义歧义的方法,它通过向某些范围有限的函数提供let表达式,这些函数被许多高级编程语言使用,如Javascript、f#和Scheme。当使用let关键字定义表达式时,将生成一个let节点,然后它将指向表达式中的操作符和变量,而不是仅仅作为DAG构建变量之间的计算关系。在基于dag的编译器中,当一个进程需要获取一个表达式的返回值时,它首先访问相应的节点并搜索相关的节点,这也称为递归下降技术。相比之下,基于let绑定的编译器计算出let表达式中变量的所有结果,并构建一个变量映射。当需要一个特定的结果时,编译器查找这个映射来决定表达式的结果。在DL编译器中,TVM的Relay IR[40]同时采用了基于dag的IR和基于let的IR,获得了两者的优点。
4.1.3Representing Tensor Computation
不同的图IRs有不同的方法来表示张量上的计算。根据这些具体的表示形式,将不同DL框架的运算符转换为图IRs。定制操作符也需要按照这种表示方式进行编程。张量计算的表示可以分为以下三类。
Function-based: 基于函数的表示只是提供封装的运算符,这被Glow、nGraph和XLA所采用。以高级优化器(HLO, XLA的IR)为例,它由符号编程中的一组函数组成,其中大多数函数都没有副作用。这些指令被组织成三个级别,包括HloModule(整个程序)、hlocomputing(一个函数)和HloInstruction(操作)。XLA使用HLO IR表示图IR和操作IR,使HLO的操作范围从数据流级到操作级。
Lambda expression: lambda表达式是一个索引公式表达式,描述通过变量绑定和替换进行计算。使用lambda表达式,程序员可以快速定义计算,而无需实现新函数。TVM表示使用张量表达式的张量计算,它是基于lambda表达式的。在TVM中,由输出张量的形状和计算规则的lambda表达式定义张量表达式中的计算算子。
Einstein notation: 爱因斯坦符号,也被称为求和约定,是一种表示求和的符号。它的编程简单性优于lambda表达式。以TC为例,临时变量的索引不需要定义。根据爱因斯坦表示法,IR可以通过未定义变量的出现求出实际表达式。在爱因斯坦表示法中,运算符必须是结合律和交换律。这个限制保证了约简运算符可以按任何顺序执行,从而使进一步的并行化成为可能。
4.2graph IR的实现
graph IR在DL编译器中的实现了对数据和操作的管理。
4.2.2数据表示
DL编译器中的数据(例如,输入、权重和中间数据)通常以张量的形式组织,也称为多维数组。DL编译器可以通过内存指针直接表示张量数据,或者以一种更灵活的方式(占位符)表示。占位符包含张量的每个维度的大小。或者,张量的维度大小可以标记为未知。为了进行优化,DL编译器需要数据布局信息。此外,迭代器的边界应该根据占位符进行推断。
Placeholder: 占位符在符号编程中被广泛使用(例如Lisp, Tensorflow)。占位符只是一个具有显式形状信息的变量(例如,每个维度的大小),它将在计算的后期用值填充。它允许程序员描述操作并构建计算图,而无需考虑确切的数据元素,这有助于将计算定义与DL编译器中的确切执行分离开来。此外,在不改变计算定义的情况下,使用占位符可以方便程序员改变输入输出和其他相应的中间数据的形状。
Unknown (Dynamic) shape representation: 在声明占位符时,通常支持未知维度大小。例如,TVM使用Any表示未知维度;XLA使用None来达到同样的目的;nGraph使用它的PartialShape类。未知形状表示是支持动态模型的必要条件。然而,要完全支持动态模型,需要放宽边界推理和维数检验。此外,还需要实现额外的机制来保证内存的有效性。
Data layout: 数据布局描述了一个张量是如何在内存中组织的,它通常是从逻辑指标到内存指标的映射。数据布局通常包括维度的顺序(例如,NCHW和NHWC)、 tiling, padding, striding,等。TVM和Glow将数据布局表示为操作符参数,需要这些信息进行计算和优化。然而,将数据布局信息与操作符(而不是张量)结合起来,可以直观地实现某些操作符,并减少编译开销。XLA将数据布局表示为与后端硬件相关的约束。Relay和MLIR将在它们的张量类型系统中添加数据布局信息。
Bound inference: 在DL编译器中编译DL模型时,应用绑定推理来确定迭代器的边界。尽管DL编译器中的张量表示很方便地描述了输入和输出,这也暴露了推断迭代器边界的挑战。根据计算图和已知的占位符,通常递归或迭代地执行绑定推理。例如,在TVM中,迭代器形成一个有向无环超图,其中图的每个节点代表一个迭代器,每个超边代表两个或多个迭代器之间的关系(例如,split, fuse或rebase)。一旦根据占位符的形状确定了根迭代器的边界,就可以根据递归的关系推断出其他迭代器。
Operators Supported :DL编译器支持的运算符负责表示DL工作负载,它们是计算图的节点。操作符通常包括代数操作符(例如加减乘除、指数等),神经网络操作符(例如卷积和池化),张量操作符(例如 reshape, resize and copy),广播和约简操作符(例如min和argmin),以及控制流操作符(例如条件和循环)。在这里,我们选择在不同的DL编译器中经常使用的三个代表性操作符作为说明。此外,我们还讨论了自定义运算符的情况。
Broadcast:广播运算符可以复制数据并生成具有兼容形状的新数据。没有广播运算符,输入张量的形状受到更多的限制。例如,对于一个加法算子,输入张量应该具有相同的形状。一些编译器,如XLA和Relay,通过提供广播操作符来放宽这种限制。例如,XLA允许对矩阵和向量进行元素相加,方法是复制它,直到它的形状与矩阵匹配为止。
Control flow:在表示复杂和灵活的模型时需要控制流。RNN和强化学习(RL)等模型依赖于循环关系和依赖数据的条件执行,这需要控制流。如果在DL编译器的图IR中不支持控制流,这些模型必须依赖宿主语言(如Python中的if和while)的控制流支持或静态展开,这会降低计算效率。Relay注意到任意控制流可以通过递归和pattern实现,这已在函数式编程中得到了演示。因此,它提供了实现控制流的if运算符和递归函数。相反,XLA通过特殊的HLO操作符(如while和conditional)表示控制流。
导数:运算符Op的导数运算符以Op的输出梯度和输入数据作为输入,然后计算运算符的梯度。虽然一些DL编译器(如TVM和TC)支持自动微分,但当应用链式法则时,它们需要高级IR中所有运算符的导数。TVM致力于提供代数算子和神经网络算子的导数算子。程序员可以使用这些导数运算符来构建自定义运算符的导数。相反,PlaidML可以自动生成派生运算符,即使是定制的运算符。值得注意的是,不能支持派生运算符的DL编译器无法提供模型训练的能力。
自定义操作符:它允许程序员为特定目的定义操作符。提供对定制操作符的支持提高了DL编译器的可扩展性。例如,在Glow中定义新操作符时,程序员需要实现逻辑和节点封装。此外,如果需要,还需要额外的努力,例如降低步骤、操作IR生成和指令生成。而TVM和TC除了描述计算实现外,需要较少的编程工作。具体来说,TVM的用户只需要描述计算和调度,声明输入/输出张量的形状。此外,定制操作符通过钩子集成Python函数,这进一步减轻了程序员的负担。
结论:几乎所有DL编译器都有自己独特的高级IRs。然而,它们有着相似的设计理念,比如使用DAG和let-binding来构建计算图。此外,它们通常为程序员提供了表示张量计算的方便方法。在高级IRs中设计的数据和操作符非常灵活和可扩展,足以支持不同的DL模型。更重要的是,高级IRs与硬件无关,因此可以应用于不同的硬件后端。
4.3.1 Low-Level IR的实现:
与高级IR相比,低级IR以更细粒度的表示方式描述DL模型的计算,它通过提供调优计算和内存访问的接口来实现与目标相关的优化。在本节中,我们将底层IRs的常见实现分为三类:基于halide的IR、基于polyhedral的IR和其他惟一IR。
基于halide的IR:Halide首次被提出用于并行化图像处理,并被证明在DL编译器(如TVM)中具有可扩展性、有效率。Halide的基本理念是 computation和schedule的分离。采用Halide的编译器不是直接给出一个具体的方案,而是尝试各种可能的schedule,然后选择最好的一个。在Halide中,内存引用和循环的边界被限制为与轴对齐的有界框。因此,Halide不能用复杂的模式(例如,非矩形)表示计算。幸运的是,DL中的计算非常规则,可以用Halide很好地表示。此外,Halide可以很容易地参数化这些边界,并将它们暴露给调优机制。当应用到DL编译器的后端时,需要修改Halide的原始IR。例如,Halide的输入形状是无限的,而DL编译器需要知道数据的确切形状,以便将操作符映射到硬件指令。一些编译器,如TC,需要固定的数据大小,以确保张量数据更好的时间局部性。TVM通过以下努力将Halide IR改进为独立的符号IR。它消除了对LLVM的依赖,重构了项目模块和Halide的IR设计的结构,追求更好的组织和图形IR和前端语言(如Python)的可访问性。可重用性也得到了提高,实现了一个运行时调度机制来方便地添加定制操作符。TVM将变量定义从字符串匹配简化为指针匹配,确保每个变量都有一个定义位置(静态单赋值,SSA)[44])。
基于polyhedral的IR:多面体模型是DL编译器采用的一种重要技术。它使用线性规划、仿射变换和其他数学方法来优化具有边界和分支静态控制流的基于循环的代码。与Halide相比,在polyhedral模型中,内存引用和 loop nests的边界可以是任意形状的多面体。这种灵活性使得多面体模型在泛型编译器中得到广泛应用。然而,这种灵活性也阻碍了与调优机制的集成。然而,由于能够处理深嵌套循环,许多DL编译器,如TC和PlaidML(作为nGraph的后端),已经采用多面体模型作为低级IR。基于多面体的IR可以很容易地应用各种多面体转换(例如,融合、平铺、下沉和映射),包括设备依赖和设备独立的优化。基于多面体的编译器可以借用许多工具链,例如isl、Omega、PIP、Polylib和PPL。
MLIR受LLVM的影响很大,是一种比LLVM更纯粹的编译器基础设施。MLIR重用LLVM中的许多思想和接口,位于模型表示和代码生成之间。MLIR有一个灵活的类型系统,并允许多个抽象级别,它引入dialects来表示这些抽象级别。每种dialects都由一组已定义的不可变操作组成。当前的MLIR dialects包括TensorFlow IR、XLA HLO IR、experimental polyhedral IR、LLVM IR和TensorFlow Lite。还支持不同方言之间的灵活转换。此外,MLIR可以创建新的dialects来连接到新的低阶编译器,这为硬件开发人员和编译器研究人员铺平了道路。XLA的HLO IR可以看作是高阶IR和低阶IR,因为HLO足够细粒度,可以表示特定于硬件的信息。此外,HLO支持特定于硬件的优化,并可用于发出LLVM IR。
大多数DL编译器采用的低阶IR最终可以降低到LLVM IR,并受益于LLVM成熟的优化器和代码生成器。此外,LLVM可以为专门的加速器显式地从头设计自定义指令集。然而,当直接传递给LLVM IR时,传统编译器可能会生成糟糕的代码。为了避免这种情况,DL编译器采用两种方法来实现硬件相关的优化:1)在LLVM的上层IR中执行目标特定的循环转换(如基于halide的IR和基于多面体的IR),2)为优化通道提供关于硬件目标的附加信息。大多数DL编译器都应用这两种方法,但重点不同。一般来说,针对前端的DL编译器(例如TC、TVM、XLA和nGraph)可能会关注1),而更倾向于后端开发人员的DL编译器(例如Glow、PlaidML和MLIR)可能会关注2)。
DL编译器中的编译方案主要分为准时(JIT)和提前(AOT)两大类。对于JIT编译器,它可以动态生成可执行代码,并且可以使用更好的运行时知识来优化代码。AOT编译器首先生成所有可执行的二进制文件,然后执行它们。此外,AOT方法可以应用于嵌入式平台的交叉编译器(例如C-GOOD[50]),也可以在远程机器上执行(TVM RPC)和自定义加速器。
结论:在DL编译器中,底层IR是DL模型的细粒度表示,它反映了DL模型在不同硬件上的详细植入。底层的IRs包括基于halide的IRs、基于多面体的IRs和其他独特的IRs。尽管它们在设计上有所不同,但它们利用了成熟的编译器工具链和基础设施,为特定于硬件的优化和代码生成提供了定制的接口。低级IRs的设计也会影响新的DL加速器的设计(例如,TVM HalideIR和Inferentia,以及XLA HLO和TPU)。
五、 前端优化
在构建计算图之后,前端应用图级优化。许多优化更容易在图级识别和执行,因为图提供了计算的全局视图。这些优化只应用于计算图,而不是后端实现。因此,它们是独立于硬件的,可以应用于各种后端目标。前端优化通常由passes定义,可以通过遍历计算图的节点并执行图转换来应用。前端提供了如下方法:1)从计算图中捕获特定的特征;2)重写图以进行优化。除了预定义的passes,开发人员还可以在前端定义定制的passes。一旦一个DL模型被导入并转换为一个计算图,大多数DL编译器可以确定每个操作的输入张量和输出张量的形状。该特性允许DL编译器根据形状信息执行优化。图2显示了一个使用Tensorflow XLA优化计算图的例子。在本节中,我们将前端优化分为三类:1)节点级优化,2)块级优化,3)数据流底层优化。
5.1 节点级优化
计算图的节点足够粗粒度,可以在单个节点内进行优化。节点级优化包括节点消除(消除不必要的节点)和节点替换(用其他低成本节点替换节点)。在通用编译器中,Nop消除删除占用少量空间但不指定操作的no-op指令。在DL编译器中,Nop消除负责消除缺少足够输入的操作。例如,只有一个输入张量的和节点可以被消除,填充宽度为零的填充节点可以被消除。Zero-dim-tensor消去负责去除输入为零维张量的不必要操作。假设A是零维张量,B是常数张量,那么A和B的和运算节点可以替换为已有的常数节点B,而不影响其正确性。假设C是一个三维张量,但一维的形状为零,如{0,2,3},因此C没有元素,可以消去argmin/argmax操作节点。
5.2 块级优化
代数简化——代数简化优化包括1)代数识别,2)强度约简,我们可以用更便宜的算子取代更昂贵的算子;3)常数折叠,用它可以用它们的值替换常量表达式。这种优化方法考虑节点序列,然后利用不同类型节点的交换性、结合性和分配性来简化计算。这里我们举例说明了代数简化可以应用的常见情况:1)优化计算顺序,在这种情况下,优化根据特定的特征发现和删除重塑/转置操作。以矩阵乘法(GEMM)为例,有两个矩阵(如A和B),两个矩阵进行转置(分别产生AT和BT),然后AT和BT相乘。然而,实现GEMM的一种更有效的方法是切换参数a和B的顺序,将它们相乘,然后对GEMM的输出进行转置,这将两个转置减少为一个;2)节点组合优化,将多个连续的转置节点组合为单个节点,消除身份转置节点,在转置节点实际不移动数据时,将转置节点优化为reshape节点;3) ReduceMean节点的优化,当reduce算子的输入是4D且最后两个维数要被约简时,优化将ReduceMean替换为AvgPool节点(例如在Glow中)。
算子融合。算子融合是DL编译器不可缺少的优化方法。它能够更好地共享计算,消除中间分配,通过结合loop nests进一步优化,并减少launch和同步开销[26]。在TVM中,运算符被分为四类:injective、reduction、complex-out-fusible和不透明。定义操作符后,将确定其对应的类别。针对上述类别,TVM设计了跨算子的融合规则。在TC中,基于自动多面体变换的不同方式进行融合。然而,如何识别和融合更复杂的图模式,如带有多个广播和减少节点的块,仍然是一个问题。最近的工作[51],[52]试图解决这个问题,并提出了一个框架来探索和优化积极的融合方案。它不仅支持基于元素的节点和简化节点,还支持其他具有复杂依赖关系的计算/内存密集型节点。
算子Sinking:这种优化将转置等操作置于批处理归一化、ReLU、sigmoid和channel shuffle等操作之下。通过这种优化,许多相似的操作彼此之间移动得更近,为代数简化创造了更多的机会。
5.3 数据流级优化
通用子表达式消除(CSE)。表达式E是一个公共子表达式,如果E的值之前已经计算过,并且E的值在之前的计算之后没有改变。在这种情况下,E的值只计算一次,并且可以使用已经计算出的E的值来避免在其他地方重新计算。DL编译器在整个计算图中搜索公共子表达式,并用之前的计算结果替换下面的公共子表达式。
死代码消除(DCE)。如果一组代码的计算结果或副作用没有被使用,那么它就是死的。DCE优化会删除死代码。死代码通常不是由程序员引起的,而是由其他图优化引起的。因此,DCE和CSE是在其他图优化之后应用的。其他优化也属于DCE,比如死库消除(DSE),它将永远不会使用的存储删除到张量中。
静态内存规划。执行静态内存规划优化以尽可能重用内存缓冲区。通常有两种方法:in-place内存共享和标准内存共享。in-place内存共享使用相同的内存作为操作的输入和输出,只在计算之前分配一个内存副本。标准内存共享重用以前操作的内存,而不重叠。静态内存规划是离线完成的,这允许应用更复杂的规划算法。最近的一项工作[54]首次设计并执行了内存感知调度,以最小化边缘设备上的峰值激活内存占用,这为内存约束设备上的内存规划提供了新的研究方向。
布局转变。布局转换试图找到最佳的数据布局来存储张量到计算图中,然后将布局转换节点插入到图中。注意,这里不执行实际的转换,而是在编译器后端计算计算图时执行转换。事实上,相同操作在不同的数据布局中的性能是不同的,最佳布局在不同的硬件上也是不同的。例如,在GPU上NCHW格式的操作通常运行得更快,所以在GPU上转换为NCHW格式(例如,TensorFlow)是有效的。一些DL编译器依赖于特定于硬件的库来实现更高的性能,这些库可能需要特定的布局。此外,一些DL加速器更喜欢更复杂的布局(例如,tile)。此外,边缘设备通常配备异构计算单元,不同计算单元可能需要不同的数据布局才能更好地利用,因此布局转换需要仔细考虑。因此,编译器需要提供一种跨各种硬件执行布局转换的方法。不仅张量的数据布局对最终性能有很大的影响,而且变换操作也有很大的开销。因为它们也消耗内存和计算资源。
最近一项基于TVM针对cpu的工作[55]改变了所有卷积操作在计算图中的布局为NCHW[x]c,其中c表示通道c的分割子维数,x表示子维数的分割大小。然后,在特定于硬件的优化过程中,在提供硬件详细信息(如缓存线大小、向量化单元大小和内存访问模式)时,通过自动调优全局探索所有x参数。
结论:前端是DL编译器中最重要的组件之一,负责从DL模型到高级IR(例如,计算图)的转换,以及基于高级IR的独立于硬件的优化。尽管前端的实现在跨DL编译器的高级IR的数据表示和操作符定义方面可能有所不同,但与硬件无关的优化在三个级别上集中:节点级别、块级别和数据流级别。每一级的优化方法都利用了特定于DL的和通用的编译优化技术,从而减少了计算冗余、在计算图级提高DL模型的性能。
六、 后端优化
DL编译器的后端通常包括各种特定于硬件的优化、自动调优技术和优化的内核库。特定于硬件的优化支持针对不同硬件目标的高效代码生成。然而,自动调优在编译器后端是必不可少的,以减少获取最佳参数配置的手工工作。此外,高度优化的内核库也广泛应用于通用处理器和其他定制的DL加速器上。
6.1特定于硬件的优化
特定于硬件的优化,也称为目标相关优化,用于获得针对特定硬件的高性能代码。应用后端优化的一种方法是将低层IR转换为LLVM IR,利用LLVM基础设施生成优化的CPU/GPU代码。另一种方法是使用DL领域知识设计定制优化,更有效地利用目标硬件。由于硬件特定的优化是为特定的硬件量身定制的,本文无法详尽地介绍,因此我们提出了五种在现有DL编译器中广泛采用的方法。这些特定于硬件的优化的概述如图3所示,详细描述如下。
硬件固有的映射。硬件内在映射可以将一定的一组低层IR指令转换为已经在硬件上高度优化的内核。在TVM中,硬件本征映射采用可扩展张量化方法实现,可声明硬件本征的行为和本征映射的降低规则。这种方法使编译器后端能够将硬件实现和高度优化的手工微内核应用到特定的操作模式中,从而获得显著的性能提升。然而,Glow支持硬件内部映射,比如量化。它可以估计神经网络每个阶段可能的数值范围,并支持配置文件引导的优化自动执行量化。此外,Halide/TVM将特定的IR模式映射到每个体系结构上的SIMD操作码,以避免遇到向量模式时LLVM IR映射的低效。
内存分配和获取。内存分配是代码生成中的另一个挑战,特别是对于gpu和定制加速器。例如,GPU主要包含共享内存空间(内存大小有限,访问时延较低)和本地内存空间(容量大,访问时延较高)。这样的内存层次结构需要有效的内存分配和获取技术来改善数据局部性。为了实现这一优化,TVM引入了内存作用域的调度概念。内存作用域调度原语可以将计算阶段标记为共享的或线程本地的。对于标记为共享的计算阶段,TVM生成具有共享内存分配和协同数据获取的代码,并在适当的代码位置插入内存屏障以保证正确性。此外,TC通过外置 PPCG[56]编译器提供类似的功能(称为内存提升)。但是,TC只支持有限的预定义规则。特别是,TVM通过内存作用域调度原语在加速器中支持特殊的缓冲。
内存延迟隐藏。通过对执行管道重新排序,内存延迟隐藏也是在后端使用的一项重要技术。由于大多数DL编译器支持CPU和GPU上的并行化,内存延迟隐藏可以自然地通过硬件实现(例如,GPU上的扭曲上下文切换)。但是对于具有解耦访问库(DAE)体系结构的tpu类加速器,后端需要执行调度和细粒度同步以获得正确和高效的代码。为了获得更好的性能和减轻编程负担,TVM引入了虚拟线程调度原语,允许用户在虚拟多线程架构上指定数据并行度。然后,TVM通过插入必要的内存屏障,并将来自这些线程的操作穿插到单个指令流中,从而降低了这些实际上并行的线程,这形成了每个线程更好的执行管道,以隐藏内存访问延迟。
面向循环优化。面向循环的优化也应用于后端,为目标硬件生成高效的代码。由于Halide和LLVM30已经集成了这样的优化技术,一些DL编译器在后端利用了Halide和LLVM。面向循环优化中应用的关键技术包括循环融合、滑动窗口、平铺、循环重排序和循环展开。
并行化。由于现代处理器通常支持多线程和SIMD并行性,编译器后端需要利用并行性来最大化硬件利用率,以实现高性能。Halide使用一个叫做parallel的调度原语来指定线程级别并行化的循环的并行化维度,并通过将标记为parallel的循环维度与块和线程的注释进行映射来支持GPU并行化。它将一个大小为n的循环替换为一个宽为n的向量语句,该向量语句可以通过硬件内在映射映射到特定于硬件的SIMD操作码。Stripe发展了一种多面体模型的变体,称为嵌套多面体模型,该模型引入了并行多面体块作为迭代的基本执行元素。在此扩展之后,一个嵌套的多面体模型可以检测平铺和跨步级别之间的层次并行性。此外,一些DL编译器依赖于手工库,如Glow或硬件供应商提供的优化数学库(在第3.4.3节中讨论)。同时,Glow将向量化工作交给LLVM,因为LLVM自向量化器在提供张量维数和环路次数信息的情况下工作良好。然而,完全通过编译器后端利用并行性允许应用更多领域特定的DL模型知识,从而以更多的工程工作为代价获得更高的性能。
6.2自动调优
由于在特定于硬件的优化中参数调优有巨大的搜索空间,因此有必要利用自动调优来确定最佳参数配置。在本次调查中研究的DL编译器中,TVM、TC和XLA支持自动调优。通常,自动调优实现包括参数化、代价模型、搜索技术和加速四个关键部分。
Parameterization:数据和目标:数据参数描述数据的规格,如输入的形状。目标参数描述在优化调度和代码生成期间要考虑的特定于硬件的特征和约束。例如,对于GPU目标,需要指定共享内存和寄存器大小等硬件参数。2)优化选项:优化选项包括优化调度和相应的参数,如面向循环的优化和瓷砖大小。在TVM中,既考虑了预定义调度,也考虑了自定义调度,还考虑了参数。然而,TC和XLA倾向于参数化优化,这与性能有很强的相关性,并且可以稍后以低成本进行更改。例如,小批尺寸是CUDA中通常映射到网格尺寸的参数之一,可以在自动调优期间进行优化
成本模型。自动调优中不同成本模型的比较如下。1)黑箱模型:该模型只考虑最终的执行时间,而不考虑编译任务的特征。建立黑盒模型很容易,但是如果没有任务特性的指导,很容易导致更高的开销和更少的最优解。TC采用了这种模式。2)基于ml的成本模型:基于ml的成本模型是一种使用机器学习方法预测性能的统计方法。它使模型能够随着新配置的探索而更新,这有助于实现更高的预测精度。TVM和XLA采用这种模型,分别是梯度树提升模型(GBDT)和前馈神经网络57。3)预定义成本模型:基于预定义成本模型的方法,期望根据编译任务的特点建立一个完美的模型,能够评估任务的整体性能。与基于ml的模型相比,预定义模型在应用时产生的计算开销更少,但需要在每个新的DL模型和硬件上重新构建模型,需要大量的工程工作。
后端负责基于低级IR的bare-metal和代码生成。尽管后端设计可能因各种低级IRs而有所不同,但它们的优化可以分为特定于硬件的优化:自动调优技术和优化的内核库。这些优化可以单独执行,也可以组合执行,通过利用硬件/软件特性实现更好的数据局部性和并行性。最终,DL模型的高级IR转换为不同硬件上的高效代码实现。
参考文献:
M. Li et al., "The Deep Learning Compiler: A Comprehensive Survey," in IEEE Transactions on Parallel and Distributed Systems, vol. 32, no. 3, pp. 708-727, 1 March 2021, doi: 10.1109/TPDS.2020.3030548.