VeriLocc: End-to-End Cross-Architecture Register Allocation via LLM

TL;DR

VeriLocc leverages large language models combined with formal techniques to automate and verify register allocation for GPUs, achieving high accuracy and outperforming expert-tuned libraries in performance.

Contribution

The paper introduces VeriLocc, a novel framework that uses LLMs and formal verification for cross-architecture register allocation in GPUs, enabling generalization and correctness.

Findings

Achieves 85-99% single-shot accuracy in register assignment

Near-100% pass@100 correctness rate

Outperforms rocBLAS by over 10% in runtime

Abstract

Modern GPUs evolve rapidly, yet production compilers still rely on hand-crafted register allocation heuristics that require substantial re-tuning for each hardware generation. We introduce VeriLocc, a framework that combines large language models (LLMs) with formal compiler techniques to enable generalizable and verifiable register allocation across GPU architectures. VeriLocc fine-tunes an LLM to translate intermediate representations (MIRs) into target-specific register assignments, aided by static analysis for cross-architecture normalization and generalization and a verifier-guided regeneration loop to ensure correctness. Evaluated on matrix multiplication (GEMM) and multi-head attention (MHA), VeriLocc achieves 85-99% single-shot accuracy and near-100% pass@100. Case study shows that VeriLocc discovers more performant assignments than expert-tuned libraries, outperforming rocBLAS by over 10% in runtime.