| # 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 |
| βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ |
| ``` |
| |