QiMeng-Xpiler:使用神经符号方法对深度学习系统的张量程序进行转编译
Shouyang Dong, Yuanbo Wen, Jun Bi, Di Huang, Jiaming Guo, Jianxing Xu, Ruibai Xu, Xinkai Song, Yifan Hao, Xuehai Zhou, Tianshi Chen, Qi Guo, Yunji Chen
发布日期: 5/12/2025

摘要
异构深度学习系统 (DLS),例如 GPU 和 ASIC,已广泛应用于工业数据中心,这需要为不同的平台开发多个底层张量程序。一个颇具吸引力的解决方案是将一个平台的遗留代码转编译到其他平台。然而,当前的转编译技术要么面临巨大的手动工作量,要么存在功能错误,使得张量程序的“一次编写,随处运行”成为悬而未决的问题。 我们提出了一种新颖的转编译器 QiMeng-Xpiler,它可以通过大型语言模型 (LLM) 和符号程序合成(即神经符号合成)自动跨 DLS 转换张量程序。其关键思想是利用 LLM 强大的代码生成能力,使昂贵的基于搜索的符号程序合成变得易于计算。具体而言,我们提出了多个 LLM 辅助编译过程,通过预定义的元提示符进行程序转换。在每次程序转换过程中,我们都会使用高效的符号程序合成来修复有限规模的错误代码片段。为了实现高性能,我们提出了一种分层自动调优方法,系统地探索转换过程的参数和序列。在 4 个具有不同编程接口的 DLS 上进行的实验(分别为带有 VNNI 的 Intel DL Boost、带有 CUDA 的 NVIDIA GPU、带有 HIP 的 AMD MI 以及带有 BANG 的 Cambricon MLU),结果表明,QiMeng-Xpiler 能够以平均 95% 的准确率正确翻译不同的张量程序,并且翻译后程序的性能比供应商提供的手动优化库高出 2.0 倍。因此,通过对旧版张量程序进行转编译,DLS 的编程效率最高可提升 96.0 倍。