VeriLocc: Using LLMs to Allocate GPU Registers

Compiler Optimization
By Haohui Mai, CausalFlow.ai, November 2025

Why Low-Level Compiler Work Matters

Modern GPUs devote much of their die area to compute units and use relatively simple scheduling hardware. The compiler must therefore make many low-level decisions before a kernel runs, and those decisions can significantly affect performance.

One example is register allocation: deciding where the GPU stores temporary values and when that storage can be reused. Production compilers solve this problem with hand-crafted heuristics developed through extensive testing and hardware-specific tuning. Keeping those heuristics current is difficult as workloads and GPU architectures change.

This challenge is particularly relevant to AMD, where compiler support must keep pace with changing GPU architectures and workloads. Gaps in backend tuning can prevent applications from reaching the performance available in the hardware.

In our EMNLP 2025 paper, co-authored with Lesheng Jin, Zain Ruan, and Jingbo Shang, we explore whether large language models (LLMs) can reduce this manual effort. Our system, VeriLocc, learns register allocation from existing compiler output and verifies each generated result before returning it to the compiler.

What Is Register Allocation?

Registers are small, fast, and limited storage locations inside a GPU. A kernel uses them for addresses, loop counters, matrix values, and partial results.

During compilation, each temporary value first receives a virtual register. Register allocation maps these virtual registers to the physical registers available on the GPU. Values that are needed at the same time must use different physical registers; after a value's final use, its register can be reused.

Consider this program:

v2 = v0 + v1
v3 = v2 + v1
v4 = v2 + v3

Mapping v2 and v3 to different physical registers is valid. Mapping both to the same register is not, because both values are needed to compute v4.

Real kernels contain many more values and hardware constraints. If the allocator runs out of registers, it may move values to slower memory. Poor assignments can also reduce parallelism or create register-bank conflicts. Finding the best allocation is NP-complete, so production compilers rely on heuristics rather than an exact solution.

Register Allocation as Translation

VeriLocc treats register allocation as a language translation task. NVIDIA and AMD use different intermediate representations (IRs), but both describe operations, values, and the dependencies between them. VeriLocc views these IRs as different dialects of a shared low-level language and learns to translate a kernel into a register mapping for the target GPU.

Making this work requires three components: IR normalization, fine-tuning on real compiler output, and formal verification.

VeriLocc training and checking workflow
Figure 1: VeriLocc normalizes the compiler input, generates a register mapping, and verifies it before returning it.

First, VeriLocc normalizes the IR. It removes comments, debug data, and setup code that is not relevant to allocation. It gives equivalent NVIDIA and AMD instructions the same name, shortens register identifiers, and makes dependencies between values explicit. These changes reduce input size by 80–90%, leaving the model with a shorter and more consistent representation.

Second, we fine-tune Qwen2.5-Coder-7B on real allocations from general matrix multiplication (GEMM) and multi-head attention kernels compiled for NVIDIA RTX 4090 and AMD MI210 GPUs. For each example, the model sees the normalized compiler input and the mapping produced by the target compiler.

Finally, VeriLocc checks each generated mapping. An invalid allocation can crash a kernel or silently produce the wrong result. VeriLocc expresses the allocation rules as a Satisfiability Modulo Theories (SMT) problem and checks them with the Z3 solver. If a candidate overwrites a live value or violates a hardware constraint, VeriLocc rejects it and generates another.

Results on Matrix Multiplication and Attention

We evaluate VeriLocc on 768 GEMM cases and 1,865 attention cases using NVIDIA-only, AMD-only, and mixed NVIDIA/AMD training.

For GEMM, 98.05–99.48% of the first generated allocations pass verification. Attention is more difficult because its kernels are longer and use more values at once; its first-attempt pass rate ranges from 85.84% to 96.30%.

Generating additional candidates improves coverage. With up to 100 attempts—reported as pass@100—the best settings reach 99.86% for GEMM and 99.74% for attention. Mixed-architecture attention is the hardest setting, reaching 89.76% pass@100.

Correct register mappings produced by VeriLocc for matrix multiplication and attention
Figure 2: First-attempt and pass@100 verification rates.

Normalization has its largest effect on mixed-architecture attention. It raises the first-attempt pass rate from 79.95% to 85.84% and reduces the average attempts from 9.11 to 6.37.

The paper also compares one VeriLocc allocation with rocBLAS, AMD's optimized matrix multiplication library. For an FP16 GEMM on AMD MI250x, the VeriLocc allocation reaches 111.44 TFLOPS—11.63% higher throughput than rocBLAS for the same matrix sizes and batch setting.

VeriLocc uses registers dedicated to the matrix units more often than rocBLAS. These registers may provide faster access for this operation, although the public hardware documentation does not confirm the reason. This is a single case study, not evidence that VeriLocc is generally faster than rocBLAS.

Conclusion

VeriLocc demonstrates that a model can learn register-allocation patterns from existing compiler output, apply them across NVIDIA and AMD toolchains, and use formal verification to reject invalid results.

More broadly, the work suggests a design for next-generation compilers in which LLMs participate in the optimization loop alongside static analysis and verification. Performance measurements could feed back into training as workloads and hardware evolve. For AMD, this could allow compiler knowledge from earlier architectures to be reused instead of rebuilding optimization heuristics for each GPU generation.