【文章标题】:JIT Compiling Code in 5μs 在5微秒内进行JIT编译代码

【文章正文】: 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. 历史上,快速JIT编译曾是一门玄学。要编写一个快速的JIT编译器,你需要知道如何编写汇编代码。一个典型的例子是:当今没有任何生产就绪的数据库拥有自己的JIT编译器。它们要么使用LLVM,要么生成C/C++代码。这两种选择都存在编译时间长的问题,这限制了它们的适用性。如今,借助AI,通过直接面向汇编来编写具有快速编译时间的JIT编译器比以往任何时候都容易。这也是新数据库改进旧数据库的一个机会领域。在构建pgrust时,我最初认为实现一个JIT编译器会非常困难。最终,由于AI的协助,我发现它比预期的要容易得多,并且这最终成为pgrust如此之快的部分原因。pgrust的JIT编译器大约在5微秒内编译代码,这使我们能够对每个SQL查询进行JIT编译,而不仅仅是其中的一部分。在这篇文章中,我将带你了解如何构建自己的快速JIT编译器。我们将构建一个简单的正则表达式引擎,以JIT编译为例。

Why JIT Compilation 为什么使用JIT编译

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. JIT编译是在运行时生成编译代码的实践,即“即时编译”。如果做得正确,它可以带来巨大的性能提升,通常能达到2-5倍,有时甚至更高。JIT编译的主要用例是,当你在运行时获得的信息会极大地改变程序的行为时。这在编程语言解释器中尤其常见;它们在运行时接收要执行的代码。JIT编译器在编程语言之外的领域也很有用,例如解析数据。有时直到运行时你才知道正在解析的数据的模式,而JIT可以在这方面提供帮助。

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. 首先,让我们实现一个玩具正则表达式引擎。为了简单起见,我们只支持两个特性:字面量字符串和重复(即正则表达式中的*)。我们还将跳过解析器,将正则表达式表示为已经解析好的Rust结构。这意味着我们将能够支持如下字符串:
  • apples
  • b(an)* 但不支持交替、后行断言或任何类似功能。

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: 现在这个正则表达式引擎非常简单。它不到20行代码,但让我们看看它在性能方面的表现。为了比较,我们将代码与专门为正则表达式实现的手写代码进行对比。在我们的示例中,我们将使用正则表达式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 }

(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. 当我针对这两个版本对一些示例进行基准测试时,我发现手写版本比解释器快10-20倍。显然还有很大的改进空间。

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. 现在让我们看看如何使用JIT编译来获得一个通用的正则表达式引擎,其性能与手写版本一样好。

How to JIT Compile 如何JIT编译

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 JIT编译代码有两个步骤。首先,为你要运行的代码生成汇编代码。一旦你有了代码,然后