Движок паттернов

Откройте любой продакшн ML-компилятор — и найдёте десятки оптимизационных проходов: свёртка констант, удаление мёртвого кода, фьюзинг операторов, тайлинг циклов, векторизация, оптимизация раскладки памяти. У каждого прохода свои структуры данных, своя логика обхода, свои баги.

Morok использует другой подход: один механизм для всего.

Traditional Compiler:              Morok:
┌─────────────────────────┐       ┌─────────────────────────┐
│  Constant Folding       │       │                         │
│  Dead Code Elimination  │       │   patterns! {           │
│  Loop Unrolling         │       │       Add[x, @zero] ~> x│
│  Operator Fusion        │       │       Mul[x, @zero] ~> 0│
│  Vectorization          │       │       // ...more        │
│  Memory Planning        │       │   }                     │
│  ...20 more passes      │       │                         │
└─────────────────────────┘       │   graph_rewrite(...)    │
     Custom logic each            └─────────────────────────┘
                                       One mechanism

Каждая оптимизация в Morok выражается как паттерн: «когда видишь эту структуру, замени её вот этой». Одна и та же функция graph_rewrite() применяет алгебраическое упрощение, индексную арифметику, снижение стоимости и оптимизацию диапазонов.

DSL `patterns!`

Morok предоставляет предметно-ориентированный язык для написания оптимизационных паттернов:

patterns! {
    // Identity folding: x + 0 → x
    Add[x, @zero] ~> |x| x.clone(),

    // Constant folding: 3 + 4 → 7
    Add(a @const(a_val), b @const(b_val))
        => |a, a_val, b_val| eval_add(a_val, b_val).map(|r| UOp::const_(a.dtype(), r)),

    // Self-folding: x / x → 1
    Idiv(x, x) ~> |x| UOp::one(x.dtype()),

    // Dead code elimination: if(true) { t } else { f } → t
    Where(@true, t, _f) ~> |t| t.clone(),
}

Макрос компилирует эти паттерны в эффективный Rust-код:

Синтаксис	Значение	Пример
`(x, y)`	Упорядочено. Сопоставляется в точном порядке.	`Sub(x, @zero) ~> x`
`[x, y]`	Коммутативно. Пробуются оба порядка.	`Add[x, @zero] ~> x`
`@zero`	Нулевая константа. Совпадает с 0 или 0.0.	`Mul[_, z @ @zero] ~> z`
`@one`	Единичная константа. Совпадает с 1 или 1.0.	`Mul[x, @one] ~> x`
`@const(val)`	Извлечение константы. Связывает значение.	`Add(@const(a), @const(b))`
`x, x`	Один и тот же операнд. Автогенерируется проверка ptr_eq.	`Idiv(x, x) ~> UOp::one(...)`
`~>`	Безусловный. Всегда успешен, возвращает `Arc<UOp>`.	`Add[x, @zero] ~> x`
`=>`	Условный. Может не сработать, возвращает `Option<Arc<UOp>>`.	`=> eval(...).map(...)`
`for op in binary [...]`	Шаблон. Генерация паттернов для нескольких операций.	См. ниже
`@context Type`	С состоянием. Доступ к мутабельному контексту в паттернах.	См. ниже

Раскрытие шаблонов

Вместо написания одного и того же паттерна для каждой бинарной операции используйте for-цикл:

patterns! {
    for op in binary [Add, Mul, Sub, Idiv, Fdiv, Max] {
        op(a @const(a_val), b @const(b_val))
            => |a, a_val, b_val| eval_binary(op, a_val, b_val)
                .map(|r| UOp::const_(a.dtype(), r))
    }
}

Это раскрывается в шесть отдельных паттернов во время компиляции — по одному для каждой операции.

Паттерны с состоянием

Некоторым оптимизациям нужен контекст (например, в каком ядре мы находимся, какие диапазоны активны):

patterns! {
    @context KernelContext;

    ReduceAxis { src } => |reduce, src, ctx| {
        ctx.record_reduction(reduce);
        transform_reduce(reduce, src, ctx)
    }
}

Подъём контекста

При комбинировании матчеров с разными типами контекста используется .with_context():

let mega_pass = symbolic().with_context::<PcontigConfig>()
    + reduction_simplify_patterns().with_context()
    + buffer_removal_with_pcontig();

Как работает сопоставление паттернов

Макрос patterns! генерирует SimplifiedPatternMatcher, который диспатчит паттерны в соответствующий бакет за O(1) через HashMap-поиск, а затем последовательно пробует каждый паттерн в бакете.

Индекс OpKey

У каждого UOp есть тип операции (Add, Mul, Load и т.д.). Макрос генерирует enum OpKey, отображающий операции в хэшируемые ключи:

pub struct SimplifiedPatternMatcher<C = ()> {
    indexed: HashMap<OpKey, Vec<PatternClosure<C>>>,  // O(1) lookup
    wildcards: Vec<PatternClosure<C>>,                 // patterns matching any op
}

При сопоставлении UOp:

Извлекаем OpKey из операции UOp
Ищем в HashMap — O(1)
Пробуем каждое замыкание, пока одно не сработает
Откатываемся на wildcards, если ни один индексированный паттерн не совпал

Это в 5–10 раз быстрее линейного перебора всех паттернов.

Обработка коммутативности

Для паттернов вроде Add[x, @zero] макрос генерирует код, пробующий оба порядка:

// Try (x, @zero)
if let Some(result) = try_match_ordered(&children[0], &children[1]) {
    return result;
}
// Try (@zero, x)
if let Some(result) = try_match_ordered(&children[1], &children[0]) {
    return result;
}

Обнаружение дубликатов

Когда вы пишете Idiv(x, x), паттерн срабатывает только если оба операнда — один и тот же UOp (равенство указателей через Arc::ptr_eq, а не структурное равенство). Это использует hash consing — идентичные подвыражения разделяют один указатель.

Движок перезаписи

Одного сопоставления паттернов недостаточно. Рассмотрим выражение:

WHERE(Lt(3, 5), t, f)

Чтобы его упростить, нужны два шага:

Lt(3, 5) → true (свёртка констант)
WHERE(true, t, f) → t (удаление мёртвого кода)

Но паттерн WHERE не сработает, пока его дочерний узел не упрощён. Движок перезаписи решает это двухстадийным алгоритмом.

Стадия 0: Применение паттернов

Паттерны применяются к каждому узлу. Если ни один паттерн не совпал, подаётся сигнал сначала обработать дочерние узлы.

Стадия 1: Реконструкция источников

После перезаписи дочерних узлов перестраиваем узел с новыми потомками и снова пробуем паттерны:

Stage 0: WHERE(Lt(3, 5), t, f)     → Gate (no match, process children)
         └── Lt(3, 5)              → true (constant folding matches!)

Stage 1: WHERE(true, t, f)         → t (dead code elimination matches!)

Стадия реконструкции повторно применяет паттерны, что позволяет многошаговым оптимизациям сработать за один обход.

Стратегии перезаписи

Три функции перезаписи, соответствующие graph_rewrite в Tinygrad:

Стратегия	Паттерны видят	Когда использовать
`graph_rewrite(pm)` (по умолчанию)	ОПТИМИЗИРОВАННЫХ потомков	Алгебраическое упрощение, раскрытие
`graph_rewrite_bottom_up(bpm)`	ИСХОДНЫХ потомков	Сопоставление вложенных структур, удаление буферов
`graph_rewrite_with_bpm(pm, bpm)`	Оба (bpm: исходные, pm: оптимизированные)	Разбиение ядер (gate + трансформация в одном проходе)

Движок всегда обходит снизу вверх; различие в том, когда срабатывают паттерны: на Стадии 0 (до обработки потомков — видят оригиналы) или на Стадии 1 (после потомков — видят оптимизированные результаты). Матчеры комбинируются оператором +: matcher_a() + matcher_b() объединяет наборы паттернов в один.

Ограничения безопасности

Для предотвращения бесконечных циклов:

1000 итераций максимум на узел
500 000 итераций максимум в сумме
Panic с диагностикой при превышении лимитов

На практике корректные паттерны сходятся быстро.

Почему это важно

Отладка прямая. Паттерны — это читаемый код. Добавьте println! в любой паттерн, чтобы отследить, когда он срабатывает.

Расширяемость простая. Добавление своей оптимизации — пара строк. Не нужно разбираться во внутренностях компилятора, писать визиторы или модифицировать pass manager.

Корректность локальна. Каждый паттерн — маленькая теорема: «если появляется эта структура, замена на ту структуру сохраняет семантику». Каждый паттерн верифицируется независимо. Композиция корректных паттернов даёт корректные программы.

Производительность настраиваема. O(1) диспатч паттернов быстр по умолчанию. Комбинируйте с beam search для продакшн-нагрузок.

Глубинная идея

Сопоставление паттернов обменивает общность на компонуемость.

Универсальный оптимизационный проход может делать что угодно — и в этом проблема. Его трудно верифицировать, трудно расширять, трудно компоновать с другими проходами. Порядок важен. Взаимодействия тонки.

Паттерн ограничен: он сопоставляет конкретную структуру и порождает конкретную замену. Но ограничения дают компонуемость. Для корректно спроектированных наборов паттернов применение до фиксированной точки даёт детерминированные результаты. Новые паттерны добавляются с локальным влиянием, удаляются без каскадных сбоев — хотя на практике взаимодействия паттернов следует тестировать для обеспечения сходимости.

Каждый паттерн — теорема о семантической эквивалентности. Движок перезаписи — доказыватель теорем, находящий вывод от входа к оптимизированному выходу. Корректность следует из корректности отдельных шагов.

Это философия Unix, применённая к компиляторам: маленькие, сфокусированные инструменты, которые компонуются.

DSL patterns!​

Раскрытие шаблонов​

Паттерны с состоянием​

Подъём контекста​

Как работает сопоставление паттернов​

Индекс OpKey​

Обработка коммутативности​

Обнаружение дубликатов​

Движок перезаписи​

Стадия 0: Применение паттернов​

Стадия 1: Реконструкция источников​

Стратегии перезаписи​

Ограничения безопасности​

Почему это важно​

Глубинная идея​