VeriLocc: End-to-End Cross-Architecture Register Allocation via LLM
This addresses the bottleneck of manual compiler optimization for GPU architectures, offering a generalizable and verifiable solution.
The paper tackles the problem of GPU register allocation requiring manual re-tuning for each hardware generation by introducing VeriLocc, a framework combining LLMs with formal compiler techniques. The result is 85-99% single-shot accuracy and over 10% runtime improvement compared to expert-tuned libraries.
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.