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

Оглавление-договор. В этой части мы сначала отделим прямой комплексный путь от готового метода 1m; затем покажем конфликт между чередующимся форматом памяти и раздельной арифметикой внешних произведений; после этого разберём четыре вещественных внешних произведения, расслоение через svuzp, эпилог с комплексными alpha и beta, а в конце — почему сопряжение остаётся работой упаковщика. Следующая часть вернёт нас к упаковке и к тому месту, где эта граница ответственности особенно важна.

Для комплексного матричного умножения существует проторённый обходной путь — метод 1m, и BLIS поставляет его из коробки: представьте каждое комплексное число маленьким вещественным блоком 2×2 — и комплексный GEMM превращается в хитро организованный вещественный GEMM, для которого быстрое ядро у вас уже есть. Это остроумно, это ноль дополнительного кода, и это было моей отправной точкой. Но за это удобство приходится доплачивать: вещественному ядру приходится перекачивать вдвое больше данных, чтобы синтезировать комплексную арифметику. Весь смысл движка внешних произведений — арифметическая интенсивность; пропуская комплексные числа через 1m, мы тратим часть этой интенсивности на счетоводство.

Итак, вопрос этой части: можно ли написать собственное комплексное ядро внешних произведений — выдающее те же GFLOPS, что и sgemm, — вместо того чтобы одалживать вещественное? Ответ — да, но сразу зафиксируем, что именно будет означать это сравнение: комплексная арифметика никуда не исчезает, она всё равно раскладывается на несколько вещественных внешних произведений; выигрыш в том, что расслоение данных не становится отдельным узким местом. Путь к этому оказывается настоящим маленьким приключением, потому что комплексные числа прибывают ровно в том формате памяти, который движок SME усваивает хуже всего.

Несовпадение: чередующееся хранение против раздельной математики

Комплексный вектор хранится в памяти с чередованием (interleaved): действительные и мнимые части идут попеременно, [r0 i0 r1 i1 r2 i2 …]. Так C хранит float _Complex, так BLAS определяет свои массивы, так любой вызывающий код передаёт вам данные. И это ровно противоположно тому, что нужно машине внешних произведений. Внешняя форма хранения и внутреннее устройство счёта разошлись: память диктует чередование, математика требует раздельности — и расслоение здесь есть примирение формы с содержанием.

Распишем комплексное умножение. Для C = A·B при A = A_re + i·A_im и B = B_re + i·B_im:

C_re = A_re·B_re − A_im·B_im
C_im = A_re·B_im + A_im·B_re

Каждое слагаемое — произведение чисто вещественного вектора на чисто вещественный вектор. FMOPA вычисляет вещественные внешние произведения — и только их. Чтобы его накормить, мне нужны A_re, A_im, B_re, B_im — каждый как непрерывный ряд вещественных значений в векторном регистре. Но память выдала [r i r i r i …], где действительные и мнимые части перемешаны. Движку нужен раздельный (planar) формат; мир хранит с чередованием. Навести этот мост — выполнить расслоение (de-interleaving); а где именно его выполнять — и есть главный вопрос этой части. Единое комплексное число приходится разъять на вещественные моменты, отдать их движку порознь и собрать обратно — то же число, но прошедшее сквозь собственное разложение.

Арифметика: одно комплексное умножение = четыре вещественных внешних произведения

Прежде чем заняться перестановкой — вычисления. Двум уравнениям выше, выполненным внешними произведениями в ZA, требуется четыре вещественных FMOPA на одно комплексное внешнее произведение — но структура накопления красивее, чем кажется. Выделим по два тайла ZA на каждый комплексный выходной субтайл: один для действительной части, один для мнимой:

тайл C_re += A_re ⊗ B_re     (fmopa)
тайл C_re −= A_im ⊗ B_im     (fmops — внешнее произведение с ВЫЧИТАНИЕМ)
тайл C_im += A_re ⊗ B_im     (fmopa)
тайл C_im += A_im ⊗ B_re     (fmopa)

Элегантная деталь — fmops: внешнее произведение с плавающей точкой и вычитанием. SME предоставляет оба варианта — накопление и накопление с вычитанием, — поэтому слагаемое − A_im·B_im в C_re укладывается в одну инструкцию, без отдельной смены знака. Четыре внешних произведения (три со сложением, одно с вычитанием), два тайла ZA — и комплексное обновление ранга 1 выполнено. Один FMOPA вещественного ядра превращается в четыре; этот четырёхкратный множитель — неотъемлемая цена комплексной арифметики, и более дешёвой схемы при этой точности не существует.

Теперь бюджет тайлов. Каждый комплексный субтайл занимает два тайла ZA32 (действительная и мнимая части). Всего у FP32-части ZA четыре тайла. Значит, собственный cgemm вмещает ровно два комплексных субтайла — и ядро имеет форму 2VL×1VL (MR = 32, NR = 16 комплексных элементов), с размещением двух субтайлов вдоль направления M:

                n: 0..VL-1
  m: 0..VL-1     za0 (re) / za1 (im)
  m: VL..2VL-1   za2 (re) / za3 (im)

Вот и весь аккумулятор: четыре тайла ZA32 — два для действительных частей, два для мнимых, — вычисляющих комплексный результат 32×16. Восемь FMOPA на шаг по k (по четыре на субтайл). Это прямой комплексный аналог макротайла из части 4, и он точно так же насыщает ZA.

Перестановка: svuzp на векторном блоке, даром

Вот проектное решение, ради которого собственное комплексное ядро стоит усилий. Где выполнять расслоение? Есть два варианта: при упаковке (переставить [r i r i] в раздельное [r r r r][i i i i] один раз, во время подготовки данных) или внутри цикла по k (расслаивать каждый срез на лету, на каждой итерации). Вариант с упаковкой звучит очевидно лучше — сделать один раз и переиспользовать, — и я вернусь к нему в части 7, потому что он не так чист, как кажется. Зарегистрированное ядро расслаивает в цикле по k, и сходит это с рук лишь потому, что расслоение выполняется на совершенно другом исполнительном блоке, нежели внешние произведения.

Движок SME занят инструкциями FMOPA. Расслоение выполняют svuzp1/svuzp2 — операция unzip, SVE-перестановка, разделяющая [r0 i0 r1 i1 …] на [r0 r1 r2 …] (uzp1, чётные элементы) и [i0 i1 i2 …] (uzp2, нечётные). Эти перестановки исполняются на конвейере потоковых SVE-векторов, который не является конвейером внешних произведений ZA. Пока блок SME перемалывает восемь FMOPA, векторный блок в сторонке расслаивает операнды следующего среза. Перестановка перекрывается с вычислениями и прячется в их тени. Вот один шаг по k:

svfloat32x4_t av = svld1_x4( pc, a );   // A: [r i r i ...], 2VL комплексных = 4 вектора
svfloat32x2_t bv = svld1_x2( pc, b );   // B: [r i r i ...], 1VL комплексных = 2 вектора

// Расслоение: uzp1 собирает действительные части, uzp2 — мнимые.
svfloat32_t a_re0 = svuzp1_f32( svget4( av, 0 ), svget4( av, 1 ) );
svfloat32_t a_im0 = svuzp2_f32( svget4( av, 0 ), svget4( av, 1 ) );
svfloat32_t a_re1 = svuzp1_f32( svget4( av, 2 ), svget4( av, 3 ) );
svfloat32_t a_im1 = svuzp2_f32( svget4( av, 2 ), svget4( av, 3 ) );
svfloat32_t b_re  = svuzp1_f32( svget2( bv, 0 ), svget2( bv, 1 ) );
svfloat32_t b_im  = svuzp2_f32( svget2( bv, 0 ), svget2( bv, 1 ) );

// m-блок 0: za0 = re, za1 = im.
svmopa_za32_f32_m( 0, pall, pall, a_re0, b_re );
svmops_za32_f32_m( 0, pall, pall, a_im0, b_im );   // обратите внимание: fmopS, вычитание
svmopa_za32_f32_m( 1, pall, pall, a_re0, b_im );
svmopa_za32_f32_m( 1, pall, pall, a_im0, b_re );
// m-блок 1: za2 = re, za3 = im (те же четыре инструкции, с a_re1/a_im1)

Измерения подтверждают ставку: собственный cgemm выдаёт около 1432 GFLOPS при n = 1000 — вровень с производительностью sgemm, примерно в 10,5 раза быстрее базового варианта 1m на NEON. Расслоение, во всех практических смыслах, бесплатно. Это и есть главная находка: поскольку перестановка живёт на другом конвейере, чем математика, собственное комплексное ядро не стоит ничего сверх вещественного в пересчёте на единицу комплексной работы.

В этом коде есть небольшая, но несущая деталь: загрузки выполняются через svld1_x4/svld1_x2, а перестановка — через svuzp, вместо напрашивающейся svld2 (которая расслаивает во время загрузки). Почему? Потому что svld2 и её напарница svst2 недопустимы в потоковом режиме. Потоковый режим SME — это урезанный SVE: целый ряд инструкций сборки/разброса и структурированного доступа попросту недоступен. Приходится загружать непрерывно и расслаивать вручную. Об ограничениях такого рода узнаёшь лишь тогда, когда ассемблер отвергает твою первую, самую естественную попытку.

Эпилог: альфа, бета и перемежение ответа обратно

Чтение результата — перестановка в обратную сторону. ZA хранит ответ раздельно — za0 содержит вектор действительных частей, za1 — вектор мнимых, — но C в памяти хранится с чередованием. Поэтому эпилог расслаивает C на входе и заново перемежает результат на выходе с помощью svzip1/svzip2 (zip — операция, обратная unzip: она берёт [r r r r] и [i i i i] и производит [r i r i …]).

Комплексные альфа и бета превращают эпилог в самостоятельное комплексное умножение. out = alpha · ab, где оба значения комплексные, — тот же танец из четырёх слагаемых:

out_re = ab_re·alpha_re − ab_im·alpha_im
out_im = ab_im·alpha_re + ab_re·alpha_im

а при beta ≠ 0 ещё одно комплексное умножение-сложение вплавляет старое значение C (прочитанное из памяти, расслоённое через svuzp, объединённое и заново перемежённое через svzip при сохранении). Действует та же защита от NaN при beta == 0, что и в части 4: когда бета равна нулю, C не читается вовсе, и мусор в свежей матрице до результата попросту не добирается. Одна особенность, характерная для SME: поскольку индексы тайлов ZA обязаны быть константами времени компиляции (то же правило из части 5), два m-блока в эпилоге развёрнуты макросом, параметризованным литеральными номерами тайлов, а не циклом.

Сопряжение — и куда ядро не заглядывает

Комплексный BLAS обязан также обслуживать сопряжённые операнды: conj(A)·B, A·conj(B) и так далее. Можно было бы ожидать, что ядро ветвится по флагам сопряжения. Оно не ветвится, и причина — приятное разделение обязанностей: сопряжение применяется при упаковке. К моменту, когда упакованная панель достигает ядра, любое запрошенное сопряжение уже учтено (мнимые части меняют знак во время упаковки), поэтому вычислительное ядро полностью безразлично к сопряжению. Оно видит раздельные действительные и мнимые части и перемножает их; оно не знает и не желает знать, меняли ли мнимым частям знак выше по течению. Это работа части 7 — и потому одно и то же ядро обслуживает все четыре комбинации сопряжений без единого ветвления.

zgemm: та же история в двойной точности

Комплексная двойная точность (zgemm) относится к cgemm так же, как dgemm относился к sgemm: тот же проект, удвоенные ширины, условие на FEAT_SME_F64F64. Каждый комплексный субтайл теперь занимает два тайла ZA64, а у FP64-части ZA восемь тайлов (правило из части 5), так что zgemm вмещает четыре комплексных субтайла: форма 2VL×2VL (MR = NR = 16 комплексных элементов), восемь тайлов ZA64, четыре субтайла, шестнадцать FMOPA на шаг по k. Расслоение — svuzp над f64, арифметика идентична, и итог — около 447 GFLOPS при n = 1000, вровень с dgemm, примерно в 6,7 раза быстрее базового варианта. Вывод тот же: собственное ядро побеждает 1m, а расслоение бесплатно, потому что прячется на векторном конвейере.

На полке лежит экспериментальный проект — заметка SME_CGEMM_ZA_DEINTERLEAVING_DESIGN: вынести расслоение из цикла по k целиком и выполнять его в упаковщике, используя сам ZA как машину транспонирования, чтобы вычислительный цикл избавился от своих шести svuzp на итерацию. Это настоящая идея с настоящей тонкостью (аккумулятор одинарной точности уже занимает все четыре тайла ZA, поэтому ZA свободен только во время упаковки, а не во время вычислений), и она — естественный мост к следующей части. Выигрывает ли она в сквозном измерении — открытый вопрос, который решит только эксперимент: выигрыш в одних лишь вычислениях немного стоит, если упаковка и её переходы в потоковый режим съедают его обратно.

Вывод. Комплексный GEMM — это четыре вещественных внешних произведения (три FMOPA и один FMOPS ради знака минус), накапливаемых в спаренные тайлы ZA для действительных и мнимых частей: cgemm вмещает два комплексных субтайла в четырёх тайлах ZA32 (32×16), zgemm — четыре субтайла в восьми тайлах ZA64 (16×16). Трудность в том, что память хранит комплексные числа с чередованием, а движку нужен раздельный формат, поэтому каждый срез по k расслаивается операциями svuzp — что бесплатно, поскольку выполняется на векторном конвейере в тени внешних произведений SME (и это приходится делать вручную: svld2 недопустима в потоковом режиме). Сопряжение обслуживается при упаковке, поэтому ядро к нему безразлично; альфа и бета исполняются как полноценные комплексные умножения в эпилоге, заново перемежающем результат. Собственное комплексное ядро держит производительность вещественного: примерно 10,5× для cgemm и 6,7× для zgemm — расслоение действительно даровое.

Далее: упаковщики — и приём, связывающий весь цикл воедино. В части 7 мы поворачиваем ZA набок и используем матричный движок как машину транспонирования — тот самый ход «записать горизонтально, прочитать вертикально» из руководства ARM в части 2, — чтобы быстро упаковывать панели со столбцовым порядком хранения. И я объясняю одно тонкое правило корректности (общее с комплексными TRSM и TRMM), которое превращает собственный комплексный упаковщик либо в двукратный выигрыш, либо в тихое разочарование.

Комплексный GEMM: четыре вещественных внешних произведения в одном флаконе