对 PAI 团队在 AI 编译优化技术方向上的技术投入和取得的结果做一个系统性的总结; 阐述我们对 AI 编译优化整体技术脉络的理解,期望能够对关心 AI 编译优化技术建设的同行有所启发; 对我们关于阿里 AI 编译优化的技术路径提出一些思考,期望能够抛砖引玉,吸引更多对我们工作感兴趣的同行联系加入我们,一起打造具备阿里特色的 AI 编译优化技术体系。
本系列文章会由以下内容构成:
总体介绍及行业现状分析
访存密集算子优化
计算密集算子优化
工程篇:编译解耦及全量打开
Dynamic shape compiler
离线编译系统
业务实践
总结展望
涉及到性能优化,我们有必要先对一个 AI 作业执行过程中的性能开销的分布有一个感性的认识,所谓 what you can't measure, you can't optimize it. 在这篇 paper 里,我们针对 PAI 平台训练 workload 的性能开销占比有过一个比较细致的分析。考虑到目前我们在 AI 编译优化里还主要关注单计算设备的计算效率问题,所以我们可以宏观上将单设备上的性能开销拆解为计算密集算子(比如 GEMM 和 Convolution)和访存密集算子(比如 Elementwise Add,BN 等操作)两部分,关于计算图过于灵活带来的框架开销,我们也统一归类到访存密集算子开销占比中。
Optimizations and transformations typically done on a such graph (e.g. in Grappler) ( _ 统一优化框架的野心 _ ).
Representation of kernels for ML operations in a form suitable for optimization( _ 和下面一条一起,期望将计算密集算子和访存密集算子集中在一套框架里统一打击的野心 _ ).
Ability to host high-performance-computing-style loop optimizations across kernels (fusion, loop interchange, tiling, etc) and to transform memory layouts of data.
Code generation “lowering” transformations such as DMA insertion, explicit cache management, memory tiling, and vectorization for 1D and 2D register architectures( _ 对硬件 memory/cache/register 的分层描述能力 _ ).
Ability to represent target-specific operations, e.g. accelerator-specific high-level operations.(灵活的对 AI-domain ASIC 的描述能力)
Quantization and other graph transformations done on a Deep-Learning graph( _ 这一点也很诱人,关于模型层面的优化,如果能够基于 MLIR 的中间表示层进行统一打击,无疑对于提升行业生产效率是有重要价值的 _ ).
A collection of modular and reusable software components that enables the progressive lowering of operations, to efficiently target hardware in a common way
关于安芯教育 安芯教育是聚焦AIoT(人工智能+物联网)的创新教育平台,提供从中小学到高等院校的贯通式AIoT教育解决方案。 安芯教育依托Arm技术,开发了ASC(Arm智能互联)课程及人才培养体系。已广泛应用于高等院校产学研合作及中小学STEM教育,致力于为学校和企业培养适应时代需求的智能互联领域人才。