De register-allocator van rat

rat is mijn relatief kleine compiler-backend (met een semi-werkende C99-frontend). De x86-64 code-generator vertaalt de intermediate representation (IR) naar x86-64 instructies. Deze maken gebruik van een onbeperkt aantal virtuele registers (vregs). De register-allocator koppelt elke vreg aan een fysiek register: een general-purpose register (hiervan kunnen er 12 worden gebruikt) of een xmm-register (14 op Linux). Wanneer er geen register vrij is, wordt de vreg gekoppeld aan een stack slot.

Lange tijd maakte rat gebruik van een linear scan allocator, die live ranges bezoekt in de volgorde van het programma. Dit werkte, maar de implementatie groeide met elke fix tot 1392 regels. Daarom heb ik gemeten welke onderdelen nuttig waren, de rest weggegooid en een priority bin-packing allocator geschreven in 584 regels. Deze bezoekt live ranges op basis van belangrijkheid en plaatst elke range in het eerste register waar deze past. Het behoort tot dezelfde familie als de greedy allocator van LLVM, maar zonder de meeste complexe onderdelen, en het produceert betere code.

Een waarde is "live" vanaf het moment dat deze wordt geschreven tot het moment dat deze voor het laatst wordt gelezen. Twee waarden kunnen slechts een register delen als ze nooit tegelijkertijd live zijn. Wanneer er op een bepaald punt te veel waarden live zijn, gaan sommige naar het geheugen: ze worden gespilled. Een spill kost een store en een load. Omdat de beste toewijzing NP-hard is om te vinden, maken alle praktische allocators gebruik van heuristieken.

De calling convention

De calling convention voegt twee regels toe. Een aanroep (call) kan de caller-saved registers overschrijven (rax, rcx, rdx, rsi, rdi, r8-r11 en alle xmm-registers op Linux). Een functie moet de callee-saved registers (rbx, rbp, r12-r15) herstellen voordat deze terugkeert.

Als voorbeeld: deze functie houdt y live over een aanroep heen:

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

Vóór de allocatie zijn rdi, rsi en rax vastgelegd door de calling convention, en zijn v1-v4 virtuele registers (vregs):

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

Omdat x86 add schrijft over zijn eerste operand (two-address), kopieert instructie 5 eerst t. Na allocatie bij -O1:

push rbp
mov  rbp, rsp
sub  rsp, 0x8
push rbx              ; rbx is callee-saved: bewaar deze
mov  rbx, rsi         ; y
call g                ; x staat al in rdi
add  rax, rbx         ; t blijft in rax
pop  rbx
leave
ret

Vijf van de zes kopieën zijn verdwenen en y is naar een callee-saved register gegaan. Er staat nergens in de allocator: "plaats waarden die een call overbruggen in callee-saved registers". Dit vloeit voort uit het ontwerp, en dat is mijn favoriete onderdeel.

De vijf stappen

De allocator doorloopt vijf stappen per functie:

1. Live ranges: Nummer de instructies en bepaal waar elke vreg live is.
2. Vaste registers: Markeer waar de code direct gebruikmaakt van fysieke registers.
3. Coalescing: Voeg vregs die via een kopieeractie verbonden zijn samen tot één groep (een bundle), zodat de kopieeractie kan worden verwijderd.
4. Registers kiezen: Wijs aan elke bundle een register toe, beginnend bij de belangrijkste.
5. Spilling: Wijs stack slots toe aan bundles zonder register en herschrijf vervolgens de code.

Elke bundle behoudt zijn register of stack slot voor zijn volledige levensduur. De allocator doet het volgende nooit:

• Een register terugnemen van een bundle (geen eviction).
• Een range splitsen tussen een register en het geheugen.
• Een stap twee keer uitvoeren.

Deze onderdelen maken commerciële allocators groot. Mijn metingen tonen aan dat rat hier weinig op verliest.

Live ranges

Slots

Instructie i krijgt twee slots: hij leest zijn operanden op 2i en schrijft zijn resultaten op 2i+1. Een live range is een gesorteerde lijst van [start, end] slot-segmenten.

Waar een bron eindigt, hangt af van de instructie:

• Kopieën: De bron eindigt bij het lees-slot, en de bestemming begint bij het schrijf-slot. In instructie 2 (rdi = copy v1) eindigt v1 bij slot 4 en begint rdi bij slot 5. Ze overlappen niet, dus ze kunnen een register delen en de kopie wordt een no-op.
• Overige instructies: Een bron blijft live door het schrijf-slot, zodat een resultaat nooit een andere operand overschrijft. v2 wordt geschreven door instructie 1 en laatst gelezen door de add bij instructie 6, dus is deze live in [3, 13].

Live-out sets

rat vindt de vregs die live-out zijn voor elk blok: een later blok kan ze namelijk nog lezen. Veel compilers doen dit met één bitset per blok en een fixed-point loop. rat doet dit in plaats daarvan per vreg:

1. De vreg is live in elk blok dat de vreg leest voordat deze hem schrijft.
2. Vanuit elk zodanig blok gaat een worklist terug via de voorgangers en markeert de vreg als live-out in elk van deze.
3. De wandeling stopt bij een blok dat de vreg definieert.

De kosten groeien mee met de blokken waar elke vreg live is, niet met blokken * vregs.

Segmenten en gewichten

Vervolgens loopt rat elk blok achterwaarts af vanuit de live-out set en maakt de segmenten. Tijdens dezelfde wandeling wordt per vreg een gewicht berekend: de kosten van een spill. Elke definitie (def) en elk gebruik (use) voegt 3^d toe, waarbij d de lusdiepte is (tot maximaal 11):

Locatie	Toevoeging
Lineaire code	1
In een lus	3
In een dubbel geneste lus	9

Holes (Gaten)

Een live range kan gaten hebben: intervallen waarin de vreg dood is. Omdat blokken in codevolgorde zijn genummerd, heeft een range die een blok overslaat daar een gat:

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

Het exit-blok bevindt zich tussen de lus-blokken:

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

De offset in rax is dood in het exit-blok, dus n*3 (één lea) kan rax gebruiken, wat ook het return-register is. Een gratis winst door de blokvolgorde.

Vaste registers

rat nummert registers van 1 tot 40, zodat één U64 een set registers kan bevatten. Elk slot krijgt een masker, busy[slot]. Een gezet bit betekent dat dat register op dat slot bezet is. Dezelfde achterwaartse wandeling markeert de fysieke registers die de code direct gebruikt:

Gebruik	Register	Bezet van/tot
Inkomend argument	argument register	tot de kopie die het leest
Call argument	argument register	van de kopie die het zet tot de call
Call	alle caller-saved	in de twee slots van de call
Return waarde	rax	van de call tot de kopie die het leest
Divisie	rax, rcx, rdx	leest rax, rcx; schrijft rax, rdx

De maskers en ranges van h:

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 .. .. .. .. .=========. ..

(r en w zijn de read en write slots. # is bezet, = is een live range en . is vrij. Een streep loopt over de kloof tussen instructies. "others" zijn alle overige caller-saved registers.)

Wanneer een bundle een register krijgt, zet rat het bit van dat register in elk slot van zijn live range. Daarna zijn vregs en vaste registers bits in dezelfde maskers. Elke groep van 64 slots heeft ook een samenvattingsmasker (de OR van de 64 maskers), zodat een lange range 64 slots per keer kan overslaan.

Coalescing (Samenvoegen)

Een kopie tussen twee vregs van dezelfde klasse is een kandidaat voor coalescing. Deze kopieën ontstaan door:

• Two-address instructies.
• Phi-nodes: een waarde die uit verschillende blokken komt bij een join-point.

Als de twee live ranges niet overlappen, worden de vregs één bundle. Deze heeft de samengevoegde segmenten en het opgetelde gewicht. rat verwijdert kopieën binnen één bundle. In h is v3 [9, 10] en v4 [11, 14], dus versmelten ze.

rat sorteert de kopieën op lusdiepte, diepste eerst. "Hete" kopieën versmelten voordat "koude" kopieën hen kunnen blokkeren. De bundles worden bijgehouden in een union-find. Een merge controleert eerst beide segmentlijsten op overlap. rat slaat een merge over als de twee bundles samen meer dan 256 segmenten hebben.

Een kopie tussen een vreg en een fysiek register zet in plaats daarvan een hint: de bundle geeft de voorkeur aan dat register als het vrij is.

Registers kiezen

Elke bundle krijgt een prioriteit: prioriteit = gewicht / sqrt(lengte in slots)

Korte, hete ranges komen eerst: zij zijn het belangrijkst en het makkelijkst te plaatsen. Lange, koude ranges komen als laatste en worden gespilled. De sqrt voorkomt dat een lange lus-teller te veel prioriteit verliest.

rat roept pick aan voor elke bundle in prioriteitsvolgorde:

// cls: register class, gp of 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]; // of 64 tegelijk
    if(hint[v] != kNoReg && !(blocked >> hint[v] & 1))
        return hint[v];
    // caller-saved eerst, callee-saved laatst
    return firstFree(order[cls], blocked);
}

Het kiezen in h:

Bundle	Hint	Krijgt
v1	rdi	rdi
v3+v4	rax	rax
v2	rsi	rbx

In het diagram is rdi alleen bezet vóór en na v1, dus v1 krijgt het. Beide kopieën worden mov rdi, rdi, en de peephole-pass verwijdert deze na de allocatie. v2 overbrugt de call. Elk caller-saved register is bezet in de call-slots, dus het eerste vrije register is rbx, de eerste callee-saved. De prologue bewaart alleen de callee-saved registers die rat daadwerkelijk heeft gebruikt.

Op Linux is geen enkel xmm-register callee-saved, dus een float die een call overbrugt, gaat altijd naar de stack.

Spilling

Een bundle zonder vrij register wordt voor zijn volledige levensduur gespilled. Vervolgens:

1. rat sorteert de gespilled bundles op startpunt. Een stack slot wordt hergebruikt zodra de laatste bundle in dat slot is geëindigd.
2. rat herschrijft de code. Een vreg wiens bundle een register kreeg, wordt dat register.
3. Vóór elke instructie laadt rat elke gespilled operand in een tijdelijk register. Daarna slaat rat elk gespilled resultaat op.

Het tijdelijke register is r10 of r11 (xmm14 of xmm15 voor floats). Geen enkele bundle krijgt deze registers. Als beide bezet zijn, neemt rat het eerste register dat op dat moment vrij is.

Twee gevallen hebben geen tijdelijk register nodig: een kopie tussen een register en een gespilled bundle wordt de load of store zelf. Een call leest een gespilled stack-argument direct uit het stack slot.

In functie p zijn 14 waarden tegelijkertijd live:

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;
}

16 registers minus rsp, rbp, r10, r11 en rdi (die a bevat) laat 11 registers over voor 14 waarden. x0-x6 bevatten ook de producten (two-address imul), dus hebben zij meer uses. Van x7-x13 gaan de drie met de langste ranges naar de stack.

Vóór de peephole-pass:

mov  r12, [rdi+0x30]  ; x6, in een 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

Er staat geen instructie tussen de store van x9 en de reload ervan die r10 overschrijft. De peephole-pass verwijdert daarom de reload. Omdat er vervolgens niets meer uit dat stack slot leest, wordt ook de store verwijderd.

Het kan dom zijn

Zonder eviction is een vroege beslissing definitief. Hier is het scenario dat mij het meest irriteert:

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

rat compileert dit naar:

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

De loop-waarden zijn kort en heet, dus zij gaan eerst:

1. Het adres a+8*i neemt rax.
2. s verliest zijn hint rax en neemt rsi.
3. a[i] neemt rdi.
4. a en n komen als laatste en verliezen ook hun hints.

Het resultaat zijn drie movs, allemaal buiten de lus. Een allocator met eviction zou deze keten oplossen. Ik heb besloten dat drie koude movs de extra code niet waard zijn.

Cijfers

Vergeleken met de oude allocator:

Metriek	Verandering
Source code allocator	-58%
Emitted instructies	-3.6%
Emitted stores	-22%
Allocator tijd (sqlite bij -O0)	-77%
Totale compilatietijd (sqlite bij -O0)	-43%

De waarde van elke feature: Vóór de rewrite heb ik elke oude feature één voor één uitgeschakeld en de code gemeten. Dit was het meest nuttige uur van het project:

Feature	Instructies bespaard	Nieuwe allocator
Copy coalescing & copy hints	ongeveer een derde	behouden
Live range holes	10%	behouden
Spill keuze op basis van use weight	5.5%	behouden
Hints naar fysieke registers	1%	behouden
Optimistische tweede poging bij spilled ranges	0.01%	verwijderd
Rematerialization	niet meetbaar	verwijderd
Spill slot cache	niet meetbaar	verwijderd

Afsluiting

Geen eviction, geen splitting, geen tweede pass, en toch verslaat de nieuwe allocator de oude. Het meeste kwaliteitsverbetering komt voort uit eenvoudige zaken: coalescing, hints, holes en use weights.

De les voor mij: meet de oude code voordat je deze port. Een groot deel van de oude allocator deed niets.

Referenties

• Poletto and Sarkar, Linear scan register allocation: de basis van de oude allocator.
• Max Bernstein, Linear scan register allocation on SSA en Linear scan with lifetime holes: een leesbaar paar posts.
• Jakob Stoklund Olesen, Greedy register allocation in LLVM 3.0: de grote versie van dit idee, inclusief eviction en splitting.
• Chris Fallin, Cranelift, part 4: a new register allocator: een lange, goede lees over bundles.
• Matt Keeter, The solid-state register allocator: nog kleiner, draait in één achterwaartse pass.

Notities

1. Op Windows kunnen alleen xmm0-xmm3 worden gebruikt. xmm4 en xmm5 zijn de spill-temporaries, en rat gebruikt de callee-saved xmm6-xmm15 niet.
2. Chaitin et al. toonden aan dat elke graaf de interferentiegraaf van een programma kan zijn. Register-allocatie is dus minstens zo moeilijk als graaf-kleuring.
3. De regalloc2 van Cranelift gebruikt hetzelfde woord voor hetzelfde concept.
4. Elk blok bewaart de laatste vreg die het markeerde, zodat de wandeling nooit een bezochte array hoeft leeg te maken.
5. De mov in de lus (mov rax, r9) heeft een andere oorzaak. Dit is de two-address kopie voor de add, en deze blijft altijd bestaan.