传统编译器与AI编译器对比

传统编译器与AI编译器的对比,包括AI编译器前后端的详细优化流程,内容来自B站ZOMI大佬。

基本概念

编译器与解析器

Interpreter Compiler
程序步骤 1. 创建代码
2. 没有文件链接或机器代码生成
3. 源语句在执行过程中逐步执行
1. 创建代码
2. 解析或分析所有语言语句的正确性
3. 将源码转换为机器码
4. 链接到可执行程序
5. 运行程序
Input 每次读取一行 整个程序
Output 不产生任何中间代码 生成中间目标代码
工作机制 编译和执行同时进行 编译在执行之前完成
存储 不保存任何机器代码 在机器上存储编译后的机器代码
执行 程序执行是解释过程的一部分,逐行执行 程序执行与编译是分开的,在整个输出程序编译后执行
生成程序 不生成输出程序,每次执行过程都要评估源程序 生成可以独立于源程序运行的输出程序
修改 可以直接修改 修改后需要重新编译
运行速度 较慢 较快
内存 所需内存较少,因为不创建中间代码 由于要创建目标代码,所需内存更多
错误 解释器读取一条语句并显示错误,纠正后才能解释下一行 编译器在编译时显示所有错误和警告,不修正就不能运行程序
错误监测 较为容易 较难
编程语言 PHP、Perl、Python、Ruby C、C++、C#、Scala、Java

JIT和AOT编译方式

程序主要有两种运行方式:

  • 静态编译:程序在执行前全部被编译为机器码,称为AOT(Ahead of time),也就是提前编译;
  • 动态解释:程序一边编译一边运行,称为JIT(Just in time),也就是即时编译;
JIT AOT
优点 1. 可以根据当前硬件情况实时编译生成最优机器指令
2. 可以根据当前程序的运行情况生成最优的机器指令序列
3. 当程序需要支持动态链接时,只能使用JIT
4. 可以根据进程中内存的实际情况调整代码,能够更充分的利用内存
1. 程序运行前编译,可以避免运行时的编译性能消耗和内存消耗
2. 可以在程序运行初期就达到最高性能
3. 可以显著加快程序的启动
缺点 1. 编译需要占用运行时资源
2. 编译占用运行时间,对某些代码编译优化不能完全支持,需要在流畅和时间之间权衡
3. 在编译准备和识别频繁使用的方法需要占用时间,初始编译不能达到最高性能
1. 程序运行前编译会使程序安装时间增加
2. 牺牲高级语言的一致性问题
3. 将提前编译的内容保存会占用更多的外存

一般在训练过程中使用JIT模式,离线推理过程中使用AOT模式。

Pass和中间表达IR

Pass是对源程序的一次完整扫描或处理,一次编译过程可能有多个Pass。

IR(Intermediate representation)是编译器或虚拟机内部用来表示源代码的数据结构或代码。

传统编译器

编译器基本构成

  • 前端(Front-End):负责词法和语法分析,将源代码转化为抽象语法树;
  • 优化器(Optimizer):对前端得到的中间代码进行优化,使代码更加高效;
  • 后端(Back-End):将优化过的中间代码转化为针对各自平台的机器代码。

GCC编译过程和原理

GCC是一个可移植的编译器,支持多种硬件平台。它不仅仅是本地编译器,还能跨平台交叉编译。有多种语言前端,用于解析不同的语言。采用模块化设计,可以加入新的语言和新CPU架构支持。

编译过程

  • 预处理

包括宏定义、文件包含、条件编译三个部分。预处理过程读入源代码,检查包含预处理指令的语句和宏定义,并对其进行响应和替换。预处理过程会删除进程中的注释和多余空白字符。最后生成.i文件。

1
gcc -E hello.c -o hello.i
  • 编译

编译器对.i文件进行语法分析并优化,然后生成对应的汇编代码.s文件。

1
gcc -S hello.i -o hello.s
  • 汇编

汇编器将编译器生成的.s汇编程序汇编为机器语言或指令,也就是机器可以执行的二进制程序.o文件。

1
gcc -c hello.s -o hello.o
  • 链接

连接器会链接程序运行所需要的目标文件,以及依赖的库文件,最后身成可执行文件,以二进制形式存在磁盘中。

1
gcc hello.o -o hello

优缺点

优点:

  • 支持Java、ADA、Fortran等多种语言;
  • 支持更多平台;
  • 更流行,使用范围更广,支持较为完备;
  • GCC基于C,不需要C++编译器即可编译。

缺点:

  • GCC代码耦合度高,很难独立,若集成到IDE中,难以用模块化的方式调用GCC;
  • GCC被构建成单一静态编译器,难以作为API集成到其他工具中;
  • 越后期的版本,代码质量越差,代码量约1500万行,过于庞大。

LLVM架构和原理

LLVM的一大贡献就是提出了中间表达,将所有语言前端都对接到IR,以IR这种统一的方式对接到各种硬件。

LLVM的一大理念就是Lib base,即基于库的实现。LLVM主要有以下库:

  • LLVM Core:LLVM的核心库,主要是围绕LLVM中间代码的一些工具,提供了一个“源”和“目标”无关的优化器,以及几乎所有主流CPU类型的机器码生成器;
  • Clang:LLVM的一个子项目,基于LLVM架构的轻量级编译器,负责编译C、C++、Objective-C语言,属于LLVM架构中的编译器前端;
  • Compiler-RT:为硬件不支持的低级功能提供特定于目标的支持;
  • LLDB:LLVM的原生调试器项目,提供丰富的流程控制和数据检测的调试功能;
  • LLD:Clang/LLVM内置的连接器;
  • Dragonegg:GCC插件,可以将GCC的优化和代码生成器替换为LLVM的相应工具;
  • libc:C标准库实现;
  • libcxx/libcxxabi:C++标准库实现;
  • libclc:OpenCL标准库实现;
  • OpenMP:提供OpenMP运行时,用于Clang中的OpenMP实现;
  • polly:支持高级别的循环和数据本地化优化支持的LLVM框架,使用多面体模型实现一组缓存局部优化以及自动并行和矢量化;
  • vmkit:基于LLVM的Java和.Net虚拟机实现;
  • klee:基于LLVM编译基础设施的符号化虚拟机,使用一个定理证明器来尝试评估程序中的所有动态路径,以发现错误并证明函数的属性。klee的一个主要特征是它可以在检测到错误时生成测试用例;
  • SAFECode:用于C/C++程序的内存安全编译器,通过运行时来检查代码,以便检测内存安全错误,可以用于保护软件免受安全攻击,也可用作Valgrind等内存安全错误调试工具。

编译过程

  • 预处理
1
clang -E -c hello.c -o hello.i
  • 编译

可以导出字节码文件.bc,即IR中间表示。

1
clang -emit-llvm hello.i -c -o hello.bc

也可以导出可以识别的.ll文件。

1
clang -emit-llvm hello.i -S -o hello.ll
  • 优化Pass并生成汇编代码
1
llc hello.ll -o hello.s
  • 汇编并生成目标文件
1
clang -c hello.s -o hello.o
  • 链接
1
clang hello.o -o hello

IR详解

IR并不指一种固定的形态,LLVM架构下,编译器前端会将AST形式的IR传给优化器,优化器的每一个Pass之间都会有相同或者不同形式的IR,最后以DAG形式的IR传递给后端,后端可能再转变IR的形式进行优化。

LLVM IR作为一种编译器IR,有两个基本原则指导核心库的开发:

  • SSA表示,代码组织为三地址指令序列和无限寄存器,使优化能够快速进行;
  • 整个程序的IR存储到磁盘中,让链接时优化易于实现。

SSA(Static single assignment)即静态单赋值形式,有两个重要特征:

  • 代码组织为三地址指令序列;
  • 寄存器数量无限制,是为了和硬件解耦。

SSA程序的每个变量有且只有一个赋值语句,在LLVM IR中,每个变量都必须在使用前定义,且每个变量只能被赋值一次。使用SSA的好处是,每次使用一个值,可以立即向后追溯到给出其定义的唯一指令。还可以极大的简化优化,因为SSA形式建立了平凡的use-define链,也就是一个值到达使用处的定义的列表。

LLVM IR表示形式

LLVM IR有三种具体表示形式,三种中间格式是完全等价的:

  • 在内存中的编译中间语言,例如无法通过文件的形式得到的指令类等;
  • 在硬盘上存储的二进制中间语言,格式为.bc
  • 人类可读的代码语言,格式为.ll
LLVM IR内存模型
  • LLVM IR文件的基本单位称为module;
  • 一个module可以拥有多个顶层实体,如function和global variable;
  • 一个function define中至少有一个basicblock;
  • 每个basicblock中有若干instruction,并都以terminator instruction结尾。

LLVM前端

编译器前端将源代码转化为编译器的中间表示LLVM IR。

  • 词法分析:处理源代码的文本输入,将语言结构分解为一组单词和标记,去除注释、空白、制表符等。每个单词或标记必须属于语言子集,语言的保留字被变换为编译器内部表示;
  • 语法分析:分组标记以形成表达式、语句、函数体等。检查一组标记是否有意义,考虑代码物理布局,但并不分析代码含义,只保证语法正确性,并输出抽象语法树AST;
  • 语义分析:借助符号表检验代码是否违背了语言类型系统。符号表存储标识符和其各自的类型之间的映射,以及其它内容。类型检查的一种直觉的方法是,在解析之后遍历AST的同时从符号表中收集关于类型的信息;

经过上述步骤后,生成LLVM IR,交给优化层进行优化。

LLVM优化层

LLVM优化层的输入是LLVM前端生成的IR,具体表现为AST,输出也是IR,具体表现为DAG。

优化层通过一个个的Pass对IR进行优化,Pass分为两类:

  • 分析Pass:负责发掘性能和优化的机会;
  • 转换Pass:生成必需的数据结构,为后续的Pass所使用。

在转换Pass和分析Pass之间有两种主要依赖类型:

  • 显式依赖:转换Pass需要一种分析,则Pass管理器自动地安排它所依赖的分析Pass在它之前运行;
  • 隐式依赖:转换或分析Pass要求IR代码运用特定的表达式,需要手动地以正确的顺序将这个Pass加入到Pass队列中,可以通过命令行工具(clangopt)或Pass管理器实现。

Pass类是实现优化的主要资源,但并不会直接使用Pass,而是通过清楚的子类来使用它。当实现一个Pass时应该选择合适的Pass粒度,并使用该粒度下的子类。常见的子类有ModulePassFunctionPassBasicblockPass等。

LLVM后端

整个后端流水线用到了四种不同层次的指令表示:

  • 内存中的LLVM IR;
  • SelectionDAG结点;
  • MachineInstr;
  • MCInst。

上图中白色的Pass是非必要Pass,灰色的Pass是必须的Pass,也叫Super Pass。

  • Instruction selection 指令选择

该Pass将内存中LLVM IR变换为目标特定的Selection DAG结点,每个DAG能够表示单一基本块的计算,结点表示指令,而边编码了指令间的数据流依赖。该Pass使得LLVM代码生成程序库能够运用基于树的模式匹配指令选择算法。

  • Instruction Scheduling (I) 指令调度

这是第一次指令调度,也称为前寄存器分配(RA)调度。该Pass对指令进行排序,同时尝试发现尽可能多的指令层次的并行。然后指令被变换为MachineInstr三地址表示。

  • Register Allocation 寄存器分配

LLVM IR两个重要特性之一是寄存器集无限,这个性质会一直保持,直到寄存器分配Pass。该Pass将无限的虚拟寄存器引用转换为有限的目标特定的寄存器集合。当寄存器不够时会挤出到内存。

  • Instruction Scheduling (II) 指令调度

这是第二次指令调度,也称为后寄存器分配(RA)调度。此时可以获得真实的寄存器信息,某些类型的寄存器存在延迟,可以用来改进指令顺序。

  • Code Emission 代码输出

该Pass将指令从MachineInstr表示变换为MCInst实例。这种新的表示更适合汇编器和链接器,可以输出汇编代码或者输出二进制块特定目标代码格式。

AI编译器

AI编译器与传统编译器的关系如下:

  • 目标相同:通过自动化方式进行程序优化和代码生成,从而减少对不同硬件手工优化的工作量;
  • 优化方式类似:在编译优化层通过统一IR执行不同的Pass进行优化,从而提高执行性能;
  • 软件结构栈类似:分为前端、优化、后端三段式,IR解耦前端和后端使其能够模块化表示;
  • AI编译器依赖于传统编译器:AI编译器对Graph IR优化之后,将优化后的IR转换成传统编译器IR,最后依赖传统编译器进行机器码生成。

主要的区别如下:

  • IR差异:AI编译器的IR与传统编译器IR所抽象出来的概念和意义不同;
    • AI编译器一般会有High-level IR,用来抽象描述深度学习模型中的运算,如卷积、Matmul等;
    • 传统编译器相对而言是Low-level IR,用于描述基本指令运算,如load、store等。
  • 优化策略:AI编译器面向AI领域,优化时会引入更多领域特定知识,从而进行更高级别、更激进的优化手段;
    • AI编译器在High-level执行算子融合,传统编译器执行类似操作时会更保守;
    • AI编译器可以降低计算精度,因为深度学习对精度不那么敏感,但传统编译器一般不执行改变类型和精度等优化。

AI编译器主要分为两个场景:

  • 推理场景:输入AI框架训练出来的模型文件,输出能够在不同硬件高效执行的程序;
  • 训练场景:输入高级语言表示的神经网络代码,输出能够在不同硬件高效执行的程序。

AI编译器架构发展

朴素的AI编译器

TensorFlow早期版本,基于神经网络的编程模型,主要进行了graph图和ops算子两层抽象。

  • 图层:通过声明式编程方式,以静态图方式执行,执行前进行硬件无关和硬件相关的编译优化;
    • 硬件无关的优化:表达式简化、常量折叠、自动微分等;
    • 硬件相关的优化:算子融合、内存分配等。
  • 算子层:采用手写kernel的方式,如在NVIDIA GPU上基于CUDA kernel实现大量算子或依赖cuDNN算子优化库。

这一阶段的编译器在表达上采用静态图方式,静态图的表达式非Python原生,开发者需要使用框架的API进行构图,易用性不好。

在性能上,DSA专用加速芯片的出现加剧了性能挑战;算子层提供的算子粒度和边界提前确定后,无法充分发挥硬件的性能;硬件厂商提供的算子优化库也未必最优。

专用的AI编译器

这个阶段的编译器,在表达上以PyTorch为代表的框架灵活表达API方式称为AI框架的参考标杆,图层的神经网络编译器主要考虑如何将类PyTorch的表达转换到图层的IR进行优化,并试图打开图算边界进行融合优化。

在性能上打开计算图和算子的边界,进行重新组合优化,发挥芯片算力。计算图层下发子图中的算子并打开成小算子,基于小算子组成的子图进行编译优化。

通用的AI编译器

是一个还未到来的阶段,在这一阶段希望解决的问题:

  • 图算统一表达,实现融合优化;
  • 算子实现自动Schedule、Tiling、Codegen,降低开发门槛;
  • 更泛化优化能力,实现动静统一、动态Shape、稀疏性、高阶微分、自动并行等;
  • 包括编译器、运行时、异构计算、边缘到数据中心都模块化表示和组合,并专注于可用性。

AI编译器通用架构

IR中间表达

编译器主要分为前端和后端,分别针对硬件无关和硬件相关的处理。每个部分都有自己的IR,每个部分也会进行优化:

  • High-level IR:用于表示计算图,其出现主要是为了解决传统编译器中难以表达深度学习模型中的复杂运算这一问题,为了实现更高效的优化,所以设计了一套IR;
  • Low-level IR:能够在更细粒度的层面上表示模型,从而能够针对于硬件进行优化,主要分为三类。

前端优化

构造计算图后,前端将应用图级优化。因为图提供了计算全局概述,所以更容易在图级发现和执行许多优化。前端优化与硬件无关,可以将计算图优化应用于各种后端目标。

前端优化分为三类:

  • 结点级优化,如Zero-dim-tensor elimination、Nop Elimination;
  • 块级优化,如代数简化、常量折叠、算子融合;
  • 数据流级优化,如Common sub-expression elimination、DCE。

后端优化

  • 特定硬件的优化

目标针对特定硬件体系结构获取高性能代码。低级IR转换为LLVM IR,利用LLVM基础结构生成优化的CPU、GPU代码。利用领域知识定制优化,可以更有效地利用目标硬件。

  • 自动调整

由于在特定硬件优化中用于参数调整的搜索空间巨大,因此有必要利用自动调整来确定最佳参数设置。Halide/TVM允许调度和计算表达分开,使用自动调节得到较好配置。应用多面体模型进行参数调整。

Halide:一种基于C++的DSL(Domain Specified Language),用于算法的底层加速,优化与算法本身设计无关。

TVM:端到端的机器学习编译框架,可以优化机器学习模型,使其在不同硬件平台上高效运行。

  • 优化内核库

厂商特定优化内核库,广泛用于各种硬件上的加速深度学习模型训练和推理。特定优化原语可以满足计算要求时,使用优化的内核库可以显著提高性能,否则可能受到进一步优化约束。

前端优化详解

前端优化、图层优化、计算图优化都是相同的概念(Graph Optimizer),指的是同一个过程。

图层IR

计算图用来表示深度学习网络模型在训练和推理过程中的计算逻辑与状态,由基本数据结构张量和基本运算单元算子构成。计算图中的结点表示算子,有向边表示张量状态,也描述了计算间的依赖关系。

计算图生成

生成计算图的两种方式:

  • 静态计算图:使用前端语言定义模型形成完整的程序表达后,不使用前端语言的解释器进行执行,将前端描述的完整模型交给AI框架。AI框架先对API代码分析,获取网络层之间的连接拓扑关系以及参数变量设置、损失函数等,然后使用静态数据结构来描述拓扑结构及其他神经网络模型组件。
  • 动态计算图:使用前端语言自身的解释器对代码进行解析,利用AI框架本身的算子分发功能,算子会即刻执行并输出结果。动态图模式采用用户友好的命令式编程范式,使用前端语言构建神经网络模型更加简洁。

动态图的优点是执行计算比较灵活,动态生成可以使用前端语言的原生控制流,充分发挥前端语言的编程友好特性。

但动态生成中完整的网络结构在执行前是未知的,不能使用静态图中的图优化技术来提高计算执行性能。

动态图转换为静态图
  • 基于追踪转换:以动态图模式执行并记录调度的算子,构建和保存为静态图模型;
  • 基于源码转换:分析前端代码,将动态图代码自动转为静态图代码,底层使用静态图执行。
图优化

计算图可以使AI框架在执行前看到深度学习模型定义的全局信息,可以将计算图作为AI框架中的High-level IR,通过计算图优化Pass来简化计算图或提高执行效率。

计算图对AI编译器的作用
  • 方便底层编译优化
  • 分层优化,便于扩展

算子融合 (Operator Fusion)

这样可以减少一次Kernel启动的开销,同时减少一次数据的访问。

这种情况下,融合之前B、C是并发执行的,但都需要等待A执行完成。融合之后可以在两个分支上并行执行。

这样横向融合可以减少一次Kernel启动的开销,B、C的计算结果都依赖A的结果放入内存,缓存效率会更高。

这种融合会将A的输出数据和A经过B的输出数据放入内存,从而提升内存使用效率。

算子融合目的

算子融合的目的在于解决两堵墙:

  • 内存墙:由访存瓶颈引起。
  • 并行墙:主要由于芯片核心增加与单算子多核并行度不匹配引起。

算子融合通过对计算图上存在数据依赖的“生产者-消费者”算子进行融合,从而提升中间数据的访存局部性,以此解决内存墙问题。这种融合技术统称为Buffer融合,早期AI框架主要通过手工方式实现固定模式的Buffer融合。

可以将计算图中的算子结点进行并行编排,提升整体计算并行度。特别是对于网络中存在可并行的分支结点,这种方式可以获得较好的并行加速效果。

TVM支配树

TVM整体的算子融合是基于支配树实现的,支配树 (Dominator Tree, DOM)是由各个点的支配点构成的树,支配点则是所有能够到达当前结点的路径的公共祖先结点 (Least Common Ancestors, LCA)。

支配树会检查每个结点到其支配点的结点是否符合融合条件,融合的基本规则是融合掉的结点不会对剩余的结点产生影响。

支配树的生成过程为,根据有向无环图(DAG)生成深度由先搜索树(DFS),根据DFS树及对应的边生成支配树(DOM),然后加入一个叫做Group的数据结构来描述多个结点是否可以被融合。

TVM算子融合流程
  • 通过AST转换为Relay IR,遍历Relay IR,Relay IR是TVM中的一种中间表示;
  • 建立DAG用于支配树分析;
  • 应用算子融合算法。
TVM算子融合规则
  • Injective (one-to-one map):映射函数,如Add、Pointwise;
  • Reduction:归约函数,输入到输出有降维性质,如sum、max;
  • Complex-out-fusable:计算复杂类型,如conv2d;
  • Opaque (cannot be fused):无法被融合的算子,如sort。

在编译器中,一般融合规则都是通过Pass来承载的,不同的Pass处理不同的融合规则,而融合规则主要是人工定义的。

数据布局转换 (Layout Transform)

张量数据布局
  • NCHW:同一通道的数据连续排布,更适合需要对每个通道单独运算的操作;
  • NHWC:不同通道的同一位置元素连续排布,更适合需要对不同通道的同一数据做运算的操作。
布局转换优化

布局转换的目的是将内部数据转换为后端设备友好的形式。这种转换的会试图找到在计算图中存储张量的最佳数据布局,然后将布局转换结点插入到计算图中。张量数据布局对最终性能有很大影响,而且转换操作本身也有很大的开销。

具体来讲,编译器会判断输入的张量格式和算子要求的输入格式,如果格式不同,则会在算子之前插入数据转换结点,转换为算子要求的输入格式。输出时也是同样的,在需要时会在算子后插入一个数据转换结点。

训练场景下可能会有三种操作:

  • 不改变原计算图;
  • 插入转换算子;
  • 取消转换算子;

推理场景可能还会有额外的两种操作:

  • 权重布局转换;
  • 算子替换。

内存分配优化 (Memory Allocation)

节省内存算法主要有下列四种:

  • 空间换内存:如卸载到CPU(CPU Offload);
  • 计算换内存:重计算(Gradient Checkpointing);
  • 模型压缩:如量化训练,剪枝等;
  • 内存复用:利用AI编译器对计算图中的数据流进行分析以复用内存。

下面重点介绍内存复用,内存复用主要有两个方式:

  • 替代操作:如果一块内存不再需要,且下一个操作是element-wise的,可以原地覆盖内存;
  • 内存共享:两个数据使用内存大小相同,且有一个数据参与计算后不再需要,后一个数据可以覆盖前一个数据;
内存优化方法

对计算图的数据流进行分析,初始化一个内存分配计划,然后按照数据流开始分配内存。遇到不可以内存复用的结点,则新分配一片内存,遇到可以复用的结点则复用,最终形成一个完整的内存分配计划。

如果计算图涉及并行操作,则需要在不同的并行分支上制定单独的内存分配计划。最基本的原则就是尽可能的允许更多并行。

常量折叠 (Constant Fold)

传统编译器的常量折叠

常量折叠实现的功能如下,假设有一条语句。

1
i = 320 * 200 * 32

不进行折叠的汇编代码如下。

1
2
3
4
5
%a: load 320
%b: load 200
%x: mul %a, %b
%c: load 32
%y: mul %x, %c

但可以在编译期直接将这个值算出来,从而减少访存和计算的次数。

1
%a: load 2048000
常量传播

常量传播是编译器优化技术之一,可以在一段代码中,将表达式中的常量替换为相关表达式或字面值,再使用常量折叠技术来简化代码。

1
2
3
4
def foo():
x = 14
y = 7 - x / 2
return y * (28 / x + 2)

经过常量传播后会变成如下形式。

1
2
3
4
def foo():
x = 14
y = 7 - 14 / 2
return y * (28 / 14 + 2)
AI编译器的常量折叠

将计算图中可以预先确定输出值的结点替换为常量,并对计算图进行一些结构简化的操作。

AI编译器中的常量折叠主要有三类:

  • 类似传统编译器中的常量折叠,找到输入结点均为常量的结点,提前计算出值来完整替换该结点;
  • 与数据形状相关的常量折叠,会通过计算图已有信息推断出形状,然后代替原结点;
  • 与已知常量的代数化简有关的常量折叠,会将一些固定的常量提前计算,加入化简式相关的算子中。

公共子表达式消除 (Common Subexpression Elimination)

传统编译器的CSE

CSE过程中,编译器会视情况将多个相同的表达式替换成一个变量,该变量存储着计算该表达式后得到的值。

1
2
a = b * c + g
d = b * c + e

经过CSE后会变成如下代码。

1
2
3
temp = b * c
a = temp + g
d = temp + e
CSE的基本原则

执行CSE的可能性基于表达式的定义可达性,当以下条件成立,则一个表达式b * c在程序的某个点p被定义为可达的:

  • 从初始结点到p的每条路径在到达p之前计算过b * c
  • b * c被计算后,无论bc到达p之前都没有被重新赋值过。

由编译器计算的成本效益分析可以判断出,重复计算该表达式的开销是否大于存储该表达式的计算结果,并且这个分析也要将寄存器等因素考虑在内。

CSE分为两种:

  • 本地公共子表达式消除:工作于基本块范围内;
  • 全局公共子表达式消除:工作于整个过程;
CSE算法

对于语句%s: z = x op y,若x op y在语句s之前可用,则:

  • s开始逆向搜索IR,找到距离s最近执行w = x op y的语句;
  • 建立临时变量%u
  • 将第一步找到的语句w = x op y进行替换:
    • %u = x op y
    • %w = u
  • 使用z = u替换s,重复步骤一直到遍历完整个IR。
AI编译器中的CSE

于传统编译器的区别在于,AI编译器中子表达式是基于计算图或图层IR的,通过在计算图中搜索相同结构的子图,简化计算图的结构,从而减少计算开销。

AI编译器CSE算法
  • 获取逆后续结点集Set(Reverse):对计算图Graph IR进行深度优先遍历,将结果逆转得到逆后续结点集。逆后续得到的结果就是拓扑排序,即访问到某一结点时,该结点的依赖结点都已经被访问;
  • 创建Map(MSE):存储公共子表达式候选集,遍历Set(Reverse)时,可以从候选集Map(MSE)中查找是否有可用的表达式;
  • 遍历计算图所有结点,判断是否有公共子表达式:获取结点hash值,hash的key由输出个数、输出类型、输入结点id组成,key可以保证输入相同且输出个数和类型相同时,得到的hash值相同,达到检索公共子表达式的目的;
  • 记录入候选集:根据hash值从Map(MSE)中得到候选集,若候选集为空,即第一次遇到这样的表达式,就将该结点存入Map(MSE)中。
  • 可复用公共子表达式判断:候选集非空,则说明之前遍历到了相似的表达式,进一步判断是否可以复用已保存结点表达式的结果。即结点的输入都是来自相同的Const结点,可以保证输入数据完全相同;输出个数和类型相同,可以保证输出的结果相同;满足条件下,可以复用之前保存结点的结果。
  • 删除重复结点:判断表达式可以复用后,将候选结点的输出连接到当前结点的输出结点对应的输入,最后删除当前结点。

死代码消除 (Dead Code Elimination)

传统编译器的DCE

DCE的用途是移除对程序执行结果没有任何影响的代码。移除这类代码由两个优点,可以减少程序的大小,还可以避免程序在执行中进行不相关的运算行为,减少执行时间。

死代码主要有以下表现形式。

1
2
3
4
5
6
7
def foo():
a = 24
b = 25 # 无意义的赋值
c = a << 2;
return c
b = 24; # 不会被执行
return 0
1
2
3
4
5
6
7
def foo():
a = 5;
b = 6;
c = a * (b >> 1)
if False: # 永远不会被执行
print(c)
return c
AI编译器的DCE

DCE可以优化计算图的计算和存储效率,避免为死结点分配内存和进行计算,同时简化了计算图的结构,方便进行后续的其他图优化Pass。

DCE一般不是在定义神经网络模型结构的时候引起的,而是其他图优化Pass造成的结果,因此DCE Pass通常在其他图优化Pass之后应用。

DCE算法
  • 获取逆后续结点集Set(Reverse):对计算图Graph IR进行深度优先遍历,将结果逆转得到逆后续结点集。逆后续得到的结果就是拓扑排序,即访问到某一结点时,该结点的依赖结点都已经被访问;
  • 遍历后续结点集,判断是否为死结点:获取结点的输入值,判断是否为计算图的输出结点,如果不是输出结点且无对应输出结点,则为死结点;
  • 删除死结点,重新遍历计算图:删除死结点的同时要删除该结点对应输入的边,然后重复步骤一。

代数化简 (Algebraic Reduced)

代数化简的目的是利用交换律、结合律等规律调整图中算子的执行顺序,或者删除不必要的算子,以提高图整体的计算效率。

代数化简主要有三种形式:

  • 算术化简;
  • 运行化简;
  • 广播化简。
算术化简

通过利用代数之间算术运算法则,在计算图中可以确定优化的运算符执行顺序,从而用新的运算符替换原有复杂的运算符组合。

  • 结合律化简

考虑下列计算的化简。 以计算图的角度来看就是下列优化。

  • 提取公因式、分配律化简

  • 交换律化简

运行化简
  • 对逆函数等于自身函数的对合算子化简,如取反、倒数、逻辑非、矩阵转置等:
  • 幂等算子化简,即作用在某一元素两次与一次相同:
广播化简

多个张量形状不同,需要进行广播将张量的形状拓展为相同形状再进行运算,化简为最小计算量所需的广播运算数量。

考虑下列矩阵与标量相加的操作。 经过这样的化简,广播操作从原来的两次减少到了一次。

后端优化详解

后端优化也是算子优化(Ops Optimizer)的过程。

前后端优化的区别在于:

  • 前端优化:输入计算图,关注计算图整体拓扑结构,不关心算子具体实现。前端优化对算子结点进行融合、消除、化简等操作,使计算图的计算和存储开销最小。
  • 后端优化:关注算子结点的内部具体实现,针对具体实现使得性能达到最优。重点关心结点的输入输出,内存循环方式和计算的逻辑。

后端的优化流程:

  • 生成低级IR:不同AI编译器内部低级IR的形式和定义不同,但对于同一算子,算法的原理实质相同。对于每个具体的算子,需要用AI编译器底层的接口来定义算法,再由编译器来生成内部的低级IR;
  • 后端优化:针对不同的硬件架构、微架构,不同的算法实现的方式有不同的性能,目的是找到算子的最优实现方式,达到最优性能。同一算子的不同形态都会有不同的循环优化方法;
  • 代码生成:对优化后的低级IR转化为机器指令执行,现阶段最广泛方式是借助成熟的编译工具来实现,并非AI编译器的核心内容。

对于传统编译器来说,深度学习中主要的数据结构为张量,传统编译器不擅长对张量计算的优化。同时传统编译器主要针对通用编程语言,缺少领域特定语言DSL的支持以及相关的特殊优化。所以要将大部分的优化工作交给专用的AI编译器。

算子的计算与调度

算子

深度学习算法由一个个计算单元组成,称这些计算单元为算子 (Operator, Op)。算子是一个函数空间到函数空间上的映射。从广义上讲,对任何函数进行某一项操作都可以认为是一个算子,对于AI框架而言,所开发的算子是网络模型中涉及到的计算函数。

算法:算法是指解题方案的准确而完整的描述,是一系列解决问题的清晰指令,算法代表着用系统的方法描述解决问题的策略机制。

计算与调度
  • 计算:是算子的定义;
  • 调度:是算子的执行策略和具体实现。

同一算子会有不同的实现方式,即不同的调度方式;但只会有一种计算形态,即一种具体的定义。

计算的实现(Function / Expression)和计算在硬件单元上的调度(Schedule)是分离的。

算子的调度在硬件层面上利用单指令多数据流(SIMD)、平铺(Tiling)、展开(Unrolling)和向量化(Vectorization)等技术进行极致优化,充分利用硬件的性能,但不改变算法本身的涉及,最大化提升程序执行的速度。

算子调度具体执行的所有可能的调度方式称为调度空间,AI编译器优化的目的就是为算子提供一种最优的调度,使得算子在硬件上运行时间最优。

调度树

算子计算有三个主要特征:

  • 多重循环;
  • 没有复杂控制流;
  • 以多维张量计算为主。

算子调度程序主要由三种代码结构组成:

  • 内存分配;
  • 循环;
  • 计算。

调度树有四种结点:

  • 根结点;
  • 循环结点:表示计算函数沿给定维度的遍历方式;
  • 存储结点:表示后续会用到的中间结果的存储;
  • 计算结点:是调度树的叶结点,表示被执行的计算。

调度树可以很好的表达算子的计算原理,通过遍历的方式就可以很容易的将调度树转化为对应的程序代码。AI编译器后端优化的主要目的就是优化调度树。

调度转换

调度转换的方式有:

  • Reorder:为相同的功能切换循环结点;
  • Hoist / Lower compute:更改计算中间结果的粒度;
  • Inline / Deinline:将函数内联到调用者中或从调用者中取消内联。

算子的手工优化

循环交换

在多重循环中,采用迭代次数较小的循环驱动内层迭代次数较大的循环可以减少内存的消耗。

1
2
3
4
5
for (int i = 0; i < 10000; i++) {
for (int j = 0; j < 200; j++) {
// TODO
}
}

上述循环可以改造为下面的形式。

1
2
3
4
5
for (int i = 0; i < 200; i++) {
for (int j = 0; j < 10000; j++) {
// TODO
}
}
表达式外放

循环体内重复计算的表达式放在循环体外部,避免每次都计算。

1
2
3
4
5
for (int i = 0; i < 200; i++) {
for (int j = 0; j < 10000; j++) {
j = i * j * x / (y - 1);
}
}

上述代码调整为下列形式。

1
2
3
4
5
6
int tmp = x / (y - 1);
for (int i = 0; i < 200; i++) {
for (int j = 0; j < 10000; j++) {
j = i * j * tmp;
}
}
循环终止调用

消除循环终止条件中的函数调用。

1
2
3
for (int i = 0; i < vec.size(); i++) {
// TODO
}

上述代码可以修改为下面的形式。

1
2
3
4
int size = vec.size();
for (int i = 0; i < size; i++) {
// TODO
}

AI编译器的算子优化

算子调度优化分为三大类,九小类:

  • 循环优化 (Loop Optimization)
    • 循环展开 (Loop Unrolling)
    • 循环分块 (Loop Tiling)
    • 循环重排 (Loop Reorder)
    • 循环融合 (Loop Fusion)
    • 循环拆分 (Loop Split)
  • 指令优化 (Instructions Optimization)
    • 向量化 (Vectorization)
    • 张量化 (Tensorization)
  • 存储优化 (Memory Optimization)
    • 延迟隐藏 (Latency Hiding)
    • 内存分配 (Memory Allocation)
循环展开

对循环进行展开,以便每次迭代多次使用加载的值,使得一个时钟周期的流水线上尽可能满负荷计算。在流水线中,会因为指令顺序安排不合理而导致计算单元等待空转,影响流水线效率。循环展开为编译器进行指令调度带来了更大的空间。

1
2
3
4
5
6
7
8
9
10
11
12
13
for (int i = 0; i < n; i++) {
for (int j = 0; j < m; j++) {
C[j] = A[j] + B[i];
}
}

// Unrolling
for (int i = 0; i < n; i+=2) {
for (int j = 0; j < m; j++) {
C[j] = A[j] + B[i];
C[j + 1] = A[j + 1] + B[i];
}
}
循环分块

由于内存空间有限,代码访问的数据量过大,无法一次性将所需数据加载到设备内存,循环分块能有效提高Cache的访存效率,改善数据局部性。如果Tiling应用于外部循环,可以增强计算的空间和时间局部性,分块应和缓存共同作用,以提高流水线的效率。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
for (int j = 0; j < n; j++) {
for (int i = 0; i < m; i++) {
A[i] += B[j];
}
}

// Tiling
for (int j_o = 0; j_o < n; j_o+=tile) {
for (int i = 0; i < m; i++) {
for (int j_i = j_o; j_i < j_o + tile; j_i++) {
A[i] += B[j_i];
}
}
}
循环重排

内外层循环重排,改善空间局部性,并最大限度地利用引入缓存的数据。对循环进行重新排序,以最大程度地减少跨步,并将访问模式与内存中的数据存储模式对齐。

1
2
3
4
5
6
7
8
9
10
11
12
for (int i = 0; i < n; i++) {
for (int j = 0; j < m; j++) {
C[i][j] = A[i][j] * B[i][j];
}
}

// Reorder
for (int j = 0; j < m; j++) {
for (int i = 0; i < n; i++) {
C[i][j] = A[i][j] * B[i][j];
}
}
循环融合

循环融合将相邻或紧密间隔的循环融合在一起,减少循环开销并增加计算密度,进而改善流水线排布,增加数据的缓存局部性。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
for (int i = 0; i < n; i++) {
C[i] = A[i] + B[i];
D[i] = 2 * C[i];
}
for(int i = 1; i < n - 1; i++) {
E[i] = D[i] + A[i];
}

// Fusion
C[0] = A[0] + B[0];
D[0] = 2 * C[0];
C[n - 1] = A[n - 1] + B[n - 1];
D[n - 1] = 2 * C[n - 1];
for(int i = 1; i < n - 1; i++) {
C[i] = A[i] + B[i];
D[i] = 2 * C[i];
E[i] = D[i] + C[i];
}
循环拆分

拆分主要是将循环分成多个循环,可以在有条件的循环中使用,分为无条件循环和有条件循环。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
for (int i = 0; i < n; i++) {
C[i] = A[i] + B[i];
D[i] = 2 * C[i];
if(temp[i] > data)
E[i] = A[i];
}

// Split
for (int i = 0; i < n; i++) {
C[i] = A[i] + B[i];
D[i] = 2 * C[i];
}
for (int i = 0; i < n; i++) {
if(temp[i] > data)
E[i] = A[i];
}
向量化
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
double sum = 0.0;
for (int i = 0; i < n; i++) {
sum += a[i];
}

// Vectorization
double<4> vec_sum = {0.0, 0.0, 0.0, 0.0};
for (int i = 0; i < n; i+=4) {
double<4> a_val = load<4>(a + i);
vec_sum = add<4>(a, vec_sum);
}
double sum = 0.0;
for (int i = 0; i < 4; i++) {
sum += vec_sum[i];
}
张量化

主流的CPU/GPU硬件厂商都提供了专用于张量化计算的张量指令,如英伟达的张量核指令、英特尔的VN。利用张量指令的一种方法是调用硬件厂商提供的算子库,如英伟达的cuBLAS、cuDNN,以及英特尔的oneDNN等。

但当深度学习模型出现新的算子或需要进一步提高性能时,这种方法是有局限性的。

TVM提出了一种将硬件指令的接口和调度分开,生成硬件接口的方式。利用Tensor表达式表示高层运算,还能表示硬件底层特性。张量化提供调度表和特定的硬件原语,便于AI编译器扩展支持新的硬件架构。AI编译器后端将Tensor操作映射为硬件实现或高度优化的手工微内核,从而显著提高性能。

延迟隐藏

将内存操作与计算交叠,最大限度地提高内存和计算资源利用率。

  • CPU:延迟隐藏可以通过多线程或硬件隐式数据预取来实现;
  • GPU:依赖于Warp Schedule对多线程的管理调度和上下文切换实现;
  • NPU/TPU:采用解耦访问/执行(Decoupled Asscess/Execute, DAE)架构。

DAE架构是将内存访问单元(MAU)与管道分离。执行处理器(EP)由直接寄存器访问(DRA)和直接缓存访问(DCA)组成,可以用于管理寄存器、缓存和内存之间的数据传输。

TVM将虚拟线程并行程序转换为单个指令流,这个流包含显式的低级同步,硬件可以解释这些同步,以恢复管道并行性,需要隐藏内存访问延迟。在硬件中执行解耦访问,允许内存和计算交叠。执行正确性通过低级同步来实现,同步形式依赖于Token入队、出队操作。

内存分配
  • 局部变量:开发者定义普通变量时编译器在内存中的栈空间为其分配一段内存;
  • 全局变量(静态变量和全局变量):编译器在内存中的静态存储区分配空间;
  • 堆变量:开发者用mallocnew申请的空间位于堆中。

传统编译器通过将内存逻辑划分为不同区段来提供开发者访问内存的权限,而每段空间都有各自的大小,一旦开发者在使用过程中将该内存区域耗尽,就会出现内存不足的错误。

GPU/NPU等设备上的内存分配更加复杂,设备种类也很多,对开发者的要求会很高,所以内存分配的工作都交给AI编译器来进行。

Auto Tuning

Auto Tuning的概念不是AI编译器中才出现的,传统编译器已经有了这个概念。对于给定的程序和目标架构,找到最优的编译优化方法就是Auto Tuning。

传统编译器的两个主要优化方法是:

  • 优化选择 (Optimize Selection):选择哪些优化算法;
  • 优化顺序 (Optimize Sequence):不同优化方法的顺序如何组成。

AI编译器中的Auto Tuning与传统编译器的主要区别在于:

  • 更低实验开销:
    • 聚焦算子或Kernel级别的优化,而非整个程序;
    • Cost Model在CPU上模拟AI芯片执行,训练和推理的模拟速度要求足够快;
  • 特定领域结构:
    • 针对算子和Kernel的高度循环化、张量化、并行化的特点进行优化;
    • 优化大量类似的算子计算模式。
参数化 (Parameterization)

对调度优化问题进行建模,参数化优化空间一般由可参数化变换的可能参数取值组合构成,因此需要调度原语进行参数化表示。Halide将算法和调度解耦,TVM提供调度模板。

成本模型 (Cost Model)

成本模型用来评价某一参数化下的调度性能,根据对调度额评价来指导搜索到最优的调度策略。可以从运行时间、内存占用、编译后指令数来评价。

主要实现方式:

  • 基于硬件的黑盒模型;
  • 基于模拟的预定义模型;
  • ML-Base模型,通过机器学习模型来对调度性能进行预测。
搜索算法 (Search Algorithm)

确定初始化和搜索空间后,在搜索空间找到达到性能最优的参数配置。

常用的搜索算法有:

  • 遗传算法;
  • 模拟退火算法;
  • 强化学习。

传统编译器与AI编译器对比

https://deleter-d.github.io/posts/40483/

作者

亦初

发布于

2024-03-26

更新于

2024-06-19

许可协议

评论

:D 一言句子获取中...

加载中,最新评论有1分钟缓存...