Матричное умножение давно стало самой горячей операцией в серверном железе: на нём держатся и обучение моделей, и инференс, и шифрование с обработкой сигналов. Долгое время процессоры расплачивались за такие задачи кэшем и широкими векторами, пока в 2023 году Intel не завезла в серверные Xeon семейства Sapphire Rapids отдельный матричный блок. Расширение называется Intel Advanced Matrix Extensions, сокращённо AMX, и оно переворачивает привычные представления о том, сколько арифметики способно провернуть обычное серверное ядро без единого ускорителя.
Суть новинки в том, что у ядра появились собственные двумерные регистры, плитки, и специальное устройство TMUL, которое перемножает содержимое этих плиток одной инструкцией. Классический путь через AVX-512 прокручивает данные векторами по 512 бит, а здесь за одну операцию обрабатывается целая подматрица, и арифметическая плотность вырастает в разы. При этом всё остаётся внутри обычного процесса, без библиотек драйверов, очередей и копирований, привычных для дискретных ускорителей.
Разберём устройство расширения, работу с ним на интринсиках, взаимодействие с ядром Linux и виртуализацией, а также границы, за которыми от AMX проку меньше, чем кажется.
Устройство матричного расширения в серверных Xeon
В распоряжении каждого ядра восемь плиток с именами TMM0 по TMM7. Максимальный размер одной плитки 1024 байта: не более 16 строк по 64 байта в каждой. Всего на ядро приходится до 8 килобайт матричных данных, а рядом лежит отдельный 64-байтовый регистр конфигурации, где описана геометрия каждой плитки. Набор параметров называется палитрой, палитра 1 предназначена для вычислений в формате bf16 и целых чисел, палитра 0 пустая и служит для сброса.
Сердце расширения, блок TMUL, выполняет инструкцию вида C плюс равно A умножить на B для целых подматриц. В формате bf16 плитка A может быть размером 16 на 32 элемента, плитка B 32 на 16 элементов, и результат копится в плитку C из 16 на 16 чисел с плавающей точкой одинарной точности. Одна инструкция _tile_dpbf16_ps перемалывает 16 умножить на 16 умножить на 32, то есть 8192 операций умножения с накоплением. Для int8 формата k-измерение удваивается до 64, и инструкция _tile_dpbssd перемалывает уже 16384 умножений за раз. Другой путь исполнения на процессоре не даёт такой плотности на инструкцию. Одна плиточная инструкция выполняется за несколько тактов, поэтому выигрыш раскрывается на длинных сериях умножений.
Проверка поддержки занимает секунды. В /proc/cpuinfo ищутся флаги amx_tile, amx_bf16 и amx_int8, тот же набор показывает lscpu. Без BIOS-включения на некоторых платформах флагов не будет, поэтому на серверах их проверяют до обновления прошивок, а не после.
Инструкции и первый пример на интринсиках
Работа строится по строгому циклу: сначала загрузить конфигурацию геометрии, затем наполнять плитки данными из памяти, гонять умножения с накоплением и выгружать результат обратно. Конфигурация лежит в памяти блоком ровно на 64 байта, где для каждой из восьми плиток перечислены число строк и число байт в строке. Загрузка конфигурации делается инструкцией _tile_loadconfig, выгрузка результата _tile_stored, обнуление накопителя _tile_zero, а для загрузки плиток из памяти служат _tile_loadd и транспонированный вариант _tile_loaddt.
Минимальный пример умножения одной подматрицы в формате bf16 выглядит так:
#include <immintrin.h>
typedef struct {
unsigned char palette_id;
unsigned char start_row;
unsigned char reserved_0[14];
unsigned char reserved_1[16];
unsigned short colsb[8];
unsigned short rows[8];
} tilecfg;
tilecfg cfg = { .palette_id = 1, .start_row = 0 };
cfg.rows[0] = 16; cfg.colsb[0] = 64;
cfg.rows[1] = 32; cfg.colsb[1] = 32;
cfg.rows[2] = 16; cfg.colsb[2] = 64;
_tile_loadconfig(&cfg);
_tile_zero(2);
for (int k = 0; k < K; k += 32) {
_tile_loadd(0, A + k, 64);
_tile_loadd(1, B + k * 32, 32);
_tile_dpbf16_ps(2, 0, 1);
}
_tile_stored(2, C, 64);Здесь плитка 0 держит полосу A из 16 строк по 32 числа bf16, плитка 1 полосу B из 32 строк по 16 чисел, а накопитель в плитке 2 растёт по всему k-измерению. Сборка требует флагов -mamx-tile -mamx-bf16, либо одного -march=sapphirerapids, интринсики объявлены в immintrin.h, а компиляторы GCC и Clang поддерживают их с версий 11 и 12 соответственно. Целочисленный вариант меняет лишь инструкцию на _tile_dpbssd или _tile_dpbuud и геометрию плиток, для смешанных знаков есть пары _tile_dpbusd и _tile_dpbuud с соответствующей раскладкой.
Целочисленный вариант с накоплением в int32 настраивается своей геометрией, k-измерение растёт до 64 элементов:
cfg.rows[0] = 16; cfg.colsb[0] = 64;
cfg.rows[1] = 64; cfg.colsb[1] = 32;
cfg.rows[2] = 16; cfg.colsb[2] = 64;
_tile_loadconfig(&cfg);
_tile_zero(2);
_tile_loadd(0, A, 64);
_tile_loadd(1, B, 32);
_tile_dpbssd(2, 0, 1);Сумма получается в плитке 2 из шестнадцати на шестнадцати знаковых 32-битных чисел, а итоговая плотность достигает 16384 умножений с накоплением на инструкцию.
Важная деталь, о которую спотыкаются новички: конфигурация держится в регистре до следующей загрузки, поэтому при смене геометрии её надо перезагружать, а изменение конфигурации на ходу сбрасывает содержимое плиток. Поле start_row в конфигурации позволяет продолжить вычисление с определённой строки, этим пользуются механизмы сохранения контекста.
Ленивое включение состояния и механизм XFD в ядре
Восемь килобайт плиток плюс 64 байта конфигурации это самый жирный блок состояния процессора в системе XSAVE, и сохранять его на каждом переключении задач для всех процессов было бы расточительством. Разработчики ядра применили изящный ход под названием XFD, расширенное отключение признака. По умолчанию состояние AMX у задачи подавлено через MSR IA32_XFD с номером 0x1D4, и первая же AMX-инструкция вызывает исключение #NM, в регистре IA32_XFD_ERR с номером 0x1D5 ядро видит причину, выделяет место под состояние и снимает блокировку. Дальше задача живёт с полным состоянием, а процессы, никогда не трогавшие матрицы, не платят ни байтом накладных расходов.
Хостовая поддержка этой схемы вошла в ядро 5.16, а гостевая, с пробросом AMX в виртуальные машины, в ядро 5.17. Готовый самотест amx из tools/testing/selftests/x86 в исходниках ядра проверяет ленивое включение и сохранение контекста при сигналах, полезно прогнать его после смены ядра. Быстрая проверка без компиляции делается утилитой cpuid из одноимённого пакета: вывод с фильтром по слову amx перечисляет поддержанные подпризнаки, геометрии плиток и параметры XSAVE-областей. Плиточные данные и конфигурация видны там как отдельные компоненты XSAVE с собственными смещениями. Из практики стоит помнить: процессы, реально использующие AMX, переключаются дороже обычных, потому что ядру приходится таскать восемь килобайт состояния. Поэтому для матричных воркеров имеет смысл привязка по привязке taskset или cset к стабильным ядрам, а не пинбол между NUMA-узлами.
Виртуализация и проброс AMX в гостевые системы
Клаудам без AMX жить скучно, поэтому виртуализация была в планах с самого начала. QEMU начиная с версии 7.2 умеет выставлять гостю признак amx, классическая строка запуска выглядит как -cpu SapphireRapids,+amx,+amx-bf16,+amx-int8. Гость обязан работать на ядре с поддержкой гостевого состояния, иначе запуск придет к отказу на этапе включения процессора. Внутри гостя те же флаги в cpuinfo и те же интринсики, то есть образ приложения не отличишь от физического сервера.
В libvirt достаточно режима host-passthrough в описании процессора, тогда признак amx попадает в гостя вместе со всем набором узла:
<cpu mode="host-passthrough" check="none"/>host-passthrough остаётся самым коротким путём и для строгих сценариев.
Живое мигрирование виртуальных машин с активным AMX-состоянием тоже поддержано: 8 килобайт состояния просто добавляются к передаваемому контексту процессора. Тонкость в совместимости узлов: миграция на узел без amx запрещена, и планировщики кластера учитывают это при размещении. Для проверки в госте достаточно одной команды grep amx /proc/cpuinfo, а для полноценного теста можно собрать маленький пример с _tile_dpbf16_ps и сверить результат с эталонным умножением на процессоре.
Библиотеки и фреймворки, где AMX уже работает
Вручную писать плиточный код редко нужно, всё уже собрано в библиотеках. oneDNN с версии 2.6 включает AMX-ядра для GEMM в bf16 и int8, поверх него работают PyTorch с расширением Intel Extension for PyTorch и TensorFlow с включённым oneDNN. Для PyTorch выигрыш виден уже после простого включения: подсистема матричных умножений автоматически выбирает AMX-путь, когда размеры матриц и формат данных позволяют. Тестовые пакеты oneDNN, в частности benchdnn, позволяют прогнать эталонные GEMM и сверить прирост против AVX-512 путей на том же железе.
Числа прироста зависят от формы матриц, но порядок понятен: против скалярных и базовых векторных путей счёт идёт на десятки раз, против хорошо настроенных AVX-512 VNNI ядер в bf16 на два-четыре раза, и это согласуется с кратностью плотности арифметики на инструкцию. Инференс трансформерных моделей в int8 получает особенно много, потому что там узкое место именно в умножениях, а не в памяти. Нюанс практический: для маленьких матриц, где k и n исчисляются десятками, расходы на загрузку плиток и конфигурацию съедают выигрыш, поэтому библиотеки применяют AMX по эвристикам и переключаются на векторные пути для мелочи.
Диагностика и замер прироста на практике
Прежде чем переписывать код, стоит убедиться, что AMX вообще включается в рабочий путь библиотеки. У oneDNN для этого есть подробный журнал: запуск приложения с переменной окружения DNNL_VERBOSE=1 печатает каждое примитивное ядро с типом движка, форматом данных и выбранной реализацией, и строки с amx мгновенно показывают, попали ли матричные умножения на плиточный путь или скатились в векторные ядра. Такой прогон занимает секунды и экономит часы гаданий.
Для честного замера прироста годится простой протокол. Фиксируется эталонная GEMM-задача с реальными размерами из приложения, прогрев выполняется несколько сотен повторений, замер берётся по среднему из большой серии, а привязка делается через taskset к стабильным ядрам одного NUMA-узла, чтобы планировщик не мешал. Сравнивать стоит три пути: базовый без расширений, AVX-512 через интринсик _mm512_dpbf16_ps и AMX. Разница между вторым и третьим и есть та кратная прибавка, ради которой всё затевалось.
Утилита perf здесь работает только по времени и счётчикам циклов: отдельного счётчика на TMUL в типовых конфигурациях нет, поэтому оценка сводится к сопоставлению секунд на задачу и доли системного времени. Если доля системного времени выросла, а полезная скорость нет, это сигнал о переключениях контекста с восьмикилобайтным состоянием, и лечится привязкой потоков либо снижением их числа.
Границы применимости и типичные ошибки
AMX не заменяет векторы, а дополняет их. Плиточный механизм заточен под умножение матриц с накоплением, и попытка натянуть на него поэлементные операции, свёртки малых ядер или нелинейности будет откатом в прошлое, для этого AVX-512 остаётся на месте и делает ту же работу быстрее. Второе ограничение в размере: восемь плиток и жёсткая геометрия покрывают подматрицы 16 на 16, 16 на 32 и похожие формы, крупные матрицы разбиваются на блоки, и здесь качество блокировки определяет, сколько времени уйдёт на память вместо арифметики. Третье, о чём забывают, это захват состояния при переключениях: на ядрах с тысячами короткоживущих AMX-потоков растёт доля системного времени на сохранение плиток, и здесь спасает либо привязка, либо уменьшение числа активных потоков.
Развитие расширения продолжается: в свежих поколениях Xeon, начиная с Granite Rapids, добавился формат fp16 с инструкцией _tile_dpfp16ps, расширяющий список применимых моделей. Палитра остаётся той же идеей, меняются форматы данных и доступные геометрии, так что код, аккуратно написанный на интринсиках с конфигурацией, переносится на новые процессоры с минимальными правками.
На практике внедрение AMX сводится к короткому чек-листу: убедиться в флагах на сервере, обновить библиотеки до версий с поддержкой, включить признак в гостях, если это виртуалка, и прогнать эталонный GEMM до и после. Время на всё это уходит меньше, чем на один тестовый запуск модели, а выигрыш в производительности узлов инференса измеряется кратно, и именно поэтому матричные расширения быстро стали обязательной частью серверных процессоров.