Зливаємо конвеєр 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 / JIT | 68 / 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 / JIT | 63 мс | 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×, який дає лише справжній регістровий цикл.