Додаємо Cranelift JIT, який працює у 10–23× швидше за CPython

Заглиблення у вмикуваний Cranelift JIT для risp на Rust: перевірник типів за принципом «відхиляй, не вгадуй», зведення if/cond/let до Cranelift IR, вбудоване виявлення знакового переповнення, хвостові виклики як зворотні ребра циклу та охоронець деоптимізації, що випереджає CPython у 10–42 рази.

Джерело: github.com/kissishka/risp

Після одинадцяти кроків оптимізації байткодова віртуальна машина risp обійшла CPython усюди, окрім рекурсії, обмеженої викликамиfib(30) лишався відставати в 1.55×. Жодне полірування портативного циклу з диспетчеризацією через match не переможе спеціалізований нативний шлях викликів CPython на чистій рекурсії. Тож останній крок припиняє інтерпретувати й починає генерувати машинний код за допомогою Cranelift. Результат: fib 6.3 мс проти 61 мс у CPython (~10×), tailsum 2.4 мс (13×), арифметичний цикл 4.1 мс (23×), а насичений cond/let/and/or цикл диспетчеризації 5 мс проти 208 мс (~42×).

Найскладніше в JIT не швидкість. Це збереження коректності. Відповідь risp полягає в тому, щоб компілювати лише те, що можна довести еквівалентним деревообхідному інтерпретатору, і відхиляти все інше.

Відхиляй, не вгадуй: придатність

Придатна лише вузька підмножина: верхньорівневий (def NAME (lambda …)), фіксована арність ≤ 4, тіло з єдиного виразу, стабільне зв’язування (ніколи не перевизначається й не змінюється через set!). Драйвер навіть не виводить тип повернення. Він вгадує Int, потім Bool, і приймає перший здогад, що проходить перевірку типів як узгоджена нерухома точка:

let ret = [Ty::Int, Ty::Bool]
    .into_iter()
    .find(|&guess| checker.check(&body, guess, &mut params.clone()) == Some(guess));
if let Some(ret) = ret {
    out.push(FnDef { name: name.clone(), params, body, ret, guarded });
}

Перевірник — це тотальна функція, що повертає Option<Ty>: Some, якщо вираз доказово належить підмножині {Int, Bool}, і None для будь-чого іншого. А None поширюється через ?, тож єдиний непідтримуваний листок дискваліфікує всю функцію:

fn check(&self, expr: &Value, self_ret: Ty, scope: &mut Vec<Rc<str>>) -> Option<Ty> {
    match expr {
        Value::Int(_) => Some(Ty::Int),
        Value::Bool(_) => Some(Ty::Bool),
        Value::Symbol(s) => scope.iter().any(|l| l == s).then_some(Ty::Int),
        Value::Pair(_) => {
            let (head, args) = as_call(expr)?;
            match &*head {
                "if" => {
                    if args.len() != 3 || self.check(&args[0], self_ret, scope)? != Ty::Bool {
                        return None;   // an Int condition is DECLINED (risp's 0 is truthy)
                    }
                    let t = self.check(&args[1], self_ret, scope)?;
                    let e = self.check(&args[2], self_ret, scope)?;
                    (t == e).then_some(t)   // both branches must share a type
                }
                "and" | "or" => self.check_and_or(&args, self_ret, scope),
                "cond"       => self.check_cond(&args, self_ret, scope),
                "let"        => self.check_let(&args, self_ret, scope),
                _ if self.is_primitive(head, scope) => { /* 2 Int args -> Int/Bool */ }
                _ if head == &**self.self_name => { /* self-call: Int args -> self_ret */ }
                _ => None,   // lists, closures, cross-function calls, ...
            }
        }
        _ => None,
    }
}

Подивіться, що він відхиляє: Int, використаний як умова if (істинність у risp робить 0 істинним, що сире галуження Cranelift сприйняло б хибно), cond без завершального else (значення не завжди було б типізованим), and/or над не-Bool операндами, а також усе, що стосується списків, рядків, замикань чи міжфункційних викликів. Існує рівно один спосіб сказати «так» і десяток способів сказати «ні». Якщо JIT не може довести еквівалентність, функцію без змін повертають віртуальній машині. Саме це дозволяє йому бути агресивним (сирі i64, mem::transmute на межі), ніколи не ризикуючи неправильною відповіддю: єдиний ризик, на який він іде, — не скомпілювати те, що міг би. А оскільки ціла підмножина замкнена відносно i64 (кожен листок — це i64, / виключено, бо воно підвищується до числа з рухомою комою), JIT ніколи не доводиться загортати значення посеред виконання. Це регістри від входу до повернення.

Зведення до Cranelift IR

Кожна змінна стає Variable у Cranelift (змінюваний SSA-слот), тож параметри, зв’язування let та точки злиття всі повторно використовують власний механізм use/def будівника, а Cranelift автоматично вставляє phi-вузли. Тіла компілюються в одному з двох режимів: tail (кожен шлях завершується термінатором) або value (видає один i64). if у позиції значення чітко демонструє цей патерн — обидві гілки визначають одну змінну результату й переходять до блоку злиття:

fn value_if(&mut self, args: &[Value]) -> Result<ClValue, RispError> {
    let cond = self.value(&args[0])?;
    let then_b = self.b.create_block();
    let else_b = self.b.create_block();
    let merge = self.b.create_block();
    let rv = self.b.declare_var(types::I64);
    brif(self.b, cond, then_b, else_b);

    self.b.switch_to_block(then_b);
    let tv = self.value(&args[1])?;
    self.b.def_var(rv, tv);
    jump(self.b, merge);

    self.b.switch_to_block(else_b);
    let ev = self.value(&args[2])?;
    self.b.def_var(rv, ev);
    jump(self.b, merge);

    self.b.switch_to_block(merge);
    Ok(self.b.use_var(rv))
}

Булеві значення не мають окремого IR-типу. Це iconst 0/1, а порівняння розширюють свій результат через uextend до i64, тож уся машина працює з одним класом регістрів. and/or зводяться до ланцюжків галужень з коротким замиканням, де пропущений операнд ніколи не обчислюється (точно як у деревообхідному інтерпретаторі), а оскільки перевірник типів змусив усі операнди бути Bool, галуження безпосередньо за ними коректне. Саме тому насичений cond/let/and/or бенчмарк диспетчеризації дає JIT найбільший відрив (42×): ці форми керування стають нативними галуженнями замість диспетчеризації опкодів VM.

Перевірена арифметика в нативному коді

Швидкість, що змінює результат, — це баг. + - * генерують вбудоване виявлення знакового переповнення, тож JIT спрацьовує пасткою рівно там, де checked_add у VM видав би помилку. Додавання використовує класичну тотожність знакового біта: переповнення відбувається тоді й лише тоді, коли операнди мають однаковий знак, а результат відрізняється:

"+" => {
    let r = self.b.ins().iadd(a, c);
    let t1 = self.b.ins().bxor(a, r);
    let t2 = self.b.ins().bxor(c, r);
    let t3 = self.b.ins().band(t1, t2);
    let ov = self.b.ins().icmp_imm(IntCC::SignedLessThan, t3, 0);
    (r, ov)
}
"*" => {
    // Full 128-bit product; overflow iff it doesn't fit back in i64.
    let a128 = self.b.ins().sextend(types::I128, a);
    let b128 = self.b.ins().sextend(types::I128, c);
    let p = self.b.ins().imul(a128, b128);
    let lo = self.b.ins().ireduce(types::I64, p);
    let lo128 = self.b.ins().sextend(types::I128, lo);
    let ov = self.b.ins().icmp(IntCC::NotEqual, lo128, p);
    (lo, ov)
}
// then: brif(self.b, overflow, self.trap_block, cont);

При переповненні керування переходить до блоку пастки, який записує збій у структуру контексту й повертається. Назад у Rust межа виклику перетворює це на ту саму типізовану помилку, яку видає вбудована функція:

if ctx.trapped != 0 {
    return match ctx.errkind {
        ERR_DEOPT => Ok(None),   // re-run on the VM (see below)
        _ => Err(RispError::Custom("integer overflow".into())),
    };
}

Тож (* 9223372036854775807 2) видає ідентичну помилку незалежно від того, виконувалося воно на VM чи на JIT. Воно ніколи не переповнюється з обгортанням.

Хвостові виклики — це зворотні ребра циклу; нехвостові виклики — це нативні виклики

Функція, що робить хвостовий виклик самій собі, мусить виконуватися в сталому нативному стеку, щоб відповідати оптимізації хвостових викликів у VM, тож вона компілюється не в call, а в jump назад до заголовка циклу з переприв’язаними змінними параметрів:

// Self tail-call -> rebind params and loop (constant native stack).
_ if &*head == self.self_name && args.len() == self.arity && self.lookup(&head).is_none() => {
    let mut newvals = Vec::with_capacity(args.len());
    for a in &args { newvals.push(self.value(a)?); }   // evaluate ALL args first
    for (i, nv) in newvals.into_iter().enumerate() {
        self.b.def_var(self.param_vars[i], nv);
    }
    jump(self.b, self.header);
    return Ok(());
}

Обчислення кожного аргументу до переприв’язування будь-якого параметра — це те, що робить (f b a) коректним. Результат — справжній цикл у машинному коді: жодного call, жодного зростання стека. Нехвостовий самовиклик (як два рекурсивні виклики у fib) натомість компілюється у справжній нативний call, а потім перевіряє, чи спрацювала пастка в адресата, виходячи у разі потреби. Саме на цій відмінності тримається наступний розділ.

Деоптимізація замість аварійного завершення

Нехвостова рекурсія використовує нативний C-стек, який обмежений (~8 MiB). Лінійно глибока нехвостова рекурсія — (sum n) = (+ n (sum (- n 1))) на мільйон рівнів углиб — переповнила б його й видала SIGABRT, тоді як деревообхідний інтерпретатор рекурсує в купі і повертає значення. Щоб лишатися поведінково ідентичним без аварійного завершення, кожна функція з нехвостовим самовикликом несе охоронець глибини часу виконання, який деоптимізує за межами стелі:

let depth_entry = if def.guarded {
    let d = b.ins().load(types::I64, MemFlags::trusted(), ctx_val, 16);
    let d1 = b.ins().iadd(d, b.ins().iconst(types::I64, 1));
    b.ins().store(MemFlags::trusted(), d1, ctx_val, 16);
    let over = b.ins().icmp_imm(IntCC::SignedGreaterThan, d1, DEPTH_LIMIT);
    let deopt_block = b.create_block();
    brif(&mut b, over, deopt_block, header);
    // deopt_block: set ctx.errkind = ERR_DEOPT, return sentinel 0.
    Some(d)
} else {
    jump(&mut b, header);   // no guard, no per-call overhead at all
    None
};

За межами DEPTH_LIMIT (10 000) він записує ERR_DEOPT, що не є помилкою, і повертається, сигналізуючи VM повторно виконати виклик на своєму байткодовому шляху зі стеком у купі. Охоронець рахує живу глибину, а не загальну кількість викликів: кожне нормальне повернення відновлює збережену глибину до інкременту, тож два сестринські піддерева fib не додаються одне до одного. Ось чому fib ніколи його не спрацьовує (fib(30) має глибину лише ~30 кадрів, значно нижче за 10 000) і виконується повністю нативно заради свого 10× виграшу, тоді як шлях деоптимізації — це чиста страховка для патологічної лінійної рекурсії. Яка функція потребує охоронця, вирішується статично шляхом обходу тіла на предмет нехвостового самовиклику, тож чистий хвостоворекурсивний цикл не платить нічого.

Межа виклику

ABI нативної функції — це extern "C" fn(*mut JitCtx, i64…) -> i64, де JitCtx має #[repr(C)], аби скомпільований код міг жорстко закодувати зміщення її полів:

#[repr(C)]
struct JitCtx {
    trapped: i64,   // @ 0
    errkind: i64,   // @ 8
    depth: i64,     // @ 16
}

Коли VM досягає виклику, чий адресат має нативний код і чиї аргументи всі є Int, вона перетворює (transmute) фіналізований вказівник на код до потрібної арності й переходить усередину:

let r: i64 = unsafe {
    match entry.arity {
        1 => mem::transmute::<*const u8, extern "C" fn(*mut JitCtx, i64) -> i64>(p)(cx, int(0)),
        2 => mem::transmute::<*const u8, extern "C" fn(*mut JitCtx, i64, i64) -> i64>(p)(cx, int(0), int(1)),
        // ... 0, 3, 4
        _ => unreachable!("jit arity is capped at 4 during eligibility"),
    }
};

Виклик fib тоді викликає fib нативною інструкцією call, з розподілом регістрів, ніколи не входячи повторно в інтерпретатор. Три витрати інтерпретатора зникають одночасно: жодної диспетчеризації опкодів, жодного облаштування кадрів (аргументи їдуть у регістрах) і жодного загортання. А резервний шлях — це те, що робить безпечним постачання цього рішення: функція, яка ніколи не була придатною, не має нативної точки входу, тож місця її виклику виконують звичайний байткодовий шлях, і навіть придатна функція відступає до резерву на кожному виклику тієї ж миті, щойно побачить не-Int аргумент під час виконання. JIT — це чистий прискорювач, накладений поверх незмінного рушія кроків 1–11.

Чому він вмикуваний

Cranelift — це справжня залежність, тож увесь JIT живе за #[cfg(feature = "jit")]. Типова збірка cargo build все одно розв’язує нуль залежностей і постачає незмінний інтерпретатор лише на std; --features jit — це перемикач «я візьму залежність заради нативної швидкості». Єдина форма, яку портативний інтерпретатор не міг виграти, дістається нативному коду з відривом на порядок (fib від відставання в 1.55× до випередження в ~10×), при цьому типова збірка не бере на себе жодного рядка зовнішнього коду. І кожна її частина утримується байт-у-байт ідентичною деревообхідному інтерпретатору тими ж диференційними тестами, що охороняють VM: паритет переповнення, деоптимізація-замість-аварії та коректне-відхилення — усе це перевіряється на відповідність еталонному рушію.

Часті запитання

Наскільки Cranelift JIT для Lisp швидший за CPython?
Вмикуваний Cranelift JIT у risp виконує fib приблизно у 10 разів швидше за CPython 3.14 (6.3 мс проти 61 мс), цикл хвостової суми у 13 разів, арифметичний цикл у 23 рази, а насичений потоком керування цикл диспетчеризації приблизно у 42 рази (5 мс проти 208 мс).
Як JIT зберігає коректність?
JIT у risp компілює лише ту вузьку підмножину, яку може довести еквівалентною деревообхідному інтерпретатору (tree-walker): функції фіксованої арності над цілими та булевими значеннями, а все інше відхиляє й повертає віртуальній машині. Тотальний перевірник типів повертає None для будь-якого непідтримуваного листка, тож єдиний необроблений випадок дискваліфікує всю функцію.
Що таке деоптимізація (deopt) у JIT?
Деоптимізація означає вихід із нативного коду назад до інтерпретатора замість аварійного завершення. JIT у risp виконує нехвостову рекурсію на обмеженому C-стеку; за межами ліміту глибини він записує сигнал деоптимізації та повторно виконує виклик на VM зі стеком у купі, відтворюючи результат деревообхідного інтерпретатора без переповнення стека.

← Усі дописи