MLIR-to-RTL simulation flow: от linalg.matmul до systolic array

от автора

Вводная

Привет! Хотел бы рассказать о своем MVP проекта hw-mlir-lab, где я использую MLIR для lowering операции умножения матриц (matmul) на systolic array, который я симулирую в Verilator.

1. Зачем я делаю этот проект

В современных системах-на-кристалле всё больше вычислений уходит в специализированные аппаратные блоки: для ИИ, связи, DSP, обработки изображений, научных вычислений и других задач. Часто процессорное ядро в такой системе занимается в основном управлением, а тяжёлая вычислительная нагрузка выполняется на отдельных IP-ускорителях.

Но спроектировать быстрый аппаратный блок — это только половина задачи. Его ещё нужно эффективно использовать из программного стека. На итоговую производительность начинают влиять не только частота и микроархитектура RTL, но и tiling, padding, layout данных, формат команд, runtime-интерфейс, организация памяти и то, как compiler flow решает, какие операции отправлять на ускоритель.

Здесь возникает неприятный разрыв. Аппаратная команда обычно смотрит на RTL-блок, интерфейсы, timing и ресурсы. Софтверная команда — на компилятор, runtime, API и интеграцию с приложением. Но многие решения лежат ровно между этими мирами. Например, размер тайла или формат данных нельзя хорошо выбрать, если смотреть только на compiler side или только на hardware side.

Мне стало интересно собрать минимальный проект, который позволяет посмотреть на всю цепочку целиком: от высокоуровневой операции в MLIR до выполнения этой операции на RTL-модели ускорителя.

В текущем MVP я проверяю одну конкретную вертикаль: linalg.matmul в MLIR преобразуется под фиксированный 8×8 systolic array, заменяется на custom MLIR operation, lower-ится до C ABI и выполняется через cocotb/Verilator-симуляцию RTL-блока.

MLIR здесь используется как удобный слой для compiler flow: можно начать с высокоуровневой операции, применить transform passes, выделить hardware-specific operation и затем lower-ить её до вызова конкретного ускорителя.

Далее я хочу развивать проект в сторону открытого фреймворка для прототипирования ускорительных подсистем на кристалле: с библиотекой MLIR passes, IP-блоками, runtime-интерфейсами, RTL-симуляцией, тестами, интеграцией с synthesis tools, CIRCT-направлением.

2. Что уже работает и как устроен MVP

Сейчас проект закрывает одну конкретную end-to-end вертикаль: от linalg.matmul в MLIR до выполнения матричного умножения на RTL-модели systolic array в cocotb/Verilator-симуляции.

Упрощённо flow выглядит так:

MLIR input  |  vMLIR transform pipeline  |  vcustom standalone dialect / passes  |  vLLVM lowering  |  vnative executable  |  vC interface bridge  |  vcocotb / Verilator  |  vSystemVerilog systolic array RTL

На входе находится MLIR-программа с linalg.matmul. Это обычная высокоуровневая операция линейной алгебры, ещё не привязанная к конкретному аппаратному блоку.

Дальше MLIR transform pipeline приводит её к форме, подходящей для текущего demo-ускорителя: matmul разбивается на 8×8 тайлы, а для граничных случаев применяется padding. Это нужно потому, что RTL systolic array в MVP работает с фиксированным размером входных данных.

После этого подходящие операции заменяются на custom MLIR operation:

standalone.systolic_matmul

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

Для project-specific преобразований я использую standalone-opt. Это расширенный вариант MLIR standalone example: MLIR поставляет пример out-of-tree dialect вместе с opt-like tool, который можно взять как основу для своего dialect и своих passes. В моём случае в этот шаблон добавлены операции и преобразования, необходимые для превращения подходящих linalg.matmul в вызовы systolic array.

Дальше standalone.systolic_matmul lower-ится до обычного C-callable интерфейса:

systolic_matmul_8x8(...)

Это важная граница проекта. Для сгенерированной программы ускоритель выглядит как обычная C-функция. Сейчас за этой функцией стоит simulation bridge, но в будущем на этом месте может быть другой runtime-механизм: MMIO, DMA, FPGA driver или интерфейс к реальному ASIC/FPGA-прототипу.

Текущий C bridge передаёт данные в cocotb testbench через Unix socket. Он отправляет два входных тайла i8[8][8] и аккумулятор i32[8][8], после чего получает обратно обновлённый результат. Cocotb принимает запрос, подаёт данные на SystemVerilog-модуль systolic array, ждёт завершения вычисления и возвращает результат обратно в программу.

Таким образом, в MVP явно разделены три уровня:

MLIR level     — представление и трансформация вычисленияC ABI level    — граница между generated code и accelerator callRTL level      — выполнение операции на модели аппаратного блока

3. Как linalg.matmul превращается в вызов systolic array

Теперь посмотрим на центральную часть проекта: как обычный linalg.matmul постепенно превращается в вызов конкретного аппаратного блока.

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

01_tiled.mlir02_padded.mlir03_bufferized.mlir04_entry_wrapped.mlir05_systolic_memref.mlir06_systolic_call.mlir07_llvm.mlir08_llvm.ll

На входе находится MLIR-программа с linalg.matmul. На этом уровне операция ещё не знает ничего про systolic array, RTL, cocotb или C ABI. Это просто высокоуровневое описание матричного умножения.

Первая стадия — tiling. С помощью Transform dialect schedule matmul_tile.mlir исходный linalg.matmul разбивается на 8x8x8 тайлы. Размер выбран не случайно: текущий demo systolic array работает с фиксированным 8×8 блоком, поэтому compiler pipeline должен привести вычисление к форме, которую такое железо может принять.

Следующая стадия — padding. Если размеры исходных матриц не кратны 8, появляются граничные тайлы. Для них используется transform schedule matmul_pad.mlir: маленькие или boundary tiles дополняются до фиксированного размера. В итоге runtime-часть всегда может отправлять в systolic array данные одинаковой формы.

После этого выполняется bufferization. До этого момента IR может работать с tensor-level представлением. Pass one-shot-bufferize переводит данные на memref-level: теперь операнды представлены как буферы в памяти. Это важный переход, потому что дальше нужно прийти к C ABI и обычному native executable.

Затем создаётся C interface entry wrapper. Он нужен, чтобы C driver — файл с функцией main — мог вызвать сгенерированную MLIR-функцию через стабильную внешнюю точку входа. В текущем проекте создаётся wrapper для matmul_entry, который получает входные memref-ы и output memref, вызывает основную MLIR-функцию и возвращает результат через переданный output buffer.

После подготовки memref-level IR запускается project-specific стадия: conversion в systolic operation.

Здесь custom pass:

--convert-linalg-matmul-to-systolic

ищет подходящие linalg.matmul и заменяет их на:

standalone.systolic_matmul

Сейчас pass ожидает конкретную форму операции:

lhs: memref<8x8xi8>rhs: memref<8x8xi8>acc: memref<8x8xi32>

То есть это уже не общий matmul любого размера и типа, а hardware-specific операция, подготовленная под текущий RTL-блок.

Дальше запускается следующий custom pass:

--lower-systolic-to-func-call

Он убирает standalone.systolic_matmul и заменяет её на обычный вызов функции:

systolic_matmul_8x8(...)

Это ключевая точка всего flow. До неё операция ещё находится внутри MLIR как accelerator-specific op. После неё это уже обычный C-callable function call, который можно слинковать с runtime/interface code.

Если совсем коротко, центральная часть pipeline выглядит так:

linalg.matmul  |  | tiling: 8x8x8  vtiled linalg.matmul  |  | padding  vfixed-size matmul tiles  |  | bufferization  vmemref-level IR  |  | convert-linalg-matmul-to-systolic  vstandalone.systolic_matmul  |  | lower-systolic-to-func-call  vfunc.call @systolic_matmul_8x8

После этого оставшиеся MLIR dialects lower-ятся в LLVM dialect, затем mlir-translate превращает LLVM dialect в обычный LLVM IR. Дальше compile pipeline компилирует этот LLVM IR в объектный файл, компилирует interface.c и линкует всё вместе с C driver в native executable.

На runtime-стороне функция systolic_matmul_8x8 отправляет данные в cocotb/Verilator testbench и возвращает результат обратно.

В итоге generated program выглядит как обычное native приложение, но часть matmul-операций внутри него фактически выполняется на RTL-модели ускорителя.

4. Как запустить demo и тесты

Для того чтобы было проще запустить пример и начать эксперименты, проект сделан внутри Docker: LLVM/MLIR, Clang, Verilator, cocotb и custom standalone-opt запускаются внутри подготовленного окружения.

Основной сценарий запускается одной командой:

make demo

В результате проект собирает MLIR-инструменты, прогоняет demo-программу через compiler pipeline, компилирует native executable и запускает его вместе с cocotb/Verilator-симуляцией RTL systolic array.

После запуска можно посмотреть промежуточные артефакты MLIR pipeline:

build/mlir-pipeline/01_tiled.mlirbuild/mlir-pipeline/02_padded.mlirbuild/mlir-pipeline/03_bufferized.mlirbuild/mlir-pipeline/04_entry_wrapped.mlirbuild/mlir-pipeline/05_systolic_memref.mlirbuild/mlir-pipeline/06_systolic_call.mlirbuild/mlir-pipeline/07_llvm.mlirbuild/mlir-pipeline/08_llvm.ll

Это удобно для отладки: можно увидеть, на каком этапе linalg.matmul был разбит на тайлы, где появилась standalone.systolic_matmul и где она превратилась в обычный вызов функции.

Для проверки нескольких случаев есть тестовый запуск:

make test

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

5. Ограничения текущей версии и дальнейшие планы

Текущий MVP намеренно ограничен одной вертикалью: linalg.matmul → 8×8 systolic array → RTL-симуляция. Это позволяет быстрее проверить весь путь end-to-end.

Первое ограничение — сам hardware target. Сейчас поддерживается fixed-size systolic array, который работает с 8×8 тайлами. Входные данные имеют тип i8, а accumulator/result — i32. Этого достаточно для демонстрации идеи, но пока не покрывает другие типы данных, quantization-сценарии, более сложные memory layouts.

Второе ограничение — runtime layer. Сейчас функция systolic_matmul_8x8 общается с cocotb testbench через Unix socket. Для MVP это удобный simulation transport, но в другой системе на этом месте может быть MMIO, DMA, driver layer, command queue или другой механизм взаимодействия с ускорителем.

Также в проекте пока нет полноценной модели памяти, interconnect, DMA scheduler, cost model и автоматического выбора, какие операции выгодно отправлять на ускоритель. Сейчас pipeline выбирает только заранее подходящие linalg.matmul, которые после tiling/padding соответствуют ожидаемой форме.

Дальше проект можно развивать в нескольких направлениях.

Первое направление — расширение набора поддерживаемых операций и IP-блоков. Сейчас в качестве hardware target используется только demo systolic array для matmul. В будущем можно подключать другие ускорители: например, DSP-блоки, FFT, stencil-compute, hash core и другие вычислительные IP. Для каждого такого блока можно описывать свою MLIR operation, lowering pass, runtime interface и cocotb/Verilator-тесты.

Второе направление — более развитый compiler flow. Сюда относятся hardware constraints, cost model, выбор между CPU и accelerator path, поддержка разных типов данных и layout-ов памяти.

Третье направление — runtime и интеграция с реальным железом. Socket bridge можно рассматривать как первый simulation transport. Следующий шаг — runtime-интерфейс, который можно будет заменить на MMIO, DMA или FPGA-прототип.

Четвёртое направление — synthesis и CIRCT. Сейчас RTL написан вручную и проверяется в симуляции. Дальше интересно добавить synthesis sanity check, чтобы видеть базовые resource/timing-отчёты, а также исследовать, как CIRCT может помочь в генерации wrappers, интерфейсов, управляющей логики или других частей hardware infrastructure.

В более широком смысле я хочу двигать проект к открытому фреймворку для прототипирования ускорительных подсистем на кристалле. Не как замену HLS и не как готовый ASIC compiler, а как воспроизводимую среду, где можно экспериментировать с compiler passes, IP-блоками, runtime-интерфейсами и RTL-симуляцией в одном flow.

6. Заключение

Буду рад обратной связи от людей, которые занимаются MLIR/LLVM, RTL, FPGA/ASIC, cocotb/Verilator, EDA tooling или аппаратными ускорителями. Если кому-то интересно поучаствовать в проекте, буду рад контрибьютам, issue, обсуждениям архитектуры и идеям по новым IP-блокам или compiler/runtime-сценариям.

ссылка на оригинал статьи https://habr.com/ru/articles/1045754/