一句话：提出多智能体系统 GraphGlue，用 LLM 修复导致深度学习编译器计算图断裂的代码，从而提升推理性能与显存效率。
深度学习编译器通过将 DNN 程序捕获为算子级计算图来做优化，但开发者程序中复杂的 Python 语言特性常使编译器无法把整个程序识别为完整计算图，导致性能次优。作者分析发现，实际的图捕获失败往往只涉及少数几行代码，因此认为可通过代码修复而非大规模改造编译器来解决该问题。为此提出 GraphGlue：一个利用 LLM 修复并优化 DNN 程序以满足编译器要求的多智能体系统，包含图断裂原因挖掘（GCM）以定位隐藏的断裂成因并指导 LLM 修复，以及拒绝采样自纠正（SRS）在代码调试与重新生成之间交替，避免错误初始优化策略带来的无效反馈。实验表明，相比直接使用 TorchDynamo 最高获得 2.19 倍加速（平均 1.23 倍），相比最先进的 AI 编译器前端最高节省 15.77 倍显存（平均 8.74 倍），并在 1411 个真实用户程序上成功优化 92.63%。
Deep learning compilers optimize DNN program execution by capturing them as operator-based computation graphs. However, developers’ deep learning programs often contain complex Python language features that prevent compilers from recognizing the entire program as a complete computation graph, resulting in sub-optimal performance. Our analysis reveals that actual capture failures involve only a few lines of code, we believe this problem can be addressed through code repair rather than extensive compiler improvements. To address this challenge, we introduce GraphGlue, a multi-agent system that leverages LLMs to repair and optimize DNN programs for compiler requirements, thereby maximizing the performance benefits of deep learning compilers in inference scenarios. GraphGlue employs (1) graph-break cause mining (GCM) to identify hidden causes of computation graph breaks and facilitate LLM-based repair, and (2) self-correction with reject sampling (SRS) to alternate between code debugging and regeneration, effectively avoiding ineffective feedback attempts caused by incorrect initial optimization strategies. Experimental results demonstrate that programs optimized by GraphGlue achieve up to 2.19x (1.23x on average) speedup compared to using TorchDynamo directly, and deliver up to 15.77x (8.74x on average) memory savings compared to state-of-the-art AI compiler frontends. GraphGlue exhibits strong generalization capabilities across 1,411 real-world user programs, successfully optimizing 92.63% of them. Code is available at https://github.com/Jamesswang/GraphGlue.
DOI 原文 ·
@inproceedings{paperbot9503,
title = {Fixing Broken Graphs: LLM-Powered Automatic Code Optimization for DNN Programs},
author = {Haotian Wang and Yicheng Sui and Yudong Xie and Yicong Liu and Yufei Sun and Changqing Shi and Yuzhi Zhang},
booktitle = {ASE},
year = {2025},
doi = {10.1109/ase63991.2025.00144}
}