DISC: A Dynamic Shape Compiler for Machine Learning Workloads

TL;DR

DISC编译器实现动态形状优化,速度提升达3.3倍。

cs.DC 🔴 高级 2021-03-09 4 次浏览
Kai Zhu Wenyi Zhao Zhen Zheng Tianyou Guo Pengzhan Zhao Feiwen Zhu Junjie Bai Jun Yang Xiaoyong Liu Lansong Diao Wei Lin
动态形状 AI编译器 内核融合 MLIR 性能优化

核心发现

方法论

DISC利用MLIR框架设计动态形状编译器,采用DHLO扩展HLO方言以支持动态形状。通过编译时生成运行时流,避免解释开销,并实现主机设备协同优化。内核融合通过形状传播和约束收集实现。

关键结果

  • DISC在TensorFlow/PyTorch上实现了最高3.3倍的速度提升,平均速度提升为2.27倍。与Nimble相比,DISC在Transformer模型上实现了1.8倍的速度提升。
  • DISC在6个流行模型上平均速度提升2.27倍,显著优于现有方法。
  • 通过形状约束和传播实现了更高效的内核融合。

研究意义

DISC为动态形状工作负载提供了高效的编译解决方案,显著降低了编译开销和内存使用,简化了优化管道和部署复杂性。它为学术界和工业界提供了新的优化思路,解决了长期存在的动态形状模型优化难题。

技术贡献

DISC首次展示了如何基于MLIR基础设施构建端到端动态形状编译器。它提出了完全动态IR和编译时生成的运行时流设计,解决了缺乏完整形状信息的内核融合问题,并支持多种机器学习框架。

新颖性

DISC是第一个展示如何高效支持动态形状的编译器,采用了全动态IR和编译时生成的运行时流设计,与现有方法相比具有显著创新。

局限性

  • DISC在处理动态秩时可能表现不佳,因为动态秩在实践中不常见。
  • 对于特定形状的优化可能不如静态编译器。
  • 需要进一步研究如何在更多设备上实现优化。

未来方向

未来可以探索DISC在更多设备上的应用,以及如何进一步优化动态形状处理的效率。作者建议社区继续研究动态形状问题的解决方案。

AI 总览摘要

动态形状模型在现代机器学习中越来越常见,但现有编译器在处理这些模型时面临编译开销和内存使用等问题。DISC编译器通过MLIR基础设施提供了一个解决方案,支持动态形状的优化。它采用DHLO扩展HLO方言,生成运行时流以避免解释开销,并实现主机设备协同优化。

实验结果显示,DISC在TensorFlow/PyTorch上实现了最高3.3倍的速度提升,平均速度提升为2.27倍,显著优于现有方法。它解决了内核融合问题,通过形状传播和约束收集实现了更高效的优化。

DISC的设计为动态形状模型的优化提供了新的思路,具有重要的学术和工业意义。未来可以探索其在更多设备上的应用,以及如何进一步优化动态形状处理的效率。

深度解读

原文摘要

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.

cs.DC