简介： 近日，对于机器学习访存密集计算编译优化框架的论文《AStitch: Enabling A New Multi-Dimensional Optimization Space for Memory-Intensive ML Training and Inference on Modern SIMT Architectures》被零碎畛域顶会ASPLOS 2022接管。作者：郑祯   近日，对于机器学习访存密集计算编译优化框架的论文《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等）依然卡在访存密集型计算的性能上。一方面，访存密集型计算的总耗时与计算密集型算子相当；另一方面，也是更为重要的，微小量的访存密集型算子带来了十分重大的算子调度和GPU kernel调用开销，这些开销在很多模型上曾经超过了计算耗时自身。 XLA所应用的优化办法是kernel fusion，进一步剖析XLA的fusion优化办法，咱们诧异地发现，XLA居然没有利用shared memory来反对算子之间的数据传输，这使得可反对的fusion粒度大大降低。比方对于reduce和它的消费者算子，一个HPC工作者会天然地将reduce的后果放在shared memory上，而后消费者算子从shared memory中读数据并做后续计算；XLA却抉择不将reduce和它的消费者fuse在一起，而是分成两个kernel。对于常见的LayerNorm算子，XLA会用3个kernel来实现，而高效的手写算子则只须要一个kernel。 基于上述发现，如果咱们将shared memory引入机器学习编译优化的fusion中，增大fusion粒度，能够缩小框架调度和kernel调用开销，并缩小片外存储的拜访，进一步优化访存密集型计算的性能。然而，XLA（包含TVM）为什么没这么做呢？进一步剖析后，咱们发现最大的挑战在于，编译器须要自动化地执行优化，手工优化的想法，到了编译器这里，难度往往会被放大。 挑战到了编译器IR层面，机器学习模型的计算图会变得非常复杂，波及两层的数据依赖关系：Element层面和Operator层面。 Element层面是指消费者解决的每个element和生产者生产的每个element之间的依赖关系，比方Broadcast算子会生成一对多的数据依赖关系。

Operator层面是指图层面的算子依赖关系，下图展现了一个访存密集型计算子图的拓扑关系（来自一个Transformer模型），能够看到简单的多对一、一对多依赖关系。

现在的模型类型极为丰盛，变种繁多，计算拓扑图无奈枚举，不同于HPC畛域的手工kernel fusion，让编译器自动化地对任意可能的计算子图做fusion优化是十分具备挑战性。这也是XLA和TVM目前只做激进fusion优化的一个重要起因。 除了简单的计算图之外，输出数据形态的多样性和未知性，也给主动代码生成带来了很大的挑战。不同的tensor形态须要不同的代码生成schedule，以失去较好的并行性。XLA针对惯例的tensor形态提供了代码生成的schedule，但tensor形态的可能性是无穷无尽的，咱们在实在生产中发现，一些非凡的形态会导致XLA产生极大的性能问题。（注：这里探讨的是动态shape的多样性，而非动静shape问题。） 破局艰难和挑战总是让对技术有极致谋求的人兴奋。 Fusion问题，外围是代码生成能力问题，能力越大，fusion越好。从两个层面的依赖关系登程，代码生成须要同时思考硬件存储档次和并行度：算子之间传递数据须要抉择适合的存储媒介，每个算子须要抉择适合的并行策略，这两者是互相交杂的，须要协同思考。咱们将咱们提出的Fusion技术称为stitch，这样能够更直观地表白算子通过层次化的存储媒介“缝合”在一起这一动作。 从存储媒介登程，联合并行度的思考，咱们将stitch策略形象为四种：无依赖（Independent）、本地的（Local）、区域的（Regional）和全局的（Global）。 XLA所反对的是Local的策略，只反对通过寄存器传递数据，算子的并行策略被离开来思考。一些Fusion工作能够反对Independent依赖，次要包含Kernel Packing等。AStitch拓展了fusion的优化空间，Regional策略将数据存储在shared memory中，反对GPU thread block locality，Global策略将数据存储在global memory中，反对全局的locality。其中，并行度的思考在于，上述每种stitch策略都须要特定locality的反对，locality越是部分，对并行策略的可能限度越大。从Local，到Regional，再到Global，locality的限度逐步升高，对并行策略的限度也随之逐步勾销。须要阐明的是，Global的策略须要GPU kernel所有线程的全局同步，其相当于将kernel之间的隐式的同步inline到kernel外部来进行，缩小一次CPU（包含框架和驱动）和GPU之间的切换。

对于并行度，为访存密集型计算做shape-aware的并行代码生成（Adaptive Thread Mapping）是AStitch的一个重要奉献。其根本思维是，基于SIMT的架构，合并小的（Task Packing），拆分大的（Task Splitting），以失去适合的CUDA thread block大小和数量。 值得阐明的是，对于Global的策略，为了实现全局同步，须要保障总的CUDA thread block数量不大于一个wave能够调度的最大数量，AStitch通过做纵向的Task Packing来达到这一要求，这相当于将本来由硬件调度器调度的工作转变为软件上事后调度好的工作，实际上还有机会节俭调度开销。

对于主动代码生成，穷举每个算子的stitching策略和并行策略是不事实的，咱们提出了一种先分组解决部分代码生成，再聚合生成全局代码的办法。一个重要的insight是，轻量级的element-wise的计算只须要追随其consumer的代码生成schedule即可，最终通过Local策略进行数据传输；而reduce等简单计算则须要优先思考并行度的优化，再去思考与其consumer的stitch策略。 基于此，咱们根据算子类型进行分组，reduce及跟有broadcast的expensive element-wise（Power, Sqrt等）作为每一个组的dominant，每个dominant的间接的或间接的producer被分到其对应的组中。组内的数据流传都是通过Local策略进行，dominant生成本人的并行策略schedule后，流传给组内的其余算子；最初检查组间的数据局部性，为两个组边缘的算子抉择可达到的locality最优的stitch策略。这种办法简洁无效地解决了搜寻空间爆炸的问题，在此过程中，咱们还有一些locality和parallelism之间的衡量，其中更多的细节请浏览咱们ASPLOS 2022的论文。

结语AStitch的雏形是咱们晚期的工作FusionStitching[4]，FusionStitching最早将shared memory引入了访存密集算子的编译优化主动fusion，并采纳cost model来做fusion和代码生成的一系列决策。随着深刻优化，咱们发现能够将kernel之间的隐式同步给inline到kernel外部，并通过global memory做stitch，防止CPU和GPU之间的无谓的切换，进一步的，联合Adaptive Thread Mapping办法解决并行策略问题，能够造成弱小的代码生成能力。至此，咱们曾经不须要cost model来领导fusion决策和代码生成，咱们能够stitch所有，咱们的代码生成也具备自适应能力。 最近，咱们惊喜地发现新版本的TensorRT[3]开始反对应用shared memory做数据媒介来生成fusion了。咱们期待越来越多的工作开始在访存密集型计算的性能优化方面发力，咱们也心愿AStitch以及晚期的FusionStitching的工作能够影响到宽泛的深度学习零碎和编译优化参与者！ 论文地址：https://dl.acm.org/doi/10.114... 参考文献：[1] TensorFlow XLA. https://www.tensorflow.org/xla[2] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al. TVM: An automated end-to-end optimizing compiler for deep learning. OSDI'18.[3] TensorRT. https://developer.nvidia.com/...[4] Zhen Zheng, Pengzhan Zhao, Guoping Long, Feiwen Zhu, Kai Zhu, Wenyi Zhao, Lansong Diao, Jun Yang, and Wei Lin. FusionStitching: boosting memory intensive computations for deep learning workloads. arXiv preprint.原文链接：https://click.aliyun.com/m/10...本文为阿里云原创内容，未经容许不得转载。