Het vervangen van een Rust Enum door een 64-bit Word maakte mijn interpreter 17% sneller

Plush is een dynamisch getypeerde taal, vergelijkbaar met Python, JavaScript, Ruby, Lua en Lox. Bij dit soort talen zijn types gekoppeld aan waarden in plaats van aan variabelen. Om waarden door programma's te propageren, gebruikt een interpreter doorgaans een Value-type dat elke mogelijke waarde in de taal kan representeren. In de oorspronkelijke versie van Plush gebruikte ik hiervoor een standaard Rust tagged enum. Dit is handig omdat Rust het werken met tagged enums zeer eenvoudig maakt; we kunnen via een match-statement dispatchen naar verschillende Value-subtypes:

// Het oude Rust Value-type als een tagged enum
enum Value {
    Undef,              // Niet-geïnitialiseerde var of field, lezen triggert een fout
    Nil,
    False,
    True,
    Int64(i64),
    Float64(f64),
    String(*const Str),        // Onveranderlijke string
    HostFn(&'static HostFn),   // Functie blootgesteld door de host VM
    Fun(FunId),                // Niet-closure Plush-functie
    Closure(*mut Closure),     // Closure die variabelen capturet
    Cell(*mut Value),          // Mutable variabele gecaptured door een closure
    Object(*mut Object),       // Class-instanties
    Array(*mut Array),         // Array/lijst in JS/Python stijl
    ByteArray(*mut ByteArray), // Snelle raw byte array (bijv. frame buffer)
    Dict(*mut Dict),           // Dict in JS/Python stijl
    Class(ClassId),
}

Hoewel Plush nog steeds een 'toy language' is, heeft het veel verschillende waardetypes. Er zijn objecten (class-instanties) voor efficiënte toegang, maar ook dictionaries in JS/Python-stijl om JSON-achtige syntaxis mogelijk te maken. Daarnaast zijn er twee afzonderlijke numerieke types: Int64 en Float64. Deze keuze is gemaakt omdat het storend is dat JavaScript doet alsof alles een double is, terwijl JS-engines op de achtergrond bijhouden wat een integer is.

Het grootste probleem is echter niet het aantal varianten, maar de breedte van deze enum: 16 bytes (128 bits). Elke variant heeft slechts 64 bits nodig en de enum-tag die Rust aanmaakt heeft slechts 8 bits nodig, maar vanwege geheugenuitlijningsbeperkingen (memory alignment) kan Rust 128 bits per waarde gebruiken. Bij grote arrays van waarden leidt dit tot een enorme hoeveelheid verspilde bytes.

Een efficiënt low-bit tagging-schema

Om de Value-type in 64 bits te laten passen, heb ik gekeken naar een efficiënter low-bit tagging-schema. Een klassieke truc is gebaseerd op het feit dat heap-objectadressen op een 64-bits systeem doorgaans zijn uitgelijnd op 8-byte grenzen, wat betekent dat de laagste 3 bits van het adres altijd nul zijn. Deze bits kunnen worden "gestolen" om extra informatie op te slaan. Ook kunnen de twee laagste bits van integer-waarden worden gebruikt, ervan uitgaande dat integers zelden het volledige 64-bits bereik nodig hebben ($2^{64} \approx 1,84 \times 10^{19}$ is een astronomisch groot getal).

Een geavanceerder tagging-schema vermindert het geheugengebruik, maar introduceert bitwise-operaties om te bepalen of een waarde een integer, een pointer of een float is. Dit betekent extra instructies voor de CPU. Een ervaren VM-engineer wees me erop dat het slim is om integers nullen als tag-bits te geven; het optellen of aftrekken van twee verschoven integers blijft dan een eenvoudige machine-instructie. Het inpakken en uitpakken van floats is complexer. Ik was bezorgd over de prestaties, maar die angst bleek onterecht.

Het uiteindelijke schema is geïmplementeerd als een Rust newtype dat een u64 wrapt. Om het verlies van de Rust enum te compenseren, zijn er veel convenience-methoden toegevoegd, waarvan de meeste zijn gemarkeerd als always inline voor maximale prestaties.

Het low-bit tagging-schema codeert vijf verschillende soorten waarden:

  • Fixnums: Signed integers die in een 62-bit bereik passen.
  • Flonums: Floating point waarden, gecodeerd via een self-tagging schema.
  • Immediates: Waarden zoals nil, true, false, undef, functie- en class-id's, en host-functies geschreven in Rust (aanroepbaar via het dollar-teken, bijv. $readfile(filename)). Immediates hebben een 5-bit subtag.
  • Pointers (twee soorten): Er zijn twee soorten pointers om referentie-equality snel te maken. De meeste waarden in Plush worden vergeleken op referentie-equality, maar strings worden vergeleken op structurele equality (zoals in JS).

Door één bit in de tag te gebruiken (index 1), kan de VM direct bepalen of twee waarden vergeleken kunnen worden met één enkele vergelijkingsinstructie. Dit creëert een "fast path" voor integers, pointers, handles en kleine immediate constants.

Indien Plush in de toekomst een JIT-compiler krijgt, kan de gegenereerde code voor if (p != nil) er als volgt uitzien:

; x0 = de waarde die wordt getest
; nil is de immediate 0x05
cmp     x0, #5
b.eq    .ELSE_BRANCH  ; spring naar else branch als gelijk

Efficiënte Fixnum-operaties

Integers die in het 62-bit fixnum-bereik passen, hebben hun laagste twee bits op 00 staan en worden opgeslagen als n << 2. Hierdoor kunnen optellingen, aftrekkingen, vergelijkingen en bitwise-operaties (AND/OR/XOR) direct op de tagged words worden uitgevoerd. Een 64-bit overflow is exact het scenario waarin het resultaat niet meer in 62 bits past, waardoor standaard machine-instructies gebruikt kunnen worden om overflow te detecteren (jo op x86-64, bvs op ARM64).

Voor een loop-counter vergelijking zoals i < n kan de ARM64 machinecode er als volgt uitzien:

; Fast path voor `i < n`, 5 instructies
; x0 = i, x1 = n
orr     x2, x0, x1    ; combineer beide operanden
tst     x2, #3        ; zijn beide operanden fixnums?
b.ne    .Lslow_lt     ; floats, boxed numbers en strings gaan die kant op
cmp     x0, x1        ; vergelijk de fixnums, `00` tag bits maken niet uit
b.ge    .LOOP_EXIT    ; signed compare, fixnums zijn 62-bit signed

Het fast path voor het optellen van twee fixnums:

; Fast path voor fixnum integer addition
; 5 instructies op het fast path
; x0 = a, x1 = b, twee tagged Values
; Een fixnum is n << 2, dus de onderste twee bits zijn 00
orr     x2, x0, x1      ; combineer beide tags, één test dekt het paar
tst     x2, #3          ; zijn beide operanden fixnums?
b.ne    .Lslow
adds    x0, x0, x1      ; (a << 2) + (b << 2) == (a + b) << 2
b.vs    .Loverflow      ; signed overflow is exact de 62-bit overflow

Zelf-tagging representatie van Flonums

Er zijn diverse manieren om floating-point waarden te taggen. NaN boxing is een bekende methode (gebruikt in Firefox), waarbij de 52 mantissa-bits van een NaN worden gebruikt als payload. Hoewel dit double-precision operaties "gratis" maakt, worden andere operaties beboet.

Voor Plush heb ik gekozen voor een benadering waarbij een paar exponent-bits worden "gestolen". Omdat het bereik van doubles ($\approx 1.8 \times 10^{308}$) enorm is, merken we weinig van het verlies van een paar bits. Waarden die niet in dit bereik passen, worden simpelweg op de heap gealloceerd.

Ik heb gebruikgemaakt van een schema uit het paper Float Self-Tagging van Olivier Melançon, Manuel Serrano en Marc Feeley. Dit schema roteert de bovenste bits naar de tag-positie na het toevoegen van een bias-waarde. Hierdoor zijn geen verschuivingen of maskeringen nodig. Tagged flonums kunnen zelfs subnormals, infinity en NaNs representeren zonder heap boxing. Voor Plush is de bias zo gekozen dat de laagste twee bits 10 worden.

Het optellen van twee flonums is duurder dan bij fixnums, omdat de tagging ongedaan moet worden gemaakt en daarna opnieuw moet worden toegepast:

; 18 instructies op het fast path, met BIAS al in een register
; x0 = a, x1 = b, twee tagged Values
; x9 = BIAS, 0x6810_0000_0000_0000, eenmalig gematerialiseerd
and     x2, x0, #3      ; isoleer de tag bits van a
and     x3, x1, #3      ; en de tag bits van b
cmp     x2, #2          ; een flonum heeft 10 in de laagste twee bits
ccmp    x3, #2, #0, eq  ; kijk naar b alleen als a een flonum was
b.ne    .Lslow
ror     x2, x0, #4      ; maak de rotatie ongedaan
sub     x2, x2, x9      ; maak de bias ongedaan, laat IEEE 754 bits over
fmov    d0, x2
ror     x3, x1, #4
sub     x3, x3, x9
fmov    d1, x3
fadd    d0, d0, d1      ; de eigenlijke f64 optelling
fmov    x2, d0          ; terug naar bits
add     x2, x2, x9      ; + BIAS
ror     x0, x2, #60     ; rotate left door 4: ARM64 roteert alleen right
and     x3, x0, #3      ; isoleer de tag bits van het resultaat
cmp     x3, #2          ; landde de tag op 10, of moet dit geboxt worden?
b.ne    .Lbox

Impact op het geheugengebruik

Om de impact te meten, zijn verschillende synthetische benchmarks gebruikt: fib (functie-aanroepen), binary_tree (object-toegang en pointer-chasing), mlp (floating-point operaties in een neuraal netwerk), sha256 en nbody (celestiale simulatie).

De resultaten laten een bimodale distributie zien. Sommige benchmarks vertonen een grote reductie in geheugengebruik (tot 37% bij mlp), terwijl andere nauwelijks bewegen omdat ze weinig geheugen alloceren. Er waren twee uitschieters: quicksort (een meetfout tijdens een GC-cyclus) en sha256_unfixed. In dat laatste geval zorgden integer-overflows voor veel "boxed integers", wat het geheugengebruik verhoogde. Dit is een trade-off: in sommige gevallen kan code minder geheugenefficiënt worden als we niet oppassen. Echter, door de berekening in sha256 aan te passen aan de 62-bit fixnums, kon dit worden opgelost.

Impact op de prestaties

Verrassend genoeg is elke benchmark sneller geworden, sommige zelfs aanzienlijk. Dit is opmerkelijk, omdat floating-point operaties nu meer werk moeten verrichten (unboxing en re-boxing).

Waarom is alles sneller?

  1. Cache-vriendelijkheid: In benchmarks zoals binarytree en linkedlist nemen objecten nu ongeveer de helft van de ruimte in beslag. Dit halveert het cache- en geheugenverkeer.
  2. Register-efficiëntie: In de oude versie met de Rust enum genereerde de compiler minder efficiënte code. Waarden werden vaak naar de stack "gespild" (spills), terwijl ze in de nieuwe versie in een enkel register passen.

Vergelijking van de add instructie

In de oude versie (met match op de enum) zag de disassembly voor de Int64 + Int64 fast path er massief uit, met veel geheugenoperaties en spills naar de stack:

; Versimpeld fragment van de oude versie
ldr   w9,  [x13]              ; v1: de tag
ldr   w14, [x13, #0xc]        ; v1: payload bytes 12..16
ldur  x13, [x13, #0x4]        ; v1: payload bytes 4..12
str   w9,  [sp, #0x2b8]       ; spill v1 naar een stack slot...

In de nieuwe versie is de disassembly veel korter. De waarden passen in een register, waardoor de code veel efficiënter is:

; Versimpeld fragment van de nieuwe versie
ldr   x3, [x9, x8,  lsl #3]   ; v1, een enkele load, 8-byte stride
ldr   x2, [x9, x10, lsl #3]   ; v0, een enkele load
and   x11, x2, #3             ; v0's tag bits
cmp   x11, #2                 ; een flonum?
b.ne  .Lcheck_fixnum

Zelfs de floating-point fast path is korter geworden: van 52 instructies, 24 geheugenoperaties en 4 branches naar 36 instructies, 9 geheugenoperaties en 1 branch.

Conclusie

Deze refactoring is een groot succes. De nieuwe tagged word representatie vermindert niet alleen het geheugengebruik voor arrays en objecten, maar zorgt ook voor een aanzienlijke prestwinst. De Rust compiler kon simpelweg geen efficiënte code genereren voor de oude enum-versie, maar met waarden die in één register passen, is het resultaat uitstekend.

Hoewel heap-boxing van integers in specifieke gevallen (zoals bij left shifts of overflow-afhankelijke random number generators) extra allocaties kan veroorzaken, is de prestatie-overhead hiervan in een geïnterpreteerd systeem verwaarloosbaar.

De prestaties van 3D-graphics in Plush zijn hierdoor sterk verbeterd; de taal kan nu ongeveer 10.000 flat-shaded polygonen renderen op een interactieve framerate. Als volgende stap is de overstap van een stack-based naar een register-based interpreter gepland, wat waarschijnlijk een verdere prestatieboost zal opleveren. Daarnaast is er een plan om een minimalistisch LISP-dialect, genaamd JetLISP, te bouwen op de Plush VM.