How AI Changes the Economics of JIT Compilers - malisper.me
Primary Menu
malisper.me
malisper.me
General<br>About Me
Table of Contents for Postgres Posts
RSS
Subscribe to Blog via Email
Enter your email address to subscribe to this blog and receive notifications of new posts by email.
Email Address
Subscribe
Skip to content
Home » pgrust » How AI Changes the Economics of JIT Compilers
Historically, 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. In this post, I’ll walk you through how you can build your own 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 {<br>Literal(&'static str),<br>Concatenation(Box, Box),<br>Repetition(Box),
fn literal(text: &'static str) -> Node {<br>Node::Literal(text)
fn concatenation(left: Node, right: Node) -> Node {<br>Node::Concatenation(Box::new(left), Box::new(right))
fn repetition(body: Node) -> Node {<br>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 {<br>match node {<br>Node::Literal(text) => {<br>let literal = text.as_bytes();<br>input[pos..].starts_with(literal) && next(pos + literal.len())
Node::Concatenation(left, right) => {<br>match_node(left, input, pos, &|left_end| {<br>match_node(right, input, left_end, next)<br>})
Node::Repetition(body) => {<br>match_node(body, input, pos, &|body_end| {<br>match_node(node, input, body_end, next)<br>}) || next(pos)
fn interp_match(regex: &Node, input: &str) -> bool {<br>let bytes = input.as_bytes();<br>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 {<br>let bytes = input.as_bytes();<br>let mut pos = 0;
if pos == bytes.len() || bytes[pos] != b'b' {<br>return false;<br>pos += 1;
while pos<br>(There are ways you could optimize this code and make it much faster, but for our purposes it serves as a good comparison)
When I benchmark a couple of examples against these two, I get that the handwritten version is 10-20x faster than the interpreter. Clearly a lot of room for improvement.
Now let’s take a look at how we can use JIT compilation to get a general regular expression engine that performs as well as the handwritten version.
How to JIT Compile
There are two steps to JIT compile code. First you generate the assembly for the code you want to run. Once you have the code, you then package the assembly code into a function that you can call like any other code into your program.
To generate the assembly, we will use a variant of an approach called copy-and-patch. The idea is that we have a series of templates in...