2025-11-11T09:13:09.652713

Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators

Hong, Bhatia, Cheung et al.
Hardware accelerators, especially those designed for tensor processing, have become ubiquitous in today's computing landscape. However, even with significant efforts in building compilers, programming these tensor accelerators remains challenging, leaving much of their potential underutilized. Recently, large language models (LLMs), trained on large amounts of code, have shown significant promise in code generation and optimization tasks, but generating low-resource languages, such as specialized tensor accelerator code still poses a significant challenge. We tackle this challenge with Autocomp, an approach that empowers accelerator programmers to leverage domain knowledge and hardware feedback to optimize code via an automated LLM-driven search. We accomplish this by: 1) formulating each optimization pass as a structured two-phase prompt, divided into planning and code generation phases, 2) inserting domain knowledge during planning via a concise and adaptable optimization menu, and 3) integrating correctness and performance metrics from hardware as feedback at each search iteration. Across three distinct hardware platforms, we demonstrate that Autocomp-optimized code runs 5.6x faster than the vendor-provided library (Gemmini), outperforms expert-level hand-tuned code by 1.9x (AWS Trainium), and achieves 3.8x higher performance than a machine learning-based cost model for GPUs (NVIDIA L40S). Additionally, we demonstrate that optimization schedules generated from Autocomp can be reused across similar tensor operations, improving speedups by up to 24% under a fixed sample budget.
academic

Autocomp: Мощный и портативный оптимизатор кода для тензорных ускорителей

Основная информация

  • ID статьи: 2505.18574
  • Название: Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators
  • Авторы: Charles Hong, Sahil Bhatia, Alvin Cheung, Yakun Sophia Shao (UC Berkeley)
  • Классификация: cs.PL cs.AI cs.AR cs.LG
  • Статус публикации: Препринт. На рецензии.
  • Ссылка на статью: https://arxiv.org/abs/2505.18574

Аннотация

Аппаратные ускорители, в частности специализированные для обработки тензоров, повсеместно распространены в современной вычислительной среде. Однако несмотря на значительные инвестиции в разработку компиляторов, программирование тензорных ускорителей остаётся сложной задачей, что приводит к недоиспользованию их потенциала. В данной работе предлагается Autocomp — метод оптимизации кода посредством автоматизированного поиска на основе больших языковых моделей (LLM), позволяющий программистам ускорителей использовать знания предметной области и обратную связь от оборудования. Метод реализуется через три ключевые технологии: 1) формулировка каждого процесса оптимизации как структурированного двухэтапного запроса, разделённого на этапы планирования и генерации кода; 2) внедрение знаний предметной области во время планирования через компактное и адаптивное меню оптимизаций; 3) интеграция показателей корректности и производительности от оборудования в качестве обратной связи на каждой итерации поиска.

Исследовательский контекст и мотивация

Основные проблемы

Основные вызовы при программировании тензорных ускорителей включают:

  1. Сложность программирования: В отличие от программирования универсальных процессоров, тензорные ускорители требуют явного управления перемещением данных, конфигурацией состояния и планированием операций
  2. Затраты на адаптацию компилятора: Адаптация традиционных компиляторов для новых аппаратных платформ требует значительных инженерных работ, причём затраты на разработку программного обеспечения составляют 40-50% от общих затрат на разработку нового оборудования
  3. Проблема оптимизации расписания: Комбинаторный взрыв при определении того, какие оптимизации применять и в каком порядке
  4. Вызовы языков с ограниченными ресурсами: Наборы инструкций (ISA) и предметно-ориентированные языки (DSL) специализированных ускорителей недостаточно представлены в обучающих корпусах LLM

Ограничения существующих подходов

  1. Традиционные компиляторы: XLA, TVM, Triton поддерживают только несколько аппаратных бэкендов, в основном CPU и GPU
  2. Подходы на основе DSL: Halide, Exo предоставляют примитивы для выражения тензорных вычислений, но бремя оптимизации остаётся на программисте
  3. Методы, управляемые данными: Требуют больших объёмов данных о производительности для обучения, что крайне дефицитно для специализированных аппаратных ускорителей
  4. Прямое применение LLM: Для языков с ограниченными ресурсами генерация кода с нулевым обучением крайне ненадёжна

Основные вклады

  1. Первый метод оптимизации кода тензорных ускорителей с ограниченными ресурсами, управляемый LLM: Предложена структура Autocomp, специально разработанная для специализированных аппаратных ускорителей
  2. Высокопортативная структура оптимизации: Адаптация к новым аппаратным платформам путём модификации запросов значительно снижает инженерные затраты
  3. Превосходная производительность: Значительное превосходство над существующими методами на трёх различных аппаратных платформах
  4. Механизм повторного использования расписаний: Доказано, что расписания оптимизации могут быть повторно использованы между аналогичными тензорными операциями, повышая эффективность выборки

Подробное описание метода

Определение задачи

Вход: Неоптимизированный код тензорного ускорителя Выход: Функционально эквивалентный, но оптимизированный по производительности код Ограничения: Сохранение семантической эквивалентности, обеспечение корректности путём проверки на оборудовании

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

Этап 1: Генерация плана (Plan Generation)

Структура запроса включает:

  1. Описание ISA ускорителя: Семантика инструкций, спецификации адресации памяти, описание архитектуры оборудования
  2. Текущий код: Код, подлежащий оптимизации
  3. Обратная связь о производительности: Метрики задержки (количество циклов), использование памяти и другие показатели
  4. Меню оптимизаций: Предопределённые высокоуровневые варианты оптимизации (например, тайлирование циклов, переупорядочение, слияние и т.д.)
  5. Информация об итерации поиска: Номер текущей итерации для руководства выбором оптимизации

Этап 2: Реализация кода (Implementation)

Структура запроса включает:

  1. Описание ISA ускорителя: Аналогично этапу 1
  2. Текущий код: Аналогично этапу 1
  3. Сгенерированный план: Конкретный план оптимизации, выведенный на этапе 1
  4. Примеры контекстного обучения: Примеры кода для сложных оптимизаций (например, тайлирование)
  5. Инструкции по реализации: Инструкции на естественном языке для применения плана и вывода оптимизированного кода

Стратегия поиска с ограничением ширины

  • Поиск с ограничением ширины B=6, параллельное исследование нескольких траекторий оптимизации
  • Фильтрация корректности: Проверка кандидатов через набор функциональных тестов
  • Отбор по производительности: Сохранение только кандидатов, превосходящих по производительности родительский узел
  • Итеративная оптимизация: Процесс поиска с фиксированным бюджетом T итераций

Ключевые технологические инновации

1. Методы повышения разнообразия

  • Отсев меню оптимизаций: Случайное удаление части параметров меню при каждом планировании (с вероятностью 70%)
  • Интеграция LLM: Распределение запросов между несколькими LLM для увеличения разнообразия ответов

2. Интеграция обратной связи от оборудования

  • Показатели производительности в реальном времени (задержка, использование памяти) направляют выбор следующей оптимизации
  • Моделирование с точностью до цикла или измерение производительности на уровне кристалла

3. Механизм повторного использования расписаний

  • Запись высококачественных последовательностей расписаний
  • Повторное использование известных расписаний для аналогичных тензорных операций (с одинаковым соотношением сторон или общими размерами)
  • Дальнейшая оптимизация после лёгкого поиска

Экспериментальная установка

Аппаратные платформы

  1. Gemmini: Генератор открытого исходного кода для ускорителей, поддерживающий систолические массивы и векторные стили тензорных ускорителей
  2. AWS Trainium: Коммерческий высокопроизводительный тензорный ускоритель, использующий Neuron Kernel Interface (NKI)
  3. NVIDIA L40S GPU: Современный графический процессор для центров обработки данных с выделенными Tensor Core

Наборы тестов

  • Gemmini: GEMM и свёртки ResNet-50, прогнозирование управления TinyMPC
  • Trainium: Операторы глубокого обучения уровня учебника и продвинутые (RMSNorm, LayerNorm, GEMM, Mamba и т.д.)
  • GPU: Тесты KernelBench Level 1

Методы сравнения

  1. Высокоуровневые библиотеки программного обеспечения: Библиотека Gemmini, PyTorch NeuronX, PyTorch
  2. Неоптимизированный низкоуровневый код: Неоптимизированный Exo, учебный код nki-samples
  3. Ручная оптимизация кода: Реализации, оптимизированные экспертами вручную
  4. Модели стоимости ML: TVM MetaSchedule (GPU)
  5. Аппаратный FSM: Конечный автомат оборудования Gemmini (верхняя граница)

Результаты экспериментов

Основные показатели производительности

Платформа Gemmini

  • Тесты GEMM: Ускорение в 5,6× по сравнению с библиотекой Gemmini, превосходство над ручной оптимизацией экспертов в 1,4×
  • Тесты свёртки: Ускорение в 2,6× по сравнению с библиотекой, превосходство над ручной оптимизацией в 1,1×
  • Мелкозернистая линейная алгебра: Превосходство над неоптимизированным кодом в 2,7×, даже превосходство над оптимизированной экспертами реализацией аппаратного FSM в 1,6× (прямой проход)

Платформа AWS Trainium

  • Учебные рабочие нагрузки: Превосходство над ручной оптимизацией экспертов в 1,36× (геометрическое среднее), превосходство над скомпилированным кодом PyTorch NeuronX в 13,52×
  • Продвинутые рабочие нагрузки: Превосходство над оптимизацией экспертов в 1,9× (геометрическое среднее), ускорение одномерной глубокой свёртки до 17,37×

GPU NVIDIA L40S

  • Тесты KernelBench: Превосходство над PyTorch в 2,05× (геометрическое среднее), превосходство над TVM MetaSchedule в 3,8×
  • Превосходство над PyTorch во всех тестах, TVM превосходит PyTorch только в 2 тестах

Анализ абляционных экспериментов

Детальные абляционные эксперименты подтверждают важность каждого компонента:

  • ISA ускорителя: Удаление приводит к значительному снижению производительности, но улучшение всё ещё достижимо
  • Меню оптимизаций: Полностью необходимо, удаление полностью нарушает производительность оптимизации
  • Отсев меню: Значительно влияет на производительность, предотвращая смещение модели в сторону ограниченного набора параметров
  • Интеграция LLM: Обеспечивает важное разнообразие, производительность одной модели ниже
  • Обратная связь о производительности оборудования: Полезна, но эффект ограничен, так как меню оптимизаций уже содержит соответствующие метрики

Эффект повторного использования расписаний

  • При бюджете 100 образцов: повторное использование расписаний достигает ускорения в 4,6×, без повторного использования только 3,7×
  • При бюджете 200 образцов: повторное использование расписаний достигает ускорения в 5,0×, без повторного использования только 4,2×
  • Доказано, что расписания хорошо обобщаются и могут эффективно снизить затраты поиска для аналогичных тестов

Связанные работы

Оптимизация кода тензорных ускорителей

  • Модели производительности: Timeloop, MAESTRO используют высокоуровневые модели архитектуры оборудования
  • Методы автоматизации: Машинное обучение, линейное программирование, чёрный ящик оптимизации, обучение с подкреплением
  • Ограничения: Существующие абстракции игнорируют специфику реализации и оптимизации на уровне инструкций

Оптимизация кода на основе LLM

  • Область применения: Эволюционный поиск, поиск с увеличением контекста, итеративная оптимизация, постобучение моделей
  • Системная оптимизация: CUDA, встроенные функции SIMD
  • Исследовательский пробел: Отсутствие работ по оптимизации кода на основе LLM для специализированного оборудования (не CPU/GPU)

Заключение и обсуждение

Основные выводы

  1. Эффективность оптимизации, управляемой LLM: Autocomp значительно превосходит традиционные методы на нескольких аппаратных платформах
  2. Исключительная портативность: Адаптация к новому оборудованию требует только модификации запросов, инженерные затраты минимальны
  3. Ценность повторного использования расписаний: Расписания оптимизации хорошо обобщаются и значительно повышают эффективность выборки

Технологические инсайты

  1. Необходимость двухэтапного дизайна: Разделение этапов планирования и реализации повышает успешность сложных задач оптимизации
  2. Важность знаний предметной области: Знания предметной области, предоставляемые меню оптимизаций, критичны для производительности
  3. Ценность обратной связи от оборудования: Показатели производительности в реальном времени эффективно направляют выбор направления оптимизации

Ограничения

  1. Зависимость от возможностей LLM: Производительность метода ограничена способностями базовой LLM в генерации и рассуждении о коде
  2. Затраты на поиск: Требуются многократные вызовы LLM и моделирование оборудования, что приводит к значительным вычислительным затратам
  3. Специфичность предметной области: Меню оптимизаций требует ручного проектирования для различных аппаратных платформ
  4. Ограничения оценки: Основное внимание сосредоточено на рабочих нагрузках тензорных вычислений, применимость к другим типам вычислений неизвестна

Направления будущих исследований

  1. Автоматическая генерация меню: Исследование методов автоматического построения меню оптимизаций
  2. Миграция расписаний между платформами: Изучение передачи знаний о расписаниях между различными аппаратными платформами
  3. Оптимизация экономической эффективности: Снижение количества вызовов LLM и моделирования оборудования в процессе поиска
  4. Более широкое применение: Расширение на другие специализированные ускорители, не связанные с тензорными вычислениями

Глубокая оценка

Преимущества

  1. Высокая инновационность: Первое применение LLM к оптимизации кода тензорных ускорителей с ограниченными ресурсами, новый технический подход
  2. Высокая практическая ценность: Решение реальных инженерных проблем, значительное снижение затрат на разработку программного обеспечения для нового оборудования
  3. Полная экспериментальная оценка: Комплексная оценка на трёх различных аппаратных платформах с убедительными результатами
  4. Универсальность метода: Структура хорошо масштабируется и портативна
  5. Превосходная производительность: Значительное превосходство над существующими лучшими методами на множестве тестов

Недостатки

  1. Вычислительные затраты: Требуются значительные вызовы LLM и моделирование оборудования, что может ограничить практическое применение
  2. Зависимость от ручного проектирования: Меню оптимизаций по-прежнему требует ручного проектирования с использованием экспертных знаний, степень автоматизации ограничена
  3. Ограничения оценки: Основное внимание сосредоточено на конкретных типах тензорных вычислений, универсальность требует проверки
  4. Недостаток теоретического анализа: Отсутствуют теоретические гарантии сходимости и оптимальности метода

Оценка влияния

  1. Академическая ценность: Открывает применение LLM в оптимизации компиляции для специализированного оборудования, имеет важное академическое значение
  2. Промышленное влияние: Может значительно снизить затраты на разработку программного стека для нового оборудования, имеет важное промышленное значение
  3. Воспроизводимость: Авторы обещают открыть исходный код реализации и запросы, что способствует дальнейшим исследованиям
  4. Вдохновляющий характер: Предоставляет новый технический путь для оптимизации компиляции других специализированных ускорителей

Применимые сценарии

  1. Разработка прототипов нового оборудования: Быстрая генерация оптимизированного кода для вновь разработанных тензорных ускорителей
  2. Построение компиляторов DSL: Дополнение или замена традиционных компиляторов
  3. Инструменты оптимизации производительности: Помощь разработчикам в оптимизации существующего кода ускорителя
  4. Исследования и обучение: Предоставление автоматизированных инструментов для программирования и оптимизации ускорителей

Библиография

Статья ссылается на обширный объём связанных работ, включая:

  • Проектирование аппаратных ускорителей (Gemmini, TPU, Trainium и т.д.)
  • Компиляторы и DSL (XLA, TVM, Halide, Exo и т.д.)
  • Генерация кода на основе LLM (CodeGen, Codex и т.д.)
  • Методы автоматической оптимизации (обучение с подкреплением, эволюционные алгоритмы и т.д.)

Общая оценка: Это высококачественная исследовательская работа, вносящая важный вклад в новую междисциплинарную область применения LLM к оптимизации компиляции для специализированного оборудования. Метод отличается высокой инновационностью, полной экспериментальной оценкой и значительной практической ценностью. Хотя остаётся место для улучшения в области вычислительных затрат и степени автоматизации, работа открывает новые направления развития в этой области и имеет важное академическое и промышленное значение.