Снижение стоимости и поздние паттерны перезаписи
Снижение стоимости заменяет дорогие операции более дешёвыми эквивалентами. Эти паттерны выполняются поздно в пайплайне (Стадии 18–19), поскольку более ранним проходам нужна видимость исходной структуры операций. Например, Add(Mul(a, b), c) должен оставаться видимым для алгебраического упрощения, прежде чем будет слит в MULACC(a, b, c).
Исходники Tinygrad: tinygrad/uop/decompositions.py:438-480 (get_late_rewrite_patterns).
Исходники Morok: schedule/src/rangeify/patterns.rs (группа поздней декомпозиции) + schedule/src/symbolic/fast_div.rs.
Оценки тактов на этой странице приблизительны для современного x86-64. Реальные задержки зависят от микроархитектуры и состояния конвейера.
Все паттерны объединены в один проход перезаписи до фиксированной точки (PM_FINAL) вместе с symbolic_simple() (алгебраическая очистка) и pm_render() (векторизация CONST, CAT-в-VECTORIZE).
1. Оптимизация степеней двойки
Наиболее эффективное снижение стоимости. Целочисленное деление и взятие остатка от константных делителей крайне распространены в индексации тензоров — вычисление шагов, тайлинг и восстановление координат из плоских индексов всё это порождают.
| Паттерн | До | После | Экономия тактов |
|---|---|---|---|
x % 2^n | idiv + imul + isub (~25 тактов) | and (1 такт) | ~24x |
x * 2^n | imul (~3–4 такта) | shl (1 такт) | ~3x |
x // 2^n (без знака) | idiv (~20–40 тактов) | shr (1 такт) | ~20–40x |
Оптимизация модуля работает потому, что 2^n - 1 является битовой маской нижних n бит. Пример: x % 8 = x & 0b111.
Tinygrad: decompositions.py:448-454. Morok: pm_mod_to_and, pm_mul_to_shl, pm_div_to_shr в rangeify/patterns.rs.
:::caution Знаковое деление
Для знаковых целых x // 2^n — это НЕ просто x >> n. Арифметический сдвиг вправо округляет к отрицательной бесконечности, а целочисленное деление — к нулю.
Исправление: (x + (x < 0 ? 2^n - 1 : 0)) >> n
Поправка 2^n - 1, добавляемая для отрицательных значений, корректирует направление округления. Это соответствует тождеству:
floor(x / 2^n) = (x + 2^n - 1) >> n when x < 0
x >> n when x >= 0
Morok проверяет vmin >= 0 через анализ диапазонов (VminVmaxProperty), чтобы пропустить поправку, когда делимое доказуемо неотрицательно. Tinygrad использует принадлежность к типу (dtypes.uints) для той же цели.
Tinygrad: decompositions.py:452-454. Morok: pm_div_to_shr в rangeify/patterns.rs.
:::
Сгенерированный C-код для знакового деления на степень двойки:
// Before: x / 8
int result = x / 8;
// After: strength reduction (signed path)
int result = (x + ((x >> 31) & 7)) >> 3;
// bias for negatives ^^^ ^shift
Когда x доказуемо неотрицателен (типично для индексных вычислений), знаковый путь полностью устраняется:
// After: strength reduction (unsigned path, vmin >= 0)
int result = x >> 3;
2. Быстрое целочисленное деление (Hacker's Delight)
Для неcтепенных константных делителей x / d заменяется на умножение со сдвигом: (x * M) >> S.
Математика
Для положительной константы d и диапазона значений [0, max_val] находятся магическое число M и сдвиг S такие, что:
(x * M) >> S == x / d for all 0 <= x <= max_val
Почему это работает: Деление на d эквивалентно умножению на 1/d. Мы приближаем 1/d как M / 2^S, где M и S выбраны так, чтобы приближение было точным на диапазоне значений. Ключевая идея — целочисленное усечение делает точное представление возможным: нам нужно лишь floor(x * M / 2^S) == floor(x / d), а не вещественное равенство.
Алгоритм
Из Hacker's Delight, глава 10 (magicgu в Tinygrad, decompositions.py:272-280):
- Вычислить
nc = floor((max_val + 1) / d) * d - 1(критический порог) - Вычислить
nbits = bit_length(max_val) - Для
sот 0 до2 * nbits:- Если
2^s > nc * (d - 1 - (2^s - 1) mod d): найден валидный сдвиг - Вычислить
M = ceil((2^s + d - 1 - (2^s - 1) mod d) / d)
- Если
- Вернуть
(M, s)— наименьшую валидную пару(множитель, сдвиг)
Цикл находит наименьшее s, дающее валидное магическое число. Меньшее s означает меньшее M, что критически важно для размещения промежуточного произведения x * M в узких целочисленных типах.
Реализация в Morok: magic_unsigned() в schedule/src/symbolic/fast_div.rs.
Трёхэтапная стратегия
Соответствует decompositions.py:282-300 Tinygrad (fast_idiv):
| Этап | Условие | Преобразование | Пример |
|---|---|---|---|
| 1. Тот же dtype | M * vmax помещается в диапазон dtype | (x * M) >> S | x / 3 при x в i32 |
| 2. Факторизация степени 2 | d = 2^k * d', где d' > 1 | (x >> k) / d', затем magic на d' | x / 6 становится (x >> 1) / 3 |
| 3. Расширение до i64 | Переполнение Int32 в x * M | приведение к i64, умножение, сдвиг, приведение обратно | Запасной вариант для больших M |
Этап факторизации (2) важен: деление на 12 (= 4 * 3) превращается в сдвиг вправо на 2, за которым следует магическое деление на 3, что часто помещается в исходный dtype, тогда как прямое магическое деление на 12 вызвало бы переполнение.
Для знаковых значений добавляется коррекция: ((x * M) >> S) + (x < 0 ? 1 : 0). Это учитывает семантику усечения к нулю — без неё отрицательные делимые округляются в неправильном направлении.
Конкретный пример
x / 7 where x in [0, 255]:
magic_unsigned(255, 7) → M = 293, S = 11
Verify: (100 * 293) >> 11 = 29300 >> 11 = 14 = floor(100 / 7)
Verify: ( 7 * 293) >> 11 = 2051 >> 11 = 1 = floor( 7 / 7)
Verify: (255 * 293) >> 11 = 74715 >> 11 = 36 = floor(255 / 7)
Generated: (x * 293) >> 11 instead of x / 7
Cost: 1 imul + 1 shr (~4-5 cycles) vs 1 idiv (~20-40 cycles)
Сгенерированный LLVM IR
; Before: x / 7
%result = sdiv i32 %x, 7
; After: fast integer division (unsigned path)
%mul = mul i32 %x, 293
%result = lshr i32 %mul, 11
3. Деление float в умножение
x / c превращается в x * (1/c) для вещественной константы c.
Вещественное умножение занимает 1–2 такта (полностью конвейеризовано), тогда как деление — 10–20 тактов (не конвейеризовано на большинстве аппаратуры). Это прямолинейное ускорение в 5–10 раз для распространённого паттерна.
Защиты:
- Пропуск при
c == 0.0— деление на ноль должно сохраняться для соблюдения семантики IEEE 754 (x / 0.0даёт+/-infилиNaN) - Пропуск если
1/cне конечно (переполнение доinfозначает, чтоcслишком мало) - Только для вещественных типов
Tinygrad: decompositions.py:477-479 (backend-ы на основе FDIV генерируют RECIP как 1/x). Morok: pm_fdiv_to_mul в rangeify/patterns.rs.
// Before
float result = x / 3.14159f;
// After
float result = x * 0.31831f; // 1/pi
4. Слияние FMA (Fused Multiply-Add)
a * b + c превращается в MULACC(a, b, c).
Это отображается на аппаратные инструкции FMA (vfmadd на x86 AVX, fmadd на ARM NEON, fma.rn на CUDA). Одна инструкция заменяет две, с одним шагом округления вместо двух — что делает FMA и быстрее, и точнее, чем раздельные умножение + сложение.
Почему применяется поздно: Более ранним проходам нужно видеть структуру Add(Mul(a, b), c) для алгебраического упрощения. Если слить рано, паттерны вроде (x*2 + x*3) не смогли бы упроститься до x*5, поскольку узлы Mul были бы скрыты внутри MULACC.
Слияние сдвиг-сложение (только Tinygrad): Tinygrad также сливает (x << n) + c в MULACC(x, 2^n, c), перехватывая случаи, когда MUL-в-SHL сработал первым в том же проходе фиксированной точки. Этот паттерн пока не портирован в Morok.
Защиты: Совпадает только когда все три операнда (a, b, c) имеют одинаковый вещественный dtype. Целочисленный FMA не сливается, поскольку аппаратные инструкции FMA работают только с вещественными числами.
Tinygrad: decompositions.py:472-475. Morok: pm_fma_decomposition в rangeify/patterns.rs.
5. Извлечение отрицания
x * -1 превращается в NEG(x).
NEG — это одна инструкция (переключение бита знака для float через xorps, отрицание для int через neg). Умножение на -1 излишне занимает конвейер умножителя на 3–4 такта.
Срабатывает только когда backend поддерживает NEG как нативную операцию. Tinygrad: decompositions.py:458-459. Morok: pm_neg_from_mul.
6. Отрицания сравнений
Поздние перезаписи для отрицённых и составных сравнений целых чисел. Эти паттерны упрощают последовательности инструкций, возникающие из оптимизаций булевой логики на более ранних проходах.
| Паттерн | До | После | Экономия |
|---|---|---|---|
!(x < c) | NOT + CMP | (c-1) < x | Устранение NOT |
!(c < x) | NOT + CMP | x < (c+1) | Устранение NOT |
(c1 < x) & (x < c2) где c2 == c1+2 | 2 CMP + AND | x == (c1+1) | 2 операции устранены |
x * -1 < c | MUL + CMP | -c < x | Устранение MUL |
x * -1 < y * c | 2 MUL + CMP | y * (-c) < x | Устранение 1 MUL |
Сжатие диапазона (строка 3) особенно ценно. Когда открытый интервал (c1, c2) содержит ровно одно целое значение, два сравнения и логическое AND сворачиваются в одну проверку равенства. Это естественно возникает в тайлированных индексных вычислениях, где переменная диапазона выбирает ровно один тайл.
:::caution Целочисленное переполнение в константах
Паттерны отрицания защищены от переполнения: !(x < c) превращается в (c-1) < x только если c-1 не вызывает underflow, а !(c < x) превращается в x < (c+1) только если c+1 не вызывает overflow. В обоих случаях используются checked_sub / checked_add с возвратом None (без преобразования) при переполнении.
:::
Tinygrad: decompositions.py:461-470. Morok: pm_comparison_negations в rangeify/patterns.rs.
7. Законы де Моргана (поздние)
!a & !b --> !(a | b)
!a | !b --> !(a & b)
Эти паттерны появляются в двух местах пайплайна:
-
Ранние (Стадия 4–5):
boolean_dsl_patterns()вschedule/src/symbolic/patterns.rs, часть полного матчераsymbolic(). Перехватывают возможности де Моргана в исходной структуре выражений. -
Поздние (Стадия 18–19):
symbolic_simple()включает булевы паттерны и выполняется совместно с паттернами снижения стоимости вPM_FINAL. Перехватывает новые возможности де Моргана, созданные паттернами отрицания сравнений — например, после перезаписи!(x < 3)и!(x < 7)в2 < xи6 < x, любое AND/OR, комбинирующее их, может получить новые возможности устранения NOT.
Morok: boolean_dsl_patterns() в schedule/src/symbolic/patterns.rs.
8. Декомпозиция ERF
erf(x) заменяется полиномиальным приближением (Abramowitz & Stegun 7.1.26):
erf(x) = sign(x) * (1 - t * P(t) * exp(-x^2))
where t = 1 / (1 + 0.3275911 * |x|)
P(t) = Horner(t, [1.061405429, -1.453152027, 1.421413741, -0.284496736, 0.254829592])
Зачем: @llvm.erf — это intrinsic с вызовом библиотеки (требует линковки libm), а не нативная аппаратная инструкция. JIT-backend LLVM не линкует libm, поэтому erf должен быть декомпозирован перед кодогенерацией. Tinygrad декомпозирует erf на уровне тензоров (elementwise.py), поэтому он никогда не достигает рендера; Morok сохраняет Erf как UOp до этого позднего прохода.
Максимальная погрешность: ~1.5e-7 (достаточно для float32 ML-нагрузок).
Morok: pm_erf_decomposition в rangeify/patterns.rs.
Композиция паттернов: когда выполняется каждый паттерн
Все паттерны снижения стоимости скомпонованы в единый матчер PM_FINAL, выполняемый как перезапись графа до фиксированной точки:
PM_FINAL = symbolic_simple() + get_late_rewrite_patterns() + pm_render()
Где get_late_rewrite_patterns() объединяет:
Stage 18-19 (PM_FINAL fixed-point rewrite):
symbolic_simple() -- algebraic cleanup (identities, constant folding)
+ pm_fma_decomposition -- a*b+c -> MULACC(a,b,c)
+ pm_erf_decomposition -- erf(x) -> polynomial approx
+ pm_mod_to_and -- x % 2^n -> x & (2^n-1)
+ pm_mul_to_shl -- x * 2^n -> x << n
+ pm_div_to_shr -- x // 2^n -> x >> n
+ pm_fdiv_to_mul -- x / c -> x * (1/c)
+ pm_neg_from_mul -- x * -1 -> NEG(x)
+ pm_comparison_negations -- !(x<c) -> (c-1)<x, etc.
+ fast_division_patterns -- x // d -> (x * M) >> S
+ pm_render() -- CONST vectorization, CAT->VECTORIZE
Поскольку перезапись выполняется до фиксированной точки, паттерны могут подпитывать друг друга. Например:
pm_mul_to_shlпреобразуетx * 4вx << 2- На следующей итерации
pm_fma_decompositionсливает(x << 2) + cвMULACC(x, 4, c) symbolic_simple()убирает любые тождества, созданные преобразованиями
После завершения прохода фиксированной точки выполняется merge_sibling_ends для слияния новых узлов-сиблингов END, которые перезапись могла создать.
Перекрёстная ссылка: Обзор пайплайна кодогенерации — полный список стадий.