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

Речь о бэкенде для IBM Spyre, ускорителя с 32 ядрами, соединёнными в кольцо. У каждого ядра 2 МБ scratchpad-памяти и нет аппаратного кэша, который сглаживал бы промахи. Загрузка данных и вычисления идут на разных движках и общаются через FIFO. Видеокартовая модель компиляции, где кэш и динамический планировщик сами разбираются с перемещением данных, здесь не работает. Что должно лежать в быстрой памяти и в какой момент, приходится решать заранее.

В блоге PyTorch вышла третья часть истории о том, как фреймворк работает с этим железом. Первая была про расширения Inductor, вторая про запуск скомпилированного артефакта, а эта про середину: как решения Inductor превращаются в программу, которую рантайм может запустить. Между ними стоит MLIR-конвейер из двух половин, и граница между половинами проведена по одному конкретному промежуточному представлению.

Что такое KTIR

KTIR (Kernel Tile IR) это представление в MLIR, которое фиксирует, что делает каждый вычислительный тайл ускорителя и где лежат его данные. Оно записывает раскладку, тайлинг, деление работы между ядрами и резидентность в scratchpad, выбранные Inductor, но намеренно молчит о том, как данные попадут на место. Логика доставки остаётся планировщику.

Такой контракт понадобился потому, что ускоритель плохо описывается привычными терминами. KTIR строится на простой картинке: набор ядер, у каждого вычислительный движок и собственная быстрая память, всё это соединено внутрикристальным интерконнектом, который дотягивается и до внешней памяти. Главное свойство картинки в том, что любой вычислительный тайл видит любую память: свою scratchpad, scratchpad соседей и память устройства. Детали вроде FIFO, блоков загрузки и decoupled-исполнения из абстракции убраны, потому что это уже работа планировщика.

Из этого свойства вырастает главное отличие от GPU-ориентированных tile-IR. Triton описывает программу, которая загружает блоки тензора из глобальной памяти, считает на них и пишет результат назад, а промежуточная раскладка по shared memory и регистрам остаётся делом компилятора и железа. Хранилище инстанса программы приватно и живёт только пока он работает. Сказать, что тензор должен остаться в быстрой памяти конкретного ядра до следующей операции, или что один инстанс прочитает данные, лежащие у другого, в такой модели обычно нельзя.

На Spyre это ограничение оказалось принципиальным. Кэша, который подхватил бы промах, нет, а удержание тензора на кристалле между операциями стало одним из главных рычагов производительности. Поэтому KTIR делает размещение в быстрой памяти полноценной частью представления: аллокатор Inductor решает, какие буферы остаются резидентными, по какому адресу и на каком ядре, а IR просто записывает это решение.

Один тензор на четырёх scratchpad

Память в KTIR именованная и физическая. Пространство global означает память устройства, а ct_local с номером ядра означает scratchpad этого ядра. Из таких видов собирается конструкция, которой нет прямого аналога в GPU-стеке: distributed memory view. Несколько видов памяти, каждый со своим набором координат, объединяются в один логический тензор, разложенный по нескольким кристаллическим буферам.

Пример из статьи делает работу наглядной. Тензор 128×1024 в fp16 разложен по 32 строки на четырёх ядрах, каждое по смещению, которое выбрал аллокатор. Потребитель, которому нужны строки с 48 по 79, читает половину из своей scratchpad, половину из scratchpad второго ядра. Для него это всё равно одна операция загрузки, потому что границы быстрой памяти для вычислительного тайла прозрачны.

Такая конструкция закрывает сразу три задачи. Промежуточные тензоры остаются на кристалле между операциями: производитель пишет в distributed-вид, потребитель читает из него, и никто не упоминает внешнюю память. Несовпадение сеток деления между производителем и потребителем превращается в обычный IR: тайлы доступа просто пересекают границу scratchpad, а перекладка на кристалле не требует отдельной операции. И размещение остаётся обязывающим: пространство памяти плюс набор координат и есть контракт, на который опирается скомпилированный код.

Арифметика и управляющие конструкции берутся из апстримных диалектов MLIR, а собственный диалект ktdp добавляет ровно то, чего в них нет: явные памяти и явные обращения. Это получение идентификатора тайла, построение вида памяти по адресу без выделения и перемещения данных, построение distributed-вида, тайлы доступа с аффинным порядком обхода, косвенные тайлы для сборок по индексам вроде paged KV-cache, загрузка и запись как единственные операции перемещения, и межтайловая передача с редукцией для частичных сумм.

Второй диалект, spyreop, гораздо меньше и выставляет собственные инструкции чипа как обычные операции MLIR: экспоненту, gelu, сигмоиду, silu, softplus, корни, шаги LayerNorm, целочисленную и адресную арифметику, включая idx32toaddr для сборок по базе и шагу. Внутри linalg.generic они ведут себя как любое скалярное вычисление и фьюзятся с окружением. Как именно каждая инструкция разворачивается в последовательность команд железа, в планировщике не написано: это лежит в описании архитектуры.

Там же определён формат df16, 16-битное представление IBM с раскладкой битов 1-6-9 между знаком, экспонентой и мантиссой вместо IEEE half precision с его 1-5-10. Разница в два бита мантиссы и один бит экспоненты определяет, какие значения на этом железе представимы, и компилятор обязан это учитывать при понижении типов.

Dataflow Scheduler: три IR и шесть шагов

Планировщик это единый конвейер проходов MLIR, который понижает KTIR через ScheduleIR в DataflowIR. Вся модель исполнения держится на трёх конструкциях: вложенные циклы scf.for, ktdf.pipeline со стадиями, перекрывающимися между итерациями, и ktdf.parallel с одинаковыми инстансами, делящими пространство итераций. Любой анализ и любое преобразование рассуждают в этих терминах.

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

Шаг первый: нарезать ядро на программы. Ядро в KTIR это последовательность операций, а программа Spyre это одна конфигурация конвейера, загруженная на устройство. Планировщик идёт от каждой записи назад по вычислениям до загрузок, которые её питают, и всё связанное таким образом становится вычислительной группой, то есть цепочкой загрузка, вычисление, запись. Каждая группа уезжает в собственный вложенный модуль, а исходное ядро превращается в оркестратор, который хранит управляющий поток и адресную арифметику и вызывает программы по порядку.

Побочный эффект деления стоит проговорить: циклы в KTIR, пришедшие из грубого тайлинга Inductor, остаются в оркестраторе и становятся последовательностью запусков по тайлам, а циклы внутри программы создаёт уже планировщик, нарезая фьюзнутое вычисление под векторную ширину железа. Внутри каждой программы элементарные операции сводятся в один linalg.generic, и дальше вся работа идёт с ним.

Шаг второй: самый простой конвейер. Фьюзнутый generic нарезается под ширину SIMD вычислительного движка, взятые из описания архитектуры (на тестовом устройстве это 64 полосы fp16), циклы помечаются как параллельные или редукционные, и в самом внутреннем цикле строится конвейер из трёх стадий: загрузка в FIFO, вычисление, запись из FIFO в память.

Функционально это уже полная программа, но легальной она не является: в железе нет пути из памяти устройства прямо в вычислительный движок через FIFO. Исправляет это следующий проход.

Шаг третий: расширение пути. Планировщик строит граф маршрутизации из описания архитектуры, где узлы это памяти, вычислительные движки и блоки загрузки-записи, а рёбра это тракты передачи с определённой гранулярностью. Для каждого конвейера берутся концы, подразумеваемые трёхстадийной версией, между соседними парами ищется кратчайший путь, и результат сшивается. Где реальный путь длиннее, в конвейер вставляются стадии передачи для промежуточных переходов, в scratchpad выделяются буферы под staging, FIFO перетипизируются так, чтобы соединять действительно соседние блоки, а каждая стадия получает метку того движка, который её исполнит.

Три стадии становятся пятью: блок DMA заводит тайл в scratchpad, блок загрузки scratchpad отдаёт его в FIFO, вычислительный движок считает, блок записи scratchpad сливает результат, блок DMA возвращает его в память устройства. Здесь же начинает работать информация о размещении: если операнд уже является ct_local-видом, путь до вычислителя начинается прямо в scratchpad, и проход не добавляет для него ни одной DMA-стадии. Тензор ни разу не уходит во внешнюю память, а резидентность, выбранная в Inductor, превращается в структуру конвейера без дополнительной машинерии.

Шаг четвёртый: вложенные конвейеры. После расширения пути каждая стадия всё ещё работает на 64 элемента, что легально, но расточительно: DMA эффективен на больших объёмах, а смысл scratchpad в том, чтобы данные ждали вычисление заранее. Планировщик перестраивает гнездо циклов в иерархию конвейеров, каждый на своей гранулярности. Внешний конвейер работает тайлами, внутренний векторами. Переносы, не меняющиеся от итерации к итерации, выносятся в конвейер-сосед, который выполняется один раз перед циклом. Редукции режутся на куски, помещающиеся в scratchpad, внешняя часть становится явным циклом с накоплением, а внутренняя оставляется под редукционную поддержку железа.

Шаг пятый: пространство циклов. Буферы, которые одна стадия пишет, а другая читает внутри цикла, дублируются, и конвейер вращается между двумя копиями, так что DMA заполняет следующий тайл, пока вычисление доедает текущий. Работа делится между субъядерными инстансами, и параллельный цикл превращается в распределение по N инстансам, где N снова берётся из описания архитектуры. Размеры тайлов подставляются последними, чтобы учитывать уже выбранный параллелизм.

В сумме это настоящее пространство поиска: структура циклов, гранулярность конвейеров, глубина буферизации, число инстансов и размер тайла заданы в IR явно, и каждое решение ограничено тем, что разрешает описание устройства. Авторы честно оговаривают границу: в первой версии эти решения принимаются локальными эвристиками, без модели стоимости и без перебора. Явность задумана как задел, чтобы будущий автотюнер или модель стоимости можно было добавить, не трогая структуру конвейера.

Шаг шестой: адреса, движки и DataflowIR. Каждый буфер получает адрес в своей памяти с проверкой по ёмкостям из описания, scratchpad выделяется на программу, а память устройства живёт дольше, потому что через неё одна программа передаёт результат другой. Абстрактные виды движков конкретизируются по ядрам и субъядрам, зависимости по токенам становятся явными сигналами, и программа распадается на блоки по видам движков, каждый со своей частью работы.

Символические значения проходят через всё это нетронутыми. Стартовый адрес, вычисленный из аргумента ядра, становится символом с идентификатором, а выражение, которое его производит (аргумент плюс доля этого элемента сетки в тензоре, приведённая к гранулярности адресации движка), записывается уровнем выше. Если элементы сетки стартуют с разных смещений, каждый получает свой символ и выбирается через отображение по движку. Тело DataflowIR при этом не принимает аргументов вовсе, оно только ссылается на символы.

Архитектурное описание вместо хардкода

Всё, что планировщик знает о железе, приходит из одного файла MLIR на диалекте ktdf_arch. Устройство описано как граф: узлы это памяти с видами и размерами и исполняющие блоки с признаками, которые важны для планирования (число SIMD-полос по типам, разрядность и гранулярность доступа блоков загрузки и записи, глубина каждого FIFO), а рёбра это тракты передачи с гранулярностью. Группы ограничивают ресурсы так же, как это делает железо: ядро владеет scratchpad вместе с блоками, которые им пользуются.

Ядро подключает описание одной строкой и отображает свои пространства памяти на памяти устройства. Дальше почти каждое решение читает что-то из описания: число полос задаёт ширину первого конвейера, тракты образуют граф, по которому ищет расширение пути, гранулярности и глубины FIFO ограничивают форму каждой передачи, размеры памятей ограничивают размещение адресов, а блоки паттернов дают собственные понижения устройства. Паттерны написаны на PDL и применяются в четырёх именованных точках: до отображения, до планирования, после планирования и после понижения. Именно там каждый spyreop-intrinsic заменяется шаблоном инструкций железа, а редукция по внутреннему измерению становится SIMD-редукцией.

Практический вывод из такой схемы простой: перенос планировщика на другой чип это написание нового описания, а не правка проходов. У Spyre своё описание, тесты планировщика гоняются на игрушечном устройстве, и одни и те же проходы работают на обоих без изменений.

От DataflowIR к SpyreCode

Дальше начинается устройство-специфичная часть. Каждая программа понижается в SpyreBinary, двоичный артефакт, который конфигурирует и запускает блоки на кристалле. Любое значение, зависящее от тензоров вызывающего кода, будь то адрес или форма, остаётся в бинарнике символом. Следующая стадия записывает, как каждый символ вычисляется из аргументов ядра, а последняя упаковывает бинарник, метаданные и план работ в две части.

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

Вся эта конструкция нужна ради одной экономии: подготовка оплачивается один раз, а запуск остаётся дешёвым. Для инференса, где одно и то же ядро вызывается тысячи раз с разными тензорами, разница между «пересобрать всё» и «подставить адреса» и есть разница между работоспособным и неработоспособным стеком.

Что это значит за пределами Spyre

Первая версия планировщика умеет поэлементные операции с broadcast и редукционными измерениями в многопрограммных KTIR-ядрах и проверена сквозняком, от KTIR через планирование и генерацию кода, на представительных операциях трансформера. Поддержка матричного умножения ещё дописывается, следом идут оптимизации на основе модели стоимости с целью запустить LLM целиком.

Цифры по скорости ввода новых операций (вчетверо быстрее на первой, в восемь раз на последующих) интересны не сами по себе, а как индикатор: если добавление операции не требует хирургии в проходах, стоимость поддержки нового железа падает на порядок. Это ровно та проблема, из-за которой зоопарк ускорителей в 2026 году такой неудобный: каждый чип приносит свой компилятор, и фреймворк вынужден тянуть десяток бэкендов.

Подход с границами в виде IR отвечает на это по-другому. Inductor фиксирует размещение в KTIR, планировщик превращает его в конвейер в ScheduleIR, устройство-специфичные стадии доводят дело до SpyreCode, и на каждой границе следующий этап получает именно те факты, которые ему нужны, не больше и не меньше. Так dataflow-ускоритель может стоять за обычными интерфейсами PyTorch, включая стримы, аллокатор и символические значения, подставленные в момент запуска.

KTIR и планировщик открыты под Apache-2.0, есть путь и от Inductor, и от Triton-фронтенда. Для тех, кто делает компилятор под ускоритель со scratchpad или dataflow-исполнением, ценность здесь не в конкретных эвристиках, а в схеме: описание архитектуры как данные, PDL-паттерны на точках расширения и явные промежуточные представления вместо мегалитического прохода.

Часто задаваемые вопросы

Чем KTIR отличается от Triton?

Triton описывает программу, которая сама загружает блоки из глобальной памяти и сама их считает, а раскладку по быстрой памяти оставляет компилятору и железу. KTIR делает размещение в scratchpad первоклассной частью представления и разрешает тайлу читать чужие scratchpad. На железе без кэша это не роскошь, а условие работоспособности.

Нужно ли переписывать планировщик под другой чип?

Нет. Всё, что планировщик знает о железе, вынесено в файл описания на диалекте ktdf_arch: памяти, движки, тракты, глубины FIFO, число SIMD-полос. Устройство-специфичные понижения живут в PDL-паттернах. Новый чип означает новое описание и свои паттерны при неизменных проходах.

Можно ли уже запускать LLM?

Пока нет. Первая версия закрывает поэлементные операции с broadcast и редукциями и проверена на операциях трансформера сквозняком, но матричное умножение ещё в работе. Авторы называют следующие шаги: завершить matmul, добавить оптимизации по модели стоимости и выйти на сквозной запуск LLM.

Итог

История про KTIR и Dataflow Scheduler это редкий случай, когда компиляторный пост читается как инженерная стратегия, а не как отчёт о внутренностях. Один контракт в MLIR фиксирует решения фронтенда, второй описывает конвейер, третий отдаёт программу по движкам, и на каждом шаге знание о железе лежит в отдельном файле, а не в коде проходов. Именно поэтому первый оператор завёлся вчетверо быстрее, а последующие в восемь раз.

Если вы делаете компилятор под нестандартный ускоритель, самое полезное здесь это не эвристики планирования, а схема разделения: описание архитектуры как данные, точки расширения на PDL и IR на каждой границе. Начать можно с репозитория dataflow-scheduler: в нём есть тест на простейшее сложение двух тензоров и флаг --mlir-print-ir-after-all, который показывает, как ядро делится на программы, как появляется трёхстадийный конвейер и как расширение пути вставляет переходы через иерархию памяти.