Register allocation on machine code

This walkthrough runs the register allocator both ways: pinned, where it must reproduce the exact allocation ptxas chose, and free, where it recolors for real. It assumes you can already lift a kernel — see the quickstart or Capturing & lifting kernels — and uses only the regalloc, lower, and crosscheck commands from the CLI reference.

Why allocate registers on machine code

Register allocation is the one optimization ptxas bakes deepest into the binary. Every operand field, every .nv.info register count, every spill slot is already decided when the cubin ships — and the decision is frozen, because nothing downstream of ptxas can re-encode an instruction, let alone rename every register in one.

That decision has direct occupancy consequences. Per-thread registers are allocated out of a fixed 64k-register file per SM, so a kernel asking for 255 registers per thread — the hardware maximum — caps residency at 12.5% occupancy. That is exactly the situation the official FlashAttention forward kernel ships in: 255 registers, 85 spills to local memory, and an occupancy ceiling nothing above the SASS layer can relax. Being able to redo allocation after the fact means occupancy fixes stop requiring source access or a vendor toolchain update.

sass2mlir’s allocator is Chaitin-Briggs graph coloring over the recovered SSA — the same SSA that defuse, liveness, and uniform already run on — extended past the general-purpose class to the uniform register and predicate classes (UR/UP) that newer architectures use for warp-scalar values. It runs in two modes, and both matter:

• pinned — reproduce the byte-identical allocation the hardware was given. This is the correctness oracle for the whole analysis stack.
• free — real recoloring under a fresh interference graph, validated by SSA contributor-set equivalence.

The oracle: reproduce ptxas exactly

sass2mlir regalloc out/flash_fwd_kernel.mlir --fn flash_fwd_kernel --pinned
# regalloc(pinned): flash_fwd_kernel — byte-identical allocation reproduced

Pinned mode re-runs the full pipeline — liveness, interference-graph construction, Chaitin-Briggs coloring, writeback — and then compares the result against the original bytes. Not “a valid allocation”: the allocation, register for register.

This is what makes it the hard oracle. Liveness and interference are easy to get almost right — an off-by-one live range or a missed edge through a BSSY/BSYNC convergence pair still produces a plausible-looking coloring. But in pinned mode there is nowhere to hide: any error in the liveness or interference model shows up as a wrong register in at least one operand field, and the byte-compare catches it. A pinned run that reproduces every register the binary carries is proof that the recovered SSA, the per-block liveness, and the interference graph are all faithful to what was actually encoded — the same property the corpus-level MetricTable on the overview reports as 0 failures across 13/13 architectures.

The oracle only earns its keep at corpus scale, so it is built to be looped. With -q the per-function lines go away and only failures print; exit code 0 means zero mismatches, so a shell loop is a CI job:

fails=0
for f in corpus/sm_90a/*.mlir; do
  sass2mlir regalloc "$f" --pinned -q || { fails=$((fails+1)); echo "FAIL: $f"; }
done
echo "pinned regalloc: $fails failures"
# pinned regalloc: 0 failures

A non-zero count comes with the function name, the instruction offset, and the expected vs. actual register on stderr — enough to jump straight to the broken live range.

Real recoloring under register pressure

Free mode drops the byte-identity constraint and colors to reduce pressure: same SSA, same interference graph, but the allocator is allowed to make any sound choice. It writes a new module rather than touching the input:

sass2mlir regalloc out/attn_fwd_h32.mlir --fn attn_fwd_h32 --free -o reallocated.mlir
# regalloc(free): attn_fwd_h32 — 168 regs → 154 regs, contributor sets equivalent

The target here is a smaller attention forward kernel in the same bind as the flagship: compiled with the register budget maxed, spilling 14 values to local memory, and losing resident warps to its own register footprint. The recoloring frees 14 registers per thread — enough that the spills the old allocation was forced into no longer happen. The register count is the only contract free mode changes; it does not reorder, reschedule, or touch the CFG.

- FFMA.FTZ  R10, R8, R14, R10    // 168 regs live at peak
- STL      [R7.64], R12          // spill forced by pressure
+ FFMA.FTZ  R8, R4, R12, R8      // recolored: 154 regs at peak
+ // (spill gone — R12-class value stays resident)

Recoloring is also where the UR/UP classes matter: on sm_90-class hardware, values the uniformity analysis proved warp-scalar live in the uniform register file with its own allocator pressure, and the interference graph has to color both files consistently. Free mode treats them as what they are — separate register classes with their own budgets, one graph.

Because the output is no longer byte-identical, “it produced a module” is not a correctness argument. The in-IR validation is SSA contributor-set equivalence: the recoloring must preserve, for every program point, which original values each physical register can carry. The contributor sets equivalent line in the output is that check passing.

// live-ins: R2 (param base), R5 (loop bound)
%a = sass.MOV_0x2a4 R2 : (!sass.reg) -> (!sass.reg)        // C(%a) = {R2}
%b = sass.IADD3_0x1cc %a, 0x4, R5 : (!sass.reg, !sass.reg) -> (!sass.reg)
                                                           // C(%b) = {R2, R5}
// %a dies at the IADD3, so %b may legally take %a's color —
// their live ranges don't overlap, and {R2, R5} ⊇ {R2} is fine:
// equivalence is about the *partition*, not the color names.

Verifying the result

First, lower the recolored module to bytes:

sass2mlir lower reallocated.mlir -o reallocated.code

Then notice what you can no longer do. roundtrip against the original cubin is supposed to fail now — free mode changes bytes by design, so round-trip identity stops being the safety net the moment recoloring is allowed. The verification stack for a reallocated kernel is, in order:

1. Contributor-set equivalence — proved in-IR by the allocator itself, above. This is the argument that the new allocation means the same thing.

2. crosscheck — re-encode every instruction of the new module from its structural fields and compare against the $raw the writeback produced:

   sass2mlir crosscheck reallocated.mlir
   # crosscheck: attn_fwd_h32 — 1,042/1,042 instructions re-encoded byte-identical (0 mismatches)

   Writeback invalidates the $raw of every instruction whose register fields changed, so those bytes come out of nvisa_encode, not the original stream. crosscheck is the audit that the assembler path got the new encodings exactly right.

3. Run the kernel. The structural checks are necessary, not sufficient — the end-to-end proof is executing the patched cubin on the same inputs. The full capture → apply → in-process validation flow is the Recompiling FlashAttention walkthrough.

One caveat to keep straight: pinned mode’s 0-failures record is corpus-wide and regression-gated; the free-mode numbers on this page are for the specific small kernel shown, not a claim about the FlashAttention pipeline — the structural passes that attack the flagship kernel’s 255-register allocation are still report-only on the general corpus, and their measured results land on the overview as the sweeps complete.