共计 5115 个字符，预计需要花费 13 分钟才能阅读完成。
简介：近日，对于机器学习访存密集计算编译优化框架的论文《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… 本文为阿里云原创内容，未经容许不得转载。
