LEGO от MIT: компилятор, который автогенерирует быстрые пространственные ускорители для ИИ
Генерация аппаратуры без шаблонов
LEGO уходит от подходов, которые либо только анализируют датафлоу, либо рождают RTL по фиксированным шаблонам. Вместо этого фреймворк принимает высокоуровневое описание тензорных программ и генерирует синтезируемый RTL, что позволяет создавать архитектуры, адаптированные к конкретным слоям и сочетаниям операций, без ручной доработки шаблонов.
Аффинное, реляционно-центричное входное представление
Во фронтенде LEGO представляет тензорные вычисления как вложенные циклы с тремя классами индексов: временные (for-циклы), пространственные (параллельные FU) и вычислительные (область итераций). Две аффинные функции описывают отображение данных и их пространственную раскладку: f_{I→D} отображает вычислительные индексы в индексы тензора, f_{TS→I} — временные/пространственные индексы в вычислительные. Такое аффинное представление переводит задачу обнаружения повторного использования и генерации адресов в задачу линейной алгебры и избавляет анализ от сложных делений/модулей. Контроль потоков отдельно кодируется в векторе c, что позволяет разделять управление между FU и снижать логику управления.
Фронтенд: граф FU и совместный дизайн памяти
LEGO синтезирует уровень функциональных блоков, решая уравнения повторного использования по аффинным отношениям для поиска прямых и задержанных (FIFO) соединений между FU. Избыточные рёбра отсеиваются с помощью минимальных остовных арист (Chu–Liu/Edmonds) с учётом стоимости FIFO-буферов, а эвристика BFS перестраивает интерконнекты при сосуществовании нескольких датафлоу, отдавая приоритет цепочному повторному использованию и уже имеющим задержки узлам, чтобы уменьшить количество мультиплексоров и дублируемых узлов данных.
Банковая организация памяти вычисляется автоматически: по множеству FU, читающих/записывающих тензор в один такт, LEGO определяет число банков по дельтам индексов, с опцией сокращения числа банков через НОД. Коммутаторы распределения данных соединяют банки с FU, а повторное использование FU→FU реализуется через межблоковую сеть.
При объединении нескольких пространственных датафлоу LEGO сливает их интерконнекты в единый Architecture Description Graph (ADG) и планирует слияние так, чтобы избежать наивного увеличения мультиплексоров, что даёт заметную экономию энергии.
Бэкенд: компиляция и оптимизация в RTL
ADG понижается до детализированного графа примитивов (FIFO, mux, сумматоры, генераторы адресов). LEGO выполняет несколько LP- и графовых преобразований:
- Выравнивание задержек через LP: линейная программа выбирает задержки выходов узлов, минимизируя добавленные регистры и сохраняя выравнивание по таймингу.
- Перепривязка шины вещания: двухэтапная оптимизация превращает дорогостоящие вещания в последовательные цепочки, применяя виртуальную корректировку стоимости и MST-перепривязку, затем LP для ребалансировки задержек.
- Выделение деревьев редукции и повторное использование выводов: последовательные цепочки сумматоров превращаются в сбалансированные деревья, а 0-1 ILP переназначает входы редьюсеров между датафлоу, уменьшая число физических пинов и логики.
Эти оптимизации по датапасу приносят значительную экономию: около 35% площади и порядка 28% экономии энергии по сравнению с наивной генерацией.
Результаты и оценка
LEGO реализован на C++, использует HiGHS как LP-решатель и эмитирует SpinalHDL → Verilog. Оценка покрывает тензорные ядра и полноразмерные модели: AlexNet, MobileNetV2, ResNet-50, EfficientNetV2, BERT, GPT-2, DDPM/Stable Diffusion, LLaMA-7B и другие. Для всех моделей используется одна инстанция LEGO-ускорителя; маппер подбирает тилинг и датафлоу для каждого слоя.
По сравнению с Gemmini при сопоставимых ресурсах (256 MACs, 256 KB буфер, 128-bit шина при 16 GB/s) LEGO обеспечивает в среднем 3.2× ускорение и 2.4× экономию энергии. Источники выигрыша — точная модель производительности для маппинга и возможность динамической смены пространственных датафлоу через сгенерированные интерконнекты. На больших конфигурациях (1024 FU) LEGO показывает высокую загрузку для генеративных задач, тогда как огромные модели остаются ограничены пропускной способностью памяти.
Кому это полезно
Исследователям LEGO даёт формализованный путь от спецификации цикл-нейстов до аппаратуры с доказуемыми оптимизациями. Практикам это — hardware-as-code: целевые датафлоу и их слияния задаются высокоуровнево, а компилятор выводит интерконнекты, буферы и контроллеры. Продуктовым командам LEGO снижает порог для создания энергоэффективного кастомного кремния, где аппарат адаптируется под модель, а не наоборот.