AICompiler编译器介绍及访存密集算子优化

原创

于 2021-03-25 14:53:29 发布 · 3.1k 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#机器学习/深度学习 #人工智能 #自然语言处理 #并行计算 #算法 #Cloud Native #编译器 #芯片 #C++ #异构计算

介绍阿里云AICompiler在AI模型性能优化方面的进展，特别是在访存密集型算子优化方面的工作，包括FusionStitching技术和CodeGen策略，以及TAOCompilerV1和V2的关键改进。

简介：欢迎走进阿里云机器学习PAI AICompiler编译器系列。随着AI模型结构的快速演化，底层计算硬件的层出不穷，用户使用习惯的推陈出新，单纯基于手工优化来解决AI模型的性能和效率问题越来越容易出现瓶颈。为了应对这些问题，AI编译优化技术已经成为一个获得广泛关注的技术方向。这两年来，这个领域也异常地活跃，包括老牌一些的TensorFlow XLA、TVM、Tensor Comprehension、Glow，以及最近呼声很高的MLIR，能够看到不同的公司、社区在这个领域进行着大量的探索和推进。

阿里云机器学习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的比较简单的