作者:PAI团队
随着AI模型结构的快速演化,底层计算硬件的层出不穷,用户使用习惯的推陈出新,单纯基于手工优化来解决AI模型的性能和效率问题越来越容易出现瓶颈。为了应对这些问题,AI编译优化技术已经成为一个获得广泛关注的技术方向。这两年来,这个领域也异常地活跃,包括老牌一些的TensorFlow XLA、TVM、Tensor Comprehension、Glow,以及最近呼声很高的MLIR,能够看到不同的公司、社区在这个领域进行着大量的探索和推进。
在过去几年时间里,PAI团队在AI编译优化技术方向投入了比较专注的资源精力,对这个领域也建立起了一定的理解,撰写这个系列文章,旨在达到如下目的:
- 对PAI团队在AI编译优化技术方向上的技术投入和取得的结果做一个系统性的总结;
- 阐述我们对AI编译优化整体技术脉络的理解,期望能够对关心AI编译优化技术建设的同行有所启发;
- 对我们关于阿里AI编译优化的技术路径提出一些思考,期望能够抛砖引玉,吸引更多对我们工作感兴趣的同行联系加入我们,一起打造具备阿里特色的AI编译优化技术体系。
本系列文章会由以下内容构成:
- 总体介绍及行业现状分析
- 访存密集算子优化
- 计算密集算子优化
- 离线编译系统
- 业务实践
- 总结展望
本篇会先做一下整体介绍以及行业现状的分析。
为什么需要AI编译器
![64d4fab8afa7f5eddc2190cec39882f8.png](https://i-blog.csdnimg.cn/blog_migrate/f769773cf5008b790b4ec1ccc95211ba.png)
在上图中将近年的深度学习框架粗略分为三代。几代框架之间有一个趋势,在上层的用户API层面,这些框架在变得越来越灵活,灵活性变强的同时也为底层性能问题提出了更大的挑战。另外一个技术趋势则是系统底层的深度学习的编译器近一段时间也开始活跃起来,这些编译器试图去解决框架的灵活性和性能之间的矛盾。
![294ef7c5e0a9a01d7ac984ad48563968.png](https://i-blog.csdnimg.cn/blog_migrate/50adbaeb02433c4787f6d2c4d1cfdf12.png)
传统编译器是以高层语言作为输入,避免用户直接去写机器码,而用相对灵活高效的语言来工作,而深度学习编译器的作用相仿,其输入是比较灵活的,具备较高抽象度的计算图,输出包括CPU或者GPU等硬件平台上的底层机器码及执行引擎。
AI编译器的目标是针对AI计算任务,以通用编译器的方式完成性能优化。让用户可以专注于上层模型开发,降低用户手工优化性能的人力开发成本,进一步压榨硬件性能空间。
涉及到性能优化,我们有必要先对一个AI作业执行过程中的性能开销的分布有一个感性的认识,所谓what you can't measure, you can't optimize it. 在这篇paper里,我们针对PAI平台训练workload的性能开销占比有过一个比较细致的分析。考虑到目前我们在AI编译优化里还主要关注单计算设备的计算效率问题,所以我们可以宏观上将单设备上的性能开销拆解为计算密集算子(比如GEMM和Convolution)和访存密集算子(比如Elementwise Add,BN等操作)两部分,关于计算图过于灵活带来的框架开销,我们也统一归类到访存密集算子开销占比中。
针对不同的性能热点,所需要的优化手段也存在区别。
本章我们首先选取当前两个主流编译框架XLA(针对访存密集算子)和TVM(针对计算密集算子),以及新近非常活跃的MLIR框架为代表(MLIR严格来说不是一个compiler,而是compiler infrastructure,后面会再进行展开介绍),来对业界目前的技术情况做一个概述。
XLA现状简述
原理及主要收益来源
XLA(Accelerated Linear Algebra)是Google于2017年出推出的用于TensorFlow的编译器。XLA使用JIT编译技术来分析用户在运行时创建的 TensorFlow 图,将TensorFlow Op转换成为HLO(High LevelOptimizer)中间表示并在HLO层上完成包括Op Fusion在内的多种图优化,最后基于LLVM完成CPU/GPU等机器