Попробуйте найти примеры кода для SME — я подожду

Часть 2 цикла о программировании Apple Scalable Matrix Extension (SME2) — от первых принципов до промышленной реализации GEMM. Часть 1 утверждала, что SME — это тензорное ядро, которое отрастил ваш процессор; эта часть — о том, как учиться его программировать и как мало помощи вы получите.

Оглавление-договор. В этой части мы проверим, есть ли у SME учебная дорога, сравнимая с той, которую получили тензорные ядра GPU; затем разберём два реально полезных источника: Arm Learning Path и KleidiAI; после этого отделим то, чему они действительно учат, от того, где они останавливаются. К концу главы станет видно, какая именно «середина лестницы» отсутствует и почему следующая часть неизбежно приводит к BLIS.

Если часть 1 права и SME действительно стоит рассматривать как маленькое тензорное ядро внутри CPU, следующий вопрос напрашивается сам: где дорога к нему? Вспомните, как тензорные ядра появились на GPU. Что бы вы ни думали о модели программирования, вы не оставались с ней один на один. cuBLAS использовал их с первого дня. CUTLASS прибыл как полностью разобранный сборник ответов с открытым исходным кодом — каждый приём упаковки, каждая стадия конвейера, разложенные по шаблонам, через которые можно пройти отладчиком. Были примеры в SDK, доклады на GTC, обсуждения на форумах, статьи в блогах, препарирующие раскладку фрагментов mma вплоть до отдельного потока. Оборудование было экзотическим; дорога к нему — автострадой.

Теперь проделайте то же упражнение для SME. Оборудование поставляется с 2024 года — в каждом Mac начиная с M4, в каждом iPhone начиная с A18 простаивает матричный движок. Попробуйте найти примеры кода. Не торопитесь.

Вот что находится после такой попытки. Это не риторическая фигура — это действительно весь список:

1. Со стороны Arm: SME Programmer’s Guide, спецификация ACLE, горстка статей в блогах сообщества и ровно одно сквозное собираемое учебное руководство — Learning Path под названием «Accelerate Matrix Multiplication Performance with SME2».

2. Промышленный код: KleidiAI, библиотека микроядер Arm для фреймворков машинного обучения. Настоящая, поставляемая, быстрая — и читаемая примерно как дамп SASS.

А Apple — компания, отгрузившая больше кремния с SME, чем все остальные вместе взятые? Официально — ничего. Ни примеров кода, ни руководства по программированию, ни абзаца. Санкционированный интерфейс — Accelerate: вызовите cblas_sgemm, получите производительность, вопросов не задавайте. Оборудование выдаёт себя, только если вы знаете, какой лист sysctl прочитать (hw.optional.arm.FEAT_SME2; флаги компилятора, которые понадобятся после находки, — в части 1). Позиция Apple, по существу, состоит в том, что матричный блок — деталь реализации их библиотек. Наша позиция, начиная с этого цикла, — что он для этого слишком интересен.

Итак, картина проста: два источника на противоположных концах лестницы абстракций — «hello world» внизу и готовый продукт наверху, и ничего между ними. Разберём их по порядку: оба научили меня настоящим вещам, и оба останавливаются ровно там, где начинаются интересные задачи.

Источник №1: Learning Path от Arm, или «мой первый FMOPA»

Learning Path (автор — Арно де Гранмезон) — действительно хорошее учебное руководство, и это звание оно зарабатывает честно: строит наивное умножение матриц на C, показывает, почему оно медленное, а затем проводит то же вычисление через рукописный ассемблер и через встроенные функции C (intrinsics), с проверочной обвязкой, сверяющей оба варианта с эталоном. Оно даже включает эмуляционную среду на Docker/FVP, позволяющую запускать код SME2 без владения оборудованием SME2. Если вы никогда не касались потокового режима — начинайте оттуда.

Три вещи в нём являются несущими для всего дальнейшего в этом цикле.

Аргумент об арифметической интенсивности. Руководство формулирует всё упражнение в терминах отношения числа умножений-накоплений (multiply-accumulate, MAC) к числу загрузок. Скалярный внутренний цикл выполняет одно умножение-накопление на две загрузки — отношение 1:2. Внешнее произведение переворачивает картину: загрузить один столбец A из 16 чисел и одну строку B из 16 чисел (32 числа), выполнить один FMOPA, получить 256 MAC. Отношение 8:1. Это правильная оптика — оптимизация GEMM насквозь является инженерией арифметической интенсивности — и та же логика модели roofline, которую носит с собой каждый автор CUDA-программ.

Упаковка, введённая без самого слова. Функция preprocess_l из руководства переставляет левую матрицу так, чтобы столбцы каждого тайла SVL×SVL лежали в памяти подряд, с дополнением нулями по краям. Подано это как разовая правка раскладки. На самом деле это ваш первый взгляд на упаковку — важнейшую структурную идею высокопроизводительного GEMM и предмет половины этого цикла. Руководство даже признаёт, во вежливой врезке, что «в промышленных условиях» раскладку можно организовать иначе. Читатель, мы потратим несколько частей на то, чтобы организовать раскладку иначе.

ZA как машина транспонирования. Версия того же преобразования на встроенных функциях делает нечто прелестное: загружает строки исходной матрицы в срезы ZA горизонтально (svwrite_hor_za32_f32_vg4), а вычитывает их обратно вертикально (svread_ver_za32_f32_vg4). Записываем строки, читаем столбцы — транспонирование выполняет сам тайл, никакая сеть перестановок не требуется. Руководство уделяет этому приёму два абзаца и идёт дальше. Отложите его в память: в части 7 он станет двигателем собственного упаковщика, который ускорил мой сквозной sgemm в 2 раза на входных данных со столбцовым порядком хранения (column-major).

Где руководство останавливается

Теперь честное очерчивание границ — потому что знание того, чего «hello world» не делает, указывает, где живёт настоящая работа.

Оно использует четверть машины. Цикл умножения накапливает в za0 — один тайл. При 512-битном SVL у Apple FP32-часть ZA состоит из четырёх тайлов 16×16. Почему это важно? Снова посчитаем загрузки. Один тайл: 32 числа на один FMOPA, 256 MAC — 8 MAC на загруженное число. Теперь используем все четыре тайла как макротайл 2×2: загружаем два вектора A и два вектора B (64 числа), выполняем четыре FMOPA — 1024 MAC, 16 на загруженное число. Удвоение площади тайла удваивает арифметическую интенсивность — вот почему каждое промышленное ядро SME, которое вам доведётся дизассемблировать (ядро KleidiAI ниже, микроядро BLIS, которое мы построим в части 4), имеет форму 2VL×2VL. Одиночный тайл руководства оставляет половину достижимой интенсивности неиспользованной — намеренно, ради ясности изложения.

Нет блокирования под кэш. Для каждого выходного тайла цикл по k проходит всю размерность K, протягивая полную столбцовую панель B мимо одного тайла C, — а затем повторяет это для следующего тайла. На задаче размера 125×70×35, принятой в руководстве по умолчанию, всё умещается в L1, и никто ничего не замечает. Увеличьте M, N, K до тысяч — и такая схема доступа перечитывает B из оперативной памяти по разу на каждую строковую панель: подсистема памяти незаметно становится главным действующим лицом. Раздел руководства об измерениях, надо отдать ему должное, содержит врезку, сводящуюся к тезису «измерения — дело трудное, эти цифры иллюстративны». Это правда — и одновременно предзнаменование.

Предобработка не амортизируется. matmul_intr переупаковывает левую матрицу при каждом вызове. Реальные пользователи умножают одну и ту же матрицу на многие другие или вызывают GEMM в цикле — промышленной библиотеке приходится решать, где живёт упаковка, кто за неё платит и как она перекрывается с вычислениями. Именно это решение — а не FMOPA — составляет трудную часть, и ни один пример размером с учебное руководство не способен её показать.

Ничто из этого не критика. Руководство делает ровно то, что и должно: учит простому. Но между учебным пособием и пригодным к использованию GEMM стоит структура, накопленная за десятилетия, и последняя страница Learning Path лишь машет в её сторону рукой и останавливается.

Источник №2: KleidiAI, или ответы в конце учебника

На другом конце лестницы находится KleidiAI — промышленная библиотека микроядер Arm, то, что действительно исполняется, когда llama.cpp или LiteRT задействует матричный блок на телефоне. Вот как выглядит пиковый код для SME. Изучать его — как изучать партию гроссмейстера: все ходы записаны, но не объяснены. Ответ виден весь, а путь к нему стёрт — хотя именно путь и был содержанием.

Начнём с имени. Вот микроядро GEMM для одинарной точности, прочтите вслух, быстро и с выражением:

kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa

Это плотная запись, но она поддаётся расшифровке, и расшифровка стоит труда, потому что имя и есть проект. Читаем справа налево и снаружи внутрь: mopa — построено на внешних произведениях. sme2 — уровень архитектуры. f32p2vlx1biasf32 — правая матрица приходит предварительно упакованной (p) панелями шириной в две длины вектора, со строкой смещения (bias) в формате f32, вшитой в упакованный буфер. f32p2vlx1 — левая матрица также приходит упакованной. matmul_clamp_f32 — вычисляется матричное произведение с вплавленным в эпилог ограничением значений (clamp) по минимуму и максимуму. А форма тайла подтверждает то, что предсказала арифметика подсчёта загрузок: это ядро 2VL×2VL — все четыре FP32-тайла ZA, полный аккумулятор 32×32, максимальная арифметическая интенсивность. (Когда я выбрал ту же форму 2VL×2VL для микроядра BLIS в части 4, это имя послужило подтверждением: гроссмейстеры встали на ту же клетку.)

В этой библиотеке скрыта настоящая учебная программа, если читать структуру каталогов как кофейную гущу. Процедуры упаковки — полноправные граждане: каждое микроядро умножения сопровождается собственными спутниками pack, потому что промышленный код никогда не перемножает неупакованные матрицы. Есть даже семейство ядер 1x16vl_sme2_mla: для умножения матрицы на вектор они полностью отказываются от внешнего произведения и используют обычные потоковые SVE-инструкции слитого умножения-сложения — напоминание о том, что FMOPA есть игра на амортизацию, а одной выходной строке амортизировать нечего.

Где KleidiAI перестаёт помогать

Невозможно прочитать «как». Ядра написаны на ассемблере вручную, и значительная часть — даже не мнемоники, а шестнадцатеричные коды:

KAI_ASM_INST(0xa0840100)  // smopa za0.s, p0/m, p0/m, z8.b, z4.bKAI_ASM_INST(0xc00800ff)  // zero {za}

Сырые машинные коды с инструкцией в комментарии — потому что поставляемые ассемблеры отстают от архитектуры. Это в точности опыт чтения дампа SASS: можно проверить, что делает ядро, инструкция за инструкцией, но проектные соображения, породившие такое расписание — почему эти загрузки здесь, почему такая глубина развёртки, — невидимы. CUTLASS для SME не существует: нет открытого, шаблонного, прокомментированного среднего слоя, который учит, пока работает. Этой ступеньки лестницы попросту нет.

Это не GEMM. Посмотрите на семантику: matmul_clamp — то есть C = clamp(A·B + bias). Ни альфы, ни беты, ни накопления в существующую матрицу C, ни вариантов транспонирования, ни столбцового порядка хранения, ни комплексных чисел. И просмотрите каталог ещё раз: из 36 семейств умножения подавляющее большинство — квантованные: qai8dxp, qsi4c32p, четырёхбитные веса с поблочными масштабами — весь бестиарий инференса больших языковых моделей. KleidiAI отвечает на вопрос, который действительно задают клиенты Arm: «сделайте инференс трансформеров быстрым на телефонах». Это библиотека операторов для фреймворков машинного обучения, а не библиотека линейной алгебры. Если ваша матрица хранится по столбцам, или она комплексная, или вам нужно C = αAB + βC — контракт, который BLAS соблюдает с 1979 года, — здесь для вас ничего нет, а код слишком непрозрачен и слишком специализирован, чтобы его расширять.

Недостающая средняя ступенька

Таково состояние мира. Внизу — первоклассный «hello world», обучающий атому: один тайл, без блокирования, упаковка на правах примечания — и честно сообщающий, что на этом он заканчивается. Наверху — поставляемые ядра, использующие машину целиком: каждый тайл, дисциплинированная упаковка, вплавленные эпилоги — написанные в шестнадцатеричных кодах, ограниченные задачами инференса, не обучающие ничему. А между ними — там, где GEMM собственно строится: где решается, как тайлы складываются в панели, панели в кэш-блоки, блоки в потоки; где упаковка находит своё место в гнезде циклов; где альфа, бета, оба порядка хранения и четыре типа данных обслуживаются без сорока дублирующих ядер, — нет ничего. Ни одно руководство не тянется вверх; ни один продукт не тянется вниз. Два конца — две абстракции: внизу атом без машины, наверху машина без объяснения. Ни один не опосредует другого, оттого оба и остаются полуправдой; истина сидит в отсутствующей середине, которая их связывает.

Но вот в чём дело: эта средняя ступенька существует уже тридцать лет. Просто её строили не для SME. Есть корпус работ — статьи Гото и выросший из них фреймворк BLIS, — решивший ровно эту задачу: как структурировать GEMM, чтобы вся независимая от архитектуры машинерия была уже построена, а добавление новой машины сводилось к написанию одного микроядра и заполнению таблицы размеров блоков. Атом из Learning Path защёлкивается в гнездо, выточенное за десятилетия до появления этого оборудования.

Вывод. SME вышла в свет без дороги к ней. Вся публичная учебная программа — одно отличное руководство на один тайл и одна нечитаемая промышленная библиотека, а лестница между ними отсутствует: стратегия упаковки, блокирование под кэш, полный контракт BLAS. Но недостающие ступеньки оказываются самой старой и лучше всего изученной частью высокопроизводительных вычислений; они просто ждали, пока оба конца лестницы не понадобятся друг другу.

Далее: фреймворк, который уже решил середину лестницы. В части 3 я доказываю, что BLIS — одно из великих произведений системного программирования, и показываю, как его архитектура из пяти циклов сводит задачу «перенести GEMM на новый матричный движок» к упражнению «заполните пропуски». Пропуски всё равно интересны. Им и посвящён остаток цикла — sgemm, dgemm, комплексные ядра и собственные упаковщики.