Фаза 1: Rangeify
Цель: Преобразовать высокоуровневые операции перемещения (movement ops) в явные циклические структуры и оптимизировать RANGE.
Стадия 1: Ранние Movement Ops
Стадия кратко
Цель: Вычистить movement ops перед назначением RANGE Ключевые паттерны: Movement на INDEX, movement через обёртки, упрощение вложенных INDEX Эффект: Предотвращает пропущенные оптимизации далее по пайплайну
Что делает: Эта стадия вычищает movement ops, перемещая манипуляции с индексами туда, где они реально используются. Считайте, что это наведение порядка на столе перед сортировкой документов — инструкции сдвигаются ближе к месту использования данных.
Зачем это нужно: Movement ops (RESHAPE, PERMUTE и т.д.) — удобные абстракции, но железо требует конкретных вычислений индексов. Вычищая их рано, мы обеспечиваем корректное срабатывание паттернов на последующих стадиях.
Паттерн: pm_mops + pm_syntactic_sugar (bottom-up)
| Паттерн | Трансформация | Визуально | Расположение |
|---|---|---|---|
| Movement на INDEX | Применить movement к выражениям индексов | INDEX(PERMUTE(arr), [i, j]) → INDEX(arr, [j, i]) | movement_op_patterns() |
| Movement через AFTER | Протянуть RESHAPE через timing-обёртку (Tinygrad-специфично) | AFTER(RESHAPE(x, arg), [dep1, dep2]) → RESHAPE(AFTER(x, [dep2]), arg) | Только Tinygrad |
| Movement через END | Развернуть movement из END-обёртки (Tinygrad-специфично) | END(RESHAPE(x), ranges) → END(x, ranges) | Только Tinygrad |
| Упрощение вложенных INDEX | Удалить избыточные вложенные INDEX (Morok) | INDEX(INDEX(ptr, [i]), [i]) → INDEX(ptr, [i]) | movement_op_patterns() |
| Конкатенация вложенных INDEX | Свернуть вложенные INDEX для PtrDType | INDEX(INDEX(ptr, i), j) → INDEX(ptr, i, j) | pm_syntactic_sugar |
Почему bottom-up? Дочерние узлы должны быть вычищены до обработки родительских. Movement ops глубоко вложены; очистка снизу предотвращает пропущенные паттерны.
Примечание: Tinygrad и Morok подходят к этому по-разному. Tinygrad перемещает movement ops через обёртки (AFTER, END), потому что повторно применяет их во время буферизации. Morok полностью удаляет movement ops, трансформируя индексы во время буферизации, поэтому паттерны для AFTER/END не нужны.
Morok: movement_op_patterns() в rangeify/patterns.rs
Стадия 2: Свёртка нагрузок (Load Collapse)
Стадия кратко
Цель: Устранить REDUCE-операции, обнаруживая вычисления, независимые от RANGE Ключевые паттерны: Ограниченная сумма, свёртка gated load, общее устранение reduce Эффект: Замена итераций цикла на арифметические операции
Что делает: Устраняет REDUCE-операции, распознавая случаи, когда вычисление можно выполнить без итерации. Использует обнаружение вычислений, независимых от RANGE, и символическое упрощение.
Зачем это нужно: Замена итераций на арифметику убирает оверхед цикла. Вместо того чтобы прогонять цикл 1000 раз, вычислить ответ напрямую.
Паттерн: pm_load_collapse
// Before: Sum with bounds check
sum(1 for k in 0..64 if k >= length)
// After: Compute count directly (NO LOOP!)
count = clamp(64 - length, 0, 64)
Механизм работает так:
- Находит подвыражения, не зависящие от RANGE REDUCE
- Создаёт DEFINE_VAR для них (считает loop-инвариантными)
- Подставляет RANGE на DEFINE_VAR и запускает символическое упрощение
- Если упрощённое выражение больше не содержит RANGE, REDUCE устраняется
Примечание: Перемещение WHERE через INDEX (pm_move_where_on_load) — отдельная оптимизация, которая ставит условия перед загрузками, чтобы пропустить обращения к памяти, но не устраняет REDUCE.
Morok: pm_load_collapse() в rangeify/patterns.rs
Стадия 3: Разделение RANGE
Стадия кратко
Цель: Улучшить оптимизацию через декомпозицию divmod Ключевые паттерны: Разделение RANGE по модулю, свёртка RANGE Эффект: Внутренние RANGE можно векторизовать, внешние — распараллелить
Что делает: Обрабатывает паттерны с модулем, разделяя RANGE на внешнюю и внутреннюю компоненты.
Зачем это нужно: Разделение RANGE — как распределение большой задачи между членами команды. Если у вас 12 элементов и каждый делает 4, получается 3 человека x 4 элемента. Внутренние циклы (4 элемента одного человека) могут быть быстрыми; внешние (3 человека) могут работать параллельно.
Паттерн: pm_split_ranges + pm_flatten_range
Before: RANGE(end=12) % 4 // One loop with modulo (slow)
↓ [Split into outer × inner]
After: RANGE(end=3) * 4 + RANGE(end=4)
↑outer ↑inner
Parallel Sequential
Это позволяет:
- Внутренние RANGE могут векторизоваться (SIMD)
- Внешние RANGE могут параллелиться (GPU-блоки / CPU-потоки)
pm_flatten_range объединяет вложенные RANGE на REDUCE/STORE/END, когда это выгодно.
Контекст: Требует контекста-словаря (ctx={}) для отслеживания подстановок на уровне SINK.
Примечание: Разделение применяется только когда end % mod == 0 (проверка делимости).
Morok: pm_split_ranges() + pm_flatten_range() в rangeify/transforms.rs
Стадия 4: Начальная символика
Стадия кратко
Цель: Упростить выражения, используя правила алгебры Ключевые паттерны: Свёртка констант, удаление тождеств, рекомбинация div-mod Эффект: Устранение дорогих операций, уменьшение объёма кода
Что делает: Применяет 100+ правил свёртки констант и алгебраических упрощений.
Зачем это нужно: Компьютеры быстры в простой арифметике. Деление и остаток — медленные операции. Эта стадия использует правила алгебры, чтобы убрать медленные операции где возможно.
Паттерн: symbolic() + pm_flatten_range
Примечание: symbolic() — подмножество sym, используемое на стадии 8. Включает алгебраические правила, но без паттернов поздних стадий.
Свёртка констант:
ADD(CONST(2), CONST(3)) → CONST(5)
MUL(x, CONST(1)) → x
ADD(x, CONST(0)) → x
Рекомбинация div-mod:
(x / c) * c + (x % c) → x
Зачем? Вычисляет то же значение, что и x, но за 3 операции вместо 1. Этот паттерн находит и устраняет избыточность (часто встречается в вычислениях страйдов).
Булева алгебра:
x AND x → x
x OR FALSE → x
NOT(NOT(x)) → x
Дополнительные категории:
- Удаление тождеств (самосвёртка, избыточные операции)
- Упрощение сравнений
- Оптимизация Cast
- Проталкивание GEP (перенос вычислений адресов через ALU)
- Свёртка WHERE (объединение WHERE с одинаковыми условиями)
- Цепочка reduce mul (вынос умножений за reduce)
Morok: symbolic() в symbolic/patterns.rs
Стадия 5: Упрощение RANGE
Стадия кратко
Цель: Объединить соседние RANGE для снижения оверхеда циклов Ключевые паттерны: Слияние RANGE с анализом стоимости Эффект: Меньше циклов = меньше оверхеда
Что делает: Объединяет соседние RANGE, когда это выгодно.
Зачем это нужно: Объединение RANGE — как объединение нескольких мелких поездок в одну большую. Вместо 4 походов в магазин за 4 вещами — один поход за всем. Экономит затраты на запуск и завершение.
Паттерн: pm_flatten_range() + pm_simplify_ranges()
// Before: two separate ranges
RANGE(0..4), RANGE(0..8)
// After: merged (if compatible)
RANGE(0..32)
Критерии слияния:
- Типы осей должны быть совместимы (оба выходные, оба reduce и т.д.)
- Скоуп REDUCE должен остаться консистентным
- На основе стоимости: принимать, только если количество divmod-операций не увеличивается
Компилятор объединяет, только если это экономит операции. Слияние может потребовать деления/остатка для пересчёта индексов. Если это обходится дороже, чем экономит, слияние пропускается.
Morok: simplify_merge_adjacent() в rangeify/transforms.rs
Стадия 6: Разделение STORE
Стадия кратко
Цель: Разделить граф на границах STORE в отдельные ядра Ключевая функция:
split_all_stores()+split_store()Эффект: Даёт возможность оптимизировать каждое ядро отдельно
Что делает: Разделяет UOp-граф на границах STORE, создавая отдельные ядра для каждого выхода.
Зачем это нужно: После буферизации граф может содержать несколько STORE-операций. Каждый STORE становится отдельным ядром со своим набором буферов, RANGE и зависимостей.
Функция: run_kernel_split_pipeline() в schedule/src/rangeify/kernel.rs
Перед разделением выполняется пре-проход pm_flatten_range на всём графе целиком (bottom-up). Он объединяет вложенные RANGE для всех ядер за один обход, избегая избыточной работы на перекрывающихся подграфах. Этот пре-проход — ключевая оптимизация скорости компиляции: без него каждый split_store для каждого ядра повторно обходил бы общие подграфы.
После пре-прохода split_all_stores выполняет разделение на границах STORE, а fix_assign отвечает за нумерацию буферов (через счётчик LocalAddBufferContext.dg) и отслеживание зависимостей.
Стадия 7: Применение оптимизаций
Стадия кратко
Цель: Найти оптимальную комбинацию векторизации, развёртки, использования памяти Ключевой алгоритм: Beam search или эвристики Эффект: Может значительно улучшить производительность
Что делает: Оптимизационный поиск — beam search или эвристика — исследует разные комбинации оптимизационных действий.
Зачем это нужно: Компилятор пробует разные комбинации оптимизаций (векторизовать тут? развернуть там?) и выбирает самую быструю. Правильная комбинация может ускорить код в 10 раз.
Функция: apply_opts(sink, renderer)
Оптимизационные действия:
| Действие | Эффект | Целевое железо |
|---|---|---|
| TC | Задействовать тензорные ядра | NVIDIA GPU |
| UPCAST | Векторизовать измерение | Все (SIMD) |
| LOCAL | Использовать локальную/shared-память | GPU (LDS) / CPU (L1) |
| UNROLL | Развернуть измерение цикла | Все (избежать оверхеда цикла) |
| GROUP | Группировка операций для кэша | Все |
| GROUPTOP | Группировка для reduce-операций | GPU-тензорные ядра |
| THREAD | Параллелизм на потоках | CPU |
| NOLOCALS | Отключить использование локальной памяти | Все (ограничение, запрещает дальнейшие LOCAL-действия) |
| SWAP | Поменять местами назначения RANGE | Все (попробовать другой тайлинг) |
| PADTO | Паддинг для выравнивания | Все (выравнивание памяти) |
Как устроен поиск оптимизаций:
Компилятор ищет лучшую комбинацию:
- Эвристический режим (BEAM=0): Быстрые захардкоженные паттерны оптимизации, без компиляции
- Beam search (BEAM>=1): Компилирует и запускает кандидатов, чтобы замерить реальную производительность
Optimization Search:
├── Heuristic mode (BEAM=0): Hand-coded optimizations
└── Beam search (BEAM≥1):
├── Generate all possible actions (~162 base actions, workload-dependent)
├── Apply to all top-K candidates in parallel
├── Filter based on constraints
├── Compile and run each candidate → Measure actual time
└── Pick fastest
Примечание: NOLOCALS — ограничение, которое устанавливает dont_use_locals = True, запрещая дальнейшие LOCAL-действия и влияя на решения по использованию shared-памяти.
Morok: optimizer/mod.rs, optimizer/opts.rs