kaori02/arm-compiler-optimization-wiki

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
2. Landscape & Key Papers
3. Recipe 1: SuperCoder — RL Assembly Superoptimization (SOTA)
4. Recipe 2: Meta LLM Compiler — SFT on LLVM IR
5. Recipe 3: Compiler Feedback — Iterative Refinement
6. Recipe 4: CUDA-L1 — Contrastive RL (3-Stage Pipeline)
7. Recipe 5: StepCoder — Fine-Grained RL for Code
8. ARM Adaptation Guide
9. Datasets
10. Model Selection
11. Reward Function Design
12. Reward Hacking & Mitigations
13. Training Infrastructure: TRL GRPO Implementation
14. Full Training Script
15. Program Transformation Taxonomy
16. Results Benchmarks & Ablations
17. Citation Graph & Future Directions
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	2025	PPO/GRPO on assembly with compiler reward	95% correct, 1.46× speedup over gcc -O3
🥈	CUDA-L1	2025	3-stage SFT→Self-supervised→Contrastive RL	3.12× avg speedup on KernelBench
🥉	Meta LLM Compiler	2023	SFT from scratch on LLVM IR for pass ordering	3.0% instruction reduction over -Oz
4	Compiler Feedback	2024	Iterative refinement with compiler oracle	+0.53% over base; sampling > feedback
5	StepCoder	2024	Fine-Grained Optimization masking for code RL	+8% pass@1 on APPS+
6	MLGO	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 | 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	§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 | 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 | 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 | 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 | 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

# 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	C/C++ source + test I/O	8M+ submissions	✅ Recompile with ARM GCC	Used by SuperCoder
deepmind/code_contests	HF Hub	C/C++ solutions + tests	~2GB train	✅ Filter C, cross-compile	Has public/private/generated tests
llvm-ml/ComPile	HF Hub	LLVM bitcode IR	602GB (2.7TB source)	✅ Retarget to AArch64 via llc	C/C++/Rust/Swift from Spack
APPS+	GitHub	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.

{
    "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

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

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)

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

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

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:

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)

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

"""
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	2025
CUDA-L1	2507.14111	2025
Meta LLM Compiler	2309.07062	2023
Compiler Feedback	2403.14714	2024
StepCoder	2402.01391	2024
MLGO	2101.04808	2021
ProGraML	2012.01470	2020
DeepSeekMath (GRPO)	2402.03300	2024
VeriReason	2505.11849	2025
ACECoder	2502.01718	2025

Code & Frameworks

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

Models & Datasets

Resource	URL
Qwen2.5-Coder-7B-Instruct	huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct
deepmind/code_contests	huggingface.co/datasets/deepmind/code_contests
llvm-ml/ComPile	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
═══════════════════════════════════════════════════════════════