JIT-компиляция кода за 5 микросекунд

Долгое время быстрая JIT-компиляция считалась чем-то вроде чёрной магии. Чтобы написать быстрый JIT-компилятор, нужно было уметь писать ассемблер. Показательный факт: сегодня нет ни одной production-ready базы данных с собственным JIT-компилятором — все либо используют LLVM, либо генерируют код на C/C++. У обоих подходов есть общая проблема — высокое время компиляции, которое ограничивает область применения. Теперь, с использованием AI, написать JIT-компилятор с быстрой компиляцией, напрямую генерирующий ассемблер, стало значительно проще. Это же открывает возможность новым базам данных обойти старые. При разработке pgrust изначально казалось, что реализовать JIT-компилятор будет очень сложно. В итоге всё оказалось намного проще благодаря помощи AI, и это одна из причин, почему pgrust работает так быстро. JIT-компилятор в pgrust компилирует код примерно за 5 микросекунд, что позволяет JIT-компилировать каждый SQL-запрос, а не только часть из них. Дальше — пошаговый разбор того, как построить собственный быстрый JIT-компилятор, на примере простого движка регулярных выражений, использующего JIT-компиляцию.

Зачем нужна JIT-компиляция

JIT-компиляция — это генерация скомпилированного кода во время выполнения программы, «точно вовремя» (Just In Time). При правильной реализации это даёт заметный прирост производительности — часто в 2-5 раз, а иногда и больше. Основной сценарий применения JIT-компиляции — ситуации, когда информация, полученная во время выполнения, кардинально меняет поведение программы. Это особенно характерно для интерпретаторов языков программирования: код для выполнения они получают именно во время работы. JIT-компиляторы полезны и за пределами языков программирования — например, при разборе данных. Иногда схема разбираемых данных неизвестна до момента выполнения, и здесь JIT тоже может помочь.

Для начала реализуем игрушечный движок регулярных выражений. Для простоты будет поддерживаться всего две возможности: литеральные строки и повторение (то есть регекс-звёздочка *). Парсер тоже опустим — регулярное выражение будет представлено уже разобранными структурами Rust. Это значит, что получится обрабатывать строки вроде:

• apples
• b(an)*

но не будет поддержки альтернации, lookbehind и подобного.

В коде всё довольно просто. Будет три типа узлов: узел литеральной строки, узел повторения и узел конкатенации — комбинация двух узлов. Выглядит это так:

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))
}

Написать интерпретатор для этого движка регулярных выражений тоже несложно:

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 строк кода. Но как обстоят дела с производительностью? Для сравнения возьмём код, написанный вручную специально под конкретный 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
}

(Этот код можно было бы оптимизировать и сделать значительно быстрее, но для наших целей он служит хорошей точкой сравнения.)

При бенчмарке нескольких примеров ручная версия оказалась в 10-20 раз быстрее интерпретатора. Явно есть куда расти.

Дальше рассмотрим, как с помощью JIT-компиляции получить универсальный движок регулярных выражений, который работает так же быстро, как ручная реализация.

Как выполнить JIT-компиляцию

JIT-компиляция кода состоит из двух шагов. Сначала генерируется ассемблерный код для нужной операции. Затем этот код упаковывается в функцию, которую можно вызывать из программы как обычную.

Для генерации ассемблера используется вариант подхода copy-and-patch. Идея в том, что для разных операций, которые нужно JIT-компилировать, заранее готовятся шаблоны ассемблерного кода — их называют «трафаретами» (stencils). Когда нужно JIT-скомпилировать операцию, берётся соответствующий трафарет и вносятся небольшие правки в зависимости от специфики операции — очень похоже на заполнение настоящего трафарета. Соединяя несколько таких заполненных трафаретов, можно во время выполнения собрать программу, по производительности близкую к написанной вручную.

Путь такой: сначала рассматривается ARM64-код, который нужно сгенерировать для b(an)*. Затем повторяющиеся последовательности инструкций превращаются в переиспользуемые трафареты, пишется эмиттер, который заполняет и комбинирует эти трафареты на основе AST регулярного выражения, и, наконец, сгенерированные инструкции копируются в исполняемую память, чтобы Rust мог вызывать их как обычную функцию.

Проще всего разобраться в этом, начав с готового сгенерированного кода и двигаясь назад к самому JIT-компилятору. Снова используем регулярное выражение «b(an)*». Ключевые решения по дизайну:

• Для бэктрекинга используется стек. Стек хранит состояние, к которому нужно вернуться при попадании в тупик регулярного выражения
• Строка, с которой идёт сопоставление, заканчивается нулевым байтом. Это значит, что любое сравнение символов автоматически провалится при достижении конца строки — и не придётся нигде сравнивать длину

Для состояния программы используются следующие регистры:

• x0 — текущая позиция в строке и возвращаемое значение
• x1 — вершина стека для бэктрекинга
• x2 — дно стека для бэктрекинга (нужно, чтобы определить, пуст ли стек)
• x9 — используется как временная переменная

Во входные параметры программы передаются:

• x0 — указатель на начало строки
• x1 — указатель на область памяти под стек

Сгенерированный ARM64-код

Теперь, разобравшись с этим, пройдёмся по сгенерированному ассемблеру блок за блоком. Пример приведён конкретно для macOS на ARM64. Сначала идёт пролог, инициализирующий программу. Он просто устанавливает вершину и дно стека равными переданному значению:

0:  aa0103e2   mov   x2, x1

Дальше — код, проверяющий символ b. Если встречается символ, отличный от b, происходит переход к блоку обработки отката. В противном случае позиция в строке продвигается вперёд:

; CHAR 'b'
  4:  39400009   ldrb  w9, [x0]                  ; load current input byte
  8:  7101893f   cmp   w9, #0x62                 ; is it 'b'?
  c:  54000281   b.ne  0x5c                      ; no -> fallback block
 10:  91000400   add   x0, x0, #1                ; yes -> advance input

Далее — повторение (an)*. Для повторения нужен бэктрекинг. Если откат происходит здесь, это значит переход сразу к концу цикла. Поэтому в стек нужно сохранить и адрес инструкции после цикла, и текущую позицию в строке:

 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

После этого можно выполнять тело повторения. Оно проверяет символы «a» и «n» и, если они найдены, возвращается к началу повторения, но уже с новой позицией в строке.

; 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

Теперь цикл пройден. Именно сюда попадёт выполнение после отката. По завершении повторения достигается конец регулярного выражения. Остаётся проверить, достигнут ли конец строки. Если да — возвращается 1 (успех). Если нет — регулярное выражение не совпало, и нужно выполнить логику отката (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

И наконец, логика отката. Она проверяет, пуст ли стек. Если да — возвращается 0. Если нет — со стека снимаются адрес возврата и позиция в строке, после чего выполняется переход на этот адрес.

 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

Построение трафаретов

Теперь, увидев скомпилированный код, легче понять принцип работы copy-and-patch компилятора. Есть общие наборы инструкций с небольшими различиями между собой. Для каждого такого блока функций можно написать функцию, генерирующую соответствующий код. Каждая функция принимает значения, которые нужно вставить в код. Например, одним из аргументов stencil_char будет символ регулярного выражения, с которым идёт сравнение — этот символ вставляется прямо в машинный код.

Пролог реализуется просто, поскольку это статичный блок кода:

const PROLOGUE_WORDS: usize = 1;

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

Для сравнения символов нужно вставить сравниваемый символ и адрес перехода для логики отката:

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
    ]
}

Для повторения — начало цикла, помещающее данные в стек, и переход в конец:

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
}

И, наконец, блоки match и fail — они получаются довольно чистыми:

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
    ]
}

Для полноты картины — вспомогательные функции, которые помогают вставлять конкретные данные в инструкции:

// 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) << 5) as u32
}

// Compute the branch-offset field for an unconditional branch (b):
// same idea, but stored in bits 0..26.
fn 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) & 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) << 5
}

Генерация кода

Теперь код, управляющий этим процессом:

// Computes how many instructions a node compiles to.
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, // 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 + exit as u64 * 4));
                self.emit_node(body);
                self.emit(&stencil_jmp(self.pos(), split_at));
            }
        }
    }
}

// Generates the complete program: prologue, the compiled AST, MATCH, fail block.
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 помогает больше всего. Основной опыт с ассемблером до этого ограничивался прохождением CTF microcorruption — сам ассемблер писать не приходилось никогда. Без AI было бы очень сложно разобраться, какие именно инструкции нужны и как их модифицировать для получения нужного результата. С AI можно объяснить кодинг-агенту общую идею того, как должен работать JIT-компилятор, и он берёт на себя множество таких деталей.

Загрузка машинного кода

Чтобы завершить компилятор, нужно фактически загрузить код. Для этого используется mmap — выделяется блок памяти, доступный для чтения, записи и исполнения. Затем код копируется в эту память, а сам блок памяти превращается в функцию, которую можно вызывать:

const BSTACK_MAX: usize = 4096;

// These functions are included in the mac system library
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); // make the region writable (Apple W^X)
            std::slice::from_raw_parts_mut(buf, code.len()).copy_from_slice(&code);
            pthread_jit_write_protect_np(1); // back to executable
            sys_icache_invalidate(buf as *mut libc::c_void, nbytes);

            Jit { buf, nbytes, bstack: vec![0; BSTACK_MAX * 2] }
        }
    }

    // Runs the generated code. Input must end with a NUL byte.
    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);
        }
    }
}

Результаты

После реализации всех вариантов можно сравнить их производительность:

JIT-версия и написанная вручную реализация показывают почти одинаковую производительность — иногда быстрее оказывается JIT, иногда ручной код.

В последнее время распространён мем о том, что AI не особо помогает, потому что «код никогда не был самой сложной частью». Это верно для некоторых областей, но в других написание кода как раз и было самым сложным. JIT-компиляторы — отличный тому пример. Для многих программ JIT-компилятор мог бы значительно ускорить работу. Редкость JIT-компиляторов говорит о том, что исторически их реализация была слишком сложной, чтобы игра стоила свеч. LLM снизили порог входа и заметно упростили написание JIT-компилятора. Именно эта идея лежит в основе pgrust. Базы данных исторически были одним из самых сложных видов программного обеспечения, и это их ограничивало. Теперь, с AI, можно быть куда амбициознее в выборе того, какое программное обеспечение создавать.