CompilerDream: Learning a Compiler World Model for General Code Optimization

2024-04-24

AI Summary

Existing compiler optimization techniques suffer from inflexible fixed pipelines, low search efficiency, and poor generalization—hindering cross-language and cross-program transfer. Method: We propose the first generalizable “compiler world model,” framing optimization as a dynamic system of state evolution and transition behavior, decoupling environment simulation from policy learning. Our approach integrates model-based reinforcement learning, program representation learning, and large-scale compilation trajectory modeling, underpinned by an LLVM IR–based simulation environment. Contribution/Results: On the CompilerGym benchmark, our method achieves zero-shot superiority over LLVM’s default optimizations and prior state-of-the-art approaches. It delivers substantial end-to-end performance gains and reduces value prediction error by 37%. Crucially, it is the first to demonstrate zero-shot generalization across both programming languages and diverse programs—marking a foundational advance in adaptive, portable compiler optimization.

Technology Category

• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP
• Machine Learning: Optimization

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

Abstract

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.

Problem

• Finding optimal compiler optimization pass sequences efficiently
• Generalizing learned optimizations to unseen code effectively
• Enhancing zero-shot performance across diverse code datasets

Innovation

• Model-based reinforcement learning for optimization
• Compiler world model simulates optimization passes
• Large-scale training enables zero-shot generalization