CompilerDream：基于世界模型的通用代码优化方法学术报告

作者及机构
本研究的核心作者包括清华大学软件学院BNRIST的Chaoyi Deng、Jialong Wu（共同一作）、Ningya Feng、Jianmin Wang以及通讯作者Mingsheng Long。研究成果发表于2025年8月的ACM SIGKDD国际会议（KDD ‘25），会议论文标题为《CompilerDream: Learning a Compiler World Model for General Code Optimization》。

学术背景
编译器优化是计算机与软件工程的核心课题，其性能提升的关键在于优化阶段（optimization passes）的选择与顺序排列。传统编译器依赖固定优化序列（如-O1、-O2），而现有方法（如搜索算法或机器学习）面临两大瓶颈：搜索耗时过长或泛化能力不足。为此，研究团队提出CompilerDream——首个基于世界模型（world model）的通用代码优化框架，旨在通过模拟编译器行为与奖励平滑技术，实现高效、可泛化的代码优化。研究领域涉及机器学习（模型开发与分析）与编译器优化（LLVM阶段排序问题）。

研究流程与方法
1. 问题建模与空间设计
- 将编译器优化过程建模为部分可观测马尔可夫决策过程（POMDP），定义状态空间（IR中间表示）、动作空间（124个LLVM优化阶段）和奖励函数（基于IR指令数减少的归一化改进）。
- 观测空间融合56维Autophase特征（统计IR代码特性）与42维动作直方图（记录阶段调用频率），通过归一化提升一致性。

1. 编译器世界模型构建

   • 核心组件：
     
     • 表示模型：将真实编译器状态编码为神经编译器状态（neural compiler state）。
       
     • 转移模型：预测下一状态动态。
       
     • 观测解码器与奖励解码器：分别重构IR特征与优化收益。
       
   • 奖励平滑技术：针对优化收益稀疏性（多数阶段无即时收益），引入指数衰减平滑（α=0.6），使非零奖励覆盖更多步骤，提升模型训练稳定性。
     

2. 优化智能体设计

   • 价值预测智能体：通过世界模型模拟优化序列的累积收益，指导搜索策略。在固定核心集（core set）中评估50组序列，选择最优前缀。
     
   • 强化学习智能体（RL）：基于Dreamer算法框架，通过演员-评论家架构在模拟环境中学习优化策略，支持端到端序列生成。
     

3. 数据准备与训练

   • 使用CodeContests数据集（11万C++程序），筛选高质量、多样化代码，避免合成数据（如Csmith）的泛化缺陷。
     
   • 训练流程：
     
     • 世界模型通过真实编译器交互数据学习动态与奖励预测。
       
     • 智能体在世界模型模拟中训练，大幅降低真实编译器调用开销。
       

主要结果
1. 自动调优性能
- 在CompileGym排行榜任务中，CompilerDream以单次试验实现1.068×代码体积缩减，超越随机搜索（1.062×）与PPO（1.047×）。结合引导搜索后，以60.8秒平均耗时达到1.073×优化，刷新榜单记录。

1. 价值预测准确性

   • 在未见过的程序集（如NPB、Cbench）上，世界模型预测的优化序列收益接近理论上限（Oracle），平均提升1.038–1.140倍，优于Coreset-NVP等现有方法。
     

2. 零样本泛化能力

   • RL智能体在跨领域测试中表现优异：
     
     • 生产级代码：TensorFlow内核优化达1.032×，OpenCV达1.092×。
       
     • 新语言：Objective-C程序平均缩减1.027×，部分案例达2.87×。
       
     • AI生成代码：在FormAI数据集（10.9万程序）上，超越LLVM内置优化（-Oz）的比例显著高于PPO基线。
       

3. 模型分析验证

   • 奖励平滑必要性：去除平滑后，优化性能下降5–10%，尤其在Mibench等稀疏奖励场景。
     
   • 数据质量影响：CodeContests训练模型在人类编写代码（Cbench）上泛化能力比Csmith高20%。
     

结论与价值
1. 科学价值
- 首次将世界模型引入编译器优化，通过动态模拟与奖励设计解决了传统方法的泛化瓶颈。
- 提出的奖励平滑技术为稀疏奖励环境下的模型训练提供了通用解决方案。

1. 应用价值
   
   • 支持多种优化代理（价值预测/RL），适配自动调优、实时优化等场景。
     
   • 在嵌入式（代码体积敏感）、AI生成代码优化等新兴领域具有直接实用性。
     

研究亮点
1. 方法创新：
- 世界模型作为编译器模拟器，替代昂贵真实交互，训练效率提升10倍。
- 多智能体框架统一了监督学习与强化学习的优化范式。

1. 性能突破：

   • 零样本泛化能力覆盖从算法库（BLAS）到工业级代码（TensorFlow）。
     
   • 在CompileGym排行榜和跨语言任务中均达到SOTA。
     

2. 开源贡献：

   • 代码与数据集公开于GitHub（https://github.com/thuml/compilerdream），推动AI与编译技术的交叉研究。
     

局限与展望
当前研究聚焦代码体积优化，未来可扩展至执行时间等多目标优化。此外，结合大语言模型（LLM）解析代码语义、支持更长序列预测是潜在方向。