Часть 1 цикла о программировании Apple Scalable Matrix Extension (SME2) — от первых принципов до промышленной реализации GEMM.
Оглавление-договор. В этой части мы сначала разберёмся, почему отсутствие SVE на Apple Silicon не обязательно означает архитектурную отсталость; затем введём главную операцию SME — FMOPA; после этого сопоставим её с тензорными ядрами GPU и, наконец, зафиксируем цену этой конструкции: потоковый режим, отсутствие непотокового SVE и несколько практических ловушек для компилятора. Следующая часть будет о том, почему такой мощный механизм почти не сопровождается хорошими примерами.
Начну, пожалуй, с претензии.
Я зарабатываю на жизнь написанием векторного кода, и я люблю это занятие. Меня искренне радуют масштабируемые длины векторов. Поэтому, когда ARM выпустила SVE, по-настоящему современную векторную архитектуру, не зависящую от длины вектора, мне хотелось видеть её повсюду. А Apple, создатель самого передового кремния на планете, выпускала чип за чипом с… NEON. С фиксированным 128-битным NEON. Той же ширины, что и десять лет назад. Без SVE. На кристалле M-серии, способном трассировать лучи и исполнять трансформеры, универсальный векторный блок был — и остаётся — демонстративно архаичным.
Меня это, признаться, задевало.
А потом я перестал беспокоиться и просветлился, потому что понял, что Apple поступила очень в духе Apple. Они пропустили современную векторную архитектуру не потому, что отстали. Они пропустили её, потому что уже прыгнули в будущее. Пока остальная экосистема ARM (включая вашего покорного слугу) спорила о длинах векторов, Apple тихо возвела следующую ступеньку лестницы прогресса. Они дали нам не улучшенные векторы. Они дали нам матрицы.
Посчитаем размерности
Чтобы увидеть, почему это не просто красивая формулировка, нужно изменить вопрос. Не «какой ширины вектор?», а «какой формы сама операция?». Вот самый ясный из известных мне способов объяснить, что изменилось. Возьмём фундаментальную операцию совмещённого умножения-сложения (fused multiply-add, FMA) и спросим: в скольких измерениях эта операция на самом деле работает?
Скалярное FMA — ранг 0. Операция происходит “в точке”. Перемножить два числа, прибавить третье, получить результат: c = a*b + c. Ноль измерений. Это атом.
Скалярное произведение векторов — ранг 1. Перемножить два вектора поэлементно и просуммировать элементы в скаляр: s = Σ a[i]*b[i]. Теперь операция простирается в одном измерении — вдоль оси, по которой выполняется редукция. Это рабочая лошадка любого рукописного SIMD-умножения матриц, и одновременно источник вечной возни: свёртывание вектора в скаляр означает горизонтальные редукции, сети перестановок и элементы вектора, вечно занятые обменом данными друг с другом.
Внешнее произведение векторов — ранг 2. Взять вектор-столбец a длины VL и вектор-строку b длины VL, построить их полное внешнее произведение — целую плоскость результатов размером VL × VL — и накопить её:
ZA[i][j] += a[i] * b[j] для всех i, j
Два вектора на входе, матрица на выходе. Одна инструкция, охватывающая сразу два измерения. (Да, внешнее произведение — это матрица ранга 1 в линейно-алгебраическом смысле; но как операция это первая из перечисленных, строящая двумерный результат, и в этом скачке вся суть.) В реализации Apple VL составляет 16 чисел одинарной точности, так что одна операция внешнего произведения выполняет блок 16×16 = 256 совмещённых умножений-сложений за один приём.
Ранг 0, ранг 1, ранг 2. Apple не стала возиться с векторами. Она добавила следующее измерение. Прибавили одно измерение — арифметически пустяк; но это тот самый узел, где накопленное количество опрокидывается в новое качество: не ускоренные векторы, а другой род операции.
Добавив всего одно слово — FMOPA (не рифмовать!). Эта аббревиатура означает Fused Matrix Outer Product Accumulate, то есть совмещённая операция внешнего произведения с накоплением.
Похоже, что у процессора выросло маленькое тензорное ядро
Если эта история кажется знакомой, так и должно быть: ту самую верхнюю ступеньку вы уже видели, на совершенно другом куске кремния.
Каждый программист GPU носит в голове одну и ту же картину: где-то на кристалле находится плотный блок, перемножающий маленькие матрицы и накапливающий результат. NVIDIA называет такие блоки тензорными ядрами. Управляющие ими инструкции — wmma, mma, wgmma — берут два операнда-плитки и текущий аккумулятор и вкладывают в него матричное произведение за один шаг. Никто не думает о тензорном ядре как о цикле умножений-сложений. О нём думают как об одной матричной операции.
FMOPA — это та самая операция, а ZA — тот самый аккумулятор, и теперь они находятся внутри центрального процессора. Это действительно выглядит так, будто процессору позволили отрастить маленькое тензорное ядро — причём такое, до которого можно дотянуться из обычного компилируемого кода, без запуска вычислительных ядер и без границы между хостом и устройством. Вы вызываете функцию; внешние произведения происходят; результат лежит в вашем массиве.
Формулировка следует немедленно, и она та же, что и в любой программе для тензорных ядер. Матричное произведение C = A·B есть не что иное, как сумма внешних произведений, по одному на каждый элемент общей размерности K:
for k in 0..K: ZA += A[:, k] ⊗ B[k, :] // один FMOPA на каждый k
Этот цикл и есть GEMM. Подавайте по одному столбцу A и одной строке B за шаг, выполняйте один FMOPA — и редукция по K происходит даром, внутри ZA. Никаких горизонтальных редукций. Измерение редукции свёртывается вдоль цикла, а не вдоль элементов вектора — именно поэтому перед нами другое упражнение, а не ускоренная версия SIMD-умножения матриц.
ZA — это настоящий архитектурный регистровый файл: квадратный массив размером SVLb × SVLb байт, где SVLb — потоковая длина вектора в байтах, разделённый на нумерованные «плитки» — тайлы (tiles). При 512-битной потоковой длине вектора у Apple это массив 64×64 байта — четыре тайла 16×16 для FP32 (za0–za3), либо восемь тайлов для FP64, либо шестнадцать тайлов со 128-битными элементами. Значительная часть этого цикла статей посвящена решению, на что потратить эти тайлы, — ровно так же, как автор CUTLASS решает, на что потратить фрагменты тензорных ядер. Та же задача, другой чип.
Претензия не снята: потоковый режим — лишь половина ответа
Теперь можно вернуться к исходной претензии, но уже без прежней простоты. Она не опровергнута и не отброшена — она возвращается определённее: вопрос теперь не «почему Apple отстала», а «какой ценой оплачен этот прыжок». Если Apple действительно дала нам матрицы, почему же повседневный векторный блок всё ещё выглядит как старый NEON? Первая половина ответа — потоковый режим.
SME вводит особый режим процессора — потоковый режим (streaming mode), — в который входят инструкцией smstart и выходят инструкцией smstop. Внутри него доступен ZA и потоковая разновидность SVE с длиной вектора SVL, в словаре ARM она так и называется: Streaming SVE. Само имя обещает иерархию: сначала SVE, поверх него — потоковая надстройка. Apple эту иерархию проигнорировала и построила кривую систему: Streaming SVE есть, а самого SVE — нет. Вне потокового режима векторный блок — обычный 128-битный NEON. Но Apple не просто придержала SVE ради экономии транзисторов: она вложила идею масштабируемых векторов в матричный сопроцессор, а повседневный векторный путь оставила за NEON. Главное действующее лицо — матричный движок; векторы лишь подают ему данные. NEON никуда не делся, он понижен до служебной роли: сохранён, но лишён самостоятельности как момент чужого движения.
Вторая половина ответа чуть менее радостна и уже не философская, а практическая. Пропустив ступеньку, Apple оставила в лестнице дыру, и я в неё угодил — на первой же сборке. Укажите «очевидный» флаг -march=armv9-a+sme2 — и компилятор решит, что непотоковый SVE существует, и вставит обычную SVE-инструкцию cntd в пролог вашей матричной функции. На ядре Apple эта инструкция недоступна до входа в собственно потоковый режим, ведь не-потокового SVE нет, и программа завершается с ошибкой недопустимой инструкции в момент запуска. Лечение — перестать обещать компилятору оборудование, которого Apple никогда не строила:
# Неправильно на Apple: подразумевает непотоковый SVE, порождает cntd -> SIGILL при выполнении-march=armv9-a+sme2# Правильно: базовый armv8, только SME2, непотоковый SVE не обещан-march=armv8.7-a+sme2
Я обычно концентрируюсь на производительности одного ядра, но промышленные системы задействуют и другие механизмы параллелизма, поэтому ещё один архитектурный факт стоит усвоить уже сейчас: ZA и потоковый блок — разделяемые ресурсы; у Apple — один на кластер производительных ядер, а не на каждое ядро. Матричный движок — общий, поэтому подход «просто добавь потоков» насыщается на числе кластеров — точно так же, как «просто добавь блоков» насыщает фиксированный пул тензорных ядер. Интуиция, принесённая с GPU, продолжает окупаться.
Остаток цикла статей посвящён тому, чтобы научиться использовать потоковый сопроцессор на всю катушку. И по-настоящему трудное здесь, подчеркну заранее, не умножение. Четырёхстрочный цикл с FMOPA, приведённый выше, напишет кто угодно. Превратить его в код, обгоняющий библиотеку производителя на любой форме матриц, любой точности и любой раскладке памяти, — упражнение во всю глубину стека: упаковка данных, блокирование под кэш (cache blocking) и дотошные измерения. Этим мы и займёмся, и именно там живут хорошие истории.
Вывод. Apple не стала модернизировать векторы, она перепрыгнула через ступеньку и выпустила матрицы. FMOPA — операция ранга 2, накопление внешнего произведения в массиве ZA, и во всех практических смыслах это тензорное ядро, программируемое на C. Цена входного билета — отказ от предположения, что перед нами «SVE с дополнениями»: у Apple всё это поставляется вовсе без SVE, через сопроцессор с потоковым режимом, и этот единственный факт помогает понять всё дальнейшее изложение.
Далее: набор инструкций такой мощности должен сопровождаться целой кучей примеров кода. Не сопровождается. В части 2 я отправляюсь на поиски и нахожу ровно два источника, две составные части.
ссылка на оригинал статьи https://habr.com/ru/articles/1063752/