rat's register allocator

rat is my smallish compiler backend (with a semi-working C99 frontend). Its x86-64 code generator translates the intermediate representation (IR) into x86-64 instructions. These use an unlimited number of virtual registers (vregs). The register allocator maps each vreg to a physical register: a general-purpose register (12 can be used) or an xmm register (14 on Linux¹). When no register is free, it maps the vreg to a stack slot.

For a long time rat used a linear scan allocator, which visits live ranges in program order. It worked, but it grew one fix at a time, to 1392 lines. So I measured which of its parts helped, threw the rest away and wrote a priority bin-packing allocator in 584 lines. It visits live ranges by importance and puts each one in the first register where it fits. It is the same family as LLVM's greedy allocator, minus most of the hard parts, and it makes better code.

long g(long);
long h(long x, long y) {
	long t = g(x);
	return t + y;
}

0  v1 = copy rdi      ; x
1  v2 = copy rsi      ; y
2  rdi = copy v1      ; argument of g
3  call g             ; clobbers caller-saved
4  v3 = copy rax      ; t
5  v4 = copy v3
6  v4 = add v4, v2
7  rax = copy v4
8  ret

push rbp
mov  rbp, rsp
sub  rsp, 0x8
push rbx              ; rbx is callee-saved: save it
mov  rbx, rsi         ; y
call g                ; x is already in rdi
add  rax, rbx         ; t stays in rax
pop  rbx
leave
ret

Five steps

1. Live ranges: number the instructions and find where each vreg is live.
2. Fixed registers: mark where the code uses physical registers directly.
3. Coalescing: join vregs that a copy connects into one group (a bundle³), so the copy can go away.
4. Picking registers: give each bundle a register, most important first.
5. Spilling: give stack slots to bundles with no register, then rewrite the code.

• takes a register back from a bundle (no eviction)
• splits a range between a register and memory
• runs a step two times

Live ranges

Slots

• Copies: the source ends at the read slot, and the destination starts at the write slot. In instruction 2, rdi = copy v1, v1 ends at slot 4 and rdi starts at slot 5. They do not overlap, so they can share a register and the copy becomes a no-op.
• Other instructions: a source stays live through the write slot, so a result never overwrites a different operand. v2 is written by instruction 1 and last read by the add at instruction 6, so it lives in [3, 13].

Live-out sets

1. The vreg is live into each block that reads it before it writes it.
2. From each such block, a worklist goes back through the predecessors and marks the vreg live-out in each.
3. The walk stops at a block that defines the vreg.

Segments and weights

def or use in	adds
straight-line code	1
a loop	3
a doubly nested loop	9

Holes

long f(long* a, long n) {
	for(long i = 0; i < n; ++i)
		if(a[i] < 0)
			a[i] = 0;
	return n * 3;
}

mov  eax, 0x0         ; offset 8*i, rax in the loop
cmp  rdx, rdi
jl   loop
exit:
lea  rax, [rdi+rdi*2] ; n*3 in the hole of rax
ret
loop:
mov  rcx, r8
add  rcx, rax
...
add  rax, 0x8
cmp  rdx, rdi
jl   loop
jmp  exit

Fixed registers

use	register	busy
incoming argument	argument register	until the copy that reads it
call argument	argument register	from the copy that sets it to the call
call	all caller-saved	in the two slots of the call
return value	rax	from the call to the copy that reads it
division	rax rcx rdx	reads rax rcx, writes rax rdx

instr   0  1  2  3  4  5  6  7  8
slot    rw rw rw rw rw rw rw rw rw
rdi     #. .. .#### .. .. .. .. ..
rsi     ####. .. ## .. .. .. .. ..
rax     .. .. .. ####. .. .. .####
others  .. .. .. ## .. .. .. .. ..
v1 x    .======. .. .. .. .. .. ..
v2 y    .. .================ .. ..
v3+v4 t .. .. .. .. .=========. ..

Coalescing

• two-address instructions
• phi nodes: a value that comes from different blocks at a join point

• rat sorts the copies by loop depth, deepest first. Hot copies merge before cold copies can block them.
• The bundles are kept in a union-find.
• A merge first walks both segment lists to check for overlap. rat skips a merge when the two bundles together have more than 256 segments.

Picking registers

priority = weight / sqrt(length in slots)

• Short, hot ranges come first: they matter most and are the easiest to place.
• Long, cold ranges come last and get spilled.
• sqrt keeps a long loop counter from losing too much priority.

// cls: register class, gp or xmm
PhysReg pick(VReg v) {
	U64 blocked = ~allocatable[cls];
	for(auto [start, end] : segs[v])
		for(I32 s = start; s <= end; ++s)
			blocked |= busy[s]; // or 64 at a time
	if(hint[v] != kNoReg && !(blocked >> hint[v] & 1))
		return hint[v];
	// caller-saved first, callee-saved last
	return firstFree(order[cls], blocked);
}

Picking in h

bundle	hint	gets
v1	rdi	rdi
v3+v4	rax	rax
v2	rsi	rbx

Spilling

1. rat sorts the spilled bundles by start. It reuses a stack slot when the last bundle in it has ended.
2. rat rewrites the code. A vreg whose bundle got a register becomes that register.
3. Before each instruction, rat loads each spilled operand into a temporary register. After it, rat stores each spilled result.

void p(long* a) {
	long x0 = a[0], x1 = a[1], ..., x13 = a[13];
	a[0] = x0 * x13; a[1] = x1 * x12; a[2] = x2 * x11;
	a[3] = x3 * x10; a[4] = x4 * x9;  a[5] = x5 * x8;
	a[6] = x6 * x7;
}

mov  r12, [rdi+0x30]  ; x6, in a register
mov  r10, [rdi+0x38]  ; x7, spilled
mov  [rbp-0x8], r10
mov  r10, [rdi+0x40]  ; x8, spilled
mov  [rbp-0x10], r10
mov  r10, [rdi+0x48]  ; x9, spilled
mov  [rbp-0x18], r10
...
mov  r10, [rbp-0x18]  ; reload x9
imul r9, r10
mov  r10, [rbp-0x10]  ; reload x8
imul rbx, r10
mov  r10, [rbp-0x8]   ; reload x7
imul r12, r10

It can be dumb

long sum(long* a, long n) {
	long s = 0;
	for(long i = 0; i < n; ++i)
		s += a[i];
	return s;
}

mov  r9, rdi          ; a: rdi was taken by a[i]
mov  r8, rsi          ; n: rsi was taken by s
...
exit:
mov  rax, rsi         ; s: rax was taken by a+8*i
ret
loop:
mov  rax, r9
add  rax, rcx         ; rax = a + 8*i
mov  rdi, [rax]       ; rdi = a[i]
add  rsi, rdi
...

1. The address a+8*i takes rax.
2. s loses its hint rax and takes rsi.
3. a[i] takes rdi.
4. a and n come last and lose their hints too.

Numbers

metric	change
allocator source	-58%
instructions emitted	-3.6%
stores emitted	-22%
allocator time, sqlite at -O0	-77%
total compile time, sqlite at -O0	-43%

What each feature was worth

feature	instructions saved	new allocator
copy coalescing and copy hints	about a third	kept
live range holes	10%	kept
spill choice by use weight	5.5%	kept
hints to physical registers	1%	kept
optimistic second try at spilled ranges (37% of allocator time)	0.01%	dropped
rematerialization (recompute instead of reload)	not measurable	dropped
spill slot cache	not measurable	dropped

Wrapping up

References

• Poletto and Sarkar, Linear scan register allocation: the base of the old allocator.
• Max Bernstein, Linear scan register allocation on SSA and Linear scan with lifetime holes: a readable pair of posts.
• Jakob Stoklund Olesen, Greedy register allocation in LLVM 3.0: the big version of this idea, with eviction and splitting.
• Chris Fallin, Cranelift, part 4: a new register allocator: a long, good read on bundles.
• Matt Keeter, The solid-state register allocator: even smaller, it runs in one backward pass.

Notes

1. On Windows, only xmm0-xmm3 can be used. xmm4 and xmm5 are the spill temporaries, and rat does not use the callee-saved xmm6-xmm15. [back]
2. Chaitin et al. showed that any graph can be the interference graph of some program. So register allocation is at least as hard as graph coloring. [back]
3. Cranelift's regalloc2 uses the same word for the same idea. [back]
4. Each block stores the last vreg that marked it, so the walk never clears a visited array. [back]
5. The mov inside the loop, mov rax, r9, has a different cause. It is the two-address copy for the add, and it always stays. [back]