Часть 5 цикла о программировании Apple Scalable Matrix Extension (SME2). Часть 4 разобрала микроядро одинарной точности по винтикам. Эта часть спрашивает, что меняется при переходе к двойной точности, — и интересный ответ таков: в коде меньше, чем вы думаете, а в одном архитектурном решении, принимаемом до исполнения какого-либо кода, — больше.

Оглавление-договор. В этой части мы используем sgemm как шаблон и проверим, что действительно меняется при переходе к dgemm: сначала пройдём скучную часть с удвоенными ширинами, затем разберём прямоугольный тайл 16×32, после этого — отдельный бит возможностей FEAT_SME_F64F64 и динамическую регистрацию ядра. Следующая часть будет устроена иначе: там меняется не только ширина элемента, но и сама форма данных.

Если вы только что прочли часть 4, большая часть ядра двойной точности вызовет у вас дежавю — и именно это я хочу подчеркнуть в первую очередь. Весь довод этого цикла состоит в том, что SME встраивается в BLIS как одно маленькое ядро плюс таблица чисел. Лучшее свидетельство того, насколько это ядро маленькое: переход от sgemm к dgemm — каждый элемент вдвое шире — почти не затрагивает структуру. Та же форма цикла редукции, те же мультивекторные загрузки, тот же предикатный постолбцовый эпилог, та же забота о NaN при beta == 0, тот же приём со стековым кадром для быстрого и медленного путей. Замените типы, заново выведите геометрию тайла — и вы почти у цели.

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

Скучная часть: двойная точность — это просто более широкие числа

На уровне инструкций внешние произведения двойной точности — это ядро одинарной точности с удвоенными ширинами. FMOPA накапливает в тайлы ZA64 вместо ZA32; встроенная функция — svmopa_za64_f64_m вместо svmopa_za32_f32_m; счётчик длины вектора — svcntd() (чисел двойной точности на вектор) вместо svcntw() (слов на вектор), что при SVL 512 даёт 8 вместо 16. Везде, где часть 4 говорила f32, читайте f64; везде, где b32, — b64. Цикл редукции всё так же расшит по два вдоль K, всё так же загружает каждый операнд одним svld1_x4 на два шага, всё так же работает на сплошь истинных предикатах над дополненными нулями упакованными панелями, всё так же начинает с аккумулятора, обнулённого атрибутом __arm_new("za"). Эпилог всё так же читает ZA вертикально, применяет альфу и ветвится на beta == 0, чтобы не читать неинициализированную матрицу C. Если вы поняли ядро одинарной точности, вы уже понимаете 90% ядра двойной. Это ровно та переиспользуемость, которую обещал BLIS: выстраданная структура не зависит от точности; меняются лишь ширины арифметики.

Посвятим же остаток части тем 10%, которые действительно другие.

Сюрприз первый: тайл имеет размер 16×32, а не 16×16

Вот что подводит интуицию. При SVL 512 один тайл ZA64 — это 8×8 чисел двойной точности (VL = 8). Ядро одинарной точности использовало четыре тайла ZA32 как макротайл 2×2. Симметрия подсказывает, что ядро двойной точности должно поступить так же: четыре тайла ZA64, макротайл 8×8 → 16×16. Но у FP64-части ZA доступно восемь тайлов, а не четыре. Вот правило, которое этим управляет: ZA — фиксированный массив размером SVLb×SVLb байт, где SVLb — потоковая длина вектора в байтах, адресуемый как набор нумерованных тайлов, чьё количество равно ширине элемента в байтах: четыре тайла 16×16 при одинарной точности (4-байтовые элементы), восемь тайлов 8×8 при двойной (8-байтовые), шестнадцать при 128-битных элементах. Удвоение ширины элемента не сокращает бюджет тайлов — оно удваивает число тайлов (меньшего размера). «Просто более широкие числа» по видимости чистое количество; но здесь пройден порог, за которым количество меняет качество — другое число тайлов, другая форма макротайла, другая геометрия эпилога. Использовать лишь четыре из восьми означало бы оставить половину аккумулятора простаивать, а простаивающий аккумулятор — это арифметическая интенсивность, оставленная неиспользованной: единственный грех, от которого предостерегала часть 2.

Поэтому ядро двойной точности использует все восемь тайлов ZA64, организованных как макротайл 2×4 — 2VL в высоту и 4VL в ширину, что при VL = 8 даёт выходной тайл 16×32:

            n: 0..VL-1  VL..2VL-1  2VL..3VL-1  3VL..4VL-1
  m: 0..VL-1    za0        za2        za4         za6
  m: VL..2VL-1  za1        za3        za5         za7

Имя ядра, bli_dgemm_armsme_2vlx4vl, кодирует ровно это: MR = 2VL = 16, NR = 4VL = 32. Асимметрия — не вопрос эстетики; это форма, заполняющая все восемь тайлов, а их заполнение и удерживает блок двойной точности накормленным. Правда, она меняет баланс загрузок: для построения макротайла 2×4 нужны два VL-столбца A, но четыре VL-строки B, поэтому каждый шаг по k выполняет восемь FMOPA (два блока A × четыре блока B), а данных B загружается вдвое больше, чем данных A. Вот один шаг основного цикла:

svfloat64x4_t av  = svld1_x4( pc, a );        // 2 столбца A на 2 шага по k
svfloat64x4_t bv0 = svld1_x4( pc, b );        // первые 4 строки B...
svfloat64x4_t bv1 = svld1_x4( pc, b + 4*vl ); // ...для B нужны две загрузки x4

svmopa_za64_f64_m( 0, pall, pall, svget4( av, 0 ), svget4( bv0, 0 ) );
svmopa_za64_f64_m( 1, pall, pall, svget4( av, 1 ), svget4( bv0, 0 ) );
svmopa_za64_f64_m( 2, pall, pall, svget4( av, 0 ), svget4( bv0, 1 ) );
svmopa_za64_f64_m( 3, pall, pall, svget4( av, 1 ), svget4( bv0, 1 ) );
svmopa_za64_f64_m( 4, pall, pall, svget4( av, 0 ), svget4( bv0, 2 ) );
svmopa_za64_f64_m( 5, pall, pall, svget4( av, 1 ), svget4( bv0, 2 ) );
svmopa_za64_f64_m( 6, pall, pall, svget4( av, 0 ), svget4( bv0, 3 ) );
svmopa_za64_f64_m( 7, pall, pall, svget4( av, 1 ), svget4( bv0, 3 ) );
// ...затем те же восемь для шага l+1 с av2/av3 и bv1

Поскольку тайлов восемь, логика эпилога «какой тайл содержит столбец j» вырастает из одиночного if/else в четырёхвариантный switch по j / VL: столбцовые блоки 0, 1, 2, 3 соответствуют парам тайлов (0,1), (2,3), (4,5), (6,7). Мелочь, но прямое следствие более широкого тайла — и она наталкивается на особенность SME, которую стоит отметить: индекс тайла ZA обязан быть константой времени компиляции, поэтому проиндексировать массив тайлов переменной j нельзя. Оператор switch с литеральными номерами тайлов в каждой ветви — способ удовлетворить это ограничение, всё же выбирая тайл во время выполнения.

Вывод из первого сюрприза: форму регистрового тайла диктует правило «используй каждый выданный тебе тайл ZA», а число выданных тайлов зависит от ширины элемента так, что тайл двойной точности получается прямоугольным. Структура ядра следует за формой тайла; форма тайла следует за оборудованием.

Сюрприз второй: ядра может не существовать

Теперь о решении, которое принимается до исполнения всего этого кода. Аппаратные внешние произведения двойной точности не входят в базовый SME2. Им требуется отдельная архитектурная возможность, FEAT_SME_F64F64. На кристалле, где есть SME2, но нет FEAT_SME_F64F64, функция svmopa_za64_f64_m не медленная — её просто нет. Поэтому ядро двойной точности условно в том смысле, в каком ядро одинарной не бывает никогда: оно существует, только когда оборудование заявляет о поддержке f64f64.

Это проявляется в двух местах. Во-первых, на этапе компиляции весь файл обёрнут условием:

#if defined(__ARM_FEATURE_SME_F64F64) || defined(__ARM_FEATURE_SME2)

Условие || __ARM_FEATURE_SME2 выглядит избыточным — казалось бы, при наличии f64f64 должен быть определён его собственный макрос? — но это причуда инструментария, и о ней стоит знать каждому, кто соберёт код SME на Mac. Apple clang 21 порождает корректный код f64f64 при флаге -march=...+sme-f64f64, но не определяет макрос __ARM_FEATURE_SME_F64F64. Генерация кода работает; макрос проверки возможности попросту отсутствует. Оградите свой код двойной точности одним этим макросом — и он исчезнет при компиляции на кремнии Apple, который прекрасно его поддерживает. Поэтому мы принимаем __ARM_FEATURE_SME2 в качестве заместителя, а настоящее решение оставляем проверке во время выполнения.

Каковая и составляет второе место: при инициализации контекста applesme опрашивает процессор и регистрирует ядро двойной точности только при наличии FEAT_SME_F64F64:

const bool has_f64f64 = bli_cpuid_has_features( features, FEATURE_SME_F64F64 );
...
if ( has_f64f64 ) {
    bli_cntx_set_ukrs( cntx, BLIS_GEMM_UKR, BLIS_DOUBLE,
                       bli_dgemm_armsme_2vlx4vl, ... );
}

Если признака нет, BLIS молча сохраняет своё эталонное ядро двойной точности — корректное, переносимое, медленное — и библиотека продолжает работать. Это страховочная сетка динамической диспетчеризации из части 3 в действии: один и тот же двоичный файл работает и на машине с f64f64, и на машине без него, используя ядро SME на первой и эталонное на второй, без каких-либо видимых пользователю различий, кроме скорости. Таблица размеров блоков кодирует ту же условность: записи MR/NR/MC/KC/NC для двойной точности установлены в значение-метку -1 («сохранить эталонное значение»), когда f64f64 отсутствует, — чтобы BLIS не пытался подбирать размеры для ядра, которое не зарегистрировано.

На M5 Pro, аппаратно располагающем и SME2, и FEAT_SME_F64F64, всё это чистый выигрыш: dgemm выдаёт около 444 GFLOPS в один поток (n = 2000) против примерно 70 у базовой реализации на NEON — ускорение примерно в 6,3 раза. Множитель меньше, чем 9,6 у sgemm, и это ожидаемо: числа двойной точности вдвое длиннее, поэтому подсистема памяти переносит вдвое больше данных на каждую операцию, и арифметика интенсивности по своей природе менее щедра. Но тот же принцип KC = 2048 из части 4 действует без изменений: аккумулятор 16×32 целиком живёт в ZA64 на протяжении всей редукции, поэтому разбиение по K ограничено кэшем, а не регистрами, и стремится быть крупным.

Урок в том, что не изменилось

Отступим на шаг и посмотрим на разность между ядрами одинарной и двойной точности в целом. Что изменилось: имена типов, геометрия тайла (16×16 → 16×32, четыре тайла → восемь), выбор тайла оператором switch в эпилоге и проверка бита возможностей во время выполнения. Что не изменилось: вся управляющая структура, контракт упаковки, стратегия мультивекторных загрузок, горячий цикл на сплошь истинных предикатах, эпилог с альфой и бетой и его безопасная по отношению к NaN ветвь beta == 0, изоляция медленного пути через noinline — и каждая строка BLIS вокруг ядра. Тождество каркаса не разрушается различием ширин, а сохраняется сквозь него; оттого и урок формулируется через то, что осталось прежним.

Это соотношение — много одинаковости, немного зависящей от точности геометрии, одна архитектурная проверка — и есть весь тезис цикла, ставший осязаемым. Когда структурой владеет фреймворк, добавление типа данных сводится в основном к пересчёту того, сколько тайлов вам выдано и какой они формы, плюс к проверке, что оборудование действительно умеет нужную арифметику. Такой перенос занимает вечер, а не месяц, — именно потому, что вся мебель из части 3 стоит на своих местах.

Вывод. Ядро двойной точности — это ядро одинарной точности с удвоенными ширинами: те же циклы, те же загрузки, та же логика эпилога, — за исключением двух вещей, которые стоит запомнить. FP64-часть ZA предоставляет восемь тайлов, поэтому тайл — прямоугольный, 16×32 (2VL×4VL), чтобы задействовать их все, с восемью FMOPA на шаг по k и четырёхвариантным switch в эпилоге. А аппаратным внешним произведениям двойной точности нужен FEAT_SME_F64F64 — возможность, отдельная от SME2, — поэтому всё ядро условно: скомпилировано за макросом (с обходным приёмом для Apple clang) и зарегистрировано лишь тогда, когда проверка во время выполнения подтверждает поддержку кремнием, с откатом к эталонному ядру в противном случае.

Далее: самое трудное. В части 6 мы берёмся за комплексные числа, где чередующийся формат хранения (действительная часть, мнимая часть) сталкивается с движком внешних произведений, понимающим лишь раздельные векторы действительных чисел, — и всё приключение вращается вокруг расслоения данных, четырёх вещественных FMOPA вместо одного комплексного умножения и того, как заставить svuzp/svzip выполнить перестановку, которую оборудование выполнять отказывается.

Двойная точность: в основном «найти и заменить», плюс два настоящих сюрприза