Путь UOp: 22-стадийный пайплайн кодогенерации

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

Пайплайн — проверенный дизайн для тензорной компиляции. Понять его — значит понять, как тензорные выражения превращаются в машинный код.

---

Как читать эту главу

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

Ключевые концепции

UOp (Micro-Operation)

• Думайте об этом как об узле в блок-схеме, представляющем одно вычисление
• Пример: ADD(a, b) означает «сложить a и b»

Паттерн

• Правило «найти и заменить» для структур кода (не текста)
• Пример: «Если видишь ADD(x, 0), замени на x»
• Паттерны срабатывают повторно, пока не перестанут находить совпадения (fixpoint)

Range

• Итерация цикла: RANGE(0..10) означает «for i from 0 to 10»

AxisType

• Какого типа этот цикл?
  • Global: параллельный по GPU-блокам / CPU-потокам
  • Local: параллельный внутри воркгруппы
  • Reduce: аккумулятор (сумма, макс. и т.д.)
  • Loop: последовательная итерация

Стадия

• Один проход трансформации по коду
• Паттерны срабатывают до fixpoint, затем начинается следующая стадия

Стратегия чтения

1. Первый проход: прочитайте только секции «Что делает» и «Зачем это нужно»
2. Второй проход: посмотрите диаграммы и примеры
3. Третий проход (если хотите деталей): прочитайте описания паттернов

Вопросы, которые стоит задавать

По каждой стадии спросите:

• Что эта стадия делает? (Высокоуровневая цель)
• Зачем она нужна? (Мотивация)
• Что сломается без неё? (Последствия)

---

Обзор

22 стадии разбиты на четыре фазы:

Tensor Expression
       │
       ▼
┌─────────────────────────────────────┐
│ RANGEIFY (Stages 1-7)               │
│ Movement ops → Explicit loops       │
│                                     │
│ [Make iteration explicit,           │
│  optimize ranges]                   │
└─────────────────────────────────────┘
       │
       ▼
┌─────────────────────────────────────┐
│ EXPANDER (Stages 8-10)              │
│ UNROLL/UPCAST → Explicit vectors    │
│                                     │
│ [Expand optimization primitives]    │
└─────────────────────────────────────┘
       │
       ▼
┌─────────────────────────────────────┐
│ DEVECTORIZER (Stages 11-15)         │
│ Vector ops → Scalar code            │
│                                     │
│ [Lower to hardware-specific ops]    │
└─────────────────────────────────────┘
       │
       ▼
┌─────────────────────────────────────┐
│ LINEARIZER (Stages 16-22)           │
│ IR → Linear instruction sequence    │
│                                     │
│ [Serialize to executable code]      │
└─────────────────────────────────────┘
       │
       ▼
  Machine Code

Каждая стадия применяет перезаписи на основе паттернов. Паттерны срабатывают до fixpoint, затем начинается следующая стадия.

Дополнительные проходы

Несколько проходов выполняются между пронумерованными стадиями и не имеют собственного номера:

Проход	Между стадиями	Назначение
linearize_multi_index	До стадии 8	Свёртка многомерных индексов в линейные смещения
pm_bool_devectorize	14–15	Обработка паттернов булевых векторов
pm_reduce_devectorize	14–15	Обработка векторных редукций (K-vec, bool, горизонтальные)
merge_sibling_ends	14–15	Слияние соседних END-операций
pm_float_decomp	Post-opt	Декомпозиция операций с плавающей точкой
bool_storage_patterns	Post-opt	Конвертация bool ↔ uint8 для операций с памятью