JIT Compiling Code in 5μs

Historically, fast JIT compilation was a black art. To write a fast JIT compiler, you would need to know how to write assembly. Case in point: there is no production-ready database today that has its own JIT compiler. They all either use LLVM or generate C/C++ code. Both of these options suffer from high compile times, which limits their applicability. Now, with the use of AI, it’s easier than ever to write a JIT compiler with fast compile times by directly targeting assembly. This is also one area of opportunity for new databases to improve on old ones. When building pgrust, I initially thought it would be really hard to implement a JIT compiler. In the end, I found it much easier than I expected due to AI assistance and it ends up being part of the reason why pgrust is so fast. The pgrust JIT compiler compiles code in around 5μs, which enables us to JIT compile every SQL query, not just a subset of them. In this post, I’ll walk you through how you can build your own fast JIT compiler. We’ll build a simple regular expression engine that uses JIT compilation as an example. Why JIT Compilation JIT compilation is the practice of generating compiled code at runtime or “Just In Time”. When done right, it can result in big performance wins, often on the order of 2-5x and sometimes even more. The main use case for JIT compilation is when there’s information you gain at runtime that drastically alters the behavior of your program. This is particularly common with programming language interpreters; they receive the code to execute at runtime. JIT compilers are also useful in domains beyond programming languages, such as parsing data. Sometimes you don’t know the schema of the data you’re parsing until runtime, and a JIT can help with that. To kick things off, let’s implement a toy regular expression engine. To keep things simple, we’ll support only two features: literal strings and repetition (i.e. the regex ). We’ll also skip the parser and represent the regular expression as already parsed Rust structures. This means we’ll be able to support strings such as: - apples - b(an) but no alternation or lookbehind or anything like that. In code this is pretty simple. We’ll have 3 types of Nodes: a literal string node, a repetition node, and a concatenation node, which is the combination of two nodes. This ends up looking like this: enum Node { Literal(&'static str), Concatenation(Box , Box ), Repetition(Box ), } 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)) } Writing an interpreter for our regular expression engine is also straightforward: 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()) } Now this regular expression engine is pretty simple. It’s under 20 lines of code, but let’s see how it does in terms of performance. For comparison, we’ll compare the code against handwritten code implemented specifically for the regex. For our example we’ll use the regex b(an). The handwritten code ends up looking like: 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 fallback block 10: 91000400 add x0, x0, #1 ; yes -> advance input Next up, we have the repetition (an). For the repetition, we need to do the backtracking. If we backtrack here, that means we jump immediately to the end of the loop. That means we need to store both the address of the instruction after the loop and our position in the string on the stack. 14: d2800989 movz x9, #0x004c ; build resume address 18: f2a00009 movk x9, #0x0000, lsl #16 ; = 0x1_0000_004c 1c: f2c00029 movk x9, #0x0001, lsl #32 ; (the loop exit) 20: f2e00009 movk x9, #0x0000, lsl #48 ; 24: a8810029 stp x9, x0, [x1], #16 ; push (exit, pos) onto stack With that in place, we can now execute the body of the repetition. This will check for the characters ‘a’ and ‘n’ and, if it sees them, go back to the top of the repetition, but at a new string location. ; CHAR 'a' 28: 39400009 ldrb w9, [x0] 2c: 7101853f cmp w9, #0x61 ; 'a'? 30: 54000161 b.ne 0x5c ; no -> fallback block 34: 91000400 add x0, x0, #1 ; CHAR 'n' 38: 39400009 ldrb w9, [x0] 3c: 7101b93f cmp w9, #0x6e ; 'n'? 40: 540000e1 b.ne 0x5c ; no -> fallback block 44: 91000400 add x0, x0, #1 ; JMP 48: 17fffff3 b 0x14 ; back to top of loop Now we’re past the loop. This is where the backtracking will jump once we backtrack. Once we finish the repetition, we’re at the end of the regex. All we have to do now is check if we’re at the end of the string. If we are at the end of the string, we return 1 for success. If we are not, that means the regex failed to match, and we need to run the fail logic to do a fallback. 4c: 39400009 ldrb w9, [x0] 50: 35000069 cbnz w9, 0x5c ; not at NUL -> fallback block 54: d2800020 mov x0, #1 ; success 58: d65f03c0 ret And then finally, we have the fallback logic. This checks if the stack is empty. If it is, we return 0. If it’s not empty, we pop both the fallback address and the fallback string position off the stack, and then jump to the fallback address. 5c: eb02003f cmp x1, x2 ; any frames left? 60: 54000060 b.eq 0x6c ; no -> give up 64: a9ff0029 ldp x9, x0, [x1, #-16]! ; pop (resume, pos) 68: d61f0120 br x9 ; jump there 6c: d2800000 mov x0, #0 ; no match 70: d65f03c0 ret Building the Stencils Now that you’ve had the chance to see the compiled code, you should start to get a sense of how the copy-and-patch compiler would work. We have common sets of instructions with only minor differences between them. For each of these blocks of functions, we can write a function to generate the respective code. Each function will take in values to use to modify the code. For example, one of the arguments to stencil_char will be the char in the regex to compare against. We’ll insert that char directly into the machine code. The prologue is straightforward since it’s just a block of code: const PROLOGUE_WORDS: usize = 1; fn stencil_prologue() -> [u32; PROLOGUE_WORDS] { [0xAA0103E2] // mov x2, x1 } For character comparison, we need to insert the character we’re comparing against and where to jump for the fallback logic: 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) [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 } And then we have the match and fail blocks which are pretty clean: 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 ] } For completeness, here’s the helper functions we used which just help us insert specific data into the instructions: // Compute the branch-offset field for a conditional branch (b.ne / cbnz): // the instruction count from branch to target, stored in bits 5..24. fn cond_branch_offset(branch_pos: usize, target_pos: usize) -> u32 { let instr_count = target_pos as i64 - branch_pos as i64; // may be negative (((instr_count as u64) & 0x7FFFF) u32 { let instr_count = target_pos as i64 - branch_pos as i64; // may be negative ((instr_count as u64) & 0x3FF_FFFF) as u32 } // Extract 16 bits of an absolute address, positioned for a movz/movk immediate. fn addr_bits(addr: u64, part: usize) -> u32 { (((addr >> (16 * part)) & 0xFFFF) as u32) 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 , fail: usize, // word offset of the shared fail block base: u64, // runtime address of code[0], for absolute-address holes } impl Emitter { // Returns the offset where the next instruction will be placed. fn pos(&self) -> usize { self.code.len() } // Appends a filled stencil to the code buffer. fn emit(&mut self, stencil: &[u32]) { self.code.extend_from_slice(stencil); } // Emits the code for one node, recursing into children. 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 +