Сквозной пример и справочник
Сквозной пример: прослеживаем все 22 стадии
Проследим путь c = a + b (где a, b — тензоры [100, 100]) через весь пайплайн.
Начальный граф тензоров
[ADD]
├── [BUFFER(a)] : Float32
└── [BUFFER(b)] : Float32
После стадии 1: Ранние Movement Ops
(Без изменений — в этом примере нет movement ops)
После стадии 2: Свёртка нагрузок
(Без изменений — в этом примере нет редукций)
После стадии 3: Разделение RANGE
(Без изменений — нет операций с модулем)
После стадии 4: Начальная символика
(Без изменений — упрощение не требуется)
После стадии 5: Упрощение RANGE
(Без изменений — соседних RANGE ещё нет)
После стадии 6: Разделение STORE
(Неприменимо — GPU-бэкенд)
После стадии 7: Применение оптимизаций
Применённые оптимизации:
- UPCAST измерения j на 4 (векторизация)
- LOCAL для входных буферов (если выгодно)
После стадии 8: Пост-оптимизационная символика
Без изменений — символика уже чистая.
После стадии 9: Expander
UPCAST → UNROLL → CONTRACT (упрощённо — реальный IR содержит CONTRACT-обёртку):
[VECTORIZE]
├── [ADD]
│ ├── [LOAD(a)]
│ │ └── [INDEX]
│ │ ├── [BUFFER(a)]
│ │ ├── [RANGE(i, Global, 0..100)]
│ │ └── [UNROLL(VCONST([0,1,2,3]))] // Converted from RANGE(j, UPCAST)
│ └── [LOAD(b)]
│ └── [INDEX]
│ ├── [BUFFER(b)]
│ ├── [RANGE(i)] // Same RANGE via hash consing
│ └── [UNROLL(VCONST([0,1,2,3]))] // Same UNROLL via hash consing
После стадии 10: Добавление локальных буферов
(Если была выбрана оптимизация LOCAL)
После стадии 11: Удаление Reduce
(Без изменений — нет редукций)
После стадии 12: Добавление GPU-измерений
[SPECIAL(gidx0)] : Index // replaces RANGE(i)
После стадии 13: Добавление LOAD
(Без изменений — загрузки уже на месте)
После стадии 14: Devectorize
Векторная структура после devectorize (показан эффект, а не точная структура UOp):
[VECTORIZE] : <4 x Float32>
├── [ADD(a[0], b[0])]
├── [ADD(a[1], b[1])]
├── [ADD(a[2], b[2])]
└── [ADD(a[3], b[3])]
После стадии 15: Понижение типа Index
[SPECIAL(gidx0)] : i32 // concrete type
После стадии 16: Пост-индексная символика
Изменений не требуется.
После стадии 17: Pre-Matcher
(Нет паттернов для стандартных бэкендов)
После стадии 18: Декомпозиции
Декомпозиции не требуются — все операции поддерживаются.
После стадии 19: Финальная перезапись
Изменений не требуется.
После стадии 20: Добавление управления потоком
Зависимости отслежены — проблем нет.
После стадии 21: Линеаризация
Линейная последовательность инструкций (упрощённо):
1. DEFINE_GLOBAL(0) // Output buffer c
2. DEFINE_GLOBAL(1) // Input buffer a
3. DEFINE_GLOBAL(2) // Input buffer b
4. RANGE(i, 0..100, Global) // gidx0
5. LOAD(a, i*4+0..i*4+3) // Vector load (vec4)
6. LOAD(b, i*4+0..i*4+3) // Vector load (vec4)
7. ADD(vec_a, vec_b) // Vector add (vec4)
8. STORE(c, i*4+0..i*4+3, result) // Vector store
9. END(RANGE(i))
Примечание: UPCAST был поглощён стадией 9 (expander), поэтому отдельного цикла RANGE(j) нет. Векторизация неявно заложена в vec4-операциях.
После стадии 22: Чистка IF/ENDIF
Изменений не требуется — gated stores отсутствуют.
Результат: Готово к кодогенерации! LLVM/CUDA/другой бэкенд скомпилирует это в реальный машинный код.
Стратегия применения паттернов
Каждая стадия использует одну из двух стратегий перезаписи:
Top-down (по умолчанию): Обработка родителей до потомков. Используется, когда трансформации создают новые подтермы для сопоставления.
Bottom-up: Обработка потомков до родителей. Используется, когда состояние потомков влияет на сопоставление родителей (стадии 1, 20).
Обе стратегии итерируют до fixpoint — паттерны срабатывают, пока не перестанут находить совпадения.
Отладка пайплайна
Когда ядро выдаёт неверные результаты, баг живёт на одной из этих 22 стадий. Используйте переменные окружения для извлечения IR на каждой стадии:
# See IR after each transformation
MOROK_DEBUG=ir cargo test failing_test
Краткий справочник
| Симптом | Вероятные стадии | Что проверять |
|---|---|---|
| Неверные значения на выходе | 4, 9, 11, 18 | Символическое упрощение, расширение, девекторизация |
| Медленная производительность | 7, 9, 14, 21 | Оптимизация, расширение, девекторизация, линеаризация |
| Краши/паники | 11, 12 | Reduce, GPU-измерения |
| Неверное число итераций | 3, 5, 12 | Разделение RANGE, упрощение RANGE, GPU-измерения |
| Отсутствие векторизации | 9, 14 | Expander, devectorizer |
Частые проблемы
- Стадии 3-4: Разделение RANGE / символика могут потерять ограничения
- Стадия 9: Порядок расширения влияет на корректность векторизации
- Стадия 11: Инициализация аккумулятора должна совпадать с нейтральным элементом редукции
- Стадия 14: Несоответствие аппаратной ширине — проверьте vector fold length
- Стадия 18: Отсутствующая декомпозиция — проверьте список supported_ops бэкенда
- Стадия 21: Баги приоритетов вызывают гонки данных — верифицируйте зависимости
Итого
22-стадийный пайплайн трансформирует тензорные выражения в машинный код через систематическое уточнение:
- Стадии 1-7: Сделать итерацию явной, оптимизировать RANGE
- Стадии 8-10: Расширить оптимизационные примитивы
- Стадии 11-15: Опустить до аппаратно-специфичных операций
- Стадии 16-22: Сериализовать в исполняемые инструкции
У каждой стадии — одна ответственность. Каждая строится на предыдущей. Результат: высокоуровневый тензорный код работает с близкой к оптимальной скоростью на разном железе.
Tinygrad vs Morok: архитектурные различия
Эта глава описывает «идеальный» 22-стадийный пайплайн, основанный на реализации Tinygrad. Morok сейчас близко следует этому дизайну с минимальными отличиями.
Оставшиеся архитектурные различия
| Стадия | Tinygrad | Morok | Примечания |
|---|---|---|---|
| 1: Ранние Movement Ops | Перемещает movement ops через обёртки AFTER/END через 3 специфичных паттерна (movement через INDEX, AFTER, END) | Удаляет movement ops во время буферизации | Оба подхода функционально эквивалентны; подход Morok чище |
Выровненные стадии (ранее различались)
Следующие стадии были выровнены с Tinygrad в рамках данной реализации:
| Стадия | Что изменилось |
|---|---|
| 15: Понижение типа Index | Morok теперь имеет pm_lower_index_dtype() с полным покрытием паттернов: Binary ops, CONST, WHERE, VECTORIZE, SPECIAL, DEFINE_VAR, RANGE, CAST cleanup |
| 18: Декомпозиции | Добавлены: fast_division_patterns(), pm_div_to_shr(), pm_fdiv_to_mul(), pm_comparison_negations(), законы Де Моргана |
| 19: Финальная перезапись | pm_render() перенесён из кодогенерации в стадию 19 пайплайна schedule |
Паттерны только Tinygrad
Morok намеренно не реализует эти Tinygrad-специфичные паттерны:
| Паттерн | Назначение | Почему Morok это не нужно |
|---|---|---|
to_bufferview | Избежать копирования буферов с диска для DISK/TINYFS-устройств | Morok не поддерживает DISK/TINYFS; in-memory бэкенды не нуждаются в этом |
| Паттерны перемещения AFTER/END | Перемещение movement ops через timing-обёртки | Morok удаляет movement ops во время буферизации |
Расширения Morok
У Morok есть паттерны/расширения, отсутствующие в Tinygrad:
| Расширение | Расположение | Назначение |
|---|---|---|
| Свёртка вложенных INDEX с идентичными индексами | movement_op_patterns() | Удаление избыточных INDEX(INDEX(ptr, [i]), [i]) |
| CAT → VECTORIZE | pm_render | Преобразование CAT в явный VECTORIZE (CAT нельзя отрендерить напрямую) |
| Развёртка PTRCAT([x]) | pm_render | Удаление одноэлементных обёрток PTRCAT |
| GEP через CAST/BITCAST | gep_pushing_patterns() | Проталкивание GEP через приведения типов для лучшей оптимизации |
| Защита Image DType | pm_add_loads() | Пропуск обёртки LOAD для Image DType (обрабатывается в кодогенерации) |
Глоссарий
| Термин | Определение | Пример |
|---|---|---|
| Accumulator | Переменная с текущей суммой | acc = acc + value (в редукции) |
| Axis | Одно измерение тензора | Форма [100, 200] имеет 2 оси |
| AxisType | Как выполняется цикл | Global=параллельно, Reduce=аккумуляция |
| Buffer | Выделенная память с данными | Данные тензора хранятся в буфере |
| Bufferize | Записать результат в память вместо вычисления по требованию | Материализация промежуточного значения |
| CONTRACT | Объединить несколько значений в один вектор | [a, b, c, d] → vec4(a,b,c,d) |
| Devectorize | Разбить векторы под железо | vec8 → vec4, vec4 |
| Divmod | Операции деления и остатка | x // 7, x % 7 |
| Fixpoint | Когда применение паттернов перестаёт менять что-либо | Паттерны срабатывают до fixpoint |
| GEP | Get Element Pointer — вычисление адреса из индексов | arr[i][j] → base + i*stride + j |
| Hash consing | Переиспользование идентичных выражений | ADD(x, 0) + ADD(x, 0) разделяют память |
| Index | Целочисленный тип для индексов массивов | i32 или i64, в зависимости от устройства |
| Load | Чтение из памяти | value = arr[i] |
| Pattern | Правило «найти и заменить» для кода | ADD(x, 0) → x |
| Predicated store | Условная запись в память | Записать если валидно, иначе пропустить |
| Range | Спецификация итерации цикла | for i in 0..100 |
| Reduction | Свёртка множества значений в одно | Сумма, максимум, минимум |
| Store | Запись в память | arr[i] = value |
| Symbolic | Упрощение по правилам алгебры | (x/4)*4 → x (когда x%4=0) |
| Tensor core | Аппаратура для быстрого матричного умножения | Только NVIDIA GPU |
| Topological sort | Упорядочение узлов с учётом зависимостей | A до B, если B использует результат A |
| UNROLL | Расширить одну операцию в несколько позиций | x → [x_0, x_1, x_2, x_3] |
| UPCAST | Пометить намерение векторизовать | RANGE(0..4, UPCAST) |
| Vectorize | Обработка нескольких значений одновременно | SIMD: сложить 4 числа за раз |
| WHERE | Условный выбор | WHERE(cond, x, y) = x if cond else y |