JIT-compileren in 5μs

Tegenwoordig is het, dankzij AI, makkelijker dan ooit om een JIT-compiler met korte compilatietijden te schrijven door direct op assembly te richten. Dit is ook een kans voor nieuwe databases om bestaande systemen te verbeteren. Bij het bouwen van pgrust dacht ik aanvankelijk dat het implementeren van een JIT-compiler erg moeilijk zou zijn. Uiteindelijk bleek het, mede door AI-ondersteuning, veel eenvoudiger dan verwacht. Dit is een van de redenen waarom pgrust zo snel is: de JIT-compiler compileert code in ongeveer 5μs, waardoor we elke SQL-query kunnen JIT-compileren en niet slechts een subset daarvan.

In dit artikel leg ik uit hoe je je eigen snelle JIT-compiler kunt bouwen. Als voorbeeld maken we een eenvoudige regular expression (regex) engine die gebruikmaakt van JIT-compilatie.

Waarom JIT-compilatie?

JIT-compilatie is het proces waarbij gecompileerde code tijdens runtime ("Just In Time") wordt gegenereerd. Wanneer dit correct wordt uitgevoerd, kan dit leiden tot aanzienlijke prestatiewinst, vaak in de orde van 2 tot 5 keer, en soms zelfs meer.

De belangrijkste use case voor JIT-compilatie is wanneer er tijdens runtime informatie beschikbaar komt die het gedrag van je programma drastisch verandert. Dit komt veel voor bij interpreters van programmeertalen, die de te executeren code pas tijdens runtime ontvangen. JIT-compilers zijn echter ook nuttig in andere domeinen, zoals bij het parsen van data. Soms is het schema van de data die je parst pas tijdens runtime bekend; een JIT-compiler kan hierbij helpen.

Om te beginnen implementeren we een eenvoudige regex-engine. Om het simpel te houden ondersteunen we slechts twee functies: letterlijke strings en herhaling (de regex ). We slaan de parser over en representeren de regular expression als reeds geparsede Rust-structuren. Dit betekent dat we strings kunnen ondersteunen zoals apples of b(an), maar geen alternaties of lookbehinds.

In code ziet dit er als volgt uit:

enum Node {
    Literal(&'static str),
    Concatenation(Box<Node>, Box<Node>),
    Repetition(Box<Node>),
}

fn literal(text: &'static str) -> Node {
    Node::Literal(text)
}

fn concatenation(left: Node, right: Node) -> Node {
    Node::Concatenation(Box::new(left), Box::new(right))
}

fn repetition(body: Node) -> Node {
    Node::Repetition(Box::new(body))
}

Het schrijven van een interpreter voor deze engine is eveneens eenvoudig:

fn match_node(node: &Node, input: &[u8], pos: usize, next: &dyn Fn(usize) -> bool) -> bool {
    match node {
        Node::Literal(text) => {
            let literal = text.as_bytes();
            input[pos..].starts_with(literal) && next(pos + literal.len())
        }
        Node::Concatenation(left, right) => {
            match_node(left, input, pos, &|left_end| {
                match_node(right, input, left_end, next)
            })
        }
        Node::Repetition(body) => {
            match_node(body, input, pos, &|body_end| {
                match_node(node, input, body_end, next)
            }) || next(pos)
        }
    }
}

fn interp_match(regex: &Node, input: &str) -> bool {
    let bytes = input.as_bytes();
    match_node(regex, bytes, 0, &|pos| pos == bytes.len())
}

Hoewel deze engine simpel is en minder dan 20 regels code beslaat, is het interessant om naar de prestaties te kijken. Ter vergelijking gebruiken we handgeschreven code die specifiek is geïmplementeerd voor de regex b(an)*:

fn handwritten_b_an_star(input: &str) -> bool {
    let bytes = input.as_bytes();
    let mut pos = 0;
    if pos == bytes.len() || bytes[pos] != b'b' {
        return false;
    }
    pos += 1;
    while pos < bytes.len() {
        if bytes[pos] != b'a' {
            return false;
        }
        pos += 1;
        if pos == bytes.len() || bytes[pos] != b'n' {
            return false;
        }
        pos += 1;
    }
    true
}

Bij benchmarks blijkt dat de handgeschreven versie 10 tot 20 keer sneller is dan de interpreter. Er is dus veel ruimte voor verbetering. We gaan nu kijken hoe we JIT-compilatie kunnen gebruiken om een algemene regex-engine te bouwen die net zo goed presteert als de handgeschreven versie.

Hoe JIT te compileren

Het JIT-compileren van code bestaat uit twee stappen. Eerst genereer je de assembly voor de code die je wilt uitvoeren. Vervolgens verpak je deze assembly in een functie die je vanuit je programma kunt aanroepen.

Voor het genereren van de assembly gebruiken we een variant van een aanpak genaamd copy-and-patch. Het idee is dat we een reeks assembly-templates hebben voor de verschillende operaties die we willen JIT-compileren. Deze templates worden "stencils" genoemd. Wanneer we een operatie willen JIT-compileren, nemen we het bijbehorende stencil en passen we kleine wijzigingen toe op basis van de specifieke operatie. Door meerdere van deze ingevulde stencils achter elkaar te plaatsen, kunnen we tijdens runtime een programma construeren met prestaties die vergelijkbaar zijn met de handgeschreven versie.

Onze aanpak is als volgt:

1. Analyseren van de ARM64-code die we willen genereren voor b(an)*.
2. Herhaalde instructiesequenties omzetten in herbruikbare stencils.
3. Een emitter schrijven die deze stencils invult en combineert op basis van de regex AST.
4. De gegenereerde instructies kopiëren naar uitvoerbaar geheugen, zodat Rust ze als een normale functie kan aanroepen.

Designbeslissingen

We werken met de regex b(an)* op macOS met ARM64. De volgende ontwerpkeuzes zijn gemaakt:

• Stack voor backtracking: De stack houdt bij naar welke status we moeten terugkeren als we in de regex in een dood spoor lopen.
• Null-byte terminatie: De string die we matchen eindigt met een null-byte. Karaktervergelijkingen falen dus automatisch als we het einde van de string bereiken, waardoor lengtevergelijkingen overbodig zijn.
• Registers:
• x0: Huidige positie in de string en de returnwaarde.
• x1: Top van de stack voor backtracking.
• x2: Bodem van de stack voor backtracking (om te bepalen of de stack leeg is).
• x9: Gebruikt als tijdelijke variabele.
• Inputs:
• x0: Pointer naar het begin van de string.
• x1: Pointer naar de locatie die we gebruiken voor onze stack.

Gegenereerde ARM64-code

Laten we de gegenereerde assembly stap voor stap doornemen.

De prologue initialiseert het programma door de top en bodem van de stack gelijk te stellen aan de overgedragen waarde:

0:  aa0103e2   mov   x2, x1

Vervolgens volgt de code die controleert op het karakter 'b'. Als dit niet wordt gevonden, springt de code naar een fallback-blok. Zo ja, dan wordt de positie in de string gevorderd:

; CHAR 'b'
4:  39400009   ldrb  w9, [x0]                  ; laad huidige input byte
8:  7101893f   cmp   w9, #0x62                 ; is het 'b'?
c:  54000281   b.ne  0x5c                      ; nee -> fallback blok
10:  91000400   add   x0, x0, #1                ; ja -> advance input

Dan volgt de herhaling (an)*. Voor de herhaling is backtracking nodig. Als we hier terugkeren, springen we direct naar het einde van de loop. Daarom slaan we zowel het adres van de instructie na de loop als onze positie in de string op de stack op:

14:  d2800989   movz  x9, #0x004c               ; bouw resume adres
18:  f2a00009   movk  x9, #0x0000, lsl #16      ;   = 0x1_0000_004c
1c:  f2c00029   movk  x9, #0x0001, lsl #32      ;   (de loop exit)
20:  f2e00009   movk  x9, #0x0000, lsl #48      ;
24:  a8810029   stp   x9, x0, [x1], #16         ; push (exit, pos) op stack

Nu kunnen we de body van de herhaling uitvoeren. Deze controleert op de karakters 'a' en 'n' en gaat, indien gevonden, terug naar het begin van de herhaling op een nieuwe stringlocatie:

; CHAR 'a'
28:  39400009   ldrb  w9, [x0]
2c:  7101853f   cmp   w9, #0x61                 ; 'a'?
30:  54000161   b.ne  0x5c                      ; nee -> fallback blok
34:  91000400   add   x0, x0, #1
; CHAR 'n'
38:  39400009   ldrb  w9, [x0]
3c:  7101b93f   cmp   w9, #0x6e                 ; 'n'?
40:  540000e1   b.ne  0x5c                      ; nee -> fallback blok
44:  91000400   add   x0, x0, #1
; JMP
48:  17fffff3   b     0x14                      ; terug naar begin loop

Na de loop controleren we of we aan het einde van de string zijn. Als dat zo is, retourneren we 1 voor succes. Zo niet, dan is de regex mislukt en wordt de fail-logica uitgevoerd:

4c:  39400009   ldrb  w9, [x0]
50:  35000069   cbnz  w9, 0x5c                  ; niet NUL -> fallback blok
54:  d2800020   mov   x0, #1                    ; succes
58:  d65f03c0   ret

Ten slotte de fallback-logica. Deze controleert of de stack leeg is. Zo ja, dan retourneren we 0. Zo nee, dan poppen we het fallback-adres en de stringpositie van de stack en springen we naar dat adres:

5c:  eb02003f   cmp   x1, x2                    ; zijn er frames over?
60:  54000060   b.eq  0x6c                      ; nee -> geef op
64:  a9ff0029   ldp   x9, x0, [x1, #-16]!       ; pop (resume, pos)
68:  d61f0120   br    x9                        ; spring daarheen
6c:  d2800000   mov   x0, #0                    ; geen match
70:  d65f03c0   ret

Stencils bouwen

Met de gecompileerde code in zicht kunnen we de copy-and-patch compiler bouwen. We schrijven functies die de respectievelijke codeblokken genereren op basis van inputwaarden.

De prologue is een vast blok:

const PROLOGUE_WORDS: usize = 1;
fn stencil_prologue() -> [u32; PROLOGUE_WORDS] {
    [0xAA0103E2] // mov x2, x1
}

Voor karaktervergelijking voegen we het specifieke karakter en het sprongdoel voor de fallback-logica in:

const CHAR_WORDS: usize = 4;
fn stencil_char(byte: u8, stencil_pos: usize, fail_pos: usize) -> [u32; CHAR_WORDS] {
    [
        0x39400009,                                                 // ldrb w9, [x0]
        0x7100013F | ((byte as u32) << 10),                         // cmp  w9, #byte
        0x54000001 | cond_branch_offset(stencil_pos + 2, fail_pos), // b.ne fail
        0x91000400,                                                 // add  x0, x0, #1
    ]
}

Voor de herhaling hebben we het begin van de loop (push op de stack) en de sprong naar het einde:

const SPLIT_WORDS: usize = 5;
fn stencil_split(resume_addr: u64) -> [u32; SPLIT_WORDS] {
    [
        0xD2800009 | addr_bits(resume_addr, 0), // movz x9, #addr[0..16]
        0xF2A00009 | addr_bits(resume_addr, 1), // movk x9, #addr[16..32], lsl 16
        0xF2C00009 | addr_bits(resume_addr, 2), // movk x9, #addr[32..48], lsl 32
        0xF2E00009 | addr_bits(resume_addr, 3), // movk x9, #addr[48..64], lsl 48
        0xA8810029,                             // stp  x9, x0, [x1], #16
    ]
}

const JMP_WORDS: usize = 1;
fn stencil_jmp(stencil_pos: usize, target_pos: usize) -> [u32; JMP_WORDS] {
    [0x14000000 | branch_offset(stencil_pos, target_pos)] // b target
}

En ten slotte de match- en fail-blokken:

const MATCH_WORDS: usize = 4;
fn stencil_match(stencil_pos: usize, fail_pos: usize) -> [u32; MATCH_WORDS] {
    [
        0x39400009,                                                 // ldrb w9, [x0]
        0x35000009 | cond_branch_offset(stencil_pos + 1, fail_pos), // cbnz w9, fail
        0xD2800020,                                                 // mov  x0, #1
        0xD65F03C0,                                                 // ret
    ]
}

const FAIL_WORDS: usize = 6;
fn stencil_fail() -> [u32; FAIL_WORDS] {
    [
        0xEB02003F, // cmp  x1, x2
        0x54000060, // b.eq +3 (to the mov below)
        0xA9FF0029, // ldp  x9, x0, [x1, #-16]!
        0xD61F0120, // br   x9
        0xD2800000, // mov  x0, #0
        0xD65F03C0, // ret
    ]
}

Hieronder volgen de helperfuncties voor het invoegen van data in de instructies:

fn cond_branch_offset(branch_pos: usize, target_pos: usize) -> u32 {
    let instr_count = target_pos as i64 - branch_pos as i64;
    (((instr_count as u64) & 0x7FFFF) << 5) as u32
}

fn branch_offset(branch_pos: usize, target_pos: usize) -> u32 {
    let instr_count = target_pos as i64 - branch_pos as i64;
    ((instr_count as u64) & 0x3FF_FFFF) as u32
}

fn addr_bits(addr: u64, part: usize) -> u32 {
    (((addr >> (16 * part)) & 0xFFFF) as u32) << 5
}

Code emitteren

De volgende code stuurt het proces aan:

fn node_words(node: &Node) -> usize {
    match node {
        Node::Literal(text) => text.len() * CHAR_WORDS,
        Node::Concatenation(left, right) => node_words(left) + node_words(right),
        Node::Repetition(body) => SPLIT_WORDS + node_words(body) + JMP_WORDS,
    }
}

struct Emitter {
    code: Vec<u32>,
    fail: usize,
    base: u64,
}

impl Emitter {
    fn pos(&self) -> usize {
        self.code.len()
    }

    fn emit(&mut self, stencil: &[u32]) {
        self.code.extend_from_slice(stencil);
    }

    fn emit_node(&mut self, node: &Node) {
        match node {
            Node::Literal(text) => {
                for &byte in text.as_bytes() {
                    self.emit(&stencil_char(byte, self.pos(), self.fail));
                }
            }
            Node::Concatenation(left, right) => {
                self.emit_node(left);
                self.emit_node(right);
            }
            Node::Repetition(body) => {
                let split_at = self.pos();
                let exit = split_at + SPLIT_WORDS + node_words(body) + JMP_WORDS;
                self.emit(&stencil_split(self.base + exit as u64 * 4));
                self.emit_node(body);
                self.emit(&stencil_jmp(self.pos(), split_at));
            }
        }
    }
}

fn generate_code(regex: &Node, base: u64) -> Vec<u32> {
    let nwords = PROLOGUE_WORDS + node_words(regex) + MATCH_WORDS + FAIL_WORDS;
    let mut emitter = Emitter {
        code: Vec::with_capacity(nwords),
        fail: nwords - FAIL_WORDS,
        base,
    };
    emitter.emit(&stencil_prologue());
    emitter.emit_node(regex);
    let match_at = emitter.pos();
    emitter.emit(&stencil_match(match_at, emitter.fail));
    emitter.emit(&stencil_fail());
    assert_eq!(emitter.pos(), nwords);
    emitter.code
}

Het schrijven van assembly is waar AI mij het meest hielp. Ik heb weinig ervaring met het zelf schrijven van assembly, maar door mijn coding agent de algemene vorm van de JIT-compiler te geven, kon de AI de specifieke details en instructies voor me afhandelen.

Machinecode laden

Om de compiler te voltooien moeten we de code laden. We gebruiken mmap om een blok geheugen toe te wijzen dat leesbaar, schrijfbaar en uitvoerbaar is.

const BSTACK_MAX: usize = 4096;

unsafe extern "C" {
    fn pthread_jit_write_protect_np(enabled: libc::c_int);
    fn sys_icache_invalidate(start: *mut libc::c_void, len: libc::size_t);
}

type MatchFn = unsafe extern "C" fn(input: *const u8, bstack: *mut u64) -> u64;

struct Jit {
    buf: *mut u32,
    nbytes: usize,
    bstack: Vec<u64>,
}

impl Jit {
    fn compile(regex: &Node) -> Jit {
        let nwords = PROLOGUE_WORDS + node_words(regex) + MATCH_WORDS + FAIL_WORDS;
        let nbytes = nwords * 4;
        unsafe {
            let buf = libc::mmap(
                std::ptr::null_mut(),
                nbytes,
                libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
                libc::MAP_PRIVATE | libc::MAP_ANON | libc::MAP_JIT,
                -1,
                0,
            ) as *mut u32;
            assert!(buf as *mut libc::c_void != libc::MAP_FAILED, "mmap failed");
            let code = generate_code(regex, buf as u64);
            pthread_jit_write_protect_np(0);
            std::slice::from_raw_parts_mut(buf, code.len()).copy_from_slice(&code);
            pthread_jit_write_protect_np(1);
            sys_icache_invalidate(buf as *mut libc::c_void, nbytes);
            Jit { buf, nbytes, bstack: vec![0; BSTACK_MAX * 2] }
        }
    }

    fn is_match(&mut self, nul_terminated: &[u8]) -> bool {
        debug_assert_eq!(nul_terminated.last(), Some(&0));
        unsafe {
            let matcher: MatchFn = std::mem::transmute(self.buf);
            matcher(nul_terminated.as_ptr(), self.bstack.as_mut_ptr()) != 0
        }
    }
}

impl Drop for Jit {
    fn drop(&mut self) {
        unsafe {
            libc::munmap(self.buf as *mut libc::c_void, self.nbytes);
        }
    }
}

Resultaten

Hieronder volgt de vergelijking tussen de verschillende implementaties:

Input lengte	Interpreter	JIT	Handgeschreven	JIT versnelling	Handgeschreven versnelling
9	45 ns	3.8 ns	3.8 ns	11.7x	11.9x
33	103 ns	7.9 ns	10.5 ns	13.0x	9.8x
129	597 ns	30 ns	32 ns	19.7x	18.6x
513	1,955 ns	126 ns	120 ns	15.5x	16.2x
2,049	8,301 ns	470 ns	393 ns	17.7x	21.1x

De JIT- en de handgeschreven implementatie liggen zeer dicht bij elkaar. Soms is de JIT-versie sneller, soms de handgeschreven versie.

Er gaat een meme rond dat AI niet helpt omdat "code nooit het moeilijke deel was". Ik denk dat dit in sommige domeinen waar is, maar in andere is het schrijven van de code absoluut het moeilijkste deel. JIT-compilers zijn daar een goed voorbeeld van. De zeldzaamheid van JIT-compilers doet me vermoeden dat de implementatie ervan historisch gezien te moeilijk was om de moeite waard te zijn. LLM's hebben de drempel verlaagd. Dit is de these achter pgrust: databases waren historisch gezien de moeilijkste software om te bouwen en waren daardoor beperkt. Met AI kunnen we ambitieuzer zijn over het type software dat we bouwen.