DL compiler Survey
2023-09-15 14:08:29 0 举报
AI智能生成
对于论文The Deep Learning Compiler: A Comprehensive Survey的一个整理
作者其他创作
大纲/内容
What is AI Compiler
2个场景
推理场景
输入AI框架训练出来的模型文件,输出能够在不同硬件高效执行的程序;
训练场景
输入高级语言表示的神经网络代码,输出能够在不同硬件高效执行的程序
AI编译器是以python为主的动态解释器语言作为前端;会有多层IR的设计,包括图编译、算子编译、代码生成(这里分别对应它们自己的IR,例如图IR、算子IR、CodeGen的那套IR);主要是面向神经网络和深度学习的特定优化(domain specific 面向特定领域的);支持很多不同的芯片。
Why need AI compiler
Deep learning 被广泛应用,包括有CV、NLP、应用于HPC或者科学计算(例如遥感处理、蛋白质折叠的处理和预测、流体动力学和电磁仿真、还可以求解韦伯斯特方程)等。
所以现在AI的算法越来越多,AI的框架(如升思、tensorflow、oneflow飞浆、pytorch等)也越来越多,而且各种硬件也都在出现,并且每家都有自己的编译体系,每家都会有自己的AI编译器和AI编程体系。这就会带来两个challenge:
所以现在AI的算法越来越多,AI的框架(如升思、tensorflow、oneflow飞浆、pytorch等)也越来越多,而且各种硬件也都在出现,并且每家都有自己的编译体系,每家都会有自己的AI编译器和AI编程体系。这就会带来两个challenge:
越来越多的新算子被提出,算子库的开发、维护、优化和测试工作指数式上升
专用加速芯片爆发导致性能可移植性成为刚需。
The difference between traditional compiler and AI compiler
Likewise
目标相同:通过自动化方式进行程序优化和代码优化,从而降低对不同硬件的手工优化;
优化方式类似:在编译优化层通过统一IR执行不同的Pass进行优化,从而提高执行性能;
软件结构栈类似:分成前端、优化、后端三段式,IR解耦前端和后端使得模块化表示;
AI编译器依赖传统编译器:AI编译器对Graph IR进行优化后,将优化后的IR转换成传统编译器IR,最后依赖传统编译器生成机器码。(AI编译器是作为传统编译器的一种补充)
AI编译器依赖传统编译器:AI编译器对Graph IR进行优化后,将优化后的IR转换成传统编译器IR,最后依赖传统编译器生成机器码。(AI编译器是作为传统编译器的一种补充)
Differencies
输入不同
传统编译的主要目的是降低编程难度(面向开发者提出一个高级语言,编译器来处理低级语言),其次是优化程序性能;AI编译器的主要目的是优化程序性能,其次是降低编程难度。
IR差异:AI编译器的IR和传统编译器的IR所抽象出来的概念和意义并不相同。
AI编译器一般会有high level IR,用来抽象描述深度学习模型中的运算,如:convolution、Matmul等,甚至部分会有transformer带有图的结构。
传统编译器相对而言是low-level IR,用于描述基本指令运算,如load、store等。有了high-level IR,AI编译器在描述深度学习模型类DSL更加方便。
AI编译器一般会有high level IR,用来抽象描述深度学习模型中的运算,如:convolution、Matmul等,甚至部分会有transformer带有图的结构。
传统编译器相对而言是low-level IR,用于描述基本指令运算,如load、store等。有了high-level IR,AI编译器在描述深度学习模型类DSL更加方便。
优化策略不同:AI编译面向AI领域,优化时引入更多领域特定知识,从而处理更high-level,更加aggressive的优化手段。如:
(1)AI编译器在high-level IR执行算子融合,传统编译器执行类似loop fusion的时候,往往更加保守。缺点是可能会导致调试执行信息跟踪难。
(2)AI编译器可以降低计算精度,比如int8、fp16、bp16等,因为深度学习对计算精度不那么敏感。但传统编译器一般不执行改变变量类型和精度等优化。
(1)AI编译器在high-level IR执行算子融合,传统编译器执行类似loop fusion的时候,往往更加保守。缺点是可能会导致调试执行信息跟踪难。
(2)AI编译器可以降低计算精度,比如int8、fp16、bp16等,因为深度学习对计算精度不那么敏感。但传统编译器一般不执行改变变量类型和精度等优化。
Background
deep learning framework
文中对常见的深度学习框架进行了梳理,例如 TensorFlow、Keras、PyTorch 等等,以及许多已经消失于历史长河之中的框架。其中值得注意的是 ONNX (Open Neural Network Exchange),其定义了一个可伸缩的计算图模型,所以所有支持 ONNX 的框架之间模型可以十分轻松的转换,因此可以用 MXNet 构建模型然后用 PyTorch 运行。
deep learning hardware
通用硬件 General- purpose hardware:为通用领域设计的硬件。代表就是 GPU,其并非设计用于深度学习,但是其高度并行化和众核设计十分适合该场景
专用硬件 Dedicated hardware:专为深度学习计算定制的硬件。最知名的就是 Google 的 TPU,其中包含了专门用于矩阵乘法、激活函数等等的单元,从而加速了深度学习相关的计算。
神经形态硬件 Neuromorphic hardware:模拟生物大脑而设计的芯片。文中提到了来自 IBM 的 TrueNorth 和 Intel 的 Loihi。
hardware- specific DL code generator
FPGA、HLS
The processor architecture
The streaming architecture
Common design architeture of DL compiler
High-level IR
High-level IR 也称为 Graph IR,用于表示计算图,其出现主要是为了解决传统编译器中难以表达深度学习模型中的复杂运算这一问题,为了实现更高效的优化所以新设计了一套 IR。
DAG-based IR:是一种最为传统的方式,用点和边组织成一个有向无环图(Directed Acyclic Graph, DAG),从而充分利用现有的极其充分的 DAG 优化算法,而且实现简单,但是缺少计算域的定义会导致语义的二义性。
Let-binding-based IR:Let-binding 其实类似于变量定义,用于定义一个具有作用域的表达式。这样计算时就不必像 DAG 一样递归的计算一个个点,而是可以将每个 let 定义的结果建立一个映射表,之后查表即可。
此外不同的 IR 还有不同的方式表示 tensor 之间的运算,主要分为如下三类:
Function-based:用函数来封装操作符,XLA 的 IR 采用此种形式。
Lambda expression:通过变量绑定和替换来描述计算(?),TVM 采用这种形式。
Einstein notation:也称为爱因斯坦求和约定(summation convention),这种表达方式下无需定义中间变量,IR 可以通过未定义变量的出现找出真正的表达式,TC 采用这种形式。
在实际实现过程中还会遇到的一个问题是数据的表示,如下是常用的技术
占位符:占位符指用一个只有形状信息的变量来代表一个 tensor,直到运算时再用数据进行填充。这样再构建计算图时就不必关心真正的数据元素,从而将计算图定义和执行分离,方便调整
未知形状表示:这一技术常用于定义占位符,因为有些层的形状直到运行时才能知道形状,这也导致了对于边界和维数的检查需要放宽,也需要额外的机制来保证内存有效。
数据布局:数据布局描述了 tensor 如何在内存中组织,通常是从逻辑下标到内存下标的映射。值得注意的是数据布局信息可以和运算符放在一起,而非直觉上和 tensor 自己存放在一起,这样可以更好的实现特定的运算符,减少编译的开销。
边界推断:用于编译深度学习模型时决定迭代器的边界,主要是决定占位符的大小。
Low-level IR
Low-level IR 能够在更细粒度的层面上表示模型,从而能够针对于硬件进行优化,文中将其分为了三类
Halide-based IR:Halide 基本思想为将计算和调度分离。相较于直接给定一个特定的模式,Halide 会尝试各种可能的调度方式并选择最好的一个。
Polyhedral-based IR:Polyhedral 模型使用 linear programming, affine transformations 以及其他的数学方法来优化具有静态控制流边界和分支的以循环为基础的代码。
其他 IR:也有其他的编译器采用了除了以上两种方式之外的 IR,然后将硬件优化交给 LLVM IR 等设施,MLIR 就是一个例子。
Code generation:以上也反映了最后代码生成部分的不同思路,尽管多数编译器会最后转换成 LLVM IR 从而利用其成熟的优化器和代码生成器,但是将深度学习的代码直接交给 LLVM 之类的传统编译器很有可能会生成质量恶劣的代码,所以通常为了能够针对硬件更好的优化会 1) 在 LLVM 上层先进行优化(Halide-based IR 和 polyhedral-based IR)或者 2) 提供更多额外信息给优化 pass。
编译时也存在两种方式:JIT(just-in-time)和 AOT(ahead-of-time)。JIT 在运行时生成代码,能够根据运行情况优化代码,而 AOT 则是先生成二进制代码再运行,因此可以在编译时进行更大范围的静态分析,此外还可以交叉编译用于嵌入式平台。
The frontend
前端优化发生于计算图构建完成后,输入后端前,主要目标是优化计算图,因此是硬件无关的。优化通常是以一个称为 passes 的东西定义,通过遍历图和转化图结构来进行优化,文中将优化分为三个级别
Node-level optimizations:结点级别的优化指的是在结点内部进行优化,包括删除结点或者替换成一个更低开销的结点,例如清除 Nop 指令(Nop Elimination)和 0 维度的 tensor (Zero-dim-tensor elimination)。
Block-level optimizations:块级别的优化可以优化结点之间的关系。Algebraic simplification 可以通过调整运算顺序、结点间结合关系等等方式替换一个高开销的运算符或者是将常量间操作替换成结果(常量折叠)等等。Operator fusion 可以融合运算符,比如减少内存分配次数,结合循环,减少同步的开销等等。Operator sinking 能够将一些操作放到相近的位置,创造更多的机会进行优化。
Dataflow-level optimizations:数据流级别的优化,文中提到了四种。Common sub-expression elimination 也就是如果一个值常被计算得到但不会变化,那么后续不再计算。Dead code elimination 删除了无法抵达的代码(通常是由优化导致的)。Static memory planning 尽可能的重用内存。Layout transformation 生成更好的数据布局,不过这部分可能会依赖于硬件。
文中还针对于 TensorFlow XLA 做了一个案例介绍
The backend
后端优化首先就是针对于硬件的优化,从而生成更加高效的代码。一种方法是利用 LLVM,另一种方法是用深度学习领域定制的优化。
Hardware intrinsic mapping:将一段 low-level IR 转化成硬件上已高度优化的指令。
Memory allocation and fetching:GPU 上有专有的内存和共享的内存,两者容量和延迟均不同,因此存在不同的调度策略。
Memory latency hiding:利用流水线,通过重排指令将访存的指令尽可能靠近从而一次取出更多数据,访存同时执行之前的代码,减少内存延迟的影响。
Loop oriented optimizations:文中提到了 loop fusion(融合相同边界的循环), sliding windows(循环计算的变量仅在使用时计算,但会导致两个循环串行化), tiling(拆分循环为几部分,更好的利用局部性), loop reordering(修改循环内部代码顺序), and loop unrolling(展开循环再采用更多的优化方法)等方法优化循环。
Parallelization:将代码并行化以充分利用硬件性能。
在优化过程中许多场景下有许多参数可以调整,因此后端优化的另一方向就是自动调优(Auto-tuning),分为如下几个部分
Parameterization:将数据中特征和优化选项作为参数提取出来。
Cost model:用于评估模型的性能,分为三类 1)黑盒模型,只考虑运行时间,简单但是缺少指导优化的信息,2)ML为基础的模型,用机器学习的方法预测性能,能够在找到新配置时更新模型提升准确率,3)预定义的模型,根据编译任务的特征建立一个可以评估任务性能的模型,相比 ML 模型计算花费小,但是换到新任务上就需要重建模型。
Searching technique:搜索最佳的参数。首先需要对其初始化并决定搜索空间,之后可以采用基因算法、模拟退火算法、强化学习等方法进行搜索。
Acceleration:加速调优,主要分为两个方向,1)并行化和 2)配置重用,用之前的调优配置加速搜索。
Future directions
Dynamic shape and pre/post processing
强化学习在运行过程中,其input shape或者是模型本身都是在变化的,所以动态模型就变的很popular,特别是在NLP中,其输入会是不同的shape,所以对于DL 编译也很有挑战,因为一直到runtime之前,数据的shape都是未知的。
此外,随着未来的深度学习模型变得更加复杂,其整个contral flow可能不可避免地包括复杂的前/后处理过程。
目前,大多数深度学习编译器都使用Python作为编程语言,当Python解释器执行时,预处理/后处理可能成为性能瓶颈。现有的深度学习编译器尚未考虑到这种潜在的性能瓶颈。
在深度学习编译器中支持整个control flow,可以表达和优化深度学习模型的预处理/后处理,这为模型部署的性能加速提供了新的机会。
Advanced auto-tuning
现有的自动调整技术侧重于单个算子的优化。然而,局部最优的组合并不会导致全局最优。例如,应用于不同数据布局的两个相邻运算符可以一起调整,而无需在其间引入额外的内存转换
随着边缘计算的升起,执行时间不仅仅是DL compiler的优化目标,自动调优中还应考虑新的优化目标,例如内存占用和能耗
基于机器学习的自动调整有几个值得探索的方向
除了cost model之外,机器学习技术还可以应用于自动调整的其他阶段。例如,在选择compiler option和optimization schedule的阶段,可以使用机器学习技术直接预测可能性并开发算法来确定最终配置
基于机器学习的自动调优技术可以根据领域知识进行改进。例如,将特征工程(选择特征来表示程序)结合到自动调优技术中可能是实现更好调优结果的潜在方向。
Polyhedral model
将多面体模型和自动调优技术结合起来设计深度学习编译器以提高效率是一个有前途的研究方向。
一方面,可以应用自动调优,通过重用先前配置的多面体 JIT 编译来最小化多面体 JIT 编译的开销
另一方面,多面体模型可以用来进行自动调度,可以减少自动调整的搜索空间
在深度学习编译器中应用多面体模型的另一个挑战是支持稀疏张量。
一般来说,稀疏张量(例如 CSF [84])的格式用不再是线性的索引数组(例如,a[b[i]])表示循环索引。这种间接索引寻址会导致非仿射下标表达式和循环边界,从而阻碍多面体模型的循环优化
幸运的是,多面体社区在支持稀疏张量方面取得了进展 [94, 95],并且集成多面体模型的最新进展可以增加 DL 编译器的性能机会。
Subgraph partitioning
支持子图划分的DL编译器可以将计算图划分为多个子图,并且子图可以以不同的方式进行处理
子图划分为深度学习编译器提供了更多研究机会
首先,它提供了集成图形库以进行优化的可能性。以 nGraph 和 DNNL 为例,DNNL 是一个深度学习库,利用大量高度优化的内核进行图形优化。 DNNL 与 nGraph 的集成使 DNNL 能够加速 nGraph 生成的子图的执行速度。
其次,它开辟了异构和并行执行的可能性。一旦计算图被划分为子图,不同子图的执行就可以同时分配给异构硬件目标。以边缘设备为例,其计算单元可能由 ARM CPU、Mail GPU、DSP 以及可能的 NPU 组成。从有效利用所有计算单元的深度学习编译器生成子图可以显着加快深度学习任务的速度。
Quantization
深度学习框架中应用的传统量化策略基于一组固定方案和数据类型,几乎没有针对在不同硬件上运行的代码进行定制。
然而,支持深度学习编译器中的量化可以在编译期间利用优化机会来得出更有效的量化策略。例如,Relay[78]提供了量化重写流程,可以自动生成各种方案的量化代码。
为了支持量化,深度学习编译器需要解决几个挑战。
第一个挑战是如何在不进行大量工程工作的情况下实现新的量化算子。 AWS的尝试指出了一个可能的方向,即利用方言(dialect)的概念,在基本算子的基础上实现新的算子,从而使图级别和算子级别的优化得到复用。
第二个挑战是编译期间量化和其他优化之间的相互作用。例如,确定量化的适当阶段以及与算子融合等优化进行协作需要未来的研究调查。
Unified optimization
尽管现有的深度学习编译器在计算图优化和硬件优化方面都采用了类似的设计,但每个编译器在某些方面都有自己的优势。缺少共享最先进优化的方法,以及跨现有编译器对新兴硬件目标的支持。
我们主张统一现有深度学习编译器的优化,以便可以重用每个深度学习编译器中采用的最佳实践。此外,统一深度学习编译器的优化可以积累强大的力量来影响通用和专用深度学习加速器的设计,并为深度学习编译器和硬件的高效协同设计提供环境。
Differentiable programing
可微分编程(不懂)
Privacy Protection
关于安全的,也不会去考虑
Training Support
现在很多都是关于推理的DL compiler,很少针对training的,所以可以把视野转向training做对应的研究,这个也可以试试看
总而言之,当前的深度学习编译器主要致力于缩小将深度学习模型有效部署到不同硬件上的差距,因此他们选择推理作为主要优化目标。然而,扩展深度学习编译器支持模型训练的能力将带来大量的研究机会,例如梯度算子的优化和高阶自动微分。

收藏
0 条评论
下一页