--- title: Compiler Pass Ordering RL Environment emoji: ⚙️ colorFrom: gray colorTo: blue sdk: docker pinned: false app_port: 8000 base_path: /web tags: - openenv --- # Compiler Pass Ordering RL Environment An OpenEnv environment that simulates **compiler optimization pass ordering** — a real-world task performed by production compilers (GCC, LLVM) every time they build software. The agent must select a sequence of optimization passes to apply to a program's Intermediate Representation (IR) to minimize estimated runtime cost. This is the same class of problem addressed by DeepMind's MLGO paper (2022), which showed RL finds better pass orderings than the hand-tuned heuristics used in production compilers today. --- ## Motivation When a compiler optimizes code, it applies a sequence of transformations called **passes**: dead code elimination, loop unrolling, vectorization, function inlining, etc. There are ~50 such passes in LLVM. The order you apply them matters enormously — applying pass A before B can unlock optimizations that B before A cannot. **Why RL is necessary here (not greedy or heuristics):** - Greedy selection (always pick the pass with highest immediate benefit) achieves ~20% cost reduction. - An RL agent that learns prerequisite chains (e.g., `alias_analysis → vectorization` gives 7x effectiveness) achieves ~42-50%. - The search space is factorial (15! orderings), and the reward for an "enabling" pass only materializes 3-4 steps later — exactly the credit assignment problem RL is built for. --- ## Environment Design ### Prerequisite Gates The core mechanism that makes RL necessary. Each high-value pass has prerequisites: | Pass | Prerequisites Required | Synergy Multiplier | Without Prerequisites | |---|---|---|---| | vectorization | alias_analysis + dead_code_elimination | 7.0x | 0.3x (penalty) | | function_inlining | interprocedural_analysis | 5.0x | 0.3x | | memory_coalescing | alias_analysis | 3.5x | 0.3x | | strength_reduction | alias_analysis + constant_folding | 3.0x | 0.3x | | instruction_scheduling | register_allocation | 2.5x | 0.3x | | loop_unrolling | loop_invariant_motion | 2.0x | 0.3x | A greedy agent never applies `alias_analysis` (base effect = 0.01) and so never unlocks `vectorization`. An RL agent learns to sacrifice early reward to enable these chains. --- ## Action Space ```python class CompilerOptAction(Action): pass_id: int # Integer in [0, 14] — which pass to apply task_id: int # 1 = easy, 2 = medium, 3 = hard ``` **Available passes:** | ID | Pass Name | Base Effect | Role | |---|---|---|---| | 0 | dead_code_elimination | 0.08 | Enabler for vectorization | | 1 | constant_folding | 0.06 | Enabler for strength_reduction | | 2 | loop_unrolling | 0.10 | High value, needs LIM first | | 3 | function_inlining | 0.03 | High potential, needs interproc | | 4 | vectorization | 0.02 | Highest potential (7x with prereqs) | | 5 | loop_invariant_motion | 0.06 | Enabler for loop_unrolling | | 6 | strength_reduction | 0.05 | Good with alias + const_fold | | 7 | common_subexpr_elimination | 0.07 | Standalone value | | 8 | tail_call_optimization | 0.04 | Standalone value | | 9 | branch_prediction_hints | 0.03 | Standalone value | | 10 | register_allocation | 0.05 | Enabler for instr_scheduling | | 11 | instruction_scheduling | 0.04 | Needs reg_alloc first | | 12 | memory_coalescing | 0.06 | Needs alias first | | 13 | alias_analysis | 0.01 | **Key enabler** — low base, high unlock | | 14 | interprocedural_analysis | 0.01 | **Key enabler** — low base, high unlock | --- ## Observation Space ```python class CompilerOptObservation(Observation): # Cost tracking estimated_cost: float # Current runtime cost estimate baseline_cost: float # Cost before any optimization # Program IR features (static per episode) num_instructions: int num_loops: int num_branches: int num_functions: int loop_depth: int program_type: str # e.g. "vectorizable", "loop_heavy" # Episode progress passes_applied: List[int] # Ordered history of applied passes passes_available: List[int] step_count: int max_steps: int # 10 # Synergy state synergy_state: List[float] # Per-pass current effectiveness multiplier # Task info task_id: int task_description: str # Results done: bool reward: float improvement_pct: float # Total % cost reduction from baseline last_pass_name: str grader_score: float # 0.0–1.0, populated when done=True ``` --- ## Tasks ### Task 1 — Easy **Goal:** Optimize a vectorizable or loop-heavy program by discovering the primary synergy chain. **What the agent must learn:** Apply `alias_analysis` (low immediate value) then `dead_code_elimination`, then `vectorization` fires at 7x base effectiveness. **Grader thresholds:** - Score 1.0: ≥35% cost reduction - Score 0.7: ≥25% cost reduction - Score 0.4: ≥15% cost reduction - Score 0.0: <15% **Greedy baseline:** ~18-20% | **RL target:** ~38-45% --- ### Task 2 — Medium **Goal:** Optimize a compute-heavy or inlining-heavy program by discovering **two independent synergy chains**. **What the agent must learn:** - Chain 1: `alias_analysis + DCE → vectorization` (7x) - Chain 2: `interprocedural_analysis → function_inlining` (5x) Both chains must be activated within 10 steps. **Grader thresholds:** - Score 1.0: ≥45% cost reduction - Score 0.7: ≥35% cost reduction - Score 0.4: ≥22% cost reduction - Score 0.0: <22% **Greedy baseline:** ~20-24% | **RL target:** ~42-50% --- ### Task 3 — Hard **Goal:** Optimize any program type. Program is sampled randomly from 7 templates. The agent must generalize across program types with different optimal orderings. **What the agent must learn:** Generalized pass ordering policy. Different programs reward different chains. A `vectorizable` program rewards chain 1; a `recursive` program rewards chain 2; a `loop_heavy` program rewards `LIM → loop_unrolling`. **Grader thresholds:** - Score 1.0: ≥52% cost reduction - Score 0.7: ≥42% cost reduction - Score 0.4: ≥28% cost reduction - Score 0.0: <28% **Greedy baseline:** ~19-23% | **RL target:** ~45-52% --- ## Reward Function Reward is provided at every step (not just terminal), giving a dense signal across the full trajectory: ``` Per-step reward = marginal_improvement - 0.02 (step penalty) marginal_improvement = (cost_before - cost_after) / baseline_cost Terminal bonus: improvement >= 50%: +0.6 improvement >= 40%: +0.4 improvement >= 30%: +0.2 improvement >= 20%: +0.05 improvement <= 0%: -0.2 Penalties: Re-applying an already-used pass: -0.3 Applying pass without prerequisites: reward naturally low (0.3x effectiveness) ``` The step penalty encourages the agent to find efficient sequences (fewer passes to reach the same improvement), matching real compiler efficiency goals. --- ## Grader The grader is called at episode end and returns a score in [0.0, 1.0]: ```python score = grade_by_threshold(improvement_pct, task_id) score -= 0.05 * max(0, steps_used - minimum_steps_needed) ``` The efficiency deduction (0.05 per extra step) rewards agents that find the optimal chain quickly over those that brute-force all 10 steps. --- ## Baseline Scores Measured using `gpt-4o-mini` via the OpenAI API, 5 episodes per task: | Task | Avg Score | Avg Improvement | Best Score | |---|---|---|---| | Task 1 (Easy) | ~0.42 | ~28% | ~0.70 | | Task 2 (Medium) | ~0.28 | ~24% | ~0.40 | | Task 3 (Hard) | ~0.21 | ~21% | ~0.40 | | **Overall** | **~0.30** | **~24%** | — | The LLM agent partially discovers the alias → vectorization chain (~30% of episodes) but rarely sequences both chains correctly in Task 2/3. --- ## Setup & Usage ### Prerequisites - Python 3.10+ - Docker (for containerized deployment) ### Local Development (no Docker) ```bash # Clone / navigate to project cd compiler_opt_env # Install pip install -e . # Start server uvicorn server.app:app --reload --host 0.0.0.0 --port 8000 # In another terminal — run sanity test python test_env.py # Run LLM baseline export OPENAI_API_KEY=your_key python baseline_agent.py --base-url http://localhost:8000 --episodes 3 ``` ### Docker ```bash # Build docker build -t compiler-opt-env:latest -f server/Dockerfile . # Run docker run -p 8000:8000 compiler-opt-env:latest # Web UI open http://localhost:8000/web ``` ### Programmatic Usage ```python from compiler_opt_env import CompilerOptAction, CompilerOptEnv from compiler_opt_env.models import TASK_EASY, TASK_MEDIUM, TASK_HARD with CompilerOptEnv(base_url="http://localhost:8000").sync() as env: # Task 1: easy obs = env.reset() # Optimal sequence for Task 1: for pass_id in [13, 0, 4]: # alias → DCE → vectorization result = env.step(CompilerOptAction(pass_id=pass_id, task_id=TASK_EASY)) print(f"{result.observation.last_pass_name}: {result.observation.improvement_pct:.1f}%") print(f"Grader score: {result.observation.grader_score}") ``` --- ## API Endpoints | Endpoint | Method | Description | |---|---|---| | `/reset` | POST | Start new episode | | `/step` | POST | Apply an action | | `/state` | GET | Current episode metadata | | `/schema` | GET | Action/observation schemas | | `/health` | GET | Health check | | `/web` | GET | Interactive web UI | | `/docs` | GET | Swagger API documentation | | `/ws` | WS | WebSocket for low-latency sessions | --- ## Project Structure ``` compiler_opt_env/ ├── __init__.py # Exports: CompilerOptEnv, CompilerOptAction, CompilerOptObservation ├── models.py # Action, Observation, constants, task IDs ├── client.py # CompilerOptEnv(EnvClient) — HTTP/WebSocket client ├── baseline_agent.py # LLM baseline using OpenAI API ├── openenv.yaml # OpenEnv manifest ├── pyproject.toml # Dependencies ├── README.md # This file └── server/ ├── __init__.py ├── compiler_opt_env_environment.py # Core RL logic: passes, synergy, grader ├── app.py # FastAPI app └── Dockerfile # Container definition ``` --- ## References - [DeepMind MLGO: A Machine Learning Guided Compiler Optimizations Framework](https://arxiv.org/abs/2101.04808) - [Meta OpenEnv](https://github.com/meta-pytorch/OpenEnv) - [LLVM Optimization Passes](https://llvm.org/docs/Passes.html)