阿里云机器学习PAI是阿里云上提供的人工智能平台服务，提供一站式的机器学习解决方案。在过去的三年多时间里，阿里云机器学习PAI团队在AI编译优化技术方向投入了比较专注的资源精力，对这个领域也建立起了一定的理解。目前在内部集群上AICompiler已经实现了默认的全量打开，在阿里外部已经开始为来自多个不同行业的业务方提供服务，涉及行业领域包括在线教育，安全，娱乐，电商，直播等；涉及的业务类型包括文本分类/识别，机器翻译，语音识别/生成，图像类业务等等；支持覆盖的硬件包括云端GPU/CPU以及端上CPU等。
我们会在接下来的系列文章里介绍一下这方面的工作以及一些思考，一方面抛砖引玉，吸引更多同学来进行共同建设；另一方面也为阿里云上对训练和推理作业的性能优化有一定需求的用户同学提供一些技术方面的背景介绍。

AI编译器简介

在上图中将近年的深度学习框架粗略分为三代。几代框架之间有一个趋势，在上层的用户API层面，这些框架在变得越来越灵活，灵活性变强的同时也为底层性能问题提出了更大的挑战。另外一个技术趋势则是系统底层的深度学习的编译器近一段时间也开始活跃起来，这些编译器试图去解决框架的灵活性和性能之间的矛盾。

传统编译器是以高层语言作为输入，避免用户直接去写机器码，而用相对灵活高效的语言来工作，而深度学习编译器的作用相仿，其输入是比较灵活的，具备较高抽象度的计算图，输出包括CPU或者GPU等硬件平台上的底层机器码及执行引擎。

AI编译器的目标是针对AI计算任务，以通用编译器的方式完成性能优化。让用户可以专注于上层模型开发，降低用户手工优化性能的人力开发成本，进一步压榨硬件性能空间。

涉及到性能优化，我们有必要先对一个AI作业执行过程中的性能开销的分布有一个感性的认识，所谓what you can't measure, you can't optimize it. 在《Characterizing Deep Learning Training Workloads on Alibaba-PAI》这篇paper里，我们针对阿里云机器学习PAI平台训练workload的性能开销占比有过一个比较细致的分析。考虑到目前我们在AI编译优化里还主要关注单计算设备的计算效率问题，所以我们可以宏观上将单设备上的性能开销拆解为计算密集算子（比如GEMM和Convolution）和访存密集算子（比如Elementwise Add，BN等操作）两部分，关于计算图过于灵活带来的框架开销，我们也统一归类到访存密集算子开销占比中。

针对不同的性能热点，所需要的优化手段也存在区别。本文后续篇幅将首先介绍我们在访存密集型部分的工作，计算密集型算子的部分会在后续连载中进行介绍。

Fusion Stitching——大颗粒度访存密集型子图CodeGen
阿里云机器学习PAI团队在访存密集算子上以XLA为基点开展了大量的工作，工作内容涵盖前端、中端、后端以及执行引擎，我们将其称之为Tensor Accelerator and Optimizer(TAO) compiler。
TAO Compiler先后经历了V1 (FusionStitching: Deep Fusion and Code Generation for TensorFlow Computations on GPUs)和V2(FusionStitching: Boosting Memory Intensive Computations for Deep Learning Workloads)两代演进。

如前文所说，XLA在GPU backend上的主要收益来源是对访存密集型算子的自动Op Fusion CodeGen。催生我们做这些尝试的原因是，我们在实际业务中发现，社区XLA在最核心的CodeGen环节还有很大的问题和改进空间。

例如下图为一个LayerNorm模块的前向计算子图，手工优化的话，它可以很容易被写成一个CUDA kernel，本应该是很适合编译器通过自动op fusion来获取性能收益。但我们发现XLA实际并没有做的非常好。经过细致的分析后我们认为这里的本质问题是社区的CodeGen模版过于简单，因为模版非常简单限制了中端Fusion pass的灵活度，被迫做了非常保守的op fusion策略，图中每一条红线都有一个细节的原因导致社区XLA无法把这些算子fuse到一起，也就无法拿到节省相应访存开销所能得到的优化收益。

分析之后我们认为造成这些问题的根本原因为：
? 单一Parallel Loop的比较简单的CodeGen模版；
? 由于CodeGen模版比较简单，为保证性能而被迫保守的Op Fusion策略。

这种简单的CodeGen策略在fuse线性的Elementwise算子时足够得到理想的结果，但当计算图的连接关系变复杂，同时计算图中出现Reduce／Broadcast等复杂计算节点时，在实际业务中效果并不好，特别是对于包含了后向处理部分的训练过程计算图更是如此。

当然，我们可以直接采用复杂的CodeGen模版，通过牺牲通用性的方式换取更好的性能收益。但在Op节点类型以及计算图的拓扑结构千变万化的情况下，这种方式并不能够在编译器中通过编译的方式有效的节省人力成本。于是就引申出了TAO compiler的设计理念。

TAO Compiler V1
TAO Compiler V1的核心原理可以概括为：
? 借助于GPU硬件中低访存开销的shared memory，缝合不同schedule的计算子图，实现多个parallel loop复合
? 支持不同shape的多个tensor输出
? 保证CodeGen性能的前提下，实现更加激进的Op fusion策略

在GPU的体系结构中，SM内的shared memory的访存开销远远低于global memory的访存开销，可以帮助我们来把多个本来独立的kernel粘合在一起。这样的方案下，被粘和的每一个子图仍然可以基于自己的简单的CodeGen模版来运作，他们不必担心因为各自的pattern被fuse在一起而需要更复杂的CodeGen模版，也一定程度上不必担心因为fuse的颗粒度太大而产生过多的冗余的计算。V1的核心工作涉及中端和后端两个部分，解决了多种不同细节原因导致的细粒度算子无法fuse到一起的问题，主要是基于这个核心思路。

下面的两张图分别为社区XLA Codegen的情况和TAO Compiler V1的情况。TAO Compiler的CodeGen将社区做法中难以fuse在一起的计算pattern，通过低访存开销的shared memory作为桥接，将多个kernel缝合在一起，同时又保留被缝合的每一个部分都依然采用简单的CodeGen模版。这种方案在通用性和性能之间取得一个更好的平衡点。


TAO Compiler V1通过这样的方式，在一定程度上保证CodeGen性能的基础上，大幅提高了op fusion的颗粒度，以下面表格中的两个例子为例：

TAO Compiler V2
在做V1的推广落地的同时，我们抽取了一些业务进行了比较精细的分析，观察优化距离性能的极限还有多大的距离，结论是虽然凭借在op fusion颗粒度上的优势，在很多case上可以做到明显优于社区XLA，但离硬件性能的峰值还是有比较明显的gap的。所以我们抽取了一些有代表性的kernel，分析了这些gap的具体原因。我们发现每个case各自都有不同的原因导致没有能够接近访存的峰值性能。这些原因比较细节，在此不再展开，但宏观来说问题来自于目前的CodeGen基于Rule-Based的策略，而Rule-Based的策略本身存在其固有的局限性。
具体到compiler的各个环节：
? 在中端Op Fusion层面上，基于简单规则下局部合理的Op Fusion决策可能并非全局最优的决策；
? 在后端CodeGen层面上，基于简单规则下，同一Kernel内的不同节点的模版选择和CodeGen行为决策也可能并非整体最优的决策。

在TAO Compiler V2与V1以及社区XLA相比最大的区别可以概括为：
? 更多的考虑全局最优，而非基于贪婪的局部最优；
? Rule-Based到Cost-Based的演进。
这两点区别同时反映在中端的Op Fusion策略以及后端的CodeGen策略上。

我们引入了多个不同等级的Cost Model，其主要区别如下：

中端在Graph范围内做Op Fusion决策，由于探索范围较大，选择接入开销较低的SOL cost model；后端在Kernel范围内做tuning，可以接入开销相对较高的Proj Cost Model。

中端Op Fusion决策部分
目前XLA和TAO V1中解决HLO instruction之间的Auto Fusion问题时采用的是完全基于规则的策略。这些规则都是基于人工经验确定的，往往难以覆盖到所有的情况，导致可能会出现各种corner case以及为了解决这些corner case而打上的一系列补丁。从长期来看这种方法并不能充分的挖缺潜在的性能优化空间，也会带来越来越沉重的维护成本。为了使用更加系统性的方法解决Auto Fusion的问题，我们提出了Cost Based Fusion Strategy。
新的策略主要是对V1中以下两个大的方面进行改进：
? 兼容性度量。 兼容性度量是一种用来评估一个给定的Fusion动作是否有收益的方法。目前V1中使用的是一条简单的规则来进行兼容性度量，即只要融合之后的kernel的thread block level的并行度没有下降到1则认为融合具有收益。该策略在很多实际业务上发挥了很好的作用，但也遇到了不少Corner case。比如该规则对于大shape的pattern表现为过于激进，融合后可能因并行性大幅下降而导致性能变差，而对于小shape的pattern则表现的过于保守，丢掉了一部分融合的机会。在V2中我们基于cost model来更系统化的考虑访存量，kernel launch次数，并行度改变对于最终性能的影响等，可以更准确的评估一个融合动作的效果以指导Fusion Plan的探索。
? Fusion的探索空间。 一个计算图对应的Fusion Plan由多个Fusion Pattern组成。一个Fusion Pattern中如果只包含具有生产消费关系指令则是纵向fusion，反之如果包含有不含生产消费关系的指令否则为横向fusion。一个计算图可能存在着多种不同的Fusion Plan。目前V1中是基于一个启发式的贪心策略确定性的生成一个Fusion Plan，而再不进行其他可能的探索。这种策略好处是具有很高的效率，但缺点是会漏掉很多可能性。另一方面目前社区XLA及TAO Compiler V1中只会考虑纵向fusion而未考虑横向fusion。这里所说的横向fusion即不存在直接生产消费关系的节点之间的fusion，如下图所示。除其中近距离横向fusion可同样带来仿存节省的收益外，在更多的场景下可通过节省kernel launch开销带来收益，同时当使用空分并行的方式对横向fusion之后的kernel做CodeGen时还可能提升GPU的资源利用率。下图中展示了一些横向fusion的例子。V2的中端部分，对包含以上两种fusion的解空间进行了更多探索，以期找到更好的全局最优决策。

具体来说fusion空间探索是一个迭代式的过程，每一次迭代由以下几个步骤构成：
? 搜索高价值fusion pattern。 这里的高价值指的是基于cost model评估会带来更大收益的pattern。直接暴力枚举进行搜索是一个指数级别的算法，我们使用的是一种启发式的算法，在稀疏图场景下（一般指令间的连接图满足这个前提）是一个接近线性的算法。具体而言我们首先构造计算图指令间的一个拓扑排序，依次遍历每条指令并按照递推的方式生成以该指令结尾的top k候选fusion pattern。递推的过程中我们会进行局部的充分枚举（扩张），对整个空间进行更多的探索，同时也会及时进行剪枝（收缩），确保可以不会组合爆炸。
? 构造fusion plan。 即寻找Fusion Pattern的一个有效组合以构成一个完整的Fusion Plan。这里的有效组合指的是各个Fusion Pattern之间不相交且相应的融合动作不会引入环。 这里我们尝试了多种不同的策略，比如贪心策略，BFS+剪枝策略。目前使用的是每次选择兼容的且收益最大的fusion pattern的贪心策略。
? apply fusion plan。 也即根据Fusion Plan执行相应的融合动作。
以上的过程可以对整个优化空间进行更大程度的探索，同时又有效的控制了搜索的复杂性，在一组评测应用集合（这也是进行编译器开发经常性需要的工具支撑了）上的实验表明，该策略工作良好，符合预期。

后端CodeGen部分
在介绍V2后端CodeGen的具体工作之前，我们先对XLA和TVM做一个简单的横向比较：
? XLA是一种比较务实的方案，用基于规则的简单方式去做收益空间最大的访存密集型pattern的Automatic CodeGen，它的优势是自动化和比较小的编译开销，缺点是如果想进一步提升性能的话会非常局限；
? TVM的目标是去胜任更复杂的CodeGen，基于规则的策略无法胜任情况下，它在Halide工作的基础上，提出了一个比较先进的Schedule抽象。这个Schedule的空间非常大，如何才能找到最好的Schedule，一方面需要基于用户的经验写出比较适合硬件体系结构的Schedule，一方面需要巨大的编译开销为代价在Schedule之间进行tuning。

这里我们需要解决的问题是：需要找到一种方式，能够自动/对用户透明的，在有限的编译开销的基础上，提升性能并进一步接近硬件的性能峰值极限。

V2的工作在XLA的基础上，一定程度上试图借鉴TVM的Schedule抽象的核心想法，但又有本质的不同。根据我们过往工作的经验，对于访存密集型pattern来说，为了达到或接近硬件的峰值性能极限，所需要的Schedule的枚举空间通常并不需要非常大。因此我们试图提炼一套比TVM的schedule抽象要简化很多的CodeGen Plan抽象，希望这套抽象能够代表手工写相应Pattern的Kernel时所需要考虑的全部问题空间。同时需要保证CodeGenPlan的枚举空间是足够小的，能够满足用户对编译开销的需求。此外，我们使用Proj Cost Model而非TVM中目前基于实际profiling结果来tuning的方式，其原因也是希望能在编译开销和性能之间取得更好的平衡点。
在TAO Compiler V1的工作的基础上，V2在CodeGen部分的主要区别是：
? CodeGen Plan抽象，尽量对全部的CodeGen可能性做抽象和枚举；
? Index计算与主体计算分离，引入index_cache与value_cache，最小化冗余计算；
? 与CodeGen Plan抽象相对应的Implementation抽象，方便对不同节点增加新的算子实现模版。

篇幅原因我们不再详细介绍CodeGen Plan抽象的具体细节。下面仅试图以一个pattern为例，来说明TAO Compiler V2与V1及社区XLA之间，在CodeGen结果上的主要区别。

下图左图为社区XLA的CodeGen结果，由于简单模版的限制会生成3个简单Schedule的kernel；右图是TAO Compiler V1的结果，借助Shared Memory的粘合，可以生成一个kernel。

TAO Compiler V1距离硬件性能峰值之间可能存在差距的主要问题来自于：
? A/B/C三个子部分的CodeGen Schedule是彼此独立的，在此例中，A/B/C三个部分的计算在同一个GPU thread内可能依赖于Elementwise 4节点处的不同Element，产生冗余计算。
? A/B/C三个子部分的CodeGen过程相对独立，Elementwise 3等节点处的值和Index的计算结果无法被有效Cache，进一步增加了冗余计算指令。
? A/B/C三个子部分可能各自有多种CodeGen模版，但可能只有一部分子模版的组合可以在全局上最好的利用GPU的并发度（即GPU Occupancy）。

当冗余计算指令的数量（或者Occupancy带来的问题）达到一定临界值后，本应该是仿存Bound的计算pattern，性能热点便会迁移为计算或Occupancy，导致无法达到硬件仿存性能峰值的性能。

TAO Compiler V2经过CodeGen Plan的枚举和基于Cost Model的tuning之后，会找到一个认为相对最优的CodeGen Plan。例如，在上图这个pattern中，选择Reduce节点做为整个Kernel的Dominant节点可能是更好的方案，其它节点处，在同一个thread内需要计算的Index会由Dominant节点的Index以辐射传播的方式推导计算得到，Index的传播过程和数值的计算过程在整个kernel范围内进行Cache，通过以上机制最大化的避免冗余的Index计算和数值计算。此外，CostModel还会对枚举的CodeGen Plan中的Launch Dimension等因素对Occupancy的影响进行建模，得到对于整体性能最优的CodeGen行为。

TAO Cost Model

中端不同的fusion策略以及不同后端CodeGen生成方式会影响最终模型实现的效率。所以很自然的想到利用cost model帮助中端、后端CodeGen探索最优的实现。这里主要的挑战是：

? cost model需要高效的对成百上千种不同实现做实时的评估，需要减少其本身的overhead；
? cost model需要准确的区分不同实现的性能差异；
? 不同型号的gpu拥有不同的硬件架构（memory hierarchy，latency以及throughput指标），cost model需要对不同型号的gpu建模。

所以我们会将workload输入映射成不同的资源需求，同时对device内部流水线stage的resource、latency以及throughput做了抽象和表示。最终通过建立latency模型以及throughput模型估计当前实现的计算耗时。

Cost Model分为3个level： SOL Cost Model, Projected Cost Model以及Profiling Cost Model。
? SOL Cost Model 只考虑硬件spec，不考虑具体实现、架构特点（例如假设100% cache hit ）和编译器效率。对于CodeGen，主要用memory cycles以及device throughput（例如：处理 1000T MACs, V100 125 Tflops vs 4x TPU 180 Tflops vs new ASIC）来区分。其特点是建模速度快，误差较大， 所以可以用于粗粒度性能筛选以及不同硬件性能快速评估。因此,我们在中端OP fusion探索决策时引入SOL cost model评估不同Op fusion的cost并保留下cost 最小的Op fusion Plan。
? Projected Cost Model需要考虑具体kernel实现、架构特点（如hit rate）和编译器特点。对于CodeGen，主要用kernel latency来区分。特点是建模速度慢，但误差较小，所以可以用于细粒度性能筛选。因此,我们在做后端CodeGen时，引入Proj cost model区分同一条fused instrcution在不同CodeGen Plan下因为指令数不同以及寄存器复用情况差异造成的性能差别。
? Profiling Cost Model目的是作为cost model的ground-truth， 从profiler上得到实际运行性能结果。其特点是执行耗时长，所以可以用于构造offline cost model以及修正SOL level与Proj Level cost model。

Cost Model主要由三个模块组成： Workload Model, Device Model, 以及 Cost Evaluation。简单来说 。可以举一个简单的例子，假设任务的workload由10条硬件指令组成，芯片的处理能力是每个cycle处理2 条硬件指令，那么该workload需要5个cycle才能被处理完。

Workload model的准确性影响了整个cost model建模的准确性，因此80%以上的研发时间用于建模workload model。实践中，我们会估计某条fused instruction在当前CodeGen Plan下的寄存器消耗、各个硬件单元的硬件指令需求、memory 消耗以及shared memory消耗。

Device Model建模了GPU各个硬件单元的throughput、bandwidth以及latency。可以从公开资料或者论文中得到.在device model中，我们主要建模了芯片内部各个unit的数量、SM内各个单元的throughput和latency、L2的throughput（用于建模atom指令）以及 Memory bandwidth。
Cost Evaluation根据workload model以及device model 估计最终生成kernel的Memory cycle、Occupancy（并行度）、Wave number、Kernel Latency、Performance Limiter以及SM内部各个unit的cycles。

Cost Model主要有两种建模方式：Throughput Model以及Latency Model。
Throughput Model首先估计fused instruction所需各机器指令的数量（例如需要X条FFMA, Y条MOV, Z条IMAD等）。并对机器指令根据其对应的pipe做聚合（例如FFMA和IMAD都是发射到FMA Pipe的）。在获得当前GP各个pipe的throughput后，用pipe workload（硬件指令）的数量除以这个pipe的throughput就是其所需要的cycle数。最后在所有pipe中找cycyle最大的pipe，这个pipe就是这个fused instruction在当前gpu上的performance limiter，并且这个pipe所需要的时间就是kernel的时间。

Latency Model首先估计fused instruction所需的register数量，并计算出它的occupancy，再计算出总共需要的wave数量（波数）。然后估计fused instruction所需各硬件指令的数量（同Throughput Model中的第一条）。根据device model获得当前gpu各个硬件指令的latency。根据硬件指令数以及指令的latency计算 warp latency，用warp latency乘以wave数量得到最后的kernel latency。

对于TAO Compiler V1来说，由于大部分fused instruction已经达到或者接近GPU math或memory的理论上限，此类情况下，throughput model能很好地对fused instruction做建模。 对于TAO Compiler V2，由于fuse了更大的尺度，目前大部分fused instruction距离硬件极限还有点距离，所以用throughput model做区分是不合理的。需要用latency model。

性能数字
下图为基于本文中所述工作，在一些典型业务模型上，TAO Compiler与社区XLA在Nvidia V100上的性能数字对比，其中社区XLA与TAO的性能收益数字，均为相同计算精度下的收益数字。

下图为某业务方基于PAI的easyTransfer解决方案进行建模，在不同模型上的性能收益数字。AICompiler同样对不同应用场景给出了一致性的优化方案和相对比较明显的性能优化结果。提供的性能优化数字普遍超过了社区XLA的工作。

阿里云机器学习PAI平台面向企业客户及开发者，提供轻量化、高性价比的云原生机器学习平台，涵盖交互式建模、拖拽式可视化建模、分布式训练到模型在线部署的全流程覆盖，百余种落地场景，全面提升机器学习工程效率。
目前， PAI AICompiler已经集成在阿里云机器学习PAI的通用推理优化工具Blade敏捷版中，用户可以参考开发文档来快速体验。

作者：朱凯，赵文益，杨军