JIT编译代码在5μs内完成

# 在5微秒内即时编译代码 - malisper.me 来源：https://malisper.me/jit-compiling-code-in-5-us/ 历史上，快速的即时编译一直是一门黑魔法。要编写一个快速的即时编译器，你需要懂得如何编写汇编代码。以事实为例：目前还没有一个生产就绪的数据库拥有自己的即时编译器。它们要么使用 LLVM，要么生成 C/C++ 代码。这两种方案都因编译时间过长而受限，从而限制了它们的适用性。 如今，借助 AI 的帮助，通过直接生成汇编代码来编写编译速度快的即时编译器比以往任何时候都更容易。这也是新兴数据库超越老旧数据库的一个机遇领域。 在构建 pgrust (https://github.com/malisper/pgrust) 时，我最初认为实现一个即时编译器会非常困难。最终，得益于 AI 辅助，我发现它比我预想的容易得多，这也成为 pgrust 运行如此迅速的原因之一。pgrust 的即时编译器能在约 5 微秒内编译代码，这使得我们可以为每个 SQL 查询进行即时编译，而不仅仅是其中的一部分。 在这篇文章中，我将逐步讲解如何构建一个你自己的快速即时编译器。我们将以一个使用即时编译的简单正则表达式引擎作为示例。 ## 为何需要即时编译 即时编译是在运行时或“恰当时刻”生成编译后代码的实践。如果操作得当，它能带来显著的性能提升，通常提升幅度在 2 到 5 倍，有时甚至更高。即时编译的主要应用场景是：在运行时获得的信息会极大改变程序的行为。这在编程语言解释器中尤为常见；它们在运行时接收要执行的代码。即时编译器在编程语言以外的领域也很有用，例如数据解析。有时你直到运行时才知道要解析的数据的结构，而即时编译器可以助一臂之力。 --- 首先，让我们实现一个玩具级的正则表达式引擎。为保持简单，我们将只支持两个特性：字面字符串和重复（即正则中的 `*`）。我们还将跳过解析步骤，直接以已解析的 Rust 结构体来表示正则表达式。这意味着我们将能够支持如下的字符串： - apples - b(an)* 但不支持交替（alternation）或后顾（lookbehind）等类似功能。 代码实现相当简单。我们将有 3 种类型的节点：字面字符串节点、重复节点和连接节点（两个节点的组合）。最终结构如下： ```rust 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)) } ``` 为我们的正则表达式引擎编写一个解释器也很直接： ```rust 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()) } ``` 这个正则表达式引擎相当简单。代码不足 20 行，但让我们看看它的性能如何。为了比较，我们将针对特定正则（b(an)*）编写的、手工优化的代码与其进行对比。手工编写的代码最终看起来像这样： ```rust 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 } ``` （这段代码还有多种优化方式可以使其快得多，但就我们的目的而言，它是一个很好的比较基准。） 当我对几个例子进行基准测试后发现，手工编写的版本比解释器版本快 10-20 倍。显然，性能还有很大的提升空间。 现在，让我们看看如何使用即时编译来获得一个性能与手工版本相当的通用正则表达式引擎。 ## 如何进行即时编译 即时编译代码分为两个步骤。首先，生成你想要运行的代码对应的汇编。一旦有了代码，你就需要将这些汇编代码打包成一个函数，使其可以像程序中任何其他代码一样被调用。 为了生成汇编，我们将使用一种称为“复制并修补”的方法变体。其核心思想是：我们为一系列想要即时编译的操作准备了一组汇编模板。这些模板被称为“模板”。当我们想要即时编译某个操作时，我们取出对应的模板，并根据该操作的具体细节进行微调。这非常类似于填充一个真实的模板。通过将多个填充好的模板拼接起来，我们可以在运行时构建出一个性能与手工版本相当的程序。 我们的路径如下：首先，我们将查看为正则“b(an)*”生成的 ARM64 代码。然后，我们将重复的指令序列转换为可重用的模板，编写一个发射器从正则 AST 填充和组合这些模板，最后将生成的指令复制到可执行内存中，以便 Rust 可以像调用普通函数一样调用它们。 为了便于理解，最简单的方式是从生成的代码入手，然后反向推导到即时编译器本身。我们再次以正则“b(an)*”为例。先明确一些设计决策： - 我们将使用栈进行回溯。栈将跟踪在正则匹配失败时应回退到的状态。 - 我们要匹配的字符串将以空字节结尾。这意味着任何字符比较在遇到字符串末尾时都会自动失败。因此，我们在任何时刻都不必进行长度比较。 对于程序状态，我们将使用以下寄存器： - `x0` – 字符串中的当前位置和返回值 - `x1` – 用于回溯的栈顶 - `x2` – 用于回溯的栈底（用于判断栈是否为空） - `x9` – 用作临时变量 对于程序的输入，我们将接收： - `x0` – 指向字符串起始位置的指针 - `x1` – 指向我们将用作栈的位置的指针 ### 生成的 ARM64 代码 现在准备工作已就绪，让我们逐段分析生成的汇编代码。这是在 macOS 上针对 ARM64 的代码。 首先是序言部分，初始化程序。它只是通过将栈顶和栈底设置为传入的值来初始化栈： ```assembly 0: aa0103e2 mov x2, x1 ``` 接下来是检查字符'b'的代码。如果遇到的不是'b'，我们将跳转到处理回退逻辑的代码块。否则，我们向前推进字符串中的位置： ```assembly ; CHAR 'b' 4: 39400009 ldrb w9, [x0] ; 加载当前输入字节 8: 7101893f cmp w9, #0x62 ; 是'b'吗？ c: 54000281 b.ne 0x5c ; 不是 -> 回退块 10: 91000400 add x0, x0, #1 ; 是的 -> 前进输入 ``` 接下来是重复部分 `(an)*`。对于重复，我们需要进行回溯。如果在这里回溯，意味着我们将立即跳转到循环结束处。因此，我们需要将循环后的指令地址和当前位置都压栈保存。 ```assembly 14: d2800989 movz x9, #0x004c ; 构建恢复地址 18: f2a00009 movk x9, #0x0000, lsl #16 ; = 0x1_0000_004c 1c: f2c00029 movk x9, #0x0001, lsl #32 ; (循环出口) 20: f2e00009 movk x9, #0x0000, lsl #48 ; 24: a8810029 stp x9, x0, [x1], #16 ; 将 (退出地址, 位置) 压入栈 ``` 有了这个设置，我们现在可以执行重复体。这将检查字符'a'和'n'，如果匹配，将返回到重复的顶部，但在新的字符串位置。 ```assembly ; CHAR 'a' 28: 39400009 ldrb w9, [x0] 2c: 7101853f cmp w9, #0x61 ; 'a'? 30: 54000161 b.ne 0x5c ; 不是 -> 回退块 34: 91000400 add x0, x0, #1 ; CHAR 'n' 38: 39400009 ldrb w9, [x0] 3c: 7101b93f cmp w9, #0x6e ; 'n'? 40: 540000e1 b.ne 0x5c ; 不是 -> 回退块 44: 91000400 add x0, x0, #1 ; JMP 48: 17fffff3 b 0x14 ; 跳转到循环顶部 ``` 现在我们已经过了循环。这是我们进行回溯后将跳转到的地方。一旦完成重复，我们就到达了正则的末尾。现在我们要做的就是检查是否到达了字符串末尾。如果到达了，我们返回 1 表示成功。如果没有到达，意味着正则匹配失败，我们需要执行失败逻辑进行回退。 ```assembly 4c: 39400009 ldrb w9, [x0] 50: 35000069 cbnz w9, 0x5c ; 不是 NUL -> 回退块 54: d2800020 mov x0, #1 ; 成功 58: d65f03c0 ret ``` 最后，是回退逻辑。它检查栈是否为空。如果为空，我们返回 0。如果不为空，我们从栈中弹出回退地址和字符串位置，然后跳转到该地址。 ```assembly 5c: eb02003f cmp x1, x2 ; 还有栈帧吗？ 60: 54000060 b.eq 0x6c ; 没有 -> 放弃 64: a9ff0029 ldp x9, x0, [x1, #-16]! ; 弹出 (恢复地址, 位置) 68: d61f0120 br x9 ; 跳转到那里 6c: d2800000 mov x0, #0 ; 没有匹配 70: d65f03c0 ret ``` ### 构建模板 现在你已经看到了编译后的代码，你应该开始对复制并修补编译器的工作原理有所了解。我们有相同的指令集，彼此之间只有细微差别。对于每个这样的功能块，我们可以编写一个函数来生成相应的代码。每个函数将接收用于修改代码的值。例如，`stencil_char` 的参数之一将是正则中要比较的字符。我们将直接将该字符插入到机器码中。 序言部分很简单，因为它只是一段代码： ```rust const PROLOGUE_WORDS: usize = 1; fn stencil_prologue() -> [u32; PROLOGUE_WORDS] { [0xAA0103E2] // mov x2, x1 } ``` 对于字符比较，我们需要插入要比较的字符以及回退逻辑的跳转位置： ```rust 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 ] } ``` 对于重复，我们有循环的开始部分（压栈）和跳转到结束部分的代码： ```rust 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 } ``` 然后我们有匹配和失败块，它们非常清晰： ```rust 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 (跳转到下面的 mov) 0xA9FF0029, // ldp x9, x0, [x1, #-16]! 0xD61F0120, // br x9 0xD2800000, // mov x0, #0 0xD65F03C0, // ret ] } ``` 为完整起见，这里是我们使用的辅助函数，它们帮助我们向指令中插入特定数据： ```rust // 计算条件分支（b.ne / cbnz）的偏移字段： // 从分支指令到目标指令的指令数，存储在位 5..24 中。 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 } // 计算无条件分支（b）的偏移字段： // 同样的思路，但存储在位 0..26 中。 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 } // 从绝对地址中提取 16 位，为 movz/movk 立即数字段定位。 fn addr_bits(addr: u64, part: usize) -> u32 { (((addr >> (16 * part)) & 0xFFFF) as u32) << 5 } ``` ### 发射代码 现在是驱动代码： ```rust // 计算一个节点编译后的指令数。 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, // code[0] 的运行时地址，用于绝对地址填充 } 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)); } } } } // 生成完整的程序：序言、编译后的 AST、匹配块、失败块。 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 } ``` 这就是最困难的部分！就我个人而言，编写汇编是我认为 AI 最有帮助的地方。我关于汇编的经验主要来自完成 microcorruption CTF 挑战。我自己从未真正编写过汇编代码。如果没有 AI，我会非常难以弄清楚所需的精确指令以及如何修改它们以获得我想要的输出。有了 AI，我可以给我的编程代理描述 JIT 编译器工作的大致框架，它就能帮我处理很多细节。 ### 加载机器代码 要完成我们的编译器，我们需要实际加载代码。为此，我们将使用 `mmap` 分配一块可读、可写、可执行的内存。然后我们将代码复制到那块内存中，并将该内存块转换为一个函数，然后调用它： ```rust const BSTACK_MAX: usize = 4096; // 这些函数包含在 macOS 系统库中 unsafe extern "C" { fn pthread_jit_write_protect_np(enabled: libc::c_int); } ```