[论文解读] DISC: A Dynamic Shape Compiler for Machine Learning Workloads

Kai Zhu, Wenyi Zhao+9 more|arXiv (Cornell University)|Mar 9, 2021

Parallel Computing and Optimization Techniques参考文献 9被引用 4

一句话总结

DISC 是一种用于机器学习工作负载的动态形状编译器，通过在 MLIR 上扩展出完全动态的形状中间表示（DHLO），在编译时生成高效的运行时控制流以避免解释执行开销，并通过形状传播与约束收集实现主机-设备联合优化。在如 Transformer 等动态形状工作负载上，其性能相比 TensorFlow/PyTorch 提升高达 3.3×，相比 Nimble 提升 1.8×。

ABSTRACT

Many recent machine learning models show dynamic shape characteristics. However, existing AI compiler optimization systems suffer a lot from problems brought by dynamic shape models, including compilation overhead, memory usage, optimization pipeline and deployment complexity. This paper provides a compiler system to natively support optimization for dynamic shape workloads, named DISC. DISC enriches a set of IR to form a fully dynamic shape representation. It generates the runtime flow at compile time to support processing dynamic shape based logic, which avoids the interpretation overhead at runtime and enlarges the opportunity of host-device co-optimization. It addresses the kernel fusion problem of dynamic shapes with shape propagation and constraints collecting methods. This is the first work to demonstrate how to build an end-to-end dynamic shape compiler based on MLIR infrastructure. Experiments show that DISC achieves up to 3.3x speedup than TensorFlow/PyTorch, and 1.8x than Nimble.

研究动机与目标

• 解决动态形状机器学习工作负载中的性能瓶颈，现有编译器普遍存在高编译开销、内存膨胀和融合机会有限的问题。
• 克服 XLA 及类似系统中静态形状优化的局限性，这些系统无法高效处理动态形状，通常会为这类工作负载禁用融合。
• 提供一种原生支持动态形状计算的解决方案，无需依赖运行时解释或基于填充的变通方法。
• 通过在编译时生成控制流代码，实现在动态形状场景下的高效内核融合与主机-设备联合优化。
• 展示基于 MLIR 基础设施的端到端动态形状编译，支持多种框架，并在可能时自动回退至静态优化。

提出的方法

• 扩展 MLIR 的 HLO 语法，创建 DHLO（动态 HLO），一种完全动态的形状表示，原生支持动态形状计算，无需从零开始设计新 IR。
• 在编译时生成完整的运行时控制流（包括形状推断、缓冲区管理及主机端逻辑），避免虚拟机解释执行开销，支持主机-设备联合优化。
• 通过生产者与消费者之间的形状传播，识别相邻操作中具有兼容动态形状的融合候选。
• 在降低至 DHLO 的过程中收集形状约束，实现在编译时无需完整形状信息的前提下，通过形状兼容性推理实现更大规模的融合。
• 通过自动回退至静态编译，支持静态与动态优化的混合模式，当子图具有已知静态形状时启用。
• 通过 DHLO 中心与 TensorFlow 和 PyTorch 集成，实现动态形状模型的端到端编译。

实验结果

研究问题

• RQ1 机器学习编译器如何在不依赖运行时解释的前提下，高效表示并优化动态形状计算？
• RQ2 在编译时无法获取完整形状信息的情况下，可采用哪些技术实现动态形状工作负载中的内核融合？
• RQ3 如何实现主机-设备控制流的联合优化，以减少动态形状执行的运行时开销？
• RQ4 在保持灵活性的同时，动态形状编译器在多大程度上可实现与静态编译器相当的性能？
• RQ5 基于 MLIR 基础设施的编译系统能否原生支持端到端的动态形状编译，并将运行时开销降至最低？

主要发现

• 在动态形状工作负载上，DISC 相较于 TensorFlow 和 PyTorch 实现平均 2.27× 的性能提升，部分模型最高达 3.3×。
• 在 Transformer 中内存密集型操作上，DISC 相较于 Nimble 实现 2.61× 的性能提升，主要得益于形状提示与约束带来的更高效融合。
• 尽管内核数量略有减少，DISC 将运行时控制流的 CPU 时间降低至 Nimble 的 36.6%，主要得益于高效的代码生成与联合优化。
• DISC 在平均性能上相较 Nimble 实现 1.8× 的提升，证明了编译时生成控制流相比基于 VM 的解释执行具有显著优势。
• 当禁用回退至静态优化功能时，DISC 在平均性能上达到静态编译的 85%，范围在 74.5% 至 91.4% 之间，表明在缺乏完整形状知识时，融合机会受限导致性能差距。
• 在 Transformer 模型上，内核数量从 Nimble 的 13,924 个减少至 DISC 的 10,734 个，表明尽管未以减少内核数为主要性能目标，其融合效率仍得到显著提升。