Rat's Register Allocator

rat's register allocator 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 Live ranges: number the instructions and find where each vreg is live. Fixed registers: mark where the code uses physical registers directly. Coalescing: join vregs that a copy connects into one group (a bundle3), so the copy can go away. Picking registers: give each bundle a register, most important first. 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 The vreg is live into each block that reads it before it writes it. From each such block, a worklist goes back through the predecessors and marks the vreg live-out in each. The walk stops at a block that defines the vreg. Segments and weights 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 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 Spilling rat sorts the spilled bundles by start. It reuses a stack slot when the last bundle in it has ended. rat rewrites the code. A vreg whose bundle got a register becomes that register. 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 ... The address a+8*i takes rax. s loses its hint rax and takes rsi. a[i] takes rdi. a and n come last and lose their hints too. Numbers What each feature was worth 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 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] 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] Cranelift's regalloc2 uses the same word for the same idea. [back] Each block stores the last vreg that marked it, so the walk never clears a visited array. [back] 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]