Зливаємо конвеєр risp, щоб обігнати CPython: вбити 400 000 виділень виявилося замало

У risp, Lisp на Rust без залежностей, злиття map/filter/fold скорочує 400 000 виділень до O(1), але CPython обганяє лише мономорфний цикл i64: 4 мс проти 24 мс.

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

risp — це Lisp, який я написав на Rust із трьома рушіями виконання, складеними за одним фронтендом: деревообхідний інтерпретатор, байткодова VM з кешуванням у регістрах та вмикуваний Cranelift JIT, що перетворює самодостатні цілочислові функції на нативний машинний код. Це не «Risp» 2019 року з туторіалу Risp in (almost) 200 lines of Rust — той самий каламбур, зовсім інший звір (вихідний код на GitHub), і саме ця відмінність — суть цього допису: на цілочислових бенчмарках той JIT працює у 10–40× швидше за CPython 3.14, і все ж на списочних бенчмарках — map, filter, fold — risp програвав CPython, і щоб відіграти це назад, знадобився не більший JIT.

Канонічний конвеєр — це 400 000 виділень у плащі

Ось бенчмарк, тієї самої форми, яку використовує кожен мікробенчмарк «функційного конвеєра» — піднести діапазон до квадрата, додати квадрати:

(fold + 0 (map (lambda (x) (* x x)) (range 200000 '())))

Прочитайте його зсередини назовні й порахуйте купу. range будує cons-список на 200 000 комірок. map обходить його й будує ще один список на 200 000 комірок. fold обходить другий до єдиного числа. Це приблизно 400 000 виділень Rc<Pair> у купі, плюс увесь трафік підрахунку посилань, щоб знову їх звільнити — і це виділення і є бенчмарком. Накладні витрати інтерпретатора — похибка округлення поряд із цим.

reduce(operator.add, map(lambda x: x*x, range(1, 200001)), 0) у CPython виглядає символ-у-символ ідентично, але range і map — це ліниві ітератори: нічого не матеріалізується, конвеєр потоково подає по одному елементу за O(1) додаткової пам’яті, а reduce/operator.add написані на C. Розрив ніколи не був у повільності мого інтерпретатора — це були 400 000 виділень, яких CPython просто ніколи не виконує.

Вимірювання точно підтверджує цю історію (повний аналіз базової лінії див. у чому CPython обганяє Lisp, що компілюється в нативний код). Жадібний risp проти CPython на 2×10⁵:

конвеєр (2×10⁵)рушійчаспроти CPython
жадібний map + fold (будує ~400k cons)VM / JIT68 / 69 мс3.36× повільніше

Колонка JIT збігається з VM на 69 мс, бо код побудови списків непридатний для цілих чисел — JIT дивиться на нього, бачить cons і Pair, відмовляється й повертає його прямо назад VM як байткод. Тут йому нема чого компілювати. Мета, сформульована точно: знизити проміжне виділення з O(n) до O(1), як це роблять ітератори CPython.

Пастка, що виглядає як злиття, але виділяє пам’ять на кожен елемент

Очевидна перша ідея — лінивий потік: cons-комірка, чий хвіст — це thunk, який ви форсуєте на вимогу. Він читається як лінивий ітератор. Це катастрофа. Кожен крок виділяє cons-комірку і замикання-thunk, щоб утримати решту обчислення, тож ви перейшли від одного виділення на елемент до двох. Це гірше за жадібну версію, яку воно мало виправити.

Це несуче правило для решти допису: якщо ваше «злиття» виділяє вузол на елемент, це не злиття. Я підтвердив це на етапі дизайну й не пішов цим шляхом. Справжнє злиття мусить протягувати кожен елемент крізь усі етапи, ніколи не матеріалізуючи місце, куди його покласти.

Перетворювачі кодують конвеєр як дані, а не як замикання, що захоплюють

Форма, що насправді дає O(1) проміжного виділення, — це перетворювач (transducer): замість того, щоб кожен етап продукував колекцію, кожен етап загортає згортальну функцію, а один reduce проштовхує кожен вихідний елемент крізь увесь скомпонований перетворювач прямо в акумулятор. Жоден етап ніколи не будує колекцію.

Неочевидна деталь реалізації — як представлено етапи. У prelude.lisp risp перетворювач — це дані — список — а не замикання:

;; A transducer here is DATA, not a capturing closure: `(map-x f)` is the
;; two-element list `(map f)` and `(filter-x p)` is `(filter p)`. They are kept
;; as data on purpose — a closure that closed over `f`/`p` would capture an
;; enclosing local, which the bytecode compiler rejects, so the whole facility
;; would be tree-walker-only.
(def map-x (lambda (f) (list (list (quote map) f))))
(def filter-x (lambda (p) (list (list (quote filter) p))))

Той тег (quote map) / (quote filter) — не стилістика. Байткодовий компілятор risp не випускає upvalues для захоплених змінних — те саме обмеження, що змушує map/filter/fold явно протягувати свій функціональний аргумент, а не замикатися над ним. Перетворювач, побудований у стилі Clojure, як замикання, що захоплює f чи p, був би відхилений компілятором VM, і вся можливість застрягла б на деревообхідному інтерпретаторі. Кодування етапу як списку (map f) та протягування його крізь явні параметри тримає кожну функцію-перетворювач придатною для швидкого шляху --vm і --jit.

Власне злиття — це один хвостоворекурсивний драйвер, що протягує єдиний елемент крізь решту етапів та у згортальну функцію, без проміжного списку на жодному кроці:

;; Drive ONE source element `x` through the remaining `stages`, threading the
;; reducer `rf` and accumulator `acc`. A `map` stage rewrites the live value; a
;; `filter` stage that rejects SHORT-CIRCUITS and returns `acc` unchanged. When
;; the stages are exhausted the element has survived the whole pipeline, so it
;; is reduced in: `(rf acc x)`. No intermediate list is ever built.
(def transduce--step
  (lambda (stages rf acc x)
    (if (null? stages)
        (rf acc x)
        (let ((stage (car stages)))
          (let ((tag (car stage)) (op (car (cdr stage))))
            (if (equal? tag (quote map))
                (transduce--step (cdr stages) rf acc (op x))
                ;; filter: keep threading the SAME x only if the predicate holds.
                (if (op x)
                    (transduce--step (cdr stages) rf acc x)
                    acc)))))))

Етап map переписує живе значення й протягує (op x) далі; етап filter або протягує той самий x далі, або коротко замикається до незмінного acc. Протягується лише єдиний акумулятор. Це справжнє злиття — і воно працює:

конвеєр (2×10⁵)рушійчаспроти CPython
узагальнений transduce (без виділень, прив’язаний до замикань)VM / JIT63 мс2.78× повільніше

Без виділень — і все одно програє.

Чесний висновок: убити саме лише виділення недостатньо, щоб обігнати CPython

Усунення всіх ~400 000 виділень зрушило risp з 3.36× повільніше до 2.78× повільніше — реальне покращення, і близько не паритет. Усунення виділення — те, за що всі хапаються першим, — було великою частиною розриву, але не всім розривом.

Залишкова витрата — це поелементний виклик узагальненого замикання. transduce викликає замикання часу виконання (op, rf) на кожен окремий елемент, а узагальнене замикання — це саме те, що JIT не може спеціалізувати: це не самодостатня цілочислова функція, це міжфункційний виклик крізь захоплене значення. VM risp платить за поелементну диспетчеризацію замикань там, де C-шний reduce у CPython не платить, і JIT нічого не може з цим вдіяти — його замикання захоплюють f/p, тож вони виконуються на VM, але ніколи не придатні для JIT, з тієї самої причини «відхиляй-не-вгадуй», з якої він взагалі відмовляється від списків та міжфункційних викликів.

Чесний висновок: для динамічно типізованої VM лінивість/злиття є необхідним, щоб бути конкурентним, але недостатнім, щоб перемогти. Треба вбити і виділення, і замикання.

Прибираємо ще й вихідний список, потім мономорфізуємо весь цикл

Ще два кроки закривають це. По-перше, навіть злитий конвеєр, нагодований (range lo hi), все одно матеріалізує вхідний список. reduce-range усуває це, ведучи згортальну функцію напряму над напіввідкритим інтервалом за допомогою лічильного циклу — цілі числа існують лише як змінні циклу, потоковий аналог лінивого range у CPython:

;; Reduce `[lo hi)` with `rf` from `acc`, one integer at a time, no list built.
(def reduce-range
  (lambda (rf acc lo hi)
    (if (< lo hi)
        (reduce-range rf (rf acc lo) (+ lo 1) hi)
        acc)))

Це вбиває останнє виділення, але лишає виклик замикання. Виграш — єдине, що насправді обганяє CPython, — це припинити передавати етапи як функції часу виконання й почати передавати їх як вбудовані шаблони виразів. Це макрос злиття sum-of. Його помічник часу макророзкриття sum-of--build вклеює тіло кожного етапу прямо в цикл:

;; Build the new-accumulator expression for ONE element, folding the remaining
;; `stages` over the current-value expression `cur`, accumulating into `acc`.
;; (A macro-time helper: it returns CODE.)
(def sum-of--build
  (lambda (stages cur acc)
    (if (null? stages)
        (list (quote +) acc cur)
        (let ((stage (car stages)))
          (let ((kind (car stage))
                (var (car (car (cdr stage))))
                (form (car (cdr (cdr stage)))))
            (if (equal? kind (quote map))
                ;; (let ((var cur)) <build rest with current value = form>)
                (list (quote let) (list (list var cur))
                      (sum-of--build (cdr stages) form acc))
                ;; filter: (let ((var cur)) (if form <build rest with cur=var> acc))
                (list (quote let) (list (list var cur))
                      (list (quote if) form
                            (sum-of--build (cdr stages) var acc)
                            acc))))))))

Етап map випускає (let ((var cur)) ...), прив’язуючи елемент і протягуючи його тіло-вираз як наступне поточне значення; етап filter випускає ту саму прив’язку, загорнуту в (if form ... acc). Макрос потім загортає це у хвостоворекурсивний цикл із gensym’нутими параметрами i/acc, тож ніщо не захоплює місце виклику:

(defmacro sum-of args
  (let ((loop (gensym))
        (stages (reverse (cdr (reverse args))))   ; all but last
        (src (car (reverse args))))               ; the (range lo hi) form
    (let ((lo (car (cdr src)))
          (hi (car (cdr (cdr src)))))
      `(begin
         (def ,loop
           (lambda (i acc)
             (if (< i ,hi)
                 (,loop (+ i 1) ,(sum-of--build stages (quote i) (quote acc)))
                 acc)))
         (,loop ,lo 0)))))

Випущене тіло циклу — це чисті + - * < над i та acc, без виходу викликом — саме та самодостатня цілочислова підмножина, яку наявний JIT компілює в нативний регістровий цикл. Нема списку, який виділяти, і нема замикання, крізь яке диспетчеризувати, бо етапи вклеєно вбудовано як код. Одна зміна в компіляторі дала цьому дістатися JIT: шлях VM тепер макророзкриває замість того, щоб відкидати макроси на деревообхідний інтерпретатор, тож розкриття sum-of виконується на --vm/--jit, а його цикл на i64 стає нативним кодом.

Межа різка: перемагає лише форма, злита до цілих чисел

;; Shape:
;;   (sum-of (map (x) (* x x)) (filter (x) (> x 100000)) (range 1 N))

Перевірено, злитий map на 2×10⁵: дерево 233 мс / VM 13 мс / JIT 4 мс проти CPython 24 мс — приблизно у 6× швидше. Злитий filter на 2×10⁵: JIT 4 мс проти CPython 19 мс — приблизно у 5× швидше. Подивіться на каскад на одному навантаженні: жадібний 68 мс → узагальнено-злитий VM 13 мс → мономорфно-злитий JIT 4 мс. Злиття купує перші ~5×; вихід на JIT купує решту.

І це справді нативний код, а не випадковість згортання констант. На 20 000 000 елементах JIT виконується за 0.01 с проти 0.79 с у VM — відрив приблизно у 79×, який дає лише справжній регістровий цикл — а результат побітово ідентичний (2666686666700000) у деревообхідному інтерпретаторі, VM та JIT. Диференційні тести, що тримають JIT байт-у-байт ідентичним деревообхідному інтерпретатору, покривають і це: злиттю конвеєра не дозволено змінювати відповідь.

Межу варто сформулювати без прикрас. sum-of — це спеціалізована форма злиття — звести конвеєр map/filter над цілочисловим діапазоном до числа — а не загальне опрацювання списків. Загальні map/filter над довільними значеннями все одно будують списки. Узагальнені перетворювачі над довільними функціями обходяться без виділень, але лишаються в ~2.8× від CPython на поелементній диспетчеризації замикань. risp обганяє CPython саме на тій формі, що зливається до цілих чисел — числові згортки, підрахунки, суми, фільтровані агрегації.

Шлях до перемоги над CPython на конвеєрах пролягав крізь злиття в цілочисловий JIT, який у мене вже був, а не крізь побудову більшого. (Якщо хочете простежити, як внутрішні частини risp складаються докупи, вивчайте Rust, читаючи інтерпретатор Lisp проходить кодову базу з нуля.) Більший JIT — загорнуті значення, виділення в купі, підрахунок посилань і диспетчеризація замикань, усе в машинному коді — був би тижнями найнебезпечнішої компіляторної роботи заради виграшу, обмеженого приблизно 1.5–2.5×, бо він атакує накладні витрати інтерпретації, лишаючи виділення, а потім і диспетчеризацію замикань рівно там, де вони були. Злиття видаляє виділення; мономорфізація видаляє замикання; JIT, що вже існував, робить решту. Це лінива відповідь в обох сенсах — менше нової машинерії, і вона перемагає, не роблячи роботи: 400 000 виділень, що ніколи не були потрібні, та поелементна диспетчеризація, якої спеціалізований цикл ніколи не потребує.

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

Як злити конвеєр map/filter/fold так, щоб він виділяв пам'ять один раз?
Скористайтеся перетворювачами (transducers): компонуйте перетворення згортальної функції й запускайте єдиний reduce, що проштовхує кожен елемент крізь увесь скомпонований перетворювач прямо в акумулятор, тож на жодному етапі не будується проміжний список. У risp це знижує конвеєр приблизно з 400 000 виділень у купі до O(1) проміжних виділень.
Чи достатньо усунути виділення пам'яті, щоб обігнати CPython на map/fold?
Ні. Усунення 400 000 виділень risp через узагальнені перетворювачі зрушило його з відставання в 3.36× до 2.78× повільніше за CPython — краще, але не перемога. Залишкова витрата — це поелементний виклик узагальненого замикання, яке цілочисловий JIT не може спеціалізувати. Щоб перемогти, треба вбити і виділення, і замикання.
Наскільки швидкий злитий цілочисловий конвеєр risp проти CPython?
Макрос злиття sum-of зводить конвеєр map/filter над цілочисловим діапазоном до самодостатнього циклу на i64, який JIT компілює в нативний код: 4 мс проти 24 мс у CPython на 200 000 елементах (приблизно у 6× швидше). На 20 мільйонах елементів JIT виконується за 0.01 с проти 0.79 с у VM — відрив приблизно у 79×, який дає лише справжній регістровий цикл.

← Усі дописи