深度学习编译器关键组件

1 高层中间代码

为了克服传统编译器中采用的IR限制DL模型中复杂计算的表达的局限性,现有的DL编译器利用高层IR(称为图IR)进行高效的代码优化设计。

1.1 图表示

基于DAG的IR:基于DAG的IR是编译器构建计算图的最传统方法之一,它将节点和边组织为有向无环图(DAG)。在DL编译器中,DAG的节点表示原子DL运算符(卷积、池等),边表示张量。该图是无环图,这与泛型编译器的数据依赖图(DDG)不同。借助于DAG计算图,DL编译器可以分析各种操作符之间的关系和依赖关系,并用它们来指导优化。在DDG上已经有很多优化,比如公共子表达式消除(CSE)和死代码消除(DCE)。通过将DL的领域知识与这些算法相结合,可以对DAG计算图进行进一步的优化。

1.2 图IR的实现

数据表示:DL编译器中的数据(例如,输入、权重和中间数据)通常以张量的形式组织,张量也称为多维数组。DL编译器可以通过内存指针直接表示张量数据,也可以通过占位符以更灵活的方式表示张量数据。占位符包含张量的每个维度的大小。

占位符(Placeholder):占位符广泛应用于符号编程。占位符只是一个具有明确形状信息(例如,每个维度中的大小)的变量,它将在计算的后期用值填充。它允许程序员在不考虑精确数据元素的情况下描述操作和构建计算图,这有助于在DL编译器中将计算定义和精确执行分离开来。

动态维度表示:在声明占位符时,通常支持未知维度大小。例如,TVM使用Any表示未知维度(例如,Tensor〈(Any,3), fp32〉);XLA使用None实现相同的目的(例如,T f.placeholder(“float”, [None,3])。未知形状表示是支持动态模型的必要条件,但是要完全支持动态模型,边界推理和维度检查应该放松。

数据布局:数据布局描述了张量在内存中的组织方式,通常是从逻辑索引到内存索引的映射。数据布局通常包括尺寸顺序(如NCHW和NHWC)、tiling、padding、striding等。TVM和Glow将数据布局表示为运算符参数,并需要此类信息进行计算和优化。Relay和MLIR将把数据布局信息添加到他们的张量类型系统中。

算子支持:算子对应计算图中的节点,深度学习编译器支持的算子表示对应的深度学习工作。

2 底层中间代码

底层IR以比高层IR更细粒度的表示形式描述DL模型的计算,它通过提供接口来调整计算和内存访问,从而实现与目标相关的优化。在本节中,我们将底层IRs的常见实现分为三类:基于Halide的IR、基于polyhedral的IR和其他IR。

2.1 底层IR

基于Halide的IR:Halide的基本原理是计算和调度的分离。采用Halide的编译器没有直接给出具体的方案,而是尝试各种可能的时间表并选择最佳的。Halide不能用复杂的图案(例如,非矩形)来表示计算,但是DL中的计算非常规则,可以用Halide完美地表示。当应用到DL编译器的后端时,Halide的原始IR需要修改。例如,Halide的输入形状是无限的,而DL编译器需要知道数据的确切形状,以便将运算符映射到硬件指令。TVM将Halide IR改进为符号IR,它消除了对LLVM的依赖,重构了项目模块的结构和Halide的IR设计,追求更好的组织以及对graph-IR和前端语言(如Python)的可访问性。通过实现运行时调度机制,方便地添加自定义操作符,提高了可重用性。

基于Polyhedral的IR:Polyhedral模型是DL编译器采用的一种重要技术。它使用线性规划、仿射变换和其他数学方法来优化具有边界和分支静态控制流的基于循环的代码。与Halide不同的是,内存引用和循环嵌套的边界可以是多面体模型中任何形状的多面体。由于能够处理深度嵌套的循环,许多DL编译器,如TC和PlaidML(作为nGraph的后端)都采用了Polyhedral模型作为其底层IR。基于Polyhedral的IR使得应用各种多面体变换(例如,fusing、tiling、sinking和mapping)变得容易,包括依赖于设备和独立于设备的优化。

其它的IR:有DL编译器实现定制的底层IR没有使用Halide和Polyhedral模型。如MLIR受LLVM的影响很大,它是一个比LLVM更纯粹的编译器基础结构。MLIR重用了LLVM中的许多思想和接口,介于模型表示和代码生成之间。MLIR有一个灵活的类型系统,允许多个抽象层次,它引入方言来表示这些多个抽象层次。每个方言都由一组定义的不可变操作组成。此外,MLIR可以创建新的方言来连接到新的低级编译器,这为硬件开发人员和编译器研究人员铺平了道路。

2.2 基于底层IR的代码生成

大多数DL编译器采用的底层IR最终可以降低到LLVM IR,并受益于LLVM成熟的优化器和代码生成器。此外,LLVM可以从零开始显式地为专用加速器设计定制指令集。然而,传统的编译器在直接传递给LLVM_IR时可能会生成糟糕的代码。为了避免这种情况,DL编译器采用了两种方法来实现硬件相关的优化:1)在LLVM的上层IR(例如,基于Halide的IR和基于Polyhedral的IR)中执行特定于目标的循环变换;2)为优化过程提供有关硬件目标的附加信息。大多数DL编译器都应用这两种方法,但重点不同。一般来说,倾向于前端用户(如TC、TVM、XLA和nGraph)的DL编译器可能关注1),而更倾向于后端开发人员(如Glow、PlaidML和MLIR)的DL编译器可能关注2)。

3 前端优化

在构建计算图之后,前端应用图级优化。许多优化更容易在图级别识别和执行,因为图提供了计算的全局视图。这些优化只应用于计算图,而不是后端的实现。因此,它们独立于硬件,可以应用于各种后端目标。前端优化可以分为三类:1)节点优化、2)块优化和3)数据(全局)优化。

3.1 节点优化

计算图的节点非常粗糙,可以在单个节点内进行优化。节点级优化包括消除不必要节点的节点消除和用其他低成本节点替换节点。在通用编译器中,Nop消除删除了占用少量空间但不指定任何操作的nop指令。在DL编译器中,Nop消除负责消除缺乏足够输入的操作。例如,可以消除只有一个输入张量的和节点,消除填充宽度为零的填充节点。零维张量消除负责消除输入为零维张量的不必要运算。

3.2 块优化

代数化简:包括1)代数识别,2)强度折减,我们可以用更简便的运算符替换更复杂的运算符;3)常数折叠,我们可以用它们的值替换常数表达式。这种优化考虑了一个节点序列,然后利用不同类型节点的交换性、结合性和分布性来简化计算。除了典型的运算符(+、×等),代数简化还可以应用于特定于DL的运算符(例如,整形、转置和池)。这些操作符可以重新排序,有时甚至可以消除,这减少了冗余,提高了效率。代数简化可以应用的常见情况:1)计算顺序的优化。以矩阵乘法(GEMM)为例,存在两个矩阵(例如,A和B),两个矩阵被转置(分别产生AT和BT),然后AT和BT被相乘在一起。然而,一种更有效的实现GEMM的方法是将参数a和B的顺序进行切换,相乘,然后对GEMM的输出进行转置,从而将两个转置减少为一个;2)优化节点组合,在这种情况下,将多个连续的转置节点组合成一个节点。

算子融合:DL编译器中不可缺少的优化。它能够更好地共享计算,消除中间分配,通过组合循环嵌套促进进一步优化,并减少启动和同步开销。

运算符下沉优化:将诸如转置之类的操作下沉到诸如批处理规范化、ReLU、sigmoid和channel shuffle之类的操作之下。通过这种优化,许多相似的操作被移近彼此,为代数简化创造了更多的机会。

3.3数据优化

公共子表达式消除(CSE):如果先前计算了E的值,则表达式E是公共子表达式,并且自先前计算以来E的值不必更改。在这种情况下,E的值被计算一次,并且已经计算的E的值可以用于避免在其他地方重新计算。

死代码消除(DCE):如果不使用计算结果,则一组代码是死的。DCE优化去除了死代码。死代码通常不是由程序员引起的,而是由其他图形优化引起的。因此,DCE和CSE是在其他图优化之后应用的。其他的优化,比如死存储消除(DSE),它将存储移到永远不会使用的张量中,也属于DCE。

静态内存规划划优化:尽可能地重用内存缓冲区。通常有两种方法:就地内存共享和标准内存共享。就地内存共享将相同的内存用于操作的输入和输出,并且在计算之前只分配一个内存副本。标准内存共享重用以前操作的内存而不重叠。静态内存规划是离线完成的,这允许应用更复杂的规划算法。

布局变换:试图找到最佳的数据布局,以便在计算图中存储张量,然后将布局变换节点插入到图中。注意,这里不执行实际的转换,而是在编译器后端计算计算图时执行。事实上,相同的操作在不同的数据布局中的性能是不同的,并且最佳的布局在不同的硬件上也是不同的。例如,GPU上NCHW格式的操作通常运行得更快,因此在GPU上转换为NCHW格式(例如TensorFlow)是有效的。

4 后端优化

DL编译器的后端通常包括各种特定于硬件的优化、自动调优技术。特定于硬件的优化可以为不同的硬件目标高效地生成代码。然而,自动调优在编译器后端是必不可少的,它可以减轻人工获取最佳参数配置的工作量,本文主要涉及特定于硬件的优化技术。

特定于硬件的优化,也称为目标相关优化,用于获得针对特定硬件的高性能代码。应用后端优化的一种方法是将底层IR转换为LLVM IR,利用LLVM基础结构生成优化的CPU/GPU代码。另一种方法是利用DL领域知识设计定制的优化,从而更有效地利用目标硬件。

硬件内部映射:可以将某组低级IR指令转换为已经在硬件上高度优化的内核。在TVM中,硬件内部映射是用可扩展张量化的方法来实现的,它可以声明硬件内部映射的行为和内部映射的降低规则。此方法使编译器后端能够将硬件实现以及高度优化的手工微内核应用于特定的操作模式,从而显著提高性能。Glow支持硬件内部映射,如量化。

内存分配和获取:代码生成中的另一个挑战,特别是对于GPU和定制加速器。例如,GPU主要包含共享内存空间(较低的访问延迟和有限的内存大小)和本地内存空间(较高的访问延迟和较大的容量)。这种内存层次结构需要有效的内存分配和获取技术来改善数据的局部性。为了实现这种优化,TVM引入了内存范围的调度概念。内存范围调度原语可以将计算阶段标记为共享或线程本地。

内存延迟隐藏:通过重新排序执行管道在后端使用。由于大多数DL编译器都支持CPU和GPU上的并行化,因此内存延迟隐藏可以自然地通过硬件实现(例如,GPU上的warp上下文切换)。为了获得更好的性能和减少编程负担,TVM引入了虚拟线程调度原语,使用户可以在虚拟化的多线程体系结构上指定数据并行性。

面向循环的优化:由于Halide和LLVM(与多面体方法集成)已经结合了这样的优化技术,一些DL编译器在其后端使用Halide和LLVM。面向循环优化的关键技术包括循环融合、滑动窗口、分块、循环重排序和循环展开。

循环融合:循环融合是一种循环优化技术,可以融合具有相同边界的循环,以实现更好的数据重用。

滑动窗口:滑动窗口是Halide采用的一种循环优化技术。它的核心概念是在需要时计算值,并动态存储它们以供数据重用,直到不再需要它们为止。由于滑动窗口将两个循环的计算交织在一起并使它们串行化,这是并行性和数据重用之间的一种折衷。

分块:循环拆分为多个块,因此循环分为通过分块迭代的外部循环和在分块中迭代的内部循环。

循环重新排序:循环重新排序(也称为循环置换)改变嵌套循环中的迭代顺序,从而优化内存访问,从而增加空间局部性。它特定于数据布局和硬件功能。

循环展开:循环展开可以将特定的循环展开为循环体的固定数量的副本,这允许编译器应用严格的指令级并行性。通常,循环展开与循环拆分结合使用,循环拆分首先将循环拆分为两个嵌套循环,然后完全展开内部循环。

并行化:由于现代处理器通常支持多线程和SIMD并行,编译器后端需要利用并行性来最大限度地提高硬件利用率以获得高性能。

你可能感兴趣的:(人工智能)