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

Оглавление-договор. В этой части мы вернёмся к работе, которую предыдущие главы принимали как готовую: упаковке. Сначала найдём скалярную ловушку в упаковке B, затем используем ZA как машину транспонирования, после этого разберём 128-битный случай для комплексных чисел и правило корректности для общего PACKM_KER. В финале соберём весь цикл одной дугой: от FMOPA до работающей субконфигурации BLIS для Apple SME.

Каждая предыдущая часть считала упаковку чужой работой. Ядра предполагают, что операнды прибывают подготовленными — непрерывными, дополненными нулями, раздельными там, где это нужно, — и части с 4 по 6 опирались на это предположение, чтобы сохранить горячий цикл свободным от ветвлений. Упаковка — то место, где эта подготовка действительно происходит, и я долго относился к ней как к водопроводу: сделать корректно и идти дальше. Это было ошибкой. То, что казалось несущественным — простой обвязкой, — и было существенным: видимость водопровода скрывала само дело. Упаковка незаметно стоила мне двукратного замедления, а починка потребовала использовать движок SME для того, для чего его никогда не рекламировали.

Начну с ошибки, которая не является ошибкой.

Скалярная ловушка: BLIS упаковывает B транспонированной, и мой быстрый путь этого не заметил

BLIS упаковывает оба операнда, но упаковывает их по-разному. Матрица A укладывается так, чтобы строки микропанели были непрерывны; матрица B укладывается транспонированной — BLIS хранит B как Bᵀ, чтобы ядро могло читать строку B как непрерывный вектор. Это стандартно и корректно. Но это же означает, что два операнда приходят в упаковщик с разными схемами доступа к памяти.

Мой первый упаковщик векторизовал ровно один случай: входной шаг равен единице, то есть собираемые исходные элементы уже лежат подряд. Это панель A в GEMM со столбцовым порядком хранения — прекрасно, полностью векторизовано. Но панель B того же самого вызова из-за транспонирования приходит с inca != 1 и lda == 1: каждая исходная строка непрерывна, а записывать нужно столбцы. Мой упаковщик смотрел на это, обнаруживал inca != 1 — и проваливался прямиком в скалярный запасной путь: обычный тройной цикл, копирующий по одному элементу за раз.

Итак, каждый стандартный sgemm со столбцовым порядком — самый частый вызов на свете — упаковывал свою панель B по одному скалярному числу за раз, пока я поздравлял себя с векторизованным путём для A. Ядро летало, упаковщик полз, и, поскольку оба живут внутри одного вызова GEMM, сквозной результат был посредственным, а по одному лишь ядру причина не просматривалась. Это урок части 4 — измеряйте упаковку отдельно от вычислений, — выученный на собственном опыте.

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

Приём: записать горизонтально, прочитать вертикально — и тайл транспонирует сам

Вернитесь мысленно к части 2, к тому единственному прелестному ходу из руководства ARM, который я просил отложить в память. Его предобработка загружала строки матрицы в ZA горизонтально, а вычитывала обратно вертикально — записываем строки, читаем столбцы, — и транспонирование выполнял сам тайл, без сети перестановок. Тогда это была диковинка. Здесь это ответ.

К ZA можно обращаться двумя способами. Можно записать вектор в горизонтальный срез тайла (строку), а позже прочитать вертикальный срез (столбец) — или наоборот. Запишите матрицу строка за строкой, прочитайте её столбец за столбцом — и на выходе окажется транспонированное. Матричный движок — это ещё и машина транспонирования, и — что принципиально — во время упаковки ZA простаивает. Ядро GEMM использует ZA в более позднем, отдельном вызове; пока идёт упаковка, весь массив — свободное рабочее пространство. Транспонирование не отнимает у нас ничего, чем мы пользовались.

Вот сердцевина векторизованного транспонирующего упаковщика:

// Загружаем до VL исходных строк как горизонтальные срезы ZA (каждая строка непрерывна по k).
for ( dim_t s = 0; s < sr; ++s )
    svld1_hor_za32( 0, (uint32_t)s, pk, a + (rb + s)*inca + kb );

// Читаем тайл вертикальными срезами (= столбцами = транспонированием), масштабируем, сохраняем.
for ( dim_t q = 0; q < kk; ++q )
{
    svfloat32_t x = svread_ver_za32_f32_m( vzero, pr, 0, (uint32_t)q );
    x = svmul_f32_x( ptrue, x, vkappa );   // применяем скаляр упаковки
    svst1_f32( pst, p + (kb + q)*ldp + rb, x );
}

Загружаем строки горизонтально, читаем столбцы вертикально, по пути умножаем на скаляр упаковки каппа, сохраняем в порядке следования по k. Транспонирование, стоившее скалярного ползания элемент за элементом, превратилось в блок векторных загрузок и векторных чтений через матричный движок. Упаковщик получил трёхстороннюю диспетчеризацию: непрерывный случай (векторизован), путь с транспонированием через ZA (также векторизован) и скалярный запасной путь для действительно причудливых шагов, который теперь почти никогда не срабатывает.

Дополнение нулями получается само собой — и это элегантная сторона решения. Когда исходный блок короче полного тайла, незагруженные срезы ZA покрываются предикатными чтениями, подмешивающими значения из нулевого регистра, так что строки-заполнители оказываются нулями без явного вызова svzero_za(). Тот же предикатный механизм, что обрабатывает края матрицы, обрабатывает и заполнение тайла. Корректные края и даровое транспонирование — из одного и того же механизма.

Измерения: изолированная упаковка B ускорилась в 3,8–9,2 раза, а сквозной sgemm на данных со столбцовым порядком вырос в 1,09–2,09 раза в зависимости от формы матриц — тот самый двукратный запас, который я оставлял неиспользованным, возвращён поворотом ZA набок. dgemm использует тот же приём на ZA64 (транспонирующие тайлы 8×8).

128-битный сюрприз: транспонировать комплексное число, не разрывая его

У комплексной упаковки есть дополнительная тонкость, и её разрешение — моё любимое маленькое открытие в проекте. При транспонировании комплексной панели пара (действительная, мнимая) каждого элемента должна остаться вместе: комплексное число — одно логическое значение, и транспонирование должно перемещать его как единое целое, а не раскидывать его половины независимо.

Ход такой: считать каждое комплексное число одним широким элементом и транспонировать его. Значение scomplex — это два 32-битных числа, то есть одна 64-битная единица, поэтому я транспонирую его на ZA64: пара (r,i) проезжает сквозь транспонирование склеенной в один 64-битный элемент и выходит нетронутой. А dcomplex — это два 64-битных числа, то есть одна 128-битная единица, поэтому он транспонируется на ZA128.

ZA128 и есть сюрприз. Тайлы ZA со 128-битными элементами — полноправная часть архитектуры, но её экзотический угол, и я, честно говоря, не ожидал, что Apple clang породит работающий код для svld1_hor_za128 / svread_ver_za128. Порождает: инструкции ассемблируются, исполняются и дают корректное транспонирование, включая края. Есть одна ловушка, стоившая мне вечера: операциями над срезами ZA128 управляет предикат b64 — по два предикатных элемента на каждый 128-битный элемент данных. Поэтому, строя предикат на n значений dcomplex, отсчитывайте 2·n предикатных элементов, а не n. Пропустите это — и края окажутся незаметно наполовину неверными. Стоит это узнать — и транспонирование dcomplex упаковывается так же чисто, как float.

Итого семейство упаковщиков покрывает четыре типа данных и три ширины элемента для транспонирования: float на ZA32 (16×16), double и scomplex на ZA64 (8×8; второй — с трактовкой комплексного числа как 64-битной единицы) и dcomplex на ZA128 (4×4). Одна идея — тайл транспонирует сам себя — воплощённая поперёк всей системы типов.

Правило корректности, делающее комплексную упаковку ловушкой

Теперь выстраданный урок — самый важный абзац этой части для всякого, кто будет развивать эту работу. Сделать комплексный упаковщик быстрым — это приём с ZA, описанный выше. Сделать его корректным — задача другая и более острая, потому что упакованный формат разделяется с большим числом ядер, чем один GEMM.

Мой инстинкт после части 6 подсказывал заставить комплексный упаковщик выдавать раздельный формат — выполнять расслоение в [r r r r][i i i i] прямо при упаковке, чтобы ядро cgemm избавилось от шести svuzp на каждую итерацию по k. Чистая идея. Она же сломала trsm, trmm и gemmtrsm — все остальные комплексные операции третьего уровня — способом, на диагностику которого ушёл целый день.

Вот почему. Регистрируемое ядро PACKM_KER разделяется между всеми комплексными операциями третьего уровня, а не только GEMM. А эталонные реализации треугольных решений и умножений (bli_gemmtrsm_ref.c и родственные) вызывают зарегистрированное микроядро GEMM, но ожидают упакованный операнд в стандартном чередующемся формате. Передайте им панель, упакованную раздельно, — и они вычислят мусор; тихо, потому что буфер правильного размера, а сбой проявляется лишь неверными числами глубоко внутри решения. Собственный упакованный формат — это контракт со всей комплексной подсистемой третьего уровня, а не частная договорённость между вашим упаковщиком и вашим ядром. Момент, забывший о целом, отрицает сам себя: панель, верная лишь вашему ядру, тихо ломает всё остальное.

Отсюда правило: зарегистрированный комплексный упаковщик обязан выдавать байт в байт стандартный чередующийся формат. Мои упаковщики так и поступают. Транспонирование через ZA сохраняет пары (r,i) чередующимися (в этом весь смысл транспонирования на увеличенной ширине элемента); расслоение остаётся внутри ядра GEMM, где его разместила часть 6, на свободном векторном конвейере. Идея раздельной упаковки не ошибочна — но жить она может только за явным, включаемым по требованию контрактом упакованного формата, которого TRSM и TRMM никогда не увидят. Именно такой поэтапный, ограждённый проект и лежит в заметке SME_CGEMM_ZA_DEINTERLEAVING — не включённый в поставку, пока не докажет сквозную победу. Поставляемый упаковщик выбирает безопасный путь: стандартный формат, транспонирование ради скорости, расслоение ниже по течению.

Ещё одна тонкость из того же семейства. При комплексной упаковке с единичным скаляром и без сопряжения — а это каждый комплексный GEMM с вещественной альфой, самый частый случай — упаковщик должен полностью пропускать комплексное масштабирование и просто копировать. Когда я впервые написал упаковщик с постоянно включённым масштабированием, комплексный GEMM замедлился: безусловное комплексное умножение на каппу через uzp/zip стоило больше, чем экономило. Добавление быстрого пути noscale (каппа равна единице и сопряжения нет → простое копирование) превратило проигрыш в выигрыш. Альфа становится комплексной каппой с ненулевой мнимой частью, только когда мнимая часть у неё действительно есть; подавляюще частый случай вещественной альфы не должен платить за комплексное масштабирование, которое ему не нужно.

Весь цикл одной дугой

Вот и весь перенос. Отступим на шаг и посмотрим, чего он потребовал, — потому что его форма и есть подлинный вывод. Круг замыкается там же, где начался: результат — не отдельная находка в конце, а весь пройденный путь, собранный в форму. Метод и оказывается итогом.

Мы начали (часть 1) с процессора, отрастившего тензорное ядро, — FMOPA, инструкция внешнего произведения ранга 2, накапливающая результат в регистровом массиве ZA, — и с осознания, что Apple выпустила матрицы вместо векторов SVE. Мы обнаружили (часть 2), что публичная учебная программа для его программирования — два источника на противоположных концах лестницы с отсутствующей серединой. Мы узнали (часть 3), что средняя ступенька построена тридцать лет назад: BLIS, чей контракт микроядра и есть накопление внешнего произведения, поэтому перенос сводится к одному ядру и пяти числам. Мы написали эти ядра: одинарной точности (часть 4), двойной — с её прямоугольной геометрией на восемь тайлов и проверкой возможностей во время выполнения (часть 5) — и собственные комплексные, где расслоение даром прячется на векторном конвейере (часть 6). И замкнули круг здесь: упаковщик, кормящий их все, использует тот же матричный движок как машину транспонирования, а самые острые уроки касались контрактов корректности, а не скорости.

Соберём числа воедино: sgemm ~9,6×, dgemm ~6,3×, cgemm ~10,5×, zgemm ~6,7× относительно базовых реализаций на NEON, один поток, один и тот же двоичный файл. Полная семантика BLAS, диспетчеризация во время выполнения, все четыре типа данных, оба порядка хранения, полное комплексное семейство третьего уровня — ничего из этого мы не писали, потому что у BLIS всё это уже было. Добавленное нами — небольшое, локальное и привязанное к оборудованию: как тайлы ложатся на ZA, как их кормить, как транспонировать входные данные самим движком. Атом из руководства ARM, вставленный в гнездо, которое BLIS выточил за десятилетия до появления этого кремния.

Матричный движок в вашем ноутбуке существует, программируется из C и быстр. Дорога к нему узка, но нужный фреймворк ждал всё это время. Вот история, которую я собирался рассказать, и заканчивается она там, где должна: быстрым GEMM, честными числами и парой приёмов, которые стоит позаимствовать.

Вывод. Упаковка выглядела водопроводом, а скрывала двукратный выигрыш. BLIS упаковывает B транспонированной, из-за чего панель B каждого GEMM со столбцовым порядком уходила на скалярный путь — пока транспонирование не поручили самому ZA: записываем исходные строки в горизонтальные срезы тайла, вычитываем их вертикальными срезами, и простаивающий во время упаковки матричный движок транспонирует даром (float на ZA32, double и scomplex на ZA64, dcomplex на ZA128 — который работает под Apple clang с оговоркой о двух предикатных элементах на каждый 128-битный элемент данных). Более глубокий урок касается корректности: комплексное ядро PACKM_KER разделяется с TRSM/TRMM/GEMMTRSM и потому обязано выдавать стандартный чередующийся формат — собственный раздельный формат оборачивается тихой поломкой, — а расслоению место в ядре, не в упаковщике. Скорость дал движок; корректность — уважение к контрактам фреймворка.

На этом цикл завершается. Код — четыре микроядра GEMM для SME2 и транспонирующий упаковщик на ZA, интегрированные как субконфигурация BLIS для Apple Silicon, — есть работающее доказательство того, что матричный блок ноутбучного процессора можно довести до сравнимой с библиотекой производителя производительности GEMM средствами переносимого C, встроенного во фреймворк. Идите и поверните ZA набок.