Оптимизация Range и Reduce

Циклы — основная цель оптимизаций в тензорных компиляторах. Наивное поэлементное сложение двух тензоров [1024, 1024] порождает один цикл на 1М элементов. После оптимизации оно превращается в 1024 параллельных потока, каждый из которых обрабатывает 1024 элемента с векторизованными загрузками/записями. Оптимизация Range — это путь к такому результату.

Эти паттерны находятся в schedule/src/rangeify/ и выполняются на Стадиях 1–5 пайплайна кодогенерации.

Исходники Tinygrad: tinygrad/codegen/simplify.py.

Разбиение Range

Что: Декомпозиция одного диапазона на внешнюю и внутреннюю компоненты через divmod.

Когда: Переменная диапазона используется с модулем: RANGE(end) % c, где end % c == 0.

Before:  RANGE(end=12) % 4     One loop, modulo in body (slow)
              |
         [split: end/c outer, c inner]
              |
After:   RANGE(end=3) * 4 + RANGE(end=4)
           ^outer              ^inner
           Parallel            Sequential / Vectorize

Зачем: После разбиения внутренний диапазон может быть векторизован (UPCAST до ширины SIMD), а внешний — параллелизован (блоки GPU, потоки CPU). Без разбиения модуль препятствует обеим оптимизациям.

Механизм: Матчер паттернов pm_split_ranges собирает диапазоны с использованием модуля, но НЕ выполняет преобразование немедленно. Он ждёт узла SINK, затем выполняет все подстановки разом (избегая рассогласованных частичных перезаписей). Новым диапазонам назначаются свежие axis_id.

Защита: Срабатывает только при end % c == 0 (точная делимость). Неделимые случаи остаются без изменений.

Tinygrad: simplify.py:60-64. Morok: pm_split_ranges() в rangeify/transforms.rs.

Слияние Range

Что: Слияние двух смежных диапазонов в один, снижающее накладные расходы цикла.

Before:  RANGE(0..4), RANGE(0..8)    Two loops, 12 iterations overhead
              |
         [merge: 4 * 8 = 32]
              |
After:   RANGE(0..32)                 One loop, indices via divmod

Зачем: Накладные расходы цикла (предсказание ветвлений, инкремент счётчика) пропорциональны числу итераций. Слияние уменьшает количество циклов ценой операций divmod для восстановления исходных индексов.

Критерий решения: Слияние принимается, только если общее число операций divmod не увеличивается. Компилятор подсчитывает операции divmod до и после — если слияние вводит больше делений, чем устраняет накладных расходов цикла, оно отклоняется.

Ограничения:

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

Tinygrad: simplify.py:39-41 (simplify_merge_adjacent). Morok: pm_simplify_ranges().

Выравнивание Range

Что: Выравнивание вложенных цепочек END/REDUCE/STORE в плоские списки диапазонов.

Before:  END(END(END(comp, [r0]), [r1]), [r2])
After:   END(comp, [r0, r1, r2])

Зачем: Вложенные цепочки END возникают из последовательных преобразований. Выравнивание нормализует структуру, чтобы другие паттерны (слияние, разбиение) могли работать с чистым списком диапазонов.

Tinygrad: simplify.py:14-17. Morok: pm_flatten_range().

Свёртка загрузок

Что: Полное устранение цикла REDUCE, когда вычисление может быть выражено замкнутой арифметической формулой.

Before:  sum(1 for k in 0..64 if k >= length)    // Loop: 64 iterations
After:   clamp(64 - length, 0, 64)                // Arithmetic: 3 ops

Как работает:

Идентифицировать подвыражения, не зависящие от диапазона REDUCE
Создать DEFINE_VAR для этих подвыражений (трактовать как инварианты цикла)
Подставить диапазон через DEFINE_VAR и запустить символьное упрощение
Если упрощённое выражение не содержит оставшихся диапазонов, REDUCE устраняется

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

Tinygrad: simplify.py:145-149. Morok: pm_load_collapse().

Свёртка Reduce

Аналитическое устранение редукций ADD. Более изощрённое, чем свёртка загрузок — применяет алгебраические преобразования внутри тела редукции.

Паттерны границ

Обрабатывают условные редукции, где сравнение ограничивает, какие итерации вносят вклад:

Паттерн	До	После
Нижняя граница	`sum(r < cut ? 0 : val, r=0..N)`	`max(0, N - cut) * val`
Верхняя граница	`sum(r < cut ? val : 0, r=0..N)`	`max(0, min(N, cut)) * val`
Двусторонний	`sum(r >= lo & r < hi ? val : 0, r=0..N)`	`max(0, min(N,hi) - max(0,lo)) * val`
NE-условный (gather)	`sum(idx != r ? 0 : expr, r=0..N)`	`in_bounds ? expr[r:=idx] : 0`

Паттерн NE-условного особенно важен для операций gather — он распознаёт, что суммирование по всем индексам, где idx == r, эквивалентно одному индексированному доступу.

Преобразования подъёма

Вынос сравнений за пределы области reduce для экспонирования паттернов границ:

Преобразование	До	После
Подъём Lt	`(x + y) < c`	`x < (c - y)`
Подъём Ge	`(x + y) >= c`	`x >= (c - y)`
Подъём EQ	`(x + y) == c`	`x == (c - y)`

Дистрибутивный закон

sum(x + y) → sum(x) + sum(y) — разделение reduce по сложению. Это позволяет каждой половине быть независимо свёрнутой паттернами границ.

MUL-casted-bool

x * bool.cast() → WHERE(bool, x, 0) — преобразует умножение на приведённый bool в WHERE, который затем может быть проанализирован паттернами границ.

Tinygrad: simplify.py:82-142. Morok: pm_reduce_simplify() + reduce_collapse_inner_patterns().

Удаление буферов (Partial Contiguous)

Что: Решение, материализовать ли промежуточный результат в буфер или встроить вычисление. В кодовой базе часто называется «pcontig» (сокращение от partial contiguous — оптимизация, которая встраивает узлы BUFFERIZE, подставляя переменные диапазонов).

Когда проход rangeify создаёт узел BUFFERIZE (помечая «здесь нужен буфер»), проход удаления буферов оценивает, стоит ли реально выделять память. BUFFERIZE — это промежуточное представление Morok между «здесь нужен буфер» и финальным STORE+BUFFER+AFTER — оно позволяет этому проходу решить, действительно ли материализация нужна. Если вычисление достаточно дешёвое, подставляются переменные диапазонов и выражение встраивается напрямую.

Дерево решений

Is this an always-run op (CONTIGUOUS, COPY)?
  └─ YES → Keep buffer (always materialized)

Does inlining exceed the buffer limit?
  └─ YES → Keep buffer

Is there a reduce in scope?
  ├─ NO → Inline (cheap: just substitute ranges)
  └─ YES:
      Is pcontig level <= 2?
        ├─ YES → Keep buffer (reduce recomputation too expensive)
        └─ NO → Check input/output ratio
            ├─ Ratio low (output small relative to input) → Keep buffer
            └─ Ratio high (output >> input) → Partial inline

:::caution Унарные операции в контексте Reduce Унарные операции (вроде отрицания) НЕ встраиваются внутри области reduce, даже если они дешёвые. Причина: если argmax(-x) встроит отрицание, оно пересчитывается на каждой итерации редукции — N дополнительных отрицаний вместо одного чтения из буфера. :::

Связанные паттерны

Паттерн	Что делает
Свёртка буферов	`BUFFERIZE(CONST) → CONST` — буфер константы — это просто константа
Свёртка индексов	`INDEX(CONST) → CONST` — индексация в константу — это константа
Тождественная свёртка	`INDEX(BUFFERIZE(compute, ranges), ranges) → compute` — одинаковые диапазоны сокращаются
Выравнивание вложенности	`BUFFERIZE(BUFFERIZE(...))` — выравнивание вложенной буферизации

Morok: buffer_removal_with_pcontig() в rangeify/patterns.rs.

Удаление мёртвых осей

Что: Удаление неиспользуемых размерностей из операций BUFFERIZE.

Размерность считается «мёртвой», когда:

Имеет размер 1 (ничего не вносит)
Появляется как константа в индексе (не как переменная)
Вычислительное выражение не ссылается на неё

Мёртвые оси удаляются из BUFFERIZE, затем форма восстанавливается через RESHAPE (вставка размерностей размера 1) и EXPAND (broadcast до исходного размера). Это снижает размерность выделения буфера.

:::caution Скалярный случай Даже когда ВСЕ диапазоны мертвы (скалярный выход), BUFFERIZE должен быть сохранён с пустыми диапазонами — его полное удаление вызывает NoKernelsFound, поскольку ни один STORE не будет создан при разбиении ядер. :::

Morok: dead_axis_removal() в rangeify/patterns.rs.

Reduce Unparented

Что: Удаление диапазонов из REDUCE, на которые не ссылается тело редукции.

Операция Reduce	Неиспользуемый диапазон размера N	Преобразование
ADD	Диапазон не используется в теле	Умножить результат на N
MUL	Диапазон не используется в теле	Возвести результат в степень N
MAX / MIN	Диапазон не используется в теле	Просто удалить диапазон

Пример: sum(x, r=0..N), где x не зависит от r → x * N. Сумма константы по N итерациям равна N-кратному значению этой константы.

Tinygrad: simplify.py:82-86. Morok: pm_reduce_simplify().

Разбиение ReduceOp

Что: Разбиение крупных редукций на две стадии для лучшего параллелизма.

Когда: Отношение входа к выходу превышает 32768.

Before:  REDUCE(data, axes=[0])       // shape [65536] → scalar
After:   REDUCE(                       // shape [256] → scalar (second stage)
           CONTIGUOUS(
             REDUCE(                   // shape [65536] → [256] (first stage)
               RESHAPE(data, [256, 256]),
               axes=[1]
             )
           ),
           axes=[0]
         )

Зачем: Одна огромная редукция не может быть параллелизована. Разбиение на две стадии позволяет первой стадии выполняться параллельно (256 потоков, каждый редуцирует 256 элементов), затем вторая стадия редуцирует 256 частичных результатов.

Защита: Применяется, только когда размерность редукции может быть факторизована и отношение входа к выходу превышает порог. Нефакторизуемые размерности пропускаются.

Morok: split_reduceop() в rangeify/kernel.rs.

Разбиение Range​

Слияние Range​

Выравнивание Range​

Свёртка загрузок​

Свёртка Reduce​

Паттерны границ​

Преобразования подъёма​

Дистрибутивный закон​

MUL-casted-bool​

Удаление буферов (Partial Contiguous)​

Дерево решений​

Связанные паттерны​

Удаление мёртвых осей​

Reduce Unparented​

Разбиение ReduceOp​