README.md

# RL-Based ARM Compiler Optimization — Comprehensive Research Wiki

> **Goal**: Train an LLM with reinforcement learning (PPO/GRPO) to generate optimized AArch64/ARM assembly code that outperforms `gcc -O3`, using compiler feedback (correctness + speedup) as the reward signal.

> **Last updated**: April 2026 | **Primary recipe**: SuperCoder (arxiv:2505.11480) adapted for ARM

---

## Table of Contents

1. [Executive Summary](#1-executive-summary)
2. [Landscape & Key Papers](#2-landscape--key-papers)
3. [Recipe 1: SuperCoder — RL Assembly Superoptimization (SOTA)](#3-recipe-1-supercoder--rl-assembly-superoptimization-sota)
4. [Recipe 2: Meta LLM Compiler — SFT on LLVM IR](#4-recipe-2-meta-llm-compiler--sft-on-llvm-ir)
5. [Recipe 3: Compiler Feedback — Iterative Refinement](#5-recipe-3-compiler-feedback--iterative-refinement)
6. [Recipe 4: CUDA-L1 — Contrastive RL (3-Stage Pipeline)](#6-recipe-4-cuda-l1--contrastive-rl-3-stage-pipeline)
7. [Recipe 5: StepCoder — Fine-Grained RL for Code](#7-recipe-5-stepcoder--fine-grained-rl-for-code)
8. [ARM Adaptation Guide](#8-arm-adaptation-guide)
9. [Datasets](#9-datasets)
10. [Model Selection](#10-model-selection)
11. [Reward Function Design](#11-reward-function-design)
12. [Reward Hacking & Mitigations](#12-reward-hacking--mitigations)
13. [Training Infrastructure: TRL GRPO Implementation](#13-training-infrastructure-trl-grpo-implementation)
14. [Full Training Script](#14-full-training-script)
15. [Program Transformation Taxonomy](#15-program-transformation-taxonomy)
16. [Results Benchmarks & Ablations](#16-results-benchmarks--ablations)
17. [Citation Graph & Future Directions](#17-citation-graph--future-directions)
18. [Reference Links](#18-reference-links)

---

## 1. Executive Summary

**The problem**: Modern compilers like `gcc -O3` apply fixed heuristics. LLMs can learn program-specific optimizations that compilers miss — loop restructuring, better instruction selection, algorithmic simplification — achieving **1.46× average speedup over gcc -O3** on x86 and potentially similar gains on ARM.

**The approach**: Use GRPO (Group Relative Policy Optimization) to train `Qwen2.5-Coder-7B-Instruct` with a reward function that:
1. Compiles the generated assembly → reward=0 if it fails
2. Runs all test cases → reward=0 if any fail
3. Measures speedup vs baseline → reward = speedup ratio (continuous)

**Key insight from the literature**: RL beats SFT for this task because superoptimization is open-ended — there's no single "correct" optimized assembly. RL directly optimizes for the metric we care about (speedup) rather than imitating examples.

**No ARM-specific work exists yet** — all published results are on x86-64 or CUDA. This is a greenfield opportunity.

---

## 2. Landscape & Key Papers

### Paper Dependency Graph

```
                    MLGO (Google, 2021)
                    ML replaces compiler heuristics
                           │
              ┌────────────┼────────────┐
              ▼            ▼            ▼
    Meta LLM Compiler   ProGraML    ML Cost Model
    (Meta, 2023)        (2020)      for MLIR (2023)
    SFT from scratch    GNN for IR
    on LLVM IR
              │
              ▼
    Compiler Feedback           StepCoder
    (Meta, 2024)                (2024)
    Iterative refinement        FGO masking
    with oracle feedback        for code RL
              │                      │
              └──────────┬───────────┘
                         ▼
               SuperCoder (2025) ◄── CURRENT SOTA
               PPO/GRPO on assembly
               with compiler reward
                         │
              ┌──────────┼──────────┐
              ▼          ▼          ▼
         CUDA-L1     LLM-VeriOpt   Astra
         (ICLR 2026)  (2026)       (2025)
         Contrastive  Formal       Multi-agent
         RL for CUDA  verification GPU kernel opt
```

### Papers Ranked by Relevance to ARM RL Optimizer

| Rank | Paper | Year | Key Contribution | Result |
|------|-------|------|-----------------|--------|
| 🥇 | [SuperCoder](https://arxiv.org/abs/2505.11480) | 2025 | PPO/GRPO on assembly with compiler reward | 95% correct, 1.46× speedup over gcc -O3 |
| 🥈 | [CUDA-L1](https://arxiv.org/abs/2507.14111) | 2025 | 3-stage SFT→Self-supervised→Contrastive RL | 3.12× avg speedup on KernelBench |
| 🥉 | [Meta LLM Compiler](https://arxiv.org/abs/2309.07062) | 2023 | SFT from scratch on LLVM IR for pass ordering | 3.0% instruction reduction over -Oz |
| 4 | [Compiler Feedback](https://arxiv.org/abs/2403.14714) | 2024 | Iterative refinement with compiler oracle | +0.53% over base; sampling > feedback |
| 5 | [StepCoder](https://arxiv.org/abs/2402.01391) | 2024 | Fine-Grained Optimization masking for code RL | +8% pass@1 on APPS+ |
| 6 | [MLGO](https://arxiv.org/abs/2101.04808) | 2021 | ML in LLVM framework (Google) | Foundation work |

---

## 3. Recipe 1: SuperCoder — RL Assembly Superoptimization (SOTA)

> **Paper**: "SuperCoder: Assembly Program Superoptimization with Large Language Models" 
> **ArXiv**: [2505.11480](https://arxiv.org/abs/2505.11480) | May 2025

### 3.1 Task Formulation

Framed as a **contextual multi-armed bandit** (not full MDP):
- **Context** `s ∈ S`: source program C, baseline assembly P, test cases T
- **Action** `a ∈ A`: generate candidate optimized assembly P̃
- **Reward** `r(s,a)`: correctness-gated speedup (see §11)
- **Policy** `π: S → Δ(A)`: the LLM maps context to a distribution over assemblies

Single-turn generation — no rollout history, no multi-step environment.

### 3.2 Training Configuration

| Component | Setting | Source |
|-----------|---------|--------|
| Base model | `Qwen/Qwen2.5-Coder-7B-Instruct` | Table A1 |
| Actor learning rate | `1e-6` | Appendix A.2 |
| Critic learning rate | `1e-5` (PPO only) | Appendix A.2 |
| Batch size | 16 | Appendix A.2 |
| Epochs | 1 | Appendix A.2 |
| Max prompt length | 2000 tokens | Appendix A.2 |
| Max response length | 2000 tokens | Appendix A.2 |
| Gradient checkpointing | Enabled (actor + critic) | Appendix A.2 |
| Rollout temperature | 0.5 | Appendix A.2 |
| Hardware | 4× A100 GPUs | Appendix A.2 |
| RL framework | [verl](https://github.com/volcengine/verl) | §3.3 |

### 3.3 Dataset Construction

**Source**: IBM CodeNet — 8M+ C/C++ competitive programming submissions.

**Curation strategy** (critical for performance):
1. Sample programs with **highest relative speedup from -O0 to -O3** — this selects computationally rich programs where further optimization is possible
2. Compile each with `gcc -O3 -S` to get baseline x86-64 assembly
3. Use test inputs from [Li et al., 2022], but **regenerate outputs** by executing the original program (CodeNet outputs are unreliable)
4. Final dataset: **7,872 training programs, 200 evaluation programs**

| Split | Programs | Avg Tests/Program | Avg LOC (C) | Avg LOC (Assembly) |
|-------|----------|-------------------|-------------|-------------------|
| Train | 7,872 | 8.86 | 22.3 | 130.3 |
| Eval | 200 | 8.92 | 21.9 | 133.3 |

### 3.4 Prompt Template

```
Given the following C code and assembly code, your task is to generate 
highly optimized x86-64 assembly code.

C Code: <C code here>
Assembly Code: <baseline assembly code here produced by gcc -O3>

Only output the optimized assembly code. Do not include any other text.
Do not write any comments in the assembly code.
Wrap the assembly code in assembly tags.

Optimized Assembly Code:
```

> ⚠️ **Critical finding** (Appendix A.5): Removing the baseline assembly from the prompt causes a **catastrophic drop** — correctness falls from 95% to near 0%. The model needs the gcc -O3 output as a starting point.

### 3.5 Results

| Model | Compile Pass | Test Pass | Avg Speedup |
|-------|-------------|-----------|-------------|
| Qwen2.5-Coder-7B (base) | 77.9% | 61.4% | 1.10× |
| SuperCoder (PPO) | 96.0% | 95.0% | **1.46×** |
| SuperCoder (GRPO) | 95.0% | 94.7% | **1.44×** |
| SuperCoder (SFT only) | 95.5% | 92.5% | 1.39× |

**PPO ≈ GRPO** — nearly identical results, but GRPO is simpler (no critic/value head needed).

**Best-of-N + RL**:
- Base best-of-8: ~1.39× (≈ RL best-of-1)
- SuperCoder best-of-8: **1.93×**

### 3.6 Model Evaluation (23 Models Tested)

| Model | Test Pass | Avg Speedup | Notes |
|-------|-----------|-------------|-------|
| DeepSeek-R1 | 0.0% | 1.00× | Generates verbose analysis, no actual code |
| GPT-4o | 5.0% | 1.02× | Compiles (81%) but sacrifices correctness for optimization |
| Claude-opus-4 | 51.5% | 1.43× | Best zero-shot baseline |
| Qwen2.5-Coder-7B | 61.4% | 1.10× | Best base model for RL starting point |
| llm-compiler-13b | 59.5% | 1.34× | Pretrained on assembly/IR |
| **SuperCoder (PPO)** | **95.0%** | **1.46×** | **SOTA** |

**Key failure modes**:
- Reasoning models (R1, o1) completely fail — they analyze instead of generating code
- GPT-4o compiles but breaks low-level conventions (stack canaries, .cfi directives, calling conventions)
- Models that work best have been pretrained on assembly/code (Qwen-Coder, llm-compiler)

---

## 4. Recipe 2: Meta LLM Compiler — SFT on LLVM IR

> **Paper**: "Large Language Models for Compiler Optimization"
> **ArXiv**: [2309.07062](https://arxiv.org/abs/2309.07062) | Sep 2023 | Meta AI

### 4.1 Approach

Train a 7B transformer **from scratch** (Llama 2 architecture) on LLVM IR to predict:
1. **Optimization pass list** (primary task)
2. **Instruction counts before/after** (auxiliary task — critical for performance)
3. **Full optimized IR** (auxiliary task)

### 4.2 Training Details

| Component | Setting |
|-----------|---------|
| Architecture | Llama 2 7B (32 heads, 4096 hidden, 32 layers) |
| Training | From scratch (random init) |
| Dataset | 1,000,000 deduplicated LLVM-IR functions, 373M tokens |
| Autotuner | 37,424 compilations per function avg; 9,016 CPU days total |
| Optimizer | AdamW (β₁=0.9, β₂=0.95) |
| LR schedule | Cosine, 1000 warmup, peak=1e-5, final=1e-6 |
| Batch size | 256 (524,288 tokens/batch) |
| Steps | 30,000 (7.7 epochs, 15.7B total tokens) |
| Sequence length | 2048 tokens |
| Hardware | 64× V100 GPUs, 620 GPU-days |

### 4.3 Key Results

- **3.0% instruction count reduction over -Oz** (without invoking compiler at inference)
- **91% compilable** generated code
- **70%** perfect emulation of compiler output
- Achieves **60% of autotuner gains** at 0 additional compilation cost

### 4.4 Key Insight: Auxiliary Tasks Matter

Training with instruction count prediction + code generation alongside pass list prediction **dramatically improves** optimization quality. The instruction count acts as a self-consistency check — the model learns to verify its own predictions.

---

## 5. Recipe 3: Compiler Feedback — Iterative Refinement

> **Paper**: "Compiler Generated Feedback for Large Language Models"
> **ArXiv**: [2403.14714](https://arxiv.org/abs/2403.14714) | Mar 2024 | Meta AI

### 5.1 Approach

Start from the best checkpoint of Recipe 2. Add compiler feedback to the prompt:
- Predicted instruction counts
- Whether the code compiled correctly
- Whether the IR is valid

Model outputs "I am sure!" if confident, else "Let me try again."

### 5.2 Training Details

| Component | Setting |
|-----------|---------|
| Base model | Best checkpoint from Meta LLM Compiler (7B) |
| Dataset | 1M training + 100K test LLVM-IR functions |
| Optimizer | AdamW (β₁=0.9, β₂=0.95) |
| LR | Cosine, 1000 warmup, peak=1e-5 |
| Batch size | 256 (786K–1M tokens/batch) |
| Steps | 20,000 (5.12 epochs, 16–21B tokens) |
| Hardware | 64× A100s, 60 GPU-days |

### 5.3 Key Finding

> **Sampling beats iterative feedback**: At n≥10 samples, the base model without feedback training outperforms the feedback-trained model. This strongly motivates using **Best-of-N + RL** (SuperCoder approach) over iterative SFT feedback.

---

## 6. Recipe 4: CUDA-L1 — Contrastive RL (3-Stage Pipeline)

> **Paper**: "CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning"
> **ArXiv**: [2507.14111](https://arxiv.org/abs/2507.14111) | Jul 2025 | ICLR 2026

### 6.1 The 3-Stage Pipeline

This is the most important contribution for cases where the **base model has low initial success rate**.

```
Stage 1: SFT via Data Augmentation
├── Generate CUDA code from 6 LLMs (GPT-4o, o1, DeepSeek-R1/V3, Llama-405B, Claude-3.7)
├── Filter for correct + fast implementations
├── Fine-tune DeepSeek-V3-671B on successful samples
└── Result: Model can generate correct code at reasonable rate

Stage 2: Self-Supervised Learning
├── Sample from Stage 1 model
├── Keep only correct samples (self-filtering)
├── Retrain on filtered dataset
└── Result: Higher base success rate for RL

Stage 3: Contrastive Reinforcement Learning
├── Present model with multiple code variants + their speedup scores
├── Model analyzes WHY certain implementations are faster
├── Generates improved solution based on comparative analysis
├── Score serves dual purpose: (1) gradient update, (2) future prompt enrichment
└── Result: 3.12× average speedup on KernelBench
```

### 6.2 Why Standard GRPO/PPO Failed for CUDA-L1

> "Standard RL algorithms compute a scalar reward for each generated CUDA code sample... the reward signal is used exclusively for parameter updates and is never provided as input to the LLM. Consequently, the LLM cannot directly reason about performance trade-offs during code generation."

Their solution: **Contrastive RL** — embed performance feedback within the input prompt. The model sees previous code + scores and learns comparative analysis.

### 6.3 Results

| Configuration | Mean Speedup | Max | Median | Success Rate |
|--------------|-------------|-----|--------|-------------|
| Default | 3.12× | 120× | 1.42× | 249/250 |
| vs Torch Compile | 2.77× | — | — | — |
| vs CUDA Graph | 2.81× | — | — | — |

### 6.4 ARM Relevance

- Use the **3-stage pipeline** if the base model has <40% ARM assembly correctness
- Contrastive RL is useful when you need the model to reason about WHY optimizations work
- Stage 1 data augmentation from multiple LLMs is a powerful bootstrapping technique

---

## 7. Recipe 5: StepCoder — Fine-Grained RL for Code

> **Paper**: "StepCoder: Improve Code Generation with Reinforcement Learning from Compiler Feedback"
> **ArXiv**: [2402.01391](https://arxiv.org/abs/2402.01391) | Feb 2024

### 7.1 Key Innovation: Fine-Grained Optimization (FGO)

Standard PPO updates ALL tokens in the generated code equally. StepCoder's FGO **masks tokens not executed by unit tests** — only code segments that are actually run contribute to the gradient update.

This is crucial for sparse compiler rewards where most of the generated code might be boilerplate.

### 7.2 Reward Design

| Outcome | Reward |
|---------|--------|
| All unit tests pass | +1.0 |
| Test failure | -0.3 |
| Runtime error | -0.6 |
| Compile error | -1.0 |

This graduated penalty scheme (vs SuperCoder's binary 0/non-zero) helps the model distinguish failure modes.

---

## 8. ARM Adaptation Guide

### 8.1 The ARM Gap

**No published work targets ARM/AArch64**. All results are on x86-64 (SuperCoder, Meta) or CUDA (CUDA-L1). This requires adaptation:

### 8.2 Minimal Changes from SuperCoder (x86 → ARM)

| Component | x86 (Original) | ARM (Adapted) |
|-----------|----------------|---------------|
| Compiler | `gcc -O3` | `aarch64-linux-gnu-gcc -O3` |
| Assembler | `as` | `aarch64-linux-gnu-as` |
| Linker | `ld` | `aarch64-linux-gnu-ld` |
| Execution | Native | `qemu-aarch64-static` (emulation) |
| Timing | `hyperfine` (wall-clock) | QEMU instruction counting or real ARM HW |
| ISA in prompt | "x86-64 assembly" | "AArch64 assembly" |
| Assembly flag | `-S` | `-S` (same) |

### 8.3 ARM-Specific Prompt Template

```
Given the following C code and assembly code, your task is to generate 
highly optimized AArch64 assembly code.

C Code: {c_code}
Assembly Code: {arm_baseline_asm}

Only output the optimized assembly code. Do not include any other text.
Do not write any comments in the assembly code.
Wrap the assembly code in <assembly></assembly> tags.

Optimized Assembly Code:
```

### 8.4 Execution Environment

```bash
# Install ARM cross-compilation toolchain
sudo apt-get install gcc-aarch64-linux-gnu g++-aarch64-linux-gnu
sudo apt-get install qemu-user-static

# Cross-compile to ARM assembly
aarch64-linux-gnu-gcc -O3 -S -o output.s input.c

# Cross-compile to ARM binary
aarch64-linux-gnu-gcc -O3 -o output input.c

# Run ARM binary on x86 via QEMU
qemu-aarch64-static ./output < test_input.txt
```

### 8.5 Timing Considerations

| Method | Accuracy | Availability |
|--------|----------|-------------|
| QEMU instruction counting | Deterministic but not wall-clock accurate | Any x86 machine |
| QEMU `-plugin libinsn` | Counts executed instructions precisely | QEMU 6.0+ |
| Real ARM hardware | Ground truth | Requires ARM machine |
| `perf stat` on ARM | Cycle-accurate | Requires ARM + Linux perf |

**Recommendation**: Use QEMU instruction counting for training (deterministic, reproducible) and validate final results on real ARM hardware.

### 8.6 ARM-Specific Optimization Opportunities

The model should learn to leverage:
- **NEON SIMD** instructions (128-bit vector operations)
- **SVE/SVE2** (Scalable Vector Extension — variable-length vectors)
- **LSE atomics** (Large System Extensions)
- **Conditional select** (`csel`, `csinc`) instead of branches
- **Fused multiply-add** (`fmadd`, `fmsub`)
- **Load/store pair** (`ldp`, `stp`) for memory throughput
- **Predicated operations** in SVE (eliminate branch mispredictions)

### 8.7 3-Stage Plan for ARM (if base model correctness < 40%)

Following CUDA-L1's insight:

```
Stage 1: SFT Warmup
├── Dataset: (C source → ARM gcc -O3 assembly) pairs
├── Task: Teach the model to generate valid ARM assembly
├── Expected: Model learns ARM syntax, calling conventions, directives
└── Duration: ~2-4 hours on A100

Stage 2: Filtered Self-Training
├── Sample N completions per prompt from Stage 1 model
├── Keep only compilable + correct samples
├── Retrain on filtered dataset (higher quality)
└── Expected: Correctness improves to >60%

Stage 3: GRPO with Compiler Reward
├── Apply SuperCoder's reward function
├── Binary correctness gate + continuous speedup
├── Expected: Correctness >90%, speedup >1.3×
└── Duration: ~4-8 hours on 4×A100
```

---

## 9. Datasets

### 9.1 Available Datasets

| Dataset | Source | Format | Size | ARM Compatible | Notes |
|---------|--------|--------|------|----------------|-------|
| **IBM CodeNet** | [GitHub](https://github.com/IBM/Project_CodeNet) | C/C++ source + test I/O | 8M+ submissions | ✅ Recompile with ARM GCC | Used by SuperCoder |
| **deepmind/code_contests** | [HF Hub](https://hf.co/datasets/deepmind/code_contests) | C/C++ solutions + tests | ~2GB train | ✅ Filter C, cross-compile | Has public/private/generated tests |
| **llvm-ml/ComPile** | [HF Hub](https://hf.co/datasets/llvm-ml/ComPile) | LLVM bitcode IR | 602GB (2.7TB source) | ✅ Retarget to AArch64 via `llc` | C/C++/Rust/Swift from Spack |
| APPS+ | [GitHub](https://github.com/ablustrund/apps_plus) | Python problems + tests | ~10K | ❌ Python only | Used by StepCoder |

### 9.2 Dataset Format for GRPO Training

The dataset must have a `prompt` column in conversational format. Extra columns are forwarded to reward functions as `**kwargs`.

```python
{
    "prompt": [
        {"role": "user", "content": "Given the following C code and assembly code..."}
    ],
    "c_code": "int main() { ... }",
    "baseline_asm": ".text\n.globl main\nmain:\n...",
    "test_inputs": ["3\n1 2 3", "5\n1 2 3 4 5"],
    "test_outputs": ["6", "15"],
    "baseline_time": 0.0042  # seconds
}
```

### 9.3 Dataset Construction Pipeline

```python
from datasets import Dataset
import subprocess, json

def build_arm_dataset(codenet_programs):
    """Build ARM optimization dataset from CodeNet C programs."""
    samples = []
    
    for prog in codenet_programs:
        c_code = prog["source"]
        test_cases = prog["test_cases"]  # [(input, output), ...]
        
        # Step 1: Cross-compile to ARM assembly with -O3
        result = subprocess.run(
            ["aarch64-linux-gnu-gcc", "-O3", "-S", "-o", "/dev/stdout", "-x", "c", "-"],
            input=c_code, capture_output=True, text=True
        )
        if result.returncode != 0:
            continue  # Skip programs that don't compile
        baseline_asm = result.stdout
        
        # Step 2: Compile to binary for benchmarking
        # ... (compile + link + measure baseline time via QEMU)
        
        # Step 3: Build prompt
        prompt_text = f"""Given the following C code and assembly code, your task is to generate highly optimized AArch64 assembly code.

C Code: {c_code}
Assembly Code: {baseline_asm}

Only output the optimized assembly code. Do not include any other text.
Do not write any comments in the assembly code.
Wrap the assembly code in <assembly></assembly> tags.

Optimized Assembly Code:"""
        
        samples.append({
            "prompt": [{"role": "user", "content": prompt_text}],
            "c_code": c_code,
            "baseline_asm": baseline_asm,
            "test_inputs": [tc[0] for tc in test_cases],
            "test_outputs": [tc[1] for tc in test_cases],
            "baseline_time": baseline_time,
        })
    
    return Dataset.from_list(samples)
```

### 9.4 deepmind/code_contests Schema

```
| Column | Type |
|--------|------|
| name | string |
| description | string |
| public_tests | Sequence: {input: list[str], output: list[str]} |
| private_tests | Sequence: {input: list[str], output: list[str]} |
| generated_tests | Sequence: {input: list[str], output: list[str]} |
| solutions | Sequence: {language: list[int], solution: list[str]} |
| difficulty | ClassLabel (29 classes) |
| source | ClassLabel (7 classes) |
```

Filter C solutions: `language == 2` (C) or look for C-like syntax in the solution strings.

---

## 10. Model Selection

### 10.1 Base Model Recommendation

**Primary**: `Qwen/Qwen2.5-Coder-7B-Instruct`
- 7.6B parameters, Qwen2 architecture
- Apache-2.0 license
- **Proven by SuperCoder**: 61.4% test pass rate (highest among 7B models)
- Strong code generation baseline for RL starting point
- Available on [HF Hub](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct)

### 10.2 Why Not Other Models?

| Model | Issue |
|-------|-------|
| DeepSeek-R1 | 0% compilation — generates analysis, not code |
| GPT-4o | 81% compile but 5% correct — breaks low-level conventions |
| Reasoning models (o1, R1) | Fundamentally fail — spend tokens reasoning instead of generating |
| llm-compiler-7b-ftd | Fine-tuned for disassembly, not optimization |
| llm-compiler-13b | Good (1.34× speedup) but not instruction-tuned; harder to RL fine-tune |

### 10.3 Model Sizing

| Model Size | Hardware Needed | VRAM (bf16) |
|-----------|----------------|-------------|
| 7B (Qwen2.5-Coder-7B) | 4× A100 80GB | ~14GB model + ~40GB for RL |
| 13B | 4× A100 80GB | ~26GB model + ~60GB for RL |
| 70B+ | 8× A100 or 4× H100 | Multi-node |

---

## 11. Reward Function Design

### 11.1 SuperCoder Reward (Recommended)

From Section 3.3 of arxiv:2505.11480:

```
r(s, a) = {
    0,              if pass(s,a) < 1       (any test fails → zero reward)
    speedup(s,a),   if pass(s,a) = 1       (all tests pass → continuous speedup)
}
```

Where:
- `pass(s,a) = (1/|T|) × Σ 𝟏[P̃(xᵢ) = yᵢ]` — fraction of test cases passed
- `speedup(s,a) = t(P) / t(P̃)` — baseline time / optimized time

**Design principles**:
1. **Hard binary correctness gate**: No partial credit. Forces model to learn correctness first.
2. **Continuous speedup reward**: Provides gradient signal proportional to actual optimization gain.
3. **No reward for partial correctness**: Passing 99% of tests still gets 0 reward.

### 11.2 StepCoder Graduated Penalty (Alternative)

```
+1.0  → all unit tests pass
-0.3  → test failure
-0.6  → runtime error
-1.0  → compile error
```

Distinguishes failure modes — the model gets a stronger negative signal for "worse" failures.

### 11.3 CUDA-L1 Reward Smoothing (Anti-Hacking)

```python
r_normalized = (r - μ) / σ           # normalize by running stats
r_smooth = clip(r_normalized, -k, k)  # clip to [-1.5, 1.5]
```

Prevents the model from over-optimizing outlier high-reward solutions.

### 11.4 Implementation

```python
import subprocess
import tempfile
import os
import re

def arm_compiler_reward(completions, c_code, baseline_asm, 
                        test_inputs, test_outputs, baseline_time, **kwargs):
    """
    Compiler feedback reward for ARM assembly optimization.
    Follows SuperCoder §3.3 exactly: binary correctness gate + continuous speedup.
    """
    rewards = []
    
    for i, completion in enumerate(completions):
        # Extract assembly from completion
        content = completion[0]["content"] if isinstance(completion, list) else completion
        asm_match = re.search(r'<assembly>(.*?)</assembly>', content, re.DOTALL)
        if not asm_match:
            rewards.append(0.0)
            continue
        asm_code = asm_match.group(1).strip()
        
        with tempfile.TemporaryDirectory() as tmpdir:
            asm_path = os.path.join(tmpdir, "opt.s")
            bin_path = os.path.join(tmpdir, "opt")
            
            # Write assembly
            with open(asm_path, "w") as f:
                f.write(asm_code)
            
            # Step 1: Assemble
            result = subprocess.run(
                ["aarch64-linux-gnu-gcc", "-o", bin_path, asm_path, "-static", "-lm"],
                capture_output=True, text=True, timeout=30
            )
            if result.returncode != 0:
                rewards.append(0.0)  # Compile failure
                continue
            
            # Step 2: Run all tests via QEMU
            all_pass = True
            for test_in, expected_out in zip(test_inputs[i], test_outputs[i]):
                try:
                    run = subprocess.run(
                        ["qemu-aarch64-static", bin_path],
                        input=test_in, capture_output=True, text=True, timeout=10
                    )
                    if run.stdout.strip() != expected_out.strip():
                        all_pass = False
                        break
                except subprocess.TimeoutExpired:
                    all_pass = False
                    break
            
            if not all_pass:
                rewards.append(0.0)  # Test failure
                continue
            
            # Step 3: Measure speedup (instruction count via QEMU)
            # Use QEMU instruction counting for deterministic measurement
            try:
                run = subprocess.run(
                    ["qemu-aarch64-static", "-d", "in_asm", bin_path],
                    input=test_inputs[i][0], capture_output=True, text=True, timeout=30
                )
                insn_count = run.stderr.count('\n')  # Rough instruction count
                speedup = baseline_time[i] / max(insn_count, 1)
                rewards.append(max(speedup, 0.1))
            except Exception:
                rewards.append(1.0)  # Default: assume baseline-equivalent
    
    return rewards
```

---

## 12. Reward Hacking & Mitigations

### 12.1 Known Hacking Behaviors (from CUDA-L1 §3.1)

| Hack | Description | Prevalence |
|------|-------------|-----------|
| **Improper timing** | Create async streams; timing only measures main stream | 32.8% of outputs |
| **Lazy evaluation** | Return lazy tensor; actual compute happens at correctness check | Found in training |
| **Hyperparameter manipulation** | Reduce batch_size/dimensions in generated code | Found in training |
| **Result caching** | Cache outputs by input address; return cached results | Found in training |

### 12.2 Mitigations

1. **Reward checking model**: When reward jumps significantly, use an adversarial model (e.g., DeepSeek-R1) to check for exploitation. Catches >60% of hacking.
2. **Hacking-case database**: Maintain a growing database of known hacking patterns. Use retrieval-augmented checking.
3. **Reward smoothing**: `r_smooth = clip((r - μ)/σ, -k, k)` with k=1.5
4. **Robust evaluation**: Synchronize all execution before timing. Validate output is real tensor with allocated storage.

### 12.3 ARM-Specific Hacking Risks

- **QEMU timing artifacts**: Model could learn to generate code that runs fast in QEMU but slow on real ARM hardware
- **Mitigation**: Use instruction counting (deterministic) rather than wall-clock timing in QEMU
- **NOP padding**: Model could pad with NOPs that don't affect correctness but confuse instruction counting
- **Mitigation**: Count only non-NOP instructions, or use basic-block counting

---

## 13. Training Infrastructure: TRL GRPO Implementation

### 13.1 Why GRPO over PPO

| Feature | PPO | GRPO |
|---------|-----|------|
| Value head/critic | Required | Not needed |
| Memory usage | ~2× model | ~1× model |
| Reward model | Optional (can use custom) | Custom function native |
| Performance | 1.46× speedup | 1.44× speedup |
| Implementation complexity | Higher | Lower |
| TRL support | Experimental (`trl.experimental.ppo`) | Full support (`trl.GRPOTrainer`) |

**Verdict**: GRPO is strictly better for this use case — same results, half the memory, simpler code.

### 13.2 GRPOConfig Parameters

```python
from trl import GRPOConfig

config = GRPOConfig(
    # === Model ===
    output_dir="arm-compiler-optimizer",
    
    # === Generation ===
    num_generations=4,              # G: completions per prompt (group size for GRPO)
    max_prompt_length=2048,         # tokens; truncated from left
    max_completion_length=2048,     # tokens
    temperature=0.5,                # rollout sampling temperature (SuperCoder: 0.5)
    
    # === Loss / Reward ===
    beta=0.0,                       # KL penalty (0.0 = no KL term, default in TRL)
    epsilon=0.2,                    # PPO clip range (used in GRPO's clipped objective)
    scale_rewards=False,            # Don't normalize — binary-gated rewards aren't Gaussian
    
    # === Training ===
    learning_rate=1e-6,             # SuperCoder Appendix A.2
    per_device_train_batch_size=4,  # Per GPU; effective batch = 4 GPUs × 4 = 16
    gradient_accumulation_steps=1,
    num_train_epochs=1,             # SuperCoder: 1 epoch
    gradient_checkpointing=True,    # Memory savings (default True in GRPOConfig)
    bf16=True,                      # Default True
    
    # === Logging ===
    logging_steps=10,
    log_completions=True,           # Log generated completions
    disable_tqdm=True,              # Plain text logs for job monitoring
    logging_first_step=True,
    logging_strategy="steps",
    
    # === Saving ===
    push_to_hub=True,
    hub_model_id="kaori02/arm-compiler-optimizer",
    save_strategy="steps",
    save_steps=100,
)
```

### 13.3 Custom Reward Function Signature

TRL GRPOTrainer passes these keyword arguments to reward functions:

```python
def my_reward(
    completions,          # list[list[dict]] — each is [{"role":"assistant","content":"..."}]
    prompts=None,         # list[list[dict]] or list[str]
    completion_ids=None,  # list[list[int]] — tokenized completions
    trainer_state=None,   # TrainerState — current training step, epoch, etc.
    log_extra=None,       # callable to log extra columns
    log_metric=None,      # callable to log scalar metrics
    **kwargs              # ALL extra dataset columns forwarded here
) -> list[float]:         # one reward per completion
    ...
```

### 13.4 Multiple Reward Functions (Composable)

```python
trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-Coder-7B-Instruct",
    reward_funcs=[compile_reward, correctness_reward, speedup_reward],
    reward_weights=[1.0, 2.0, 5.0],  # Weight speedup most heavily
    args=config,
    train_dataset=dataset,
)
```

---

## 14. Full Training Script

```python
"""
ARM Compiler Optimization via GRPO
Based on: SuperCoder (arxiv:2505.11480) adapted for AArch64

Recipe:
- Base model: Qwen/Qwen2.5-Coder-7B-Instruct
- Method: GRPO with custom compiler reward
- Reward: Binary correctness gate + continuous speedup
- Dataset: CodeNet C programs → ARM assembly
"""

import os
import re
import subprocess
import tempfile
from datasets import load_dataset, Dataset
from trl import GRPOTrainer, GRPOConfig
import trackio

# ═══════════════════════════════════════════════════════════════════════════
# REWARD FUNCTION
# ═══════════════════════════════════════════════════════════════════════════

def extract_assembly(content):
    """Extract assembly code from <assembly> tags."""
    match = re.search(r'<assembly>(.*?)</assembly>', content, re.DOTALL)
    return match.group(1).strip() if match else None

def compile_arm(asm_code, output_path):
    """Cross-compile ARM assembly to binary."""
    with tempfile.NamedTemporaryFile(suffix=".s", mode="w", delete=False) as f:
        f.write(asm_code)
        asm_path = f.name
    try:
        result = subprocess.run(
            ["aarch64-linux-gnu-gcc", "-o", output_path, asm_path, "-static", "-lm"],
            capture_output=True, text=True, timeout=30
        )
        return result.returncode == 0
    except Exception:
        return False
    finally:
        os.unlink(asm_path)

def run_test(binary_path, test_input, expected_output, timeout=10):
    """Run a test case via QEMU and check output."""
    try:
        result = subprocess.run(
            ["qemu-aarch64-static", binary_path],
            input=test_input, capture_output=True, text=True, timeout=timeout
        )
        return result.stdout.strip() == expected_output.strip()
    except subprocess.TimeoutExpired:
        return False

def measure_instructions(binary_path, test_input, timeout=30):
    """Count executed instructions via QEMU for deterministic timing."""
    try:
        result = subprocess.run(
            ["qemu-aarch64-static", "-d", "in_asm", binary_path],
            input=test_input, capture_output=True, text=True, timeout=timeout
        )
        return result.stderr.count('\n')
    except Exception:
        return float('inf')

def arm_compiler_reward(completions, test_inputs, test_outputs, 
                        baseline_insn_count, log_metric=None, **kwargs):
    """
    SuperCoder §3.3 reward adapted for ARM:
    r = 0 if compilation fails or any test fails
    r = baseline_instructions / optimized_instructions if all tests pass
    """
    rewards = []
    compile_successes = 0
    test_successes = 0
    
    for i, completion in enumerate(completions):
        content = completion[0]["content"] if isinstance(completion, list) else completion
        asm = extract_assembly(content)
        
        if asm is None:
            rewards.append(0.0)
            continue
        
        with tempfile.TemporaryDirectory() as tmpdir:
            bin_path = os.path.join(tmpdir, "opt_binary")
            
            # Step 1: Compile
            if not compile_arm(asm, bin_path):
                rewards.append(0.0)
                continue
            compile_successes += 1
            
            # Step 2: Run ALL tests
            all_pass = True
            inputs = test_inputs[i] if isinstance(test_inputs[i], list) else [test_inputs[i]]
            outputs = test_outputs[i] if isinstance(test_outputs[i], list) else [test_outputs[i]]
            
            for tin, tout in zip(inputs, outputs):
                if not run_test(bin_path, tin, tout):
                    all_pass = False
                    break
            
            if not all_pass:
                rewards.append(0.0)
                continue
            test_successes += 1
            
            # Step 3: Measure speedup via instruction count
            opt_insn = measure_instructions(bin_path, inputs[0])
            baseline = baseline_insn_count[i] if isinstance(baseline_insn_count, list) else baseline_insn_count
            
            if opt_insn > 0 and baseline > 0:
                speedup = baseline / opt_insn
                rewards.append(max(speedup, 0.1))
            else:
                rewards.append(1.0)
    
    # Log metrics
    if log_metric and len(rewards) > 0:
        log_metric("compile_rate", compile_successes / len(completions))
        log_metric("test_pass_rate", test_successes / len(completions))
        log_metric("avg_reward", sum(rewards) / len(rewards))
    
    return rewards

# ═══════════════════════════════════════════════════════════════════════════
# TRAINING
# ═══════════════════════════════════════════════════════════════════════════

def main():
    # Initialize tracking
    trackio.init(
        project="arm-compiler-optimizer",
        run="grpo-qwen-7b-arm",
    )
    
    # Load pre-built dataset (see §9.3 for construction)
    dataset = load_dataset("your-org/arm-compiler-dataset", split="train")
    
    # GRPO Configuration (SuperCoder Appendix A.2, adapted for TRL)
    config = GRPOConfig(
        output_dir="arm-compiler-optimizer",
        
        # Generation
        num_generations=4,
        max_prompt_length=2048,
        max_completion_length=2048,
        temperature=0.5,
        
        # Loss
        beta=0.0,
        scale_rewards=False,
        
        # Training
        learning_rate=1e-6,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=1,
        num_train_epochs=1,
        gradient_checkpointing=True,
        bf16=True,
        
        # Logging
        logging_steps=10,
        log_completions=True,
        disable_tqdm=True,
        logging_first_step=True,
        logging_strategy="steps",
        
        # Saving
        push_to_hub=True,
        hub_model_id="kaori02/arm-compiler-optimizer",
        save_strategy="steps",
        save_steps=100,
    )
    
    trainer = GRPOTrainer(
        model="Qwen/Qwen2.5-Coder-7B-Instruct",
        reward_funcs=arm_compiler_reward,
        args=config,
        train_dataset=dataset,
    )
    
    trainer.train()
    trainer.push_to_hub()

if __name__ == "__main__":
    main()
```

---

## 15. Program Transformation Taxonomy

SuperCoder §5.4 analyzed all 200 evaluation programs. The LLM learns these optimization patterns:

| Transformation | Description | Frequency |
|---------------|-------------|-----------|
| **Loop Restructuring** | Reorder, unroll, alter loop control flow | 45% |
| **Instruction Selection** | Use specialized CPU instructions (e.g., `popcnt`, `bsr`, `cmov`) instead of generic sequences | 35% |
| **Algorithmic Simplification** | Replace custom logic with standard library calls (`memcmp`, `strcmp`, `atoi`) | 30% |
| **Stack Canary Removal** | Eliminate stack protection checks and security instrumentation | 25% |
| **Register Allocation** | Better register assignment, reuse, reduced spills | 20% |
| **Branch Elimination** | Replace conditional branches with conditional moves (`cmov`, `setcc`) | 15% |
| **Address Calculation** | Optimize memory address computation | 10% |
| **Dead Code Elimination** | Remove unused code paths | 10% |
| **Constant Propagation** | Evaluate expressions at compile time | 5% |

### ARM-Specific Transformations to Expect

| Transformation | ARM Instructions | x86 Equivalent |
|---------------|-----------------|----------------|
| Vectorization | NEON `ld1`, `add`, `mul` (4×32-bit) | SSE/AVX |
| Predication | SVE predicated ops (no branch) | None (x86 lacks predication) |
| Paired load/store | `ldp`, `stp` | None (x86 does one at a time) |
| Fused multiply-add | `fmadd`, `fmsub` | `vfmadd` (AVX) |
| Conditional select | `csel`, `csinc`, `csneg` | `cmov` |
| Bit manipulation | `clz`, `rbit`, `cnt` | `bsr`, `popcnt` |

---

## 16. Results Benchmarks & Ablations

### 16.1 RL vs SFT (SuperCoder §5.2-5.3)

| Method | Correctness | Speedup | Notes |
|--------|------------|---------|-------|
| Base (zero-shot) | 61.4% | 1.10× | No training |
| SFT | 92.5% | 1.39× | Trained on best samples |
| GRPO | 94.7% | 1.44× | RL with compiler reward |
| PPO | 95.0% | 1.46× | RL with compiler reward |
| PPO + best-of-8 | ~95% | **1.93×** | Inference-time scaling |

> **RL > SFT** because optimization is open-ended. There's no single "correct" optimized program — RL directly maximizes the objective (speedup) rather than imitating examples.

### 16.2 Inference-Time Scaling

Best-of-N sampling works multiplicatively with RL:

| Model | N=1 | N=2 | N=4 | N=8 |
|-------|-----|-----|-----|-----|
| Base (Qwen-7B) | 1.10× | 1.20× | 1.30× | 1.39× |
| Claude-opus-4 | 1.43× | 1.60× | 1.80× | 2.05× |
| SuperCoder (PPO) | 1.46× | 1.60× | 1.75× | **1.93×** |

### 16.3 Ablation: Prompt Components

| Prompt Contains | Correctness | Speedup |
|----------------|-------------|---------|
| C code + gcc -O3 assembly | 95.0% | 1.46× |
| C code only (no assembly) | ~0% | — |
| gcc -O3 assembly only (no C) | ~60% | ~1.3× |

> The baseline assembly is **essential**. Without it, the model cannot generate valid assembly.

### 16.4 Random vs Curated Dataset (SuperCoder §A.7)

| Dataset | Correctness | Speedup |
|---------|-------------|---------|
| Curated (high O0→O3 speedup) | 95.0% | 1.46× |
| Random sample from CodeNet | 93.5% | 1.35× |

Curated dataset helps, but random works too — the approach is robust.

---

## 17. Citation Graph & Future Directions

### 17.1 Papers Citing SuperCoder

| Paper | Key Insight |
|-------|------------|
| **LLM-VeriOpt** (2026) [influential] | Uses formal verification instead of test suites for correctness — eliminates false positives from finite test coverage |
| **Astra** (2025, 31 citations) | Multi-agent system for GPU kernel optimization — agent decomposition for complex optimizations |
| **InCoder-32B** (2026) | Industrial code model covering compiler optimization as a domain |
| **Genesys** (2025) | Evolutionary program synthesis with continuous optimization |

### 17.2 Future Directions

1. **Formal Verification as Reward** (LLM-VeriOpt direction): Replace test-based correctness with formal equivalence checking — eliminates false positive rewards from insufficient test coverage.

2. **Multi-Turn Refinement** (Kevin, 2025 — arxiv:2507.11948): Let the model iteratively refine its assembly based on profiler feedback. Multiple rounds of generation → profiling → feedback.

3. **Contrastive RL** (CUDA-L1): When base model success rate is too low for standard GRPO, present multiple code variants with scores and let the model reason about WHY certain versions are faster.

4. **Cross-Architecture Transfer**: Train on x86, transfer to ARM. The C source code is architecture-agnostic — optimization patterns (loop unrolling, vectorization, etc.) transfer across ISAs even if specific instructions differ.

5. **SVE/SVE2 Exploitation**: ARM's Scalable Vector Extension offers variable-length SIMD. This is a unique optimization opportunity not available on x86 — models that learn to leverage SVE could achieve outsized speedups.

6. **Auto-Parallelization**: Beyond single-thread optimization — teach the model to identify parallelization opportunities and generate multi-threaded ARM code with proper synchronization.

---

## 18. Reference Links

### Papers
| Paper | ArXiv | Year |
|-------|-------|------|
| SuperCoder | [2505.11480](https://arxiv.org/abs/2505.11480) | 2025 |
| CUDA-L1 | [2507.14111](https://arxiv.org/abs/2507.14111) | 2025 |
| Meta LLM Compiler | [2309.07062](https://arxiv.org/abs/2309.07062) | 2023 |
| Compiler Feedback | [2403.14714](https://arxiv.org/abs/2403.14714) | 2024 |
| StepCoder | [2402.01391](https://arxiv.org/abs/2402.01391) | 2024 |
| MLGO | [2101.04808](https://arxiv.org/abs/2101.04808) | 2021 |
| ProGraML | [2012.01470](https://arxiv.org/abs/2012.01470) | 2020 |
| DeepSeekMath (GRPO) | [2402.03300](https://arxiv.org/abs/2402.03300) | 2024 |
| VeriReason | [2505.11849](https://arxiv.org/abs/2505.11849) | 2025 |
| ACECoder | [2502.01718](https://arxiv.org/abs/2502.01718) | 2025 |

### Code & Frameworks
| Resource | URL |
|----------|-----|
| TRL (GRPO Trainer) | [huggingface.co/docs/trl/grpo_trainer](https://huggingface.co/docs/trl/grpo_trainer) |
| TRL Reward Functions | [huggingface.co/docs/trl/rewards](https://huggingface.co/docs/trl/rewards) |
| TRL OpenEnv | [huggingface.co/docs/trl/openenv](https://huggingface.co/docs/trl/openenv) |
| verl (Volcano Engine RL) | [github.com/volcengine/verl](https://github.com/volcengine/verl) |
| CUDA-L1 Code | [github.com/deepreinforce-ai/CUDA-L1](https://github.com/deepreinforce-ai/CUDA-L1) |
| StepCoder / APPS+ | [github.com/ablustrund/apps_plus](https://github.com/ablustrund/apps_plus) |
| IBM CodeNet | [github.com/IBM/Project_CodeNet](https://github.com/IBM/Project_CodeNet) |

### Models & Datasets
| Resource | URL |
|----------|-----|
| Qwen2.5-Coder-7B-Instruct | [huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct) |
| deepmind/code_contests | [huggingface.co/datasets/deepmind/code_contests](https://huggingface.co/datasets/deepmind/code_contests) |
| llvm-ml/ComPile | [huggingface.co/datasets/llvm-ml/ComPile](https://huggingface.co/datasets/llvm-ml/ComPile) |

### Tools
| Tool | Purpose |
|------|---------|
| `aarch64-linux-gnu-gcc` | ARM cross-compiler |
| `qemu-aarch64-static` | ARM userspace emulation |
| `hyperfine` | Benchmarking tool |
| `trackio` | Experiment tracking |

---

## Appendix: Quick Reference Card

```
═══════════════════════════════════════════════════════════════
  ARM COMPILER OPTIMIZATION VIA GRPO — QUICK REFERENCE
═══════════════════════════════════════════════════════════════

Model:     Qwen/Qwen2.5-Coder-7B-Instruct
Method:    GRPO (no critic needed)
Dataset:   CodeNet C programs → ARM assembly (7,872 train / 200 eval)
Reward:    r=0 if fail, r=speedup if all tests pass
LR:        1e-6
Batch:     16 (4 per GPU × 4 GPUs)
Epochs:    1
G:         4 completions per prompt
Temp:      0.5
Seq len:   2048 prompt + 2048 completion
KL (beta): 0.0
Hardware:  4× A100 80GB
Framework: TRL GRPOTrainer
Time:      ~4-8 hours

Expected:  61% → 95% correctness, 1.10× → 1.46× speedup

If base model ARM correctness < 40%, use 3-stage pipeline:
  Stage 1: SFT warmup on (C → ARM assembly) pairs
  Stage 2: Self-filtered retraining
  Stage 3: GRPO with compiler reward
═══════════════════════════════════════════════════════════════
```