翻译《The Deep Learning Compiler: A Comprehensive Survey》,我看了前面的部分主要是在介绍不同的ai编译器主要使用了什么技术,我对以下的内容更加关心,于是只翻译了这部分内容。
4 深度学习编译器的通用设计
4.1 设计概述
深度学习编译器的通用设计主要包含两个部分:编译器前端和编译器后端,如图2所示。中间表示(IR)横跨前端和后端。通常,中间表示是程序的抽象,用于程序优化。具体来说,深度学习模型在深度学习编译器中被转换为多级中间表示,其中高级中间表示位于前端,低级中间表示位于后端。不同深度学习编译器中的中间表示实现如表2所示。基于高级中间表示,编译器前端负责硬件无关的转换和优化。基于低级中间表示,编译器后端负责硬件特定的优化、代码生成和编译。深度学习编译器中前端、后端和多级中间表示的功能简要描述如下:

图2:DL编译器常用设计概述
高级中间表示(也称为图中间表示)表示计算和控制流,并且是硬件无关的。高级中间表示的设计挑战在于能够抽象计算和控制流,能够捕获和表达多样化的深度学习模型。高级中间表示的目标是建立控制流以及运算符和数据之间的依赖关系,并为图级优化提供接口。它还包含丰富的语义信息用于编译,并为自定义运算符提供可扩展性。高级中间表示的详细讨论在第4.2节中介绍。
低级中间表示是为在不同硬件目标上进行硬件特定优化和代码生成而设计的。因此,低级中间表示应该足够细粒度以反映硬件特征并表示硬件特定的优化。它还应该允许在编译器后端使用成熟的第三方工具链,例如Halide[77]、多面体模型[13]和LLVM[59]。低级中间表示的详细讨论在第4.3节中介绍。
前端接受来自现有深度学习框架的深度学习模型作为输入,然后将模型转换为计算图表示(即图中间表示)。为了支持不同框架中的多样化格式,前端需要实现各种格式转换。计算图优化结合了来自通用编译器和深度学习特定优化的优化技术,这些技术减少了冗余并提高了图中间表示的效率。这些优化可以分为节点级(例如,空操作消除和零维张量消除)、块级(例如,代数化简、算子融合和算子下沉)和数据流级(例如,公共子表达式消除CSE、死代码消除DCE、静态内存规划和布局转换)。在前端之后,生成优化的计算图并传递给后端。注意,一些深度学习编译器进一步将优化的计算图转换为操作中间表示。前端的详细讨论在第4.4节中介绍。
后端将高级中间表示转换为低级中间表示,同时执行硬件特定的优化。一方面,它可以直接将高级中间表示转换为第三方工具链(如LLVM IR),以利用LLVM基础设施进行通用优化和CPU/GPU代码生成。另一方面,它可以利用深度学习模型和硬件特征的先验知识,通过定制的编译pass来生成更高效的代码。常用的硬件特定优化包括硬件内在函数映射、内存分配和获取、内存延迟隐藏、并行化以及循环导向优化。为了解决上述优化引入的大规模解空间,现有深度学习编译器中广泛采用两种方法,例如自动调度(例如多面体模型)和自动调优(例如AutoTVM)。优化后的低级中间表示使用JIT或AOT编译,为不同的硬件目标生成代码。后端的详细讨论在第4.5节中介绍。

表 2. 不同 DL 编译器的 IR 实现。
4.2 高级中间表示
为了克服传统编译器中采用的中间表示限制了深度学习模型中使用的复杂计算的表达,现有的深度学习编译器利用具有专门设计的数据结构的图中间表示来进行高效的代码优化。为了更好地理解深度学习编译器中使用的图中间表示,我们从以下几个方面描述图中间表示的语义和表示。
4.2.1 图中间表示的语义
我们从编程语言(PL)的角度介绍图中间表示的语义。图中间表示的抽象语法树(AST)影响图中间表示的表达能力,也显示了深度学习编译器如何分析图中间表示代码。此外,不同的图中间表示有其描述张量计算的方法。这些方法应该是用户友好的,并且计算表示具有可扩展性。我们首先描述图中间表示的AST。
• 基于DAG的中间表示:基于DAG的中间表示是编译器构建计算图的最传统方式之一,其节点和边将被组织为有向无环图。在DAG计算图上有丰富的优化算法,例如活变量分析和变量依赖分析。基于DAG的中间表示由于其简单性便于编程和编译,但在其他方面存在不足,例如由于计算作用域未定义而导致的语义歧义[5]。
• 基于绑定的中间表示:Let绑定是一种通过为某些函数提供Let表达式来解决上述语义歧义的方法,具有受限作用域,被许多编程语言(如F#[76]和Scheme[57])使用。Let绑定将改变中间表示的AST结构。当使用Let关键字定义表达式时,将生成一个let节点并指向表达式中的运算符和变量,而不是像DAG那样仅在变量之间建立计算关系。当进程需要获取一个表达式的返回结果时,基于DAG的编译器将首先访问add节点并递归搜索与之相关的节点。相比之下,基于Let绑定的编译器将计算Let表达式中每个变量的结果并建立变量映射。当需要结果时,编译器将查找映射以确定表达式的结果。
计算表示 - 不同的图中间表示对张量计算有不同的表示。编译器将根据其特定的表示形式翻译框架运算符。它们的自定义运算符也需要用这些表示形式编程。这些形式可以分为以下三类。

图 3. Lambda 表达式和 Einstein 表示法
• 基于函数的形式:基于函数的形式只提供封装的运算符,HLO(XLA的中间表示)采用这种形式。因此我们以HLO为例来描述基于函数的形式。HLO IR是一种符号编程的编程语言。该中间表示由一组函数组成,其中大多数没有副作用。因此很难检索计算过程中使用的中间值。指令被组织成三个级别:HloModule,类似于整个程序;HloComputation,类似于函数;HloInstruction,表示基本操作。XLA使用HLO IR来表示图中间表示和操作中间表示,因此HLO的操作范围从节点级到计算级。
• Lambda表达式:TVM通过lambda表达式(如图3(a)所示)[29]表示操作中间表示,这是一个索引公式表达式,TVM通过lambda表达式提供了一种称为张量表达式的语言。Lambda表达式通过变量绑定和替换来描述计算,这在编程语言(如Python和C++11)中广泛存在。在TVM中,张量表达式中的计算运算符由两部分组成:输出张量的形状和计算规则的lambda表达式。
• 爱因斯坦记号:爱因斯坦记号(也称为求和约定)被TC采用。爱因斯坦记号是表达求和的一种记号。以矩阵向量乘法的定义(如图3(b)所示)[89]为例,不需要定义索引的临时变量。TC IR可以通过未定义变量的出现来推断实际表达式。例如,在图3(b)的第二行中,变量k只存在于右侧,因此运算符+将在k表示的维度上执行归约运算符。在爱因斯坦记号表达式中,运算符需要是结合的和可交换的,像max这样的运算符不被支持。这个限制保证了归约运算符可以按任何顺序执行,使进一步的并行化成为可能。
4.2.2 图表示
深度学习编译器中的图表示采用传统计算图的方法,由数据和运算符组成。
数据组织 - 深度学习编译器总是通过多级节点建立其计算图以进行组织。以Relay为例,图4(a)显示了用于表示计算图的数据结构。Relay将整个计算图划分为称为调度节点的较小图以进行组织。一个调度节点由许多阶段节点组成,它形成一个DAG图。每个阶段节点恰好有一个运算符节点来表示计算运算符,例如ComputeOp或TensorComputeOp。运算符节点也用于产生迭代变量之间的关系。

图 4 TVM 中图 IR 的数据组织
此外,Relay在其图数据结构中结合迭代变量以进行边界推断。Relay通过上述多级节点分层组织迭代变量,图4(b)[5]显示了每个级别的相关节点保持的迭代变量的数据结构。在调度级别,将为必要的迭代变量保持边界映射,可以由每个阶段节点获取。每个阶段节点管理其迭代变量。一个阶段节点中的迭代变量可以分为叶变量和根变量。前者确定特定运算符中的范围,后者计算出循环维度。
边界推断 - 编译器在执行优化时需要确保每个迭代变量的范围。然而,一些计算表示并不直观地指出循环变量的作用域。相反,这些计算形式只提供关于输入和输出的信息。因此,需要边界推断过程来确定变量范围。边界推断算法可以根据推断方向分为两类:
• 从输出到输入:TVM的推断是从输出到输入的。这与其表达形式——张量表达式有关,它直接需要输出的形状。Relay通过一个称为InferBound pass的pass来处理推断过程,图5显示了伪代码和示意图。输入是Schedule节点,它由一系列阶段节点组成。首先,Relay将建立其数据流图,其节点是阶段节点。然后,执行反向拓扑排序以获得一个列表,保证消费者节点在输入节点之前。这为推断过程提供了一个主要假设:当处理一个阶段时,其输出的形状已经是已知的。对于每个节点,推断pass将首先确定根迭代变量的范围,使用其输出张量的形状。然后,Relay将通过阶段节点记录的迭代变量的关系和它们的关系来计算出叶迭代变量。关系包括split、fuse和rebase。由该关系组织的变量依赖关系可以很容易地分析。考虑split节点为例。当面对给定的因子时,如果没有越界,变量的范围将被划分为(0, factor)和(factor, extend)。根变量的边界结果将保存在全局映射中供其他阶段节点使用。
• 从输入到输出:爱因斯坦记号(由TC IR使用)通过自动推断索引范围使编程更简洁,但它基于可理解的推断规则。正如上面提到的,有些运算符是爱因斯坦记号无法处理的,需要用户添加where子句来指导编译器。如果需要花费大量时间来考虑哪个表达式应该使用where子句,他宁愿显式编写循环并放弃TC IR。为了避免这种尴尬,TC在推断算法中采用了一种直接的方法,该方法受到在多面体库[95]中实现的Presburger算术的启发。推断解决的问题是在索引不能超出输入边界的约束下,尽可能大地寻找所有未定义变量的范围。该过程是迭代的。首先,算法为所有未解决的变量初始化一个列表,这些变量不受where子句约束。然后,算法将找到一个只有单个未解决变量的表达式,并使用Halide函数solve_for_inner_interval在现有条件下找到该变量的最大范围。然后将已解决的变量从未解决列表中删除。算法重复第二步,直到列表为空。如果不存在具有单个未解决变量的表达式且未解决列表不为空,编译器将停止推断并向用户请求where。推断的另一个挑战是递归定义,第一个版本不支持,使TC在实现RNN时失败。

图 5. Relay 中的绑定推理。
数据表示 - 深度学习编译器中数据的表示包括输入、权重和变量。一些编译器通过特定指针直接表示数据,而其他编译器使用占位符或其他设计来间接指向数据。数据表示不仅包含特定值,还包含其他相关信息。由于图优化的需求,主流深度学习编译器更倾向于将数据布局信息添加到数据表示中。不同编译器的数据表示如下所示。
• 占位符(Placeholder):占位符在符号编程中被广泛使用。使用占位符的编译器没有关于具体张量的信息。此外,程序员通过使用占位符而不是更改整个语义来更改输入和输出形状更方便。在编译过程中,占位符可以用不同的张量替换而不改变语义。TVM在其张量表达式语言中使用占位符作为构造空张量对象的方法。
• 未知形状表示:TVM使用Any来表示未知维度以支持动态模型。Any可以在张量类型定义中用作张量形状的维度之一。与占位符不同,包含Any的形状在编译时是未知的,因此需要放宽推断和检查维度的工作,并且需要额外的形状函数来保证内存有效性。
• 数据布局:编译器的数据布局有三个视角:从运算符视角(TVM、Glow、PlaidML)、从后端视角(XLA)和张量视角(Relay、MLIR)。前者将数据布局视为运算符的附加或必要参数,并需要这些信息进行计算和优化。将数据布局信息与运算符而不是张量结合使得某些运算符的直观实现成为可能,编译器更倾向于将数据布局作为参数而不是运算符的字段,以减少编译期间的动态消耗[4]。XLA将数据布局视为与其后端设备相关的约束:CPU和GPU。例如,当常量数组仅在CPU后端由dot运算符使用时,XLA将要求数组为列主序。Relay和MLIR正在将布局信息添加到张量的类型系统中。
支持的运算符 - 深度学习编译器提供的计算运算符可以分为三个级别:基本计算节点、高级计算节点和融合计算节点。基本节点是整个运算空间的组成部分,例如指数运算。高级节点包括神经网络模型的单元,例如卷积和池化。融合节点表示通过图优化融合的节点。
此外,一些编译器还支持自定义运算符,这也属于计算运算符。对自定义运算符的支持是可扩展性的最基本设计之一,它允许用户为新的后端指令或仅为技术需求定义其运算符。对自定义运算符的支持可以分为两个级别:节点级和运算符级。
• 节点级:Glow支持图IR节点和操作IR指令的自定义运算符。Glow建议使用现有的IR来避免实现新的运算符。对于定义新的图IR,除了完成逻辑实现外,还需要为节点封装做额外的努力,包括节点注册、布局指定和运算符加载完成。此外,如有必要,用户还需要实现降低步骤、操作IR生成和指令生成。
• 运算符级:TVM和TC在函数级别上对运算符有更好的可扩展性,几乎不需要额外的努力,只需描述计算行为。TVM将外部函数视为黑盒,并原生支持自定义运算符调用[1]。具体来说,用户在调用外部函数时只需要描述计算并声明输出张量的形状。用户还可以将python函数挂钩为外部函数,这使得TVM在面对自定义运算符的挑战时更加灵活。
深度学习编译器中支持大量的运算符。这里我们选择四个具有代表性且在不同深度学习编译器中最常用的运算符进行说明。
• 广播(Broadcasting):如果没有广播运算符,运算符的输入形状将更加严格。例如,在add运算符中,期望输入张量具有相同的形状。一些编译器(如XLA和Relay)放宽了限制,提供了一个称为broadcasting的运算符。例如,XLA允许通过复制向量直到其形状与矩阵相同来对矩阵和向量进行逐元素加法。
• 谓词(Predication):原始谓词是一种广泛使用的方法,通过在执行前判断条件来决定要执行的代码。条件总是用布尔标志指出,false意味着相关代码不应该被执行。谓词也可以用作控制流中的一种方法,例如通过非零判断来避免非法除法。编译器可能由于不同的实现而重新定义谓词的语义。Glow将谓词视为一种优化,这意味着如果谓词的执行不加速程序,后端可以忽略内容[4]。加速场景可以在RNN实现中找到,通过避免不同批次大小带来的一些计算。然而,Glow假设谓词不影响程序的主要语义,这意味着Glow中的谓词不允许用于与控制流相关的代码。TVM提供了一个assert语句,适合Halide IR和Python AST[5]。
• 自动微分(Automatic differentiation):Glow通过将梯度和高级微分运算符(如随机梯度下降SGD)降低为低级运算符(如加法和乘法)来自动微分其计算图。通过这样做,Glow可以支持没有实现复杂训练操作的后端。受反向模式[75]的启发,Relay通过特定类型表示偏导数,相关的梯度节点通过转换内部AST来重写。
• 控制流(Control flow):在图IR中,在表示像RNN这样的复杂灵活模型时需要控制流。Relay注意到任意控制流可以通过递归和模式实现,这已经被函数式编程社区证明[78]。因此,Relay为实现控制流提供了if运算符和递归函数,这意味着需要手动通过判断和跳转来实现while循环。相反,XLA通过特殊的HLO运算符while和conditional来表示控制流。XLA已经通过一系列基本HLO操作实现了控制流运算符。
4.3 低级中间表示
4.3.1 低级中间表示的实现
低级中间表示以比高级中间表示更细粒度的表示描述深度学习模型的计算,这通过提供调优计算和内存访问的接口来实现目标相关的优化。它还允许开发人员在编译器后端利用成熟的第三方工具链,例如Halide[77]和多面体模型[13]。在本节中,我们将低级中间表示分为三类:基于Halide的中间表示、基于多面体模型的中间表示和其他独特的中间表示。
基于Halide的中间表示 - Halide最初是为了并行化图像处理而提出的,TVM证明了它在深度学习编译器中具有可扩展性和高效性。Halide的基本概念是计算和调度的分离。编译器采用Halide不会直接给出具体方案,而是尝试各种可能的调度并选择最好的一个。当用于深度学习编译问题时,Halide的原始中间表示需要在设计和实现上进行更改。例如,Halide的输入形状是无限的,而深度学习编译器需要知道数据的确切形状以便将运算符映射到硬件指令。一些编译器(如TC)需要固定大小的数据,考虑到机器学习应用中的张量比一般程序具有更高的时间局部性。
TVM通过以下努力[7]将基于Halide的中间表示从原始Halide程序中分离出来并改进为独立的符号中间表示。首先,TVM移除了原始Halide IR对LLVM的依赖。Halide的项目模块结构和中间表示的整体设计都已重新组织,以追求更好的逻辑性和从图IR和前端语言(如python)的公开可访问性。还考虑了可重用性,以便通过引入运行时分发机制轻松添加自定义运算符。TVM将变量定义从字符串匹配改为指针匹配,保证每个变量有单一的定义位置,如静态单赋值(SSA)[35],这对进一步优化很有用。
基于多面体模型的中间表示 - 多面体模型是深度学习编译器的另一个重要参考。它使用线性规划和其他数学方法来优化基于循环的代码,其中边界和分支的控制流是静态的。基于多面体模型的中间表示经历多个多面体转换(例如融合、分块、下沉和映射),包括设备相关和设备无关的优化。当针对特定架构时,多面体转换涉及调度和映射策略的变化。多面体模型中的主要优化空间是循环降低,例如迭代步长和循环级别。多面体模型总是使用各种特定节点来呈现程序的语义,这将在下面介绍。进一步的优化也基于这些节点。有许多可以被基于多面体的编译器借用的工具链,例如isl[94]、Omega[56]、PIP[36]、Polylib[66]和PPL[20]。由于能够处理众多深度嵌套的循环,许多深度学习编译器(如Tensor Comprehension(TC)和PlaidML)使用多面体模型,无论是否进行修改。
TC在低级中间表示中有其独特的设计,它结合了Halide和多面体模型。它使用基于Halide的中间表示来表示计算,但采用基于多面体模型的中间表示来表示循环结构。因此循环结构可以使用多面体技术进行优化。TC通过抽象实例呈现详细的表达式并引入特定的节点类型,其中一些类似于原始多面体模型。节点及其功能可以如表3所述。简而言之,TC使用domain节点来指定索引变量的范围,使用context节点来引入与硬件相关的新迭代变量,例如GPU的块大小。band节点确定迭代的顺序。filter节点表示与语句实例结合的迭代器。set和sequence是关键字,用于指定filter的执行类型,set用于并行,sequence用于串行执行。此外,TC使用extension节点来引入代码生成所需的其他必要指令,例如内存移动语句。

表 3. TC 中使用的基于多面体模型的 IR 节点。
Stripe/PlaidML通过嵌套多面体模型表示张量操作。嵌套多面体模型通过将并行多面体块的嵌套扩展到多个级别来创建可并行化代码的层次结构。此外,它允许将嵌套多面体分配给嵌套的内存单元,提供了一种将计算与多级硬件拓扑的缓存结构匹配的方法。Stripe的硬件配置独立于内核代码完成。Stripe包含向优化pass发出信号的标签。这些标签不改变内核结构,但为优化pass提供关于硬件目标的附加信息。Stripe将机器学习操作分割成适合本地硬件资源的”块(tile)”。
其他独特的中间表示 - 有些深度学习编译器在不使用Halide和多面体模型的情况下实现了定制的低级中间表示。在定制的低级中间表示上,它们应用硬件特定的优化并直接降低到LLVM IR。
Glow中的低级中间表示是基于指令的表达式,它对通过地址引用的张量进行操作[79]。低级中间表示不仅可以实现高级中间表示无法实现的目标无关优化,还可以基于低级指令(例如异步DMA操作)表示目标特定的操作。此外,低级中间表示允许编译器创建调度来隐藏内存操作的延迟。低级中间表示中基于指令的函数有两个部分:声明和程序。第一部分声明在程序生命周期内存在的多个常量内存区域(例如输入、权重、偏置)。第二部分是局部分配区域的列表,包括函数(例如conv和pool)和临时变量。指令可以在全局内存区域(例如声明)或局部分配区域(例如程序)上运行。此外,每个操作数都用限定符之一(“@in”/”@out”/”@inout”)标注。”@In”表示操作数从缓冲区读取。”@Out”表示操作数写入缓冲区。”@Inout”表示操作数读取和写入缓冲区。这些指令和操作数限定符帮助Glow确定何时可以执行某些内存优化(例如复制消除或缓冲区共享)。完成低级中间表示优化后,Glow通过LLVM执行硬件特定的优化和代码生成。
MLIR深受LLVM的影响,并重用了其许多好的想法和接口。MLIR比LLVM更像是一个纯粹的编译器基础设施。MLIR IR位于模型表示和生成硬件特定代码的低级编译器之间。MLIR具有灵活的类型系统,并允许在同一编译单元中组合多个抽象级别。MLIR引入方言(dialect)来表示这些多个抽象级别。每个方言由一组定义的不可变操作组成。MLIR当前的方言包括TensorFlow IR、XLA HLO IR、实验性多面体IR、LLVM IR和TensorFlow Lite。通过LLVM IR的方言,MLIR可以利用LLVM类型系统来定义完全自定义的类型和操作。此外,MLIR可以创建新方言来连接到新的低级编译器,这为硬件开发人员和编译器研究人员铺平了道路。
XLA的HLO IR既是高级中间表示又是低级中间表示,因为HLO足够细粒度以表示硬件特定的信息。此外,HLO支持硬件特定的优化,并可用于生成LLVM IR。
4.3.2 基于低级中间表示的代码生成
上述提到的所有深度学习编译器最终都可以降低到LLVM IR,它们受益于LLVM成熟的优化器和代码生成器。此外,LLVM可以从头开始为专用加速器显式设计自定义指令集。然而,当直接传递给LLVM IR时,传统编译器可能生成较差的代码。为了避免这种情况,深度学习编译器应用两种方法来实现硬件相关的优化:1) 在LLVM的上层中间表示(例如基于Halide的中间表示和基于多面体模型的中间表示)中执行目标特定的循环转换;2) 为优化pass提供关于硬件目标的附加信息。大多数深度学习编译器同时应用这两种方法,但侧重点不同。一般来说,更倾向于前端用户的深度学习编译器(例如TC、TVM、XLA和nGraph)可能专注于1);更倾向于后端开发人员的深度学习编译器(例如Glow、PlaidML和MLIR)可能专注于2)。
深度学习编译器中的编译方案主要可以设计为两种类型:即时编译(JIT)和提前编译(AOT)。对于JIT编译器,它可以动态生成可执行代码,并且可以利用更好的运行时知识来优化代码。AOT是深度学习编译器的另一种方法,它首先生成所有可执行二进制文件然后执行它们。与JIT编译相比,AOT编译在静态分析方面可以有更大的范围,因此受益于更彻底的优化。AOT方法可以与目标平台的交叉编译器一起应用,以支持嵌入式平台(例如C-GOOD[55])或在远程机器上执行(TVM RPC)。
对于在CPU(X86和ARM)上的代码生成,深度学习编译器(例如TVM、PlaidML、TC和Glow)基于优化的低级中间表示生成LLVM IR,并调用LLVM进行JIT编译。对于GPU上的代码生成,XLA通过LLVM NVPTX支持NVIDIA GPU,这也允许将PTX代码JIT/AOT编译为原生GPU机器代码。然而,TVM根据基于Halide的中间表示生成CUDA代码,TC根据基于多面体模型的中间表示生成CUDA代码。对于其他定制加速器,深度学习编译器通常为AOT编译生成相应的代码,例如PlaidML为AMD GPU生成OpenCL代码。
4.4 前端优化
在构建计算图之后,前端会应用图级别的优化。许多优化在图级别更容易被发现和执行,因为图提供了计算的全局视图。这些优化仅应用于计算图,而不是后端的具体实现;因此,它们是硬件无关的,这意味着计算图优化可以应用于各种后端目标。
在传统编译器中,代码的中间表示(IR)由细粒度的三地址指令组成,基本块是连续三地址指令的最大序列。基本块成为数据流图的节点,其边表示基本块的执行顺序。因此,实现了两个范围的优化,包括局部(块级)优化和全局(数据流级)优化。
然而,在深度学习编译器中,深度学习模型的IR由粗粒度的操作节点组成。一个节点表示对张量的操作,一条边表示两个操作之间的依赖关系。节点足够粗粒度,可以在单个节点内部进行优化。几个相邻的节点可以被视为一个节点块,优化可以在块内部应用。这些优化通常由pass来定义,也可以通过遍历计算图的整个节点并执行图转换来应用。前端提供方法来:1) 从计算图中捕获特定特征;2) 重写图的部分以进行优化。除了预定义的pass之外,开发人员还可以通过前端定义自定义的pass。
在本节中,我们将前端优化分为三类:1) 节点级优化(例如,零维张量消除、空操作消除),2) 块级(窥孔、局部)优化(例如,代数化简、融合),以及3) 数据流级(全局)优化(例如,公共子表达式消除CSE、死代码消除DCE)。我们描述每个类别中的详细优化如下。
4.4.1 节点级优化
大多数深度学习编译器一旦导入深度学习模型并转换为计算图,就可以确定每个操作的输入张量和输出张量的形状。这个特性允许深度学习编译器根据形状信息执行优化。节点级优化有两种类型:节点消除(消除不必要的操作节点)和节点替换(用其他低成本节点替换操作节点)。
空操作消除(Nop Elimination) 在通用编译器中移除占用少量空间但不指定任何操作的空操作指令。但在深度学习编译器中,空操作消除负责移除缺少足够输入的操作。例如,只有一个输入张量的求和操作可以被消除,填充宽度为零的填充操作可以被消除。
零维张量消除(Zero-dim-tensor elimination) 负责移除输入为零维张量的不必要操作。假设A是零维张量,B是常量张量,A和B的求和操作节点可以用已经存在的常量节点B替换,而不影响正确性。假设C是一个三维张量,但其中一个维度的形状为零,例如{0,2,3},因此C没有元素,argmin/argmax操作节点可以被消除。
4.4.2 块级优化
代数化简(Algebraic simplification) - 代数化简优化在深度学习编译器和通用编译器中都包含三个部分的优化:(1) 代数恒等式,例如,我们可以将x × 1的计算改为x;(2) 强度削减,我们可以用更便宜的运算符替换更昂贵的运算符,例如将x/2替换为x × 0.5;(3) 常量折叠,我们可以用常量表达式的值替换它们,例如将表达式x × 2 × 3替换为最终值x × 6。它们考虑一系列操作节点,然后利用不同类型运算符的交换性、结合性和分配性来简化计算。
除了上述描述的典型运算符(+, −, ×, ÷等)之外,代数化简优化还可以应用于深度学习特定操作,例如reshape、transpose和pooling,因此运算符可以被重新排序,有时可以被消除,这减少了冗余并提高了效率。这里我们展示几个常见的情况:
- reshape/transpose节点的优化 - 这个优化根据比代数化简更具体的特征来查找和移除reshape/transpose操作。以矩阵乘法(GEMM)为例,有两个矩阵(例如A和B),两个矩阵都被转置(分别产生A^T和B^T),然后A^T和B^T相乘。然而,实现GEMM的更有效方法是切换参数A和B的顺序,将它们相乘,然后转置GEMM的输出。实际上,这种方法将两个转置操作减少到只有一个,这个优化可以进行这种重写;
- transpose节点的优化 - 这个优化将多个连续的转置节点合并为单个节点,消除恒等转置节点,并在实际上不移动数据时将转置节点优化为reshape节点;
- pool节点的优化 - 这个优化交换相邻的池化节点和ReLU节点的顺序,以便ReLU运算符可以在较小的张量上操作。因为ReLU操作占整个计算的一小部分,这个优化可能带来轻微的性能改进;
- ReduceMean节点的优化 - 如果reduce运算符的输入是4D且最后两个维度要被约简,这个优化会将ReduceMean替换为AvgPool节点(在Glow中)。
算子融合(Operator fusion) - 算子融合是深度学习编译器不可或缺的构建模块,确实是深度学习编译器中的一项关键优化。它能够更好地共享计算、消除中间分配、通过组合循环嵌套促进进一步优化[78],以及减少启动和同步开销[89]。在TVM中,运算符被分为四类:1) 注入式(injective,一对一映射),2) 归约式(reduction,复杂不可融合),3) 复杂输出可融合(complex-out-fusible,可以将逐元素映射融合到输出,例如conv2d),以及4) 不透明(opaque,无法融合,例如sort)。当定义运算符时,其对应的类别也被确定。为此,TVM设计了几条规则来执行融合:1) 多个注入式运算符可以融合成一个新的注入式运算符,2) 归约运算符可以与输入之前的注入式运算符融合,3) 复杂输出运算符可以与输出之后的注入式运算符融合。根据运算符类别和融合规则,可以方便地执行融合。在Tensor Comprehension中,融合是基于自动多面体转换以不同的方式进行的。然而,如何识别和融合越来越复杂的图模式,例如具有多个广播和归约节点的块,仍然是一个问题。
算子下沉(Operator sinking) - 这个优化将诸如转置的操作下沉到批归一化、ReLU、sigmoid、通道混洗等操作之下。通过这样做,许多相似的操作被移动并彼此靠近,这为代数化简创造了更多机会。
4.4.3 数据流级优化
公共子表达式消除(Common Sub-expression Elimination, CSE) - 如果表达式E的值之前已经计算过,并且E的值自上次计算以来没有改变,则表达式E是公共子表达式[19]。在这种情况下,E的值只计算一次,已经计算的E的值可以用来避免在其他地方重新计算。深度学习编译器在整个计算图中搜索公共子表达式,并用之前计算的结果替换后续的公共子表达式。程序员和编译器优化都可能导致公共子表达式。
死代码消除(Dead Code Elimination, DCE) - 如果一组代码的计算结果或副作用没有被使用,则该代码是死代码,DCE优化移除死代码。死代码通常不是由程序员造成的,而是由其他图优化引起的。因此,DCE以及CSE在其他图优化之后应用。其他详细的优化,例如死存储消除(DSE),它移除对张量的存储但张量永远不会被使用,也属于DCE。
静态内存规划(Static memory planning) - 静态内存规划优化旨在尽可能重用内存缓冲区。通常有两种方法:原位内存共享和标准内存共享。原位内存共享对操作的输入和输出使用相同的内存,在计算之前只分配一份内存。标准内存共享找到其他内存共享的方式,它重用先前操作的内存而不重叠。两种方法都有应该避免的专门条件。内存规划问题类似于通用编译器中的寄存器分配。静态内存规划是离线完成的,这使得可以使用更复杂的规划算法。
布局转换(Layout transformation) - 这个优化试图找到最佳的数据布局来存储计算图中的张量,然后向计算图插入布局转换节点。注意,实际的转换不在这里执行,而应该在编译器后端评估计算图时执行。例如,通道优先格式(NCHW)和通道在后格式(NHWC)之间的转换是典型的布局转换。
实际上,相同操作在不同布局中的性能是不同的,并且最佳布局在不同硬件上是不同的(例如CPU、GPU和其他定制加速器)。例如,GPU上NCHW格式的操作通常运行得更快,因此在GPU上转换为NCHW格式是高效的(实际上,这是TensorFlow的布局优化器总是试图做的)。一些深度学习编译器可能依赖调用硬件特定的库来实现更高的性能,但这些库可能需要特定的布局。此外,一些深度学习加速器可能更喜欢更复杂的布局(例如tile)。因此,深度学习编译器面向各种硬件,它们需要提供一种执行布局转换的方法。
不仅输入、输出和中间张量的数据布局对最终性能有重要影响,转换操作也有显著的开销。因为转换操作也消耗内存和计算资源。
最近基于TVM针对CPU的一项工作[64]首先将计算图中所有卷积操作的布局更改为NCHW[x]c,其中c表示通道C的分割子维度,x表示子维度的分割大小。然后在提供硬件细节时通过自动调优全局探索所有x参数,例如缓存行大小、向量化单元大小和内存访问模式,这属于硬件特定的优化。
4.4.4 Tensorflow XLA案例研究
为了具体说明计算图优化,我们在Tensorflow XLA中每个pass之前和之后转储HLO图。我们选择Alexnet模型作为XLA编译器的输入,Volta GPU作为目标硬件。优化如图6所示。为了简单起见,我们移除了一些节点的数据布局信息。代数化简包括将连续的转置和reshape节点减少为单个reshape节点,以及将reshape节点替换为bitcast节点。CSE重用broadcast节点。cuDNN转换将卷积节点转换为cuDNN的函数调用(convForward),以使图优化能够利用cuDNN库。常量折叠将相邻的卷积(convForward)和add节点转换为带偏置的卷积(convBiasActivationForward)。算子融合融合了几个bitcast节点和一个add节点。注意,深度学习编译器(例如XLA)中前端优化的实现包含几个阶段。因此,优化会执行多次,这可能每次都改变计算图,从而为进一步优化引入更多机会。

图6:从Tensorflow XLA的Alexnet转储HLO图中提取的计算图优化示例。
4.4.5 讨论
前端是深度学习编译器中最重要的组件之一,负责从深度学习模型到高级IR(即计算图)的转换以及基于高级IR的硬件无关优化。尽管高级IR的数据表示和运算符定义在不同深度学习编译器之间的实现可能不同,但硬件无关优化在节点级、块级和数据流级三个层面上趋同。每个层面的优化方法利用深度学习特定的以及通用编译优化技术,这在计算图层面减少了计算冗余并提高了深度学习模型的性能。具体来说,XLA的前端在现有深度学习编译器中包含最详尽的硬件无关优化。
4.5 后端优化

图7:深度学习编译器中应用的后端优化概述
4.5.1 硬件特定优化
在深度学习编译器的后端,应用硬件特定优化(也称为目标相关优化)来获得针对特定硬件架构的高性能代码。应用后端优化的一种方法是将低级IR转换为LLVM IR,以利用LLVM基础设施生成优化的CPU/GPU代码。另一种方法是利用深度学习领域知识设计定制优化,这可以更有效地利用目标硬件。由于硬件特定优化是为特定硬件架构或实现量身定制的,我们介绍现有深度学习编译器中广泛采用的五种方法,包括硬件内在函数映射、内存分配和获取、内存延迟隐藏、并行化和循环导向优化技术。
硬件内在函数映射(Hardware Intrinsic Mapping) - 硬件内在函数映射可以将特定的低级IR指令集转换为已经在硬件上高度优化的内核。在TVM中,硬件内在函数映射通过可扩展张量化(extensible tensorization)的方法实现,它可以声明硬件内在函数的行为和用于内在函数映射的下降规则。这种方法使编译器后端能够应用硬件实现以及高度优化的手工微内核到特定的操作模式,从而获得显著的性能提升。而Glow支持硬件内在函数映射,例如量化,这通常用于最小化内存占用和提高推理速度。Glow可以估计神经网络每个阶段的可能数值范围,并支持配置文件引导的优化来自动执行量化。此外,Halide作为深度学习编译器(如TVM)中广泛使用的低级IR,将特定的IR模式映射到每个架构上的SIMD操作码(例如x86上的SSE/AVX和ARM上的NEON)。这种方法可以避免LLVM IR映射在遇到向量模式时的低效率。
内存分配和获取(Memory allocation and Fetching) - 内存分配是后端代码生成中的另一个挑战,特别是对于GPU和定制加速器。例如,GPU主要包含两个内存空间,共享内存空间和本地内存空间,其中共享内存具有较低的访问延迟但内存大小有限,本地内存具有较高的访问延迟但容量大。这种内存层次结构需要高效的内存分配和获取技术来提高数据局部性。为了实现内存分配和获取的优化,TVM引入了内存作用域(memory scope)的调度概念。内存作用域调度原语可以将计算阶段标记为共享或线程局部。对于标记为共享的计算阶段,TVM生成具有共享内存分配以及协作数据获取的代码,它在适当的代码位置插入内存屏障以保证正确性。此外,TC也通过扩展PPCG[96]编译器提供共享和本地(又称私有)内存的类似功能。与TVM不同,TC中的内存分配和获取受到更多约束(称为内存提升),它仅支持在预定义规则下的优化内存分配和获取。除了GPU之外,其他深度学习加速器也需要在代码生成中进行高效的内存分配和获取以获得更好的性能。特别是,TVM通过其内存作用域调度原语在加速器中启用特殊缓冲。
内存延迟隐藏(Memory Latency Hiding) - 内存延迟隐藏也是后端中使用的一项重要技术,通过尽可能重新排序流水线执行来实现。由于大多数深度学习编译器支持CPU和GPU上的并行化,内存延迟隐藏可以通过硬件优化(例如GPU上的warp上下文切换)自然实现。但对于使用解耦访问-执行(DAE)架构实现的类TPU加速器,后端需要执行调度和细粒度同步以获得正确和高效的代码。为了获得更好的性能并减少编程负担,TVM引入虚拟线程(virtual threading)调度原语,它使用户能够在虚拟化多线程架构上指定数据并行性。然后TVM通过插入必要的内存屏障来下降这些虚拟并行化的线程,并将这些线程的操作交错到单个指令流中,这形成了每个线程更好的执行流水线,以尽可能隐藏内存访问延迟。
循环导向优化(Loop Oriented Optimizations) - 循环导向优化也在后端应用,以为目标硬件生成高效代码。由于Halide和LLVM(基于多面体方法)[59]已经包含了这些优化技术,一些深度学习编译器在其后端利用Halide和LLVM。循环导向优化中应用的关键技术包括:循环融合、滑动窗口、分块、循环重排序和循环展开。
• 循环融合(Loop fusion):循环融合是一种循环优化技术,可以将具有相同边界的循环融合在一起以获得更好的数据重用。对于像PlaidML、TVM、TC和XLA这样的编译器,这个优化是通过Halide调度或多面体方法执行的,而Glow通过其算子堆叠进行循环融合。
• 滑动窗口(Sliding windows):滑动窗口是Halide采用的循环优化技术。滑动窗口的核心概念是在需要时计算值并存储它们直到不再需要。当嵌套循环的输出值由前一个计算阶段(循环)计算的值计算时,滑动窗口可以动态缓存所需的值以进行数据重用。由于滑动窗口会交错两个阶段的计算并使其串行化,这是并行性和数据重用的权衡。
• 分块(Tiling):分块是高性能计算中常用的另一种循环优化技术。分块意味着将循环分割成几个块(tile),因此循环被分为遍历块的外循环和在块内部迭代的内循环。这种转换旨在通过将块适配到硬件缓存中来实现块内部更好的数据重用,这在现代处理器的内存层次结构中至关重要。由于块的大小非常特定于硬件,很难提供规则来定义分块的模式和大小。因此,许多深度学习编译器支持通过自动调优自动决定分块的模式和大小,这在4.5.2节中详细描述。基于多面体模型的方法可以通过修改其band节点的仿射函数来实现分块。
• 循环重排序(Loop reordering):循环重排序(也称为循环置换)意味着改变嵌套循环中的迭代顺序,这可以优化内存访问,从而增加空间局部性。它非常特定于数据布局和硬件特性,Halide提供了一个名为reorder的调度原语,可以通过自动调优来优化循环顺序。然而,当语句沿迭代顺序有依赖关系时,执行循环重排序是不安全的。
• 循环展开(Loop unrolling):循环展开也是编译器中常用的优化技术。循环展开将特定循环展开为固定数量的循环体副本,这可以从更少的循环控制成本(或为该循环生成更少的指令数)中受益。重用Halide工作的编译器支持循环展开的调度原语,但它的循环展开是完全循环展开,这意味着它将大小为n的循环替换为n个循环体的副本。广义循环展开通过循环分割和循环展开的组合来表达,它首先将循环分割成两个嵌套循环,然后完全展开内循环。
并行化(Parallelization) - 由于现代处理器通常支持多线程和SIMD并行性,编译器后端利用并行性以最大化硬件利用率以获得高性能非常重要。Halide和多面体模型是现有深度学习编译器中用于利用硬件并行化和SIMD向量化的两种主要技术。
• Halide:对于线程级并行化,Halide使用名为parallel的调度原语来指定循环的并行化维度,可以通过自动调优进行调整。Halide还通过将标记为parallel的循环维度映射到block和thread的注释来支持GPU并行化。对于向量化,Halide将大小为n的循环替换为n宽的向量语句,可以通过硬件内在函数映射映射到硬件特定的SIMD操作码。
• 多面体模型(Polyhedral model):在多面体模型中,可以对多面体执行转换,从而获得转换后的循环映射以获得更好的并行化。由于多面体模型实际上是一种循环转换来检测潜在的并行性,它可以应用于使用CPU线程以及GPU内核进行并行化。此外,Stripe开发了多面体模型的变体,称为嵌套多面体模型,它引入并行多面体块作为其迭代的基本执行元素。经过这个扩展,嵌套多面体模型可以在分块和跨步的层次之间检测层次并行化,而不管并行多面体块内部的复杂控制流。
除了上述两种技术之外,一些深度学习编译器依赖手工库,如Glow[79]或硬件供应商提供的优化数学库(详细讨论见4.5.3节),这可以在目标硬件上实现更高的性能。同时,Glow将向量化卸载到LLVM编译器,因为当提供张量维度和循环迭代次数的知识时,LLVM自动向量化器工作良好。与依赖优化库和LLVM基础设施的方法相比,完全由编译器后端利用并行性允许应用更多深度学习模型的领域特定知识,从而在不同硬件目标上实现更高性能,但代价是更多的工程努力。
4.5.2 自动调优
由于硬件特定优化中参数调优的搜索空间巨大,有必要利用自动调优来确定最优参数设置。Halide/TVM允许程序员首先定义硬件特定优化(调度),然后使用自动调优来推导最优参数设置。通过这种方式,Halide/TVM程序员可以通过反复检查特定参数设置的性能来更新或重新设计调度。此外,自动调优也可以应用于多面体模型进行参数调优[89]。例如,TC利用自动调优来探索配置(包括多面体调度和参数)并更新编译缓存,以最小化多面体JIT编译的开销。通常,自动调优的实现包括三个部分:参数化、代价模型、搜索技术和性能优化。
参数化(Parameterization) - 1) 数据和目标:数据参数描述数据的规格,例如输入形状。目标参数描述在优化调度和代码生成期间要考虑的硬件特定特征和约束。例如,对于GPU目标,需要指定共享内存和寄存器大小等硬件参数。2) 优化选项:优化选项包括优化调度和相应的参数,例如循环导向优化和块大小。在TVM中,预定义和用户定义的调度以及参数都被考虑在内。而在TC中,它倾向于参数化与性能有强相关性且可以稍后以低成本更改的优化。例如,minibatch维度是参数之一,它通常映射到CUDA中的网格维度,可以在自动调优期间进行优化。
代价模型(Cost model) - 自动调优中应用的不同代价模型的比较如表4所示。

表4:自动调整中应用的不同成本模型的比较。
- 黑盒模型(Black-box model):这个模型只考虑最终执行时间而不考虑编译任务的特征。构建黑盒模块很容易,但在没有任务特征指导的情况下会导致更高的开销和较不理想的解决方案。
- 预定义代价模型(Pre-defined cost model):理想情况下,基于预定义代价模型的方法期望一个完美的模型,该模型建立在编译任务的特征之上,能够评估任务的整体性能。有许多因素影响模型的准确性,例如内存访问模式和流水线依赖关系,这些与任务和硬件目标相关。与基于机器学习的模型相比,预定义模型在应用时产生更少的计算开销,但需要在每个新的深度学习模型和硬件上重新构建模型的大量工程努力。
- 基于机器学习的代价模型(ML-based cost model):基于机器学习的代价模型是一种使用机器学习方法预测性能的统计方法。使用基于机器学习的模型使模型能够随着新配置的探索而更新,这有助于实现更高的预测准确性。TVM提出了一个基于机器学习的模型和一个神经网络模型来预测生成代码在给定硬件目标上的运行时间。代价模型的输入是一个下降的循环程序,两个模型在进行回归预测时使用不同的自变量。具体来说,基于机器学习的模型使用从输入中提取的特征,神经网络模型使用程序的AST。TVM声称两个模型的准确性相似,并选择基于机器学习的模型(树提升)作为默认代价模型。
搜索技术(Searching technique) -
- 初始化和搜索空间确定:初始选项可以随机设置,也可以根据已知配置设置,例如用户给出的配置或历史最优配置。在搜索空间方面,应该在自动调优之前指定。TVM允许开发人员使用其领域特定知识指定搜索空间,并为每个硬件目标提供基于计算描述的自动搜索空间提取。而TC依赖于编译缓存和预定义规则。
- 遗传算法(Genetic algorithm):TC使用的遗传搜索算法将每个调优参数视为基因,将每个配置视为候选者。新候选者由两种方法生成:杂交和突变。对于每个新候选者,根据其适应度选择三个父代,这意味着性能越高,被选中的可能性越大,父代的基因被随机选择,然后形成新候选者。之后,基因以低概率(称为突变率)被赋予随机值,用于控制探索和利用之间的权衡。还有其他需要设置的算法参数,例如迭代界限和其他与调度选择相关的选项。
- 模拟退火算法(Simulated annealing algorithm):TVM使用这个算法。初始配置可以随机生成,然后在每一步随机预测一个类似的配置。如果给定代价模型预测的新配置具有更低的成本,则该步骤被认为是成功的。如果代价模型被修改,算法从最后一步开始,即前一个代价模型的搜索结果。
性能优化(Performance Optimization) -
- 并行化(Parallelization):改进自动调优性能的一个方向是并行化。TC在考虑到遗传算法需要在每个下一次生成步骤之前评估所有候选者的情况下,提出了一种通用的多线程、多GPU策略。首先,该策略将候选配置排队,并由多个CPU线程编译它们。生成的代码在GPU上并行评估,每个候选者拥有其由父代选择步骤使用的适应度。在整个评估完成后,由搜索算法生成新候选者,新的编译作业被排队,等待CPU编译。类似地,TVM支持交叉编译和RPC,允许用户在本地机器上编译并在多个目标上使用不同的自动调优配置运行程序。
- 配置重用(Configuration reuse):改进自动调优性能的另一个方向是重用以前的自动调优配置。TC通过编译缓存存储与给定配置对应的已知最快生成代码版本,并使用元组作为缓存条目来呈现与版本相关的必要信息。在编译期间,在每次内核优化之前查询缓存以实现持久性和重用,如果缓存未命中则触发自动调优。类似地,TVM生成一个日志文件,存储所有调度运算符的最优配置,并在编译期间查询日志文件以获取最佳配置。值得一提的是,TVM对Halide IR中的每个运算符(例如conv2d)执行自动调优,而不是将所有运算符作为一个整体对待,因此最优配置是针对运算符单独的,而不是针对低级IR。
4.5.3 优化内核库
有几个高度优化的内核库被广泛用于在各种硬件上加速深度学习训练和推理。英特尔公司的DNNL(以前称为MKL-DNN)支持英特尔CPU及其集成GPU。DNNL在运行时检测支持的ISA,并为最新支持的ISA(例如Skylake-X上最新的AVX-512 ISA)即时(JIT)部署优化代码。计算密集型原语(例如卷积、GEMM和RNN)和内存带宽受限的原语(例如批归一化、池化和shuffle)都得到支持并高度优化。此外,它支持低精度训练和推理,包括FP32、FP16和INT8(仅推理)以及非IEEE浮点格式bfloat16[97]。NVIDIA公司的cuDNN也提供了在DNN应用中频繁出现的高度调优的原语。此外,它支持可自定义的数据布局(例如4D张量的灵活维度排序、跨步和子区域),这使得cuDNN易于集成到深度学习应用中,并避免频繁的数据布局转换。cuDNN可以充分利用Volta和Turing GPU系列上新的张量核心操作。它还支持低精度训练和推理。AMD公司的MIOpen也对AMD GPU上的高性能机器学习原语进行了类似的优化。然而,与DNNL和cuDNN相比,仅支持有限的功能。例如,只有GEMM原语得到优化,只支持FP16。其他定制的深度学习加速器也维护其特定的内核库以提高深度学习计算的性能。
现有的深度学习编译器,如TVM、nGraph和TC,可以在代码生成期间(例如JIT和AOT)生成对这些库的函数调用。然而,如果深度学习编译器需要利用现有的优化内核库,它们需要首先将数据布局和融合样式转换为内核库中预定义的类型,这种转换可能会破坏最优的控制流。此外,深度学习编译器将内核库视为黑盒,因此在调用内核库时无法跨运算符应用优化(例如算子融合)。总之,当计算可以被特定的高度优化原语满足时,使用优化内核库可以实现显著的性能提升,否则可能受到进一步优化的限制并遭受较不理想的性能。
4.5.4 讨论
深度学习编译器的后端通常采用包括各种硬件特定优化、自动调优技术和优化内核库的设计。硬件特定优化能够为不同的硬件目标(如CPU、GPU和定制深度学习加速器)生成高效代码。后端广泛使用的硬件特定优化包括硬件内在函数映射、内存分配和获取、内存延迟隐藏、循环导向优化和并行化。然而,由于深度学习硬件的多样性,优化不限于本节中介绍的这些。为了解决硬件特定优化引入的大量参数调优空间,自动调优在编译器后端变得至关重要,以减轻推导最优参数设置的手工工作。自动调优的设计通常由四个组件组成,如参数化、代价模型、搜索技术和性能优化。为了进一步提高生成代码的性能,优化内核库也在深度学习编译器的后端广泛使用。当深度学习计算满足高度优化库中的内核定义时,可以实现显著的性能提升。然而,依赖优化内核库可能会浪费高级优化(如跨多个运算符的算子融合)的性能机会。
5 结论与未来方向
在本综述中,我们对现有的深度学习编译器进行了全面分析。首先,我们从各个方面对现有深度学习编译器进行了全面比较,这可以作为用户为其定制场景选择合适的深度学习编译器的指南。然后,我们深入探讨了现有深度学习编译器中采用的通用设计,包括多级中间表示、前端和后端。我们详细介绍了每个组件的设计理念和参考实现,重点关注深度学习编译器特有的独特中间表示和优化。我们总结了本综述中的发现,并突出了深度学习编译器的未来方向,我们希望这能够启发更多研究人员和从业者为这一领域做出贡献。
动态形状和控制流(Dynamic shape and control flow) - 动态模型在深度学习领域变得越来越流行,其输入形状甚至模型本身可能在执行过程中发生变化。特别是,动态形状是动态模型的主要关注点,TVM和其他深度学习编译器已部分支持。在深度学习社区,尤其是在自然语言处理(NLP)中,模型可能接受各种形状的输入,这对深度学习编译器来说是一个挑战,因为数据的形状直到运行时才是未知的。现有的深度学习编译器需要更多研究努力来高效支持新兴动态模型的动态形状。
此外,未来的深度学习模型将变得更加复杂,可能包括复杂的控制流。由于大多数深度学习框架和编译器使用Python作为其编程语言,如果模型使用控制流实现,性能将成为一个严重问题,因为这会导致模型由Python解释器执行。此外,深度学习模型通常需要复杂的数据/结果预处理和后处理。尽管预处理/后处理可能成为训练和推理中的性能瓶颈,但现有的深度学习编译器尚未考虑这一点。编译器对控制流和常用预处理/后处理的支持将进一步提高模型部署的性能增益。实际上,Relay[78]目前正在开发Relay虚拟机以支持Relay运行时中的控制流,这可以避免使用低效的Python解释器。
高级自动调优(Advanced auto-tuning) - 现有的自动调优技术专注于单个运算符的优化。然而,局部最优的组合并不能导致全局最优。例如,在不同数据布局上应用的两个相邻运算符可以一起进行性能调优,而无需在它们之间引入额外的内存转换节点。此外,随着边缘计算的兴起,执行时间不再是深度学习编译器的唯一优化目标。自动调优中还应考虑新的优化目标,例如更少的内存占用和更低的能耗。
对于基于机器学习的自动调优技术,有几个方向值得进一步探索。首先,机器学习技术应该应用于自动调优的其他阶段,而不仅仅是代价模型。例如,在选择编译器选项和优化调度的阶段,机器学习技术可以用于直接预测期望的可能性并开发算法来确定最终结果,而不是开发代价模型。此外,不同深度学习模型的各种计算特征和优化目标需要用不同的数据集和目标函数训练的特定调优模型,而不是共享一个通用的自动调优模型。其次,基于机器学习的自动调优技术也可以基于深度学习模型的领域知识进行改进。例如,将特征工程(选择特征来表示程序)[99]纳入自动调优技术可能是实现更好调优结果的潜在方向。此外,除了从中间表示收集的静态代码特征外,其他计算特征也具有表现力,并与性能强相关。例如,计算图的特征可以更直观地表示数据依赖和内存移动,并能够更好地预测执行时间。
多面体模型(Polyhedral model) - 如4.5.2节所述,自动调优可以应用于最小化多面体JIT编译的开销,通过重用以前的配置。另一方面,多面体模型可以用于执行自动调度,这可以减少自动调优的搜索空间。将多面体模型和自动调优的组合进一步应用于深度学习编译器的设计以提高效率是一个有前景的研究方向。
多面体模型的另一个挑战是支持稀疏张量。一般来说,稀疏张量的格式(如CSF[84])用索引数组(例如a[b[i]])表达循环索引,这不再是线性的。这种间接索引寻址导致非仿射下标表达式和循环边界,这禁止了多面体模型的循环优化[26, 88]。幸运的是,多面体社区在支持稀疏张量方面取得了进展[92, 93],整合多面体模型的最新进展可以为未来的深度学习编译器增加性能机会。
子图分区(Subgraph partitioning) - 支持子图分区的深度学习编译器可以将计算图分区为几个子图。通过这种方法,计算图不再被视为一个整体,子图可以以不同的方式处理。子图分区可以为深度学习编译器带来至少两个优势。
首先,它开启了集成更多应用图优化的库的可能性。以nGraph与DNNL为例,DNNL是CPU上的深度学习优化库,它可以通过利用其多样化的高度优化内核集合来应用层融合和其他图优化。这种集成使DNNL能够优化和执行大部分兼容的子图,同时将剩余的图留给nGraph。TensorFlow与TensorRT的集成采用了类似的方法。然而,其他深度学习编译器未能利用这种集成。以TVM为例,TVM支持调用DNNL库,但将其视为常规BLAS库,而不利用子图优化。类似的问题发生在专用深度学习加速器上,它们依赖定制的图优化库。总之,集成具有图优化的库可以提供另一种方法来改进深度学习编译器生成的代码的性能。此外,集成还提供了一种通过调用图优化库来支持硬件目标的新方法。
其次,它开启了异构和并行执行的可能性。目前,在两个极端规模的场景(如边缘设备和数据中心)中部署深度学习模型都表现出异构和并行执行的趋势。一旦计算图被分区为子图,不同子图的执行可以同时分配给异构硬件目标。以边缘设备为例,其计算单元可能包括ARM CPU、Mali GPU、DSP,可能还有NPU。从深度学习编译器生成能够高效利用所有计算单元的代码可以显著加速深度学习任务,例如人脸识别和语音助手。
量化(Quantization) - 量化是深度学习模型中众所周知的优化技术,可以通过降低操作数据的精度来减轻计算和内存负担。通过减少位宽,可以用更少的资源存储和计算数据,代价是模型精度略有下降。量化的主要挑战是设计在低精度的好处和模型精度损失之间权衡的策略。深度学习框架中应用的传统量化策略基于一组固定的方案和数据类型,几乎没有针对在不同硬件上运行的代码的定制。而在深度学习编译器中支持量化可以在编译期间利用更多优化信息来推导更高效的量化策略。例如,Relay[78]提供了一个量化重写流程,可以自动为各种方案生成代码。
为了支持量化,深度学习编译器中还有更多挑战需要解决。第一个挑战是如何在不需要大量工程努力的情况下实现具有降低精度的新运算符。为了在Relay中重用量化实现,[15]提出了一个新的方言,用基本指令实现新运算符,而不是从头实现新运算符,这消除了重新实现图级和运算符级优化的工程努力。编译期间量化与其他优化之间的交互是另一个挑战。例如,确定量化的适当位置以及与算子融合等优化的协作需要未来的研究调查。同时,量化也影响深度学习编译器中的硬件特定优化,可以重新设计以利用低精度运算符和数据所需资源更少的性能机会。
统一优化(Unified optimizations) - 尽管现有的深度学习编译器在计算图优化和硬件特定优化方面都采用了相似的设计,但每个编译器都有自己的实现,在某些方面具有优势。缺少一种方法来共享最先进的优化,以及跨现有编译器支持新兴硬件目标。我们倡导统一现有深度学习编译器的优化,以便可以重用每个深度学习编译器中采用的最佳实践。此外,跨深度学习编译器统一优化可以积累强大的力量来影响通用和专用深度学习加速器的设计,并为深度学习编译器和硬件的高效协同设计提供环境。目前,Google MLIR是朝着这个方向的一个有前景的倡议。它提供了多级中间表示的基础设施,包含中间表示规范和工具包以跨每个级别的中间表示执行转换。它还提供灵活的方言,以便每个深度学习编译器可以为高级和低级中间表示构建其定制的方言。通过跨方言的转换,一个深度学习编译器的优化可以被另一个编译器重用。然而,方言的转换可能需要精细的权衡和一些工程努力。
可微编程(Differentiable programming) - 可微编程是一种编程范式,可微编程范式中的程序是完全可微的。它最近吸引了深度学习编译器社区的关注。许多新的编译器项目已经用可微编程替代了计算图,例如Myia[25]和Swift for TensorFlow[18]。此外,Flux[51]是最有前景的机器学习堆栈之一,具有其微分语言Julia[23]。Julia是一种非常适合机器学习编程的语言,专为数学和数值计算设计,Flux扩展了Julia以支持可微算法和由编译器实现的加速。不幸的是,现有的深度学习编译器几乎没有对微分语言的支持。
支持微分语言对现有的深度学习编译器来说是相当具有挑战性的。困难不仅来自数据结构,还来自语言语义。例如,为了实现从Julia到XLA HLO IR的转换,项目[37]面临的挑战之一是Julia使用的命令式语言和XLA使用的符号语言之间的控制流不同。为了有效使用HLO IR,编译器还需要为Julia提供操作抽象,以支持XLA的特定语义,例如MapReduce和broadcast。此外,Julia和XLA之间微分语义的差异也导致自动微分算法和相应设计的重大变化。
图神经网络(Graph neural network, GNN) - 图神经网络近年来一直是深度学习领域的热门研究方向[22, 102, 104, 107]。传统的深度学习网络在欧几里得空间中规则结构化的数据(例如图片和语音)上有效,而在不规则结构化的数据(例如社交和电子商务)上表现不佳。以电子商务为例,用户、产品和广告可以被视为节点,用户购买商品和点击广告等操作可以被视为边。节点和边形成一个大图。图神经网络将这个大图作为输入,通过聚合邻居信息获得节点、边或子图的表示,从而实现分类、预测和推荐等任务。已经公认图神经网络的设计在非结构化数据上表现良好。
大量的深度学习框架已经提供了图神经网络的实现,例如TensorFlow、PyTorch、MXNet、PaddlePaddle和Theano。然而,现有的深度学习编译器对图神经网络几乎没有支持。只有TVM在Cora数据集上发布了使用Relay和MXNet构建图卷积网络(GCN)的原始教程。现有深度学习编译器中图神经网络支持不成熟的原因之一是其独特的设计。例如,图神经网络总是很浅,其中大多数不超过三层[107]。这是因为图神经网络将相邻节点的表示合并得更接近,这导致层数更少。然而,鲁莽地堆叠多个图神经网络层很容易导致过度平滑。对于深度学习编译器来说,如何使用图级优化来避免图神经网络的过度平滑需要更多的研究努力。
隐私保护(Privacy protection) - 随着传感器和移动手机等边缘设备在我们日常生活中的广泛使用,边缘-云系统正变得无处不在,可以运行深度学习模型来完成人脸检测和语音识别等智能任务。在边缘-云系统中,深度学习模型通常被分成两半,每个部分模型分别在边缘设备和云服务上运行,这可以提供更好的响应延迟并消耗更少的通信带宽。然而,边缘-云系统的缺点之一是用户隐私变得脆弱。原因是攻击者可以拦截从边缘设备发送到云的中间结果,然后使用中间结果训练另一个模型,该模型可以揭示与原始用户任务偏离的隐私信息[40, 69, 73]。
为了保护边缘-云系统中的隐私,现有方法[40, 69, 73]提出向中间结果添加具有特殊统计属性的噪声,这可以降低攻击者模型的准确性,而不会严重恶化原始用户模型的准确性。然而,现有方法的困难在于确定应该插入噪声的层,识别最优层是相当费力的。上述困难为深度学习编译器支持隐私保护提供了很好的机会,因为编译器维护关于所有层的计算、通信和熵的丰富信息,这可以指导跨层的噪声生成以自动实现更好的隐私保护。我们相信这可能是深度学习社区中安全和编译学科交叉的一个有趣的研究方向。