Часть 4 цикла о программировании Apple Scalable Matrix Extension (SME2). Часть 3 доказывала, что BLIS сводит задачу «перенести GEMM на новую машину» к написанию одного микроядра и заполнению пяти чисел размеров блоков. Эта часть заполняет самый важный пропуск: микроядро одинарной точности, разобранное строка за строкой.
Оглавление-договор. В этой части мы возьмём первый из пропусков BLIS — микроядро sgemm; сначала зафиксируем форму аккумулятора 32×32 в четырёх тайлах ZA, затем разберём горячий цикл по K, потом эпилог с alpha, beta и ловушкой NaN, а в конце посмотрим, какие числа производительности даёт именно это ядро. Следующая часть оставит структуру почти той же, но заменит геометрию тайла и условия существования ядра.
Всё предыдущее было разбегом; здесь наконец взлетаем. Микроядро одинарной точности — тот код, который во внутреннем цикле действительно касается матричного движка, и он достаточно мал — около 120 строк содержательной работы, — чтобы прочитать его целиком и понять каждое решение. К концу этой части вы будете точно знать, что именно исполняется от 1357 до 1450 миллиардов раз в секунду на производительном ядре M5, и почему оно устроено именно так.
Вспомним контракт из части 3. BLIS уже упаковал высокую узкую панель A и широкую узкую панель B в непрерывные, дополненные нулями буферы. Микроядру он передаёт эти две панели, длину K, альфу, бету и указатель на один тайл MR×NR выходной матрицы C. Наша работа: вычислить C := alpha·(A·B) + beta·C для этого одного тайла и вернуть управление. BLIS вызывает нас по разу на тайл и делает буквально всё остальное.
Для этого ядра MR = NR = 2VL; при 512-битной потоковой длине вектора у Apple VL составляет 16 чисел одинарной точности, так что тайл имеет размер 32×32. Запомните это число: в нём вся причина быстроты ядра.
Где живёт тайл 32×32: все четыре тайла ZA сразу
Часть 1 представила ZA как квадратный байтовый массив, разделённый на нумерованные тайлы. При SVL 512 FP32-часть ZA — четыре тайла 16×16, от za0 до za3. Один FMOPA обновляет один из них — блок 16×16 (VL×VL). Нам нужен аккумулятор 32×32 — четыре таких блока, — поэтому мы используем все четыре тайла ZA как один макротайл 2×2:
n: 0..VL-1 n: VL..2VL-1 m: 0..VL-1 za0 za2 m: VL..2VL-1 za1 za3
Это выигрыш от подсчёта загрузок из части 2. Один FMOPA загружает 32 числа (VL-столбец A, VL-строку B) и выполняет 256 умножений-сложений — 8 на загруженное число. Но чтобы заполнить макротайл 2×2, мы загружаем два VL-столбца A и две VL-строки B — 64 числа — и выполняем четыре FMOPA, то есть 1024 умножения-сложения. Это 16 умножений-накоплений на загруженное число — вдвое больше арифметической интенсивности однотайлового ядра из учебного руководства. Тайл 32×32 не произволен: это наименьший тайл, насыщающий FP32-массив ZA, а насыщение ZA отодвигает загрузки операндов достаточно далеко на задний план, чтобы узким местом стал блок SME, а не подсистема памяти.
Четыре FMOPA, строящие макротайл из одной пары векторов A (a0, a1) и одной пары векторов B (b0, b1), — попросту четыре блока внешнего произведения:
za0 += a0 ⊗ b0 za2 += a0 ⊗ b1 za1 += a1 ⊗ b0 za3 += a1 ⊗ b1
Каждая строка A встречает каждый столбец B; это обновление ранга 1 размером 32×32 — один шаг редукции по k. Четыре тайла при этом — не четыре отдельных ответа, а моменты одного целого: пока идёт редукция, макротайл живёт как единое C и распадается на четвёрку лишь при записи.
Цикл редукции: расшит по два, загружает по четыре
Теперь, когда форма аккумулятора ясна, можно смотреть на код без ощущения, что перед нами случайный набор встроенных функций. Вот горячий цикл дословно — код, исполняющийся KC раз на каждый тайл и миллионы раз на каждый вызов GEMM:
for ( ; l + 1 < k; l += 2 ) { svfloat32x4_t av = svld1_x4( pc, a ); svfloat32x4_t bv = svld1_x4( pc, b ); svmopa_za32_f32_m( 0, pall, pall, svget4( av, 0 ), svget4( bv, 0 ) ); svmopa_za32_f32_m( 1, pall, pall, svget4( av, 1 ), svget4( bv, 0 ) ); svmopa_za32_f32_m( 2, pall, pall, svget4( av, 0 ), svget4( bv, 1 ) ); svmopa_za32_f32_m( 3, pall, pall, svget4( av, 1 ), svget4( bv, 1 ) ); svmopa_za32_f32_m( 0, pall, pall, svget4( av, 2 ), svget4( bv, 2 ) ); svmopa_za32_f32_m( 1, pall, pall, svget4( av, 3 ), svget4( bv, 2 ) ); svmopa_za32_f32_m( 2, pall, pall, svget4( av, 2 ), svget4( bv, 3 ) ); svmopa_za32_f32_m( 3, pall, pall, svget4( av, 3 ), svget4( bv, 3 ) ); a += 4*vl; b += 4*vl; }
Здесь стоит остановиться на трёх вещах.
Одна мультивекторная загрузка на операнд — два шага по k. svld1_x4 — мультивекторная загрузка SME2: одна инструкция заполняет четыре последовательных векторных регистра. Поскольку упакованная панель хранит шаг l и шаг l+1 вплотную друг к другу, один svld1_x4 по A захватывает оба столбца A для обоих шагов (av0, av1 для шага l; av2, av3 для шага l+1), и то же справедливо для B. Таким образом, цикл расшит по два вдоль K, и каждый операнд обходится ровно в одну инструкцию загрузки на два шага. Это важнее, чем кажется: на коротких K меньшее число более широких загрузок измеримо ускоряет ядро, потому что механизм выборки инструкций не тонет в потоке одиночных загрузок. Первый блок из четырёх FMOPA потребляет шаг l; второй — шаг l+1; затем оба указателя продвигаются на 4·VL (два шага × два VL-вектора), и всё повторяется.
Предикаты сплошь истинные — намеренно. Каждый FMOPA маскирован предикатом pall = svptrue_b32(), истинным для каждого элемента. Никаких проверок краёв во внутреннем цикле. Это дивиденд упаковки (часть 7): упакованные панели BLIS дополнены нулями до полного тайла 32×32, поэтому даже когда настоящая матрица имеет размер 30×17, ядро перемножает полные 32×32, по большей части нулей. Дополнение вносит в аккумулятор ровно ноль и никогда не сохраняется, так что оно безвредно, — а взамен цикл редукции лишён ветвлений, работает с постоянным предикатом и тривиально планируется. Мы платим немного лишней арифметики на нулях, чтобы горячий цикл был идеально регулярным. Для ядра, ограниченного вычислениями, это прекрасная сделка.
ZA пришёл заранее обнулённым. Перед циклом нет вызова svzero_za(). Функция объявлена с атрибутом __arm_new("za"), и ARM ACLE гарантирует, что этот атрибут обнуляет состояние ZA при входе. Обнуление вручную породило бы вторую, избыточную инструкцию zero {za} — Apple clang действительно порождает две, если написать это явно; так я это и заметил. Поэтому мы просто начинаем накапливать.
После основного цикла скалярный хвост обрабатывает нечётный последний шаг по k (через svld1_x2 и четыре FMOPA), потому что K не всегда чётно. Та же операция, половинной ширины.
Читаем ответ: альфа, бета и ловушка с NaN
Когда редукция по K завершена, произведение 32×32 находится в za0–za3. Теперь применяем уравнение BLAS и записываем результат в C. Этого не делает ни одно из двух ядер, рассмотренных в части 2, и именно здесь микроядро зарабатывает звание «GEMM».
Мы читаем ZA по одному выходному столбцу, вертикально. Для выходного столбца j значения находятся в za0/za1, если j < VL (левая половина макротайла), или в za2/za3, если j ≥ VL (правая половина), с разделением на верхние VL строк и нижние VL строк:
if ( (uint64_t)j < vl ) { ab0 = svread_ver_za32_f32_m( svundef_f32(), pm0, 0, (uint32_t)j ); ab1 = svread_ver_za32_f32_m( svundef_f32(), pm1, 1, (uint32_t)j ); } else { ab0 = svread_ver_za32_f32_m( svundef_f32(), pm0, 2, (uint32_t)(j - vl) ); ab1 = svread_ver_za32_f32_m( svundef_f32(), pm1, 3, (uint32_t)(j - vl) ); }
pm0 и pm1 — строковые предикаты, маски svwhilelt, покрывающие строки 0..m и VL..m, — так что тайл, чья настоящая высота меньше 32, читает и записывает только существующие строки. Столбцовый цикл просто перебирает j от 0 до настоящего n; короткие тайлы никогда не касаются дополненных столбцов. Вот и вся история обработки краёв при записи: предикатные чтения по оси строк, укороченный цикл по оси столбцов. Никаких дополнительных ядер для подчистки.
Затем само обновление — и тонкость, ради которой стоит весь этот раздел. Уравнение BLAS — C := alpha·AB + beta·C, но в нём скрывается особый случай. Когда beta == 0, уравнение превращается в C := alpha·AB: матрица C перезаписывается, и её старое содержимое читать нельзя. Почему это существенно? Потому что свежевыделенная выходная матрица может содержать мусор, включая значения NaN, а 0·NaN равно NaN, а не нулю. Если добросовестно вычислить alpha·AB + beta·C при beta = 0 над неинициализированной C, случайный NaN испортит результат. Стандарт BLAS прямо требует, чтобы beta == 0 означало «вообще не заглядывать в C». Поэтому ядро ветвится:
if ( *beta == 0.0f ) { // перезапись: C не читается никогда svst1_f32( pm0, cj, svmul_f32_x( pm0, ab0, valpha ) ); svst1_f32( pm1, cj + vl, svmul_f32_x( pm1, ab1, valpha ) ); } else { // накопление: читаем C, вплавляем бету svfloat32_t c0 = svld1_f32( pm0, cj ); svfloat32_t c1 = svld1_f32( pm1, cj + vl ); c0 = svmla_f32_x( pm0, svmul_f32_x( pm0, ab0, valpha ), c0, vbeta ); c1 = svmla_f32_x( pm1, svmul_f32_x( pm1, ab1, valpha ), c1, vbeta ); svst1_f32( pm0, cj, c0 ); svst1_f32( pm1, cj + vl, c1 ); }
Ветвь beta == 0 умножает содержимое ZA на альфу и сохраняет — C не загружается никогда. Общая ветвь загружает C, вычисляет alpha·AB + beta·C умножением и слитым умножением-сложением и сохраняет результат. Ветвление небольшое, но именно оно отделяет корректный BLAS от неуловимо сломанного — и это ровно та деталь, о которой ядрам из учебных руководств думать не приходилось: они всегда вычисляли только C = A·B. Ноль здесь — не пустое ничто: 0·NaN не гасит мусор, а разносит его. Поэтому обнуление C — определённое отрицание, активная перезапись, а не формальное умножение на ноль.
Приём со стековым кадром: как сохранить быстрый путь крошечным
Одна структурная деталь — приятный образец инженерии, отладка которого стоила настоящих усилий. Микроядро, вызываемое BLIS, bli_sgemm_armsme_2vlx2vl, — тонкий диспетчер:
if ( rs_c == 1 ) { // C по столбцам: быстрый путь bli_sgemm_armsme_2vlx2vl_body( m, n, k, alpha, a, b, beta, c, cs_c ); return; } bli_sgemm_armsme_2vlx2vl_ct( ... ); // произвольный шаг по строкам: медленный путь
Потоковая функция _body требует единичного шага по строкам (C хранится по столбцам) — это подавляюще частый случай, и именно для него наше ядро объявляет предпочтение (часть 3). Когда у C произвольный шаг по строкам, записывать в неё напрямую нельзя, и механизм временных микротайлов BLIS проводит результат через промежуточный буфер. Но этот буфер занимает 16 КиБ, и стоит функции его упомянуть, как компилятор резервирует под него стек и вставляет проверку стека (stack probe) — на каждом вызове, включая быстрый путь, если оба пути живут в одной функции. Поэтому медленный путь изолирован в отдельной вспомогательной функции с атрибутом __attribute__((noinline)). Стековый кадр быстрого пути остаётся крошечным — ни проверки, ни напрасного резерва, — и вызов со столбцовым порядком, составляющий 99% случаев, не платит ничего за существование редкого. Такие оптимизации находят, только разглядывая дизассемблированный код и недоумевая, откуда у горячей функции проверка стека, которой там быть не должно.
Что это даёт в числах
Один и тот же двоичный файл, один поток, M5 Pro, n = 2000: это ядро разгоняет sgemm до 1357–1450 GFLOPS против примерно 141 у базовой реализации на NEON. Ускорение примерно в 9,6 раза — от одного микроядра и таблицы размеров блоков, при том что каждую прочую строку библиотеки поставил BLIS. Пиковая производительность самого ядра, измеренная изолированно при больших K, — около 2130 GFLOPS; в этой точке стеной становится сам блок SME. Разрыв между 1450 и 2130 — упаковка и эффекты кэша, предмет настройки размеров блоков, к которой я ещё вернусь.
То, что KC стремится быть огромным (2048), прямо следует из сказанного. Аккумулятор 32×32 целиком живёт в ZA и не выгружается в память, сколь бы долго ни длилась редукция по K. Значит, единственное ограничение KC — сколько упакованных данных A и B удаётся удержать в кэше, а вовсе не нехватка регистров, обычный ограничитель. Делайте KC настолько большим, насколько позволяет кэш, — и фиксированные издержки (вход в потоковый режим, чтение ZA в конце) амортизируются на всё более длинной редукции. Гигантский регистровый тайл, который даёт SME, требует гигантского блока по K себе на прокорм, и таблица размеров блоков BLIS — ровно то место, где об этом сообщают.
Вывод. Микроядро sgemm — аккумулятор 32×32, распределённый по всем четырём FP32-тайлам ZA; его питает цикл по K, расшитый по два и загружающий данные мультивекторными инструкциями SME2, с горячим путём на сплошь истинных предикатах, возможным благодаря упаковке с дополнением нулями. Эпилог читает ZA столбец за столбцом, применяет альфу и бету и тщательно избегает прикосновения к C при beta == 0, чтобы значения NaN из свежей матрицы не просочились в результат. Всё это, плюс пять чисел, — GEMM одинарной точности с ускорением в 9,6 раза.
Далее: двойная точность — и сюрприз в том, что нового кода почти нет. В части 5 мы увидим, что меняется, когда числа одинарной точности становятся числами двойной: тайл другой формы, восемь тайлов ZA вместо четырёх и один бит возможностей (FEAT_SME_F64F64), проверяемый во время выполнения и решающий, существует ли ядро двойной точности вообще, — и что остаётся ровно тем же.
