CompilerDream: Learning a Compiler World Model for General Code Optimization

📄 arXiv: 2404.16077v4 📥 PDF

作者: Chaoyi Deng, Jialong Wu, Ningya Feng, Jianmin Wang, Mingsheng Long

分类: cs.PL, cs.LG

发布日期: 2024-04-24 (更新: 2026-05-28)

备注: KDD 2025 camera-ready version with extended appendix. Code is available at https://github.com/thuml/CompilerDream. This update additionally fixes an issue in Table 6 where the dataset names in three rows were ordered incorrectly

💡 一句话要点

提出CompilerDream以解决通用代码优化问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 代码优化 编译器设计 强化学习 模型训练 自动调优

📋 核心要点

  1. 现有编译器优化方法依赖固定的优化序列，难以适应不同代码，且寻找最佳序列的算法效率低下。
  2. 提出CompilerDream，通过构建编译器世界模型，利用强化学习生成有效的优化策略，提升代码优化的灵活性和效率。
  3. 实验结果显示，CompilerDream在多种数据集上表现优异，超越LLVM优化，尤其在零-shot泛化能力上具有显著优势。

📝 摘要（中文）

有效的代码优化在编译器中至关重要，其成功主要依赖于优化过程的选择和顺序。现有编译器通常依赖固定的优化序列，而寻找最佳序列的方法往往使用速度较慢的搜索算法或难以推广到未见代码的学习方法。本文提出了CompilerDream，一种基于模型的强化学习方法，旨在实现通用代码优化。CompilerDream包含一个准确模拟优化过程内在特性的编译器世界模型，以及一个在该模型上训练的代理，用于生成有效的优化策略。通过在大规模程序数据集上训练，CompilerDream能够在各种应用场景和源代码语言中作为通用代码优化器。实验结果表明，CompilerDream在自动调优方面表现出色，领先于CompilerGym排行榜，并且在零-shot泛化能力上超越LLVM内置优化和其他先进方法。

🔬 方法详解

问题定义：本文旨在解决编译器中代码优化过程的灵活性和效率问题。现有方法通常依赖固定的优化序列，难以适应不同的代码结构，且寻找最佳序列的算法往往效率低下，无法处理未见代码。

核心思路：CompilerDream的核心思路是构建一个编译器世界模型，通过强化学习训练一个代理，使其能够生成适应不同代码的优化策略。这种设计使得优化过程不仅依赖于固定序列，而是能够根据代码特性动态调整优化策略。

技术框架：CompilerDream的整体架构包括两个主要模块：编译器世界模型和强化学习代理。编译器世界模型负责模拟优化过程的内在特性，而强化学习代理则在该模型上进行训练，以生成有效的优化策略。

关键创新：最重要的技术创新在于引入了编译器世界模型，使得优化过程能够在未见代码上进行零-shot泛化，显著提升了优化的灵活性和效果。这一创新与传统方法的固定序列优化形成了本质区别。

关键设计：在模型训练中，采用了大规模程序数据集，并设计了适应性强的损失函数，以确保代理能够在多样化的代码场景中学习到有效的优化策略。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，CompilerDream在自动调优任务中表现优异，领先于CompilerGym排行榜，并在多个数据集上实现了零-shot泛化能力，超越LLVM内置优化和其他先进方法，展示了显著的性能提升。

🎯 应用场景

CompilerDream的潜在应用领域包括各种编程语言的编译器优化、软件性能提升以及自动化代码调优等。其实际价值在于能够根据不同代码特性动态生成优化策略，未来可能对编程语言的编译器设计和软件开发流程产生深远影响。

📄 摘要（原文）

Effective code optimization in compilers is crucial for computer and software engineering. The success of these optimizations primarily depends on the selection and ordering of the optimization passes applied to the code. While most compilers rely on a fixed sequence of optimization passes, current methods to find the optimal sequence either employ impractically slow search algorithms or learning methods that struggle to generalize to code unseen during training. We introduce CompilerDream, a model-based reinforcement learning approach to general code optimization. CompilerDream comprises a compiler world model that accurately simulates the intrinsic properties of optimization passes and an agent trained on this model to produce effective optimization strategies. By training on a large-scale program dataset, CompilerDream is equipped to serve as a general code optimizer across various application scenarios and source-code languages. Our extensive experiments first highlight CompilerDream's strong optimization capabilities for autotuning, where it leads the CompilerGym leaderboard. More importantly, the zero-shot generalization ability of large-scale trained compiler world model and agent, excels across diverse datasets, surpassing LLVM's built-in optimizations and other state-of-the-art methods in both settings of value prediction and end-to-end code optimization.