QiMeng-Xpiler: Transcompiling Tensor Programs for Deep Learning Systems with a Neural-Symbolic Approach
作者: 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
分类: cs.CL, cs.LG, cs.PL
发布日期: 2025-05-04
备注: Accepted to OSDI 2025
💡 一句话要点
QiMeng-Xpiler:利用神经-符号方法转译张量程序,实现深度学习系统跨平台部署
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 张量程序 跨平台转译 深度学习系统 大型语言模型 符号程序合成
📋 核心要点
- 异构深度学习系统需要为不同平台编写底层张量程序,现有转译技术面临人工成本高或功能不准确的挑战。
- QiMeng-Xpiler利用LLM的代码生成能力辅助符号程序合成,实现跨平台张量程序的自动转译。
- 实验表明,QiMeng-Xpiler在多个DLS上实现了95%的转译准确率,并获得了高达2.0倍的性能提升。
📝 摘要(中文)
异构深度学习系统(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倍。
🔬 方法详解
问题定义:论文旨在解决异构深度学习系统(DLS)中张量程序跨平台移植的问题。现有方法要么依赖大量人工优化,成本高昂;要么自动转译的准确性不足,无法保证功能正确性。这使得“一次编写,随处运行”的目标难以实现。
核心思路:论文的核心思路是结合大型语言模型(LLM)和符号程序合成的优势,构建一个神经-符号转译器。LLM负责生成候选代码,降低符号程序合成的搜索空间,使其在计算上可行。符号程序合成则负责修复LLM生成的错误代码片段,保证功能正确性。
技术框架:QiMeng-Xpiler的整体框架包含以下几个主要阶段:1) LLM辅助的程序转换:通过预定义的元提示,利用LLM进行程序转换,生成候选代码。2) 符号程序合成:对LLM生成的代码片段进行验证和修复,确保功能正确性。3) 分层自动调优:系统地探索转换过程的参数和序列,以获得最佳性能。
关键创新:论文的关键创新在于将LLM和符号程序合成相结合,构建了一个神经-符号转译器。LLM强大的代码生成能力降低了符号程序合成的搜索空间,使其能够处理更复杂的程序。同时,符号程序合成保证了转译后的程序的功能正确性。
关键设计:论文提出了多个LLM辅助的编译过程,每个过程都包含一个预定义的元提示,用于指导LLM生成候选代码。此外,论文还提出了一种分层自动调优方法,用于系统地探索转换过程的参数和序列。具体的参数设置、损失函数、网络结构等技术细节在论文中未详细描述,属于未知信息。
🖼️ 关键图片
📊 实验亮点
实验结果表明,QiMeng-Xpiler在四个不同的DLS平台上实现了平均95%的转译准确率。与供应商提供的手动优化库相比,转译后的程序性能提升高达2.0倍。通过转译遗留张量程序,DLS的编程效率提高了96.0倍。
🎯 应用场景
QiMeng-Xpiler可应用于各种异构深度学习系统,如GPU、ASIC等,降低了深度学习模型的部署成本和开发难度。通过自动转译遗留代码,可以显著提高编程效率,加速AI应用的落地。该研究有望推动深度学习系统软件生态的繁荣。
📄 摘要(原文)
Heterogeneous deep learning systems (DLS) such as GPUs and ASICs have been widely deployed in industrial data centers, which requires to develop multiple low-level tensor programs for different platforms. An attractive solution to relieve the programming burden is to transcompile the legacy code of one platform to others. However, current transcompilation techniques struggle with either tremendous manual efforts or functional incorrectness, rendering "Write Once, Run Anywhere" of tensor programs an open question. We propose a novel transcompiler, i.e., QiMeng-Xpiler, for automatically translating tensor programs across DLS via both large language models (LLMs) and symbolic program synthesis, i.e., neural-symbolic synthesis. The key insight is leveraging the powerful code generation ability of LLM to make costly search-based symbolic synthesis computationally tractable. Concretely, we propose multiple LLM-assisted compilation passes via pre-defined meta-prompts for program transformation. During each program transformation, efficient symbolic program synthesis is employed to repair incorrect code snippets with a limited scale. To attain high performance, we propose a hierarchical auto-tuning approach to systematically explore both the parameters and sequences of transformation passes. Experiments on 4 DLS with distinct programming interfaces, i.e., Intel DL Boost with VNNI, NVIDIA GPU with CUDA, AMD MI with HIP, and Cambricon MLU with BANG, demonstrate that QiMeng-Xpiler correctly translates different tensor programs at the accuracy of 95% on average, and the performance of translated programs achieves up to 2.0x over vendor-provided manually-optimized libraries. As a result, the programming productivity of DLS is improved by up to 96.0x via transcompiling legacy tensor programs.