机器学习访存密集计算编译优化框架 AStitch，大幅提升任务执行效率

近日，关于机器学习访存密集计算编译优化框架的论文《AStitch: Enabling A New Multi-Dimensional Optimization Space for Memory-Intensive ML Training and Inference on Modern SIMT Architectures》被系统领域顶会ASPLOS 2022接收。

AStitch通过编译优化的手段来自动化地提高机器学习任务的执行效率, 提出了一种大粒度计算融合的编译优化手段，通过计算图的依赖关系特性、GPU多层次存储架构上的数据局部性、以及不同数据尺寸之下的线程并发性等三个方面的联合考虑，自动化地为大粒度的复杂访存密集算子子图生成高效的GPU代码，从而大幅减少GPU kernel调用及框架层算子调度的额外开销，避免了不必要的重复计算。大幅减少片外访存的同时，可适配各种数据尺寸以得到最佳并行效率。对比XLA[1]，AStitch最高可以取得2.73倍的性能加速。

背景

近年来，深度学习编译优化已经成为深度学习系统领域最活跃的方向之一。“编译“方法将机器学习DSL（TensorFlow、PyTorch等）自动化地翻译为底层硬件支持的程序，“优化”方法则在翻译（编译）过程中执行一系列的性能优化方案，使得最终生成的程序能够跑得更快更好。

深度学习编译器在生成最终可执行程序之前，往往需要将上层的DSL先翻译为中间层的IR，以方便编译器进行处理。比如，上层的一个LayerNorm算子，会被翻译为中间层的一个子图，该子图包含十多个算子，包括Add、Sub、Mul、Div、Reduce、Broadcast等类型。在一个好的IR定义中，这些“原子”性质的算子可以组合表达任意用户定义的上层计算。

随着算法、硬件以及深度学习系统生态的发展，深度学习系统的瓶颈在不断变化。我们发现新出现的模型的性能瓶颈更多地表现在访存密集型算子上（Element-wise、Reduce等计算）。

一方面，新模型中访存密集型算子的比重越来越多。较早的MLP和CNN以计算密集型算子为主（卷积和矩阵乘），而随着模型结构的发展（比如Transformer开始进入各个领域），LayerNorm、Softmax、GELU以及各种算法工程师自定义的部件使得访存密集型算子在算子数量和执行时间上都开始超过计算密集型算子。另一方面，新硬件的算力提升速度高于访存带宽提升速度，尤其是TensorCore这种“加速器上的加速器”的出现，让算力有了急剧的提升。算力的提升让计算密集型算子的计算时间进一步缩短，使得访存密集型计算的性能问题更加突出。

XLA是最早对访存密集型计算有所关注的SOTA之一（TVM[2]对访存密集型计算的优化方法与其类似），也是工业上最成熟的机器学习编译器之一。在GPU上对XLA进行性能评测时，我们发现其优化后的很多模型（Transformer、ASR等ÿ