一句话：提出 AStitch 编译器，通过四种算子缝合方案、层次化数据复用与自适应线程映射，在多维空间优化内存密集型 ML 计算。
论文指出现代机器学习模型中内存密集型计算正成为关键性能瓶颈：现有 ML 优化编译器在复杂的两级依赖叠加即时（just-in-time）需求下难以有效融合，只能在代价高昂的冗余计算融合与产生大量小 kernel 的不融合之间二选一；同时，面对生产中常见的不规则张量形状，并行度往往偏低。为此作者提出 AStitch，一个面向内存密集型 ML 计算的优化编译器：它系统抽象出四种算子缝合（stitching）方案以覆盖多维优化目标，借助新颖的层次化数据复用处理复杂计算图依赖，并通过自适应线程映射高效支持各种张量形状，最终以即时编译方式为训练与推理同时提供这些优化。AStitch 可作为独立编译器引擎移植到任意版本的 TensorFlow，其思想也可推广到其他框架。在五个生产负载上相较 Google XLA 平均加速 1.84 倍（最高 2.73 倍），并已部署于数千 GPU 的生产集群运行逾 10 个月，每周为 7 万个任务节省约 2 万 GPU 小时。
This work reveals that memory-intensive computation is a rising performance-critical factor in recent machine learning models. Due to a unique set of new challenges, existing ML optimizing compilers cannot perform efficient fusion under complex two-level dependencies combined with just-in-time demand. They face the dilemma of either performing costly fusion due to heavy redundant computation, or skipping fusion which results in massive number of kernels. Furthermore, they often suffer from low parallelism due to the lack of support for real-world production workloads with irregular tensor shapes. To address these rising challenges, we propose AStitch, a machine learning optimizing compiler that opens a new multi-dimensional optimization space for memory-intensive ML computations. It systematically abstracts four operator-stitching schemes while considering multi-dimensional optimization objectives, tackles complex computation graph dependencies with novel hierarchical data reuse, and efficiently processes various tensor shapes via adaptive thread mapping. Finally, AStitch provides just-in-time support incorporating our proposed optimizations for both ML training and inference. Although AStitch serves as a stand-alone compiler engine that is portable to any version of TensorFlow, its basic ideas can be generally applied to other ML frameworks and optimization compilers. Experimental results show that AStitch can achieve an average of 1.84x speedup (up to 2.73x) over the state-of-the-art Google's XLA solution across five production workloads. We also deploy AStitch onto a production cluster for ML workloads with thousands of GPUs. The system has been in operation for more than 10 months and saves about 20,000 GPU hours for 70,000 tasks per week.
@inproceedings{ZhengYZLZZZLYZS22,
title = {AStitch: enabling a new multi-dimensional optimization space for memory-intensive ML training and inference on modern SIMT architectures},
author = {Zhen Zheng and Xuanda Yang and Pengzhan Zhao and Guoping Long and Kai Zhu and Feiwen Zhu and Wenyi Zhao and Xiaoyong Liu and Jun Yang and Jidong Zhai and Shuaiwen Leon Song and Wei Lin},
booktitle = {ASPLOS},
year = {2022},
doi = {10.1145/3503222.3507723}
}