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

Оглавление-договор. В этой части мы начнём с тупика, в который приводит голый цикл FMOPA; затем покажем, какую часть GEMM BLIS уже построил за нас; после этого разберём пять циклов BLIS и место микроядра внутри них. Финал главы должен сделать дальнейший план конкретным: что именно нужно добавить для Apple SME, а что уже относится к готовой «мебели» фреймворка.

Расскажу о лучшем часе, потраченном на весь этот проект, — часе, который я провёл, не написав ни строчки кода.

У меня работал цикл с FMOPA — четырёхстрочное ядро из части 1, накапливающее внешние произведения в ZA. Оно перемножало две матрицы. В качестве BLAS оно было при этом совершенно бесполезно: ни альфы и беты, ни возможности прибавить результат к существующей матрице C, ни обработки транспонирования, один тип данных — и производительность, падавшая с обрыва, едва матрицы вырастали из кэша L1. Всё, чего, по словам части 2, не делает учебное руководство, моё ядро тоже не делало. Я стоял ровно там, где должна быть недостающая средняя ступенька, держа в руках атом — и никакой лестницы.

Инстинкт в такой момент — начать строить наружу: написать цикл блокирования под кэш, затем процедуру упаковки, затем слой многопоточности, затем эпилог с альфой и бетой, затем краевые случаи, затем второй тип данных — и так несколько месяцев, в конце которых у вас будет, если вы очень хороши и очень удачливы, посредственная повторная реализация того, что уже существует. Потому что оно действительно уже существует. Середину лестницы выточили в 1990-х и оформили во фреймворк в 2010-х, и в ней есть гнездо, по форме в точности совпадающее с моим атомом.

Этот фреймворк — BLIS. Прежде чем хвалить его вслух, покажу форму, из которой эта похвала следует: понимание его структуры — единственная причина, по которой остаток цикла короток. Стоит увидеть форму BLIS — и «перенести GEMM на SME» перестаёт быть проектом и становится анкетой для заполнения.

Озарение: одно ядро, а всё остальное — мебель

Вот наблюдение, на котором построен BLIS; оно принадлежит Кадзусигэ Гото и было формализовано Филдом Ван Зи и Робертом ван де Гейном из Техасского университета в Остине. Возьмите любую высокопроизводительную реализацию GEMM — для x86, для POWER, для ARM, для дюжины умерших архитектур — и прищурьтесь. Все они — одна и та же программа. Все упаковывают операнды в удобные для кэша панели; все разбивают три цикла по M, N и K на блоки так, чтобы нужные данные находились в нужном уровне кэша в нужное время; все опираются в основании на маленькое, вручную настроенное ядро, умножающее узкий срез A на узкий срез B. Блокирование одинаково. Упаковка одинакова. Порядок циклов одинаков. Единственное, чем великий GEMM для AVX-512 отличается от великого GEMM для NEON, — крохотное ядро в самом основании и горстка чисел, описывающих размеры кэшей. Одинаковость — это всеобщее; архитектуры — его особенные обличья; крохотное ядро в основании — единичное, в котором всё различие и умещается.

И BLIS делает очевидную — задним числом — вещь: записывает всю эту одинаковость один раз, на переносимом C, а различия выставляет в виде небольшого набора точек подключения. Переносимая часть — которую статьи о BLIS называют «макроядром», вместе с окружающими его циклами, — огромна, тонка, и трогать её никогда не приходится. Она ведает объектным API, интерфейсами CBLAS и Fortran, транспонированием, нормализацией порядка хранения, блокированием под кэш на трёх уровнях, расписанием упаковки, многопоточностью и граничными случаями, когда размеры матрицы не кратны ничему удобному. Это тысячи строк самого выстраданного знания в вычислительной математике, и перенос на новую архитектуру переиспользует его целиком.

Вы поставляете мебель, которая встаёт в гнёзда:

Микроядро: умножить одну упакованную микропанель A на одну упакованную микропанель B, накопить результат в крохотном тайле C размером MR×NR, применить альфу и бету. Вот сюда идёт ваш цикл с FMOPA.
Ядро упаковки (необязательно — есть эталонное): переставить блок A или B в непрерывную, дополненную нулями раскладку, которую микроядро хочет потреблять потоком.
Таблицу чисел: размеры регистрового тайла MR и NR и размеры кэш-блоков MC, KC, NC.

Это и есть весь перенос. Микроядро — несколько сотен строк, размеры блоков — пять целых чисел, а всё остальное делает BLIS. Когда я говорю, что остаток цикла — «заполнение пропусков», это те самые пропуски. Они по-прежнему интересны — матричный движок достаточно необычен, чтобы заполнить их хорошо стоило настоящего труда, о чём и рассказывают части с 4 по 7, — но объём работы ограничен так, как реализация GEMM с нуля не бывает ограничена никогда.

Пять циклов, снаружи внутрь

Чтобы разместить пропуски, нужна картина гнезда циклов BLIS. GEMM в BLIS — это пять вложенных циклов, обёрнутых вокруг одного микроядра, и каждый цикл существует, чтобы угодить своему уровню иерархии памяти. От внешнего к внутреннему:

Цикл 5 (NC), по N. Нарезать задачу на столбцовые панели B шириной NC. Каждая панель — то, что потребит один проход через остальную машину.
Цикл 4 (KC), по K. Нарезать размерность K на отрезки длиной KC. Столбцовая панель B высотой KC (размером NC×KC) теперь упаковывается в непрерывный, размещённый в кэше буфер: этот блок будет переиспользоваться снова и снова, поэтому его стоит реорганизовать. Упакованная панель B живёт в диапазоне кэшей L3/L2.
Цикл 3 (MC), по M. Нарезать A на строковые панели высотой MC. Блок A размером MC×KC упаковывается в собственный непрерывный буфер, рассчитанный на L2/L1. Этот упакованный блок A будет многократно проходить сквозь упакованный блок B.
Цикл 2 (NR), по упакованной панели B. Обходить упакованный блок B полосами шириной NR — по одной микропанели B.
Цикл 1 (MR), по упакованной панели A. Обходить упакованный блок A полосами высотой MR — по одной микропанели A. Пара (MR, NR) выбирает один регистровый тайл C, и мы вызываем микроядро: пропустить панель A размером MR×KC через панель B размером KC×NR, накопив KC внешних произведений в регистрах, а затем обновить C.

Прочтите это изнутри наружу — и получится рассказ о переиспользовании данных. Микроядро удерживает тайл C в регистрах (в нашем случае — в ZA) на протяжении всей редукции по KC, так что C записывается в память один раз, а не KC раз. Упакованный блок A подобран так, чтобы оставаться в L2 и переиспользоваться на каждой NR-полосе B. Упакованный блок B подобран так, чтобы оставаться в L3 и переиспользоваться на каждой MC-панели A. Каждый байт упаковывается потому, что его предстоит прочитать много раз; каждая граница цикла — уровень кэша. Здесь нет ничего специфичного для x86, ARM или SME. Это специфично для устройства памяти, а оно у всех общее.

И выигрыш для наших целей: две вещи, делающие эту схему быстрой на конкретной машине, — ровно те две вещи, которые BLIS предлагает вам подключить. Регистровый тайл (MR, NR) и микроядро определяют, как вы используете вычислительный блок; кэш-блоки (MC, KC, NC) — как вы его кормите. SME меняет то, чем регистровый тайл является — это массив ZA, а не набор векторных регистров, — но не меняет того, что регистровый тайл есть то, во что сводит результат самый внутренний цикл. Гнездо подходит.

Почему это правильный фреймворк именно для матричного движка

Есть и более глубокая причина, по которой BLIS подходит SME, помимо «это хороший фреймворк для GEMM». Центральная абстракция BLIS: микроядро накапливает обновление ранга 1 на каждом шаге по K в размещённый в регистрах тайл C. Оглянитесь на часть 1: FMOPA и есть обновление ранга 1 — внешнее произведение, — накапливаемое в размещённый в регистрах тайл по имени ZA. Контракт микроядра BLIS и инструкция SME, никогда не встречавшись, описывали одну и ту же операцию.

На SIMD-машине микроядру BLIS приходится имитировать внешнее произведение: транслировать один элемент A на весь вектор, умножать на вектор B и повторять это MR раз, синтезируя обновление ранга 1 из операций FMA ранга 0. Этот синтез — неказистое, вечно стеснённое нехваткой регистров сердце каждого SIMD-ядра GEMM. На SME синтез исчезает: оборудование выполняет обновление ранга 1 само, одной инструкцией, в регистровый файл, созданный служить тайлом C. BLIS тридцать лет притворялся, что векторные регистры — это матричный аккумулятор. SME просто делает это правдой. Тридцать лет она была правдой лишь в возможности — записанной в контракте; кремний сделал её действительной. Вставить микроядро SME в гнездо BLIS — не согласование несовместимого; это абстракция, наконец получившая оборудование, которое она всё это время описывала.

По той же причине таблица размеров блоков — то место, где живёт интересная настройка. Поскольку ZA огромен — регистровый тайл 32×32 одинарной точности против тайлов порядка 8×8, которыми располагает ядро на NEON, — числа MC/KC/NC, приводящие его в равновесие, не похожи ни на что в стандартных конфигурациях ARM. В части 4 мы увидим, что KC стремится быть огромным (2048): весь смысл гигантского аккумулятора ZA в том, что он никогда не выгружается в память, поэтому редукция по K ограничена лишь тем, сколько упакованных данных A и B удаётся удержать в кэше. Фреймворк тот же; числа, заставляющие его петь, — новые.

Что BLIS даёт вам, конкретно, для переноса на Apple SME

Сделаем утверждение «заполните пропуски» конкретным, потому что это сквозная нить всего дальнейшего. Вот полный список того, что работа над Apple SME добавляет в BLIS, — а всё, чего нет в этом списке, BLIS уже сделал:

Субконфигурация (applesme), которая регистрирует ядра и, что принципиально, распознаёт оборудование во время выполнения — SME2 для ядер одинарной точности, дополнительный признак FEAT_SME_F64F64 для ядер двойной точности — и молча оставляет эталонные ядра BLIS, если нужного кремния нет. (Сага о флагах компилятора из части 1 тоже живёт здесь.)
Четыре микроядра GEMM, по одному на тип данных: sgemm (часть 4), dgemm (часть 5) и два комплексных ядра (часть 6).
Ядро упаковки, использующее сам ZA как машину транспонирования (часть 7).
Таблица размеров блоков, в которой MR и NR выводятся во время выполнения из потоковой длины вектора, так что один и тот же двоичный файл адаптируется к любому будущему SVL.

Обратите внимание, чего в списке нет. Никакой логики альфы и беты в циклах — микроядро применяет их на уровне тайла, а BLIS дирижирует. Никакой обработки транспонирования — BLIS нормализует порядок хранения до того, как ядро увидит данные. Никакой многопоточности — BLIS распределяет внешние циклы по ядрам процессора, и (оговорка об общем движке из части 1) потолок задаётся оборудованием, а не чем-либо, что пишем мы. Никакой прослойки CBLAS, никакого интерфейса Fortran, никаких дополнительных циклов подчистки краёв. Всё это — мебель, и вся она поставлена.

Регистрируя ядро, вы также сообщаете BLIS один бит предпочтения: хочет ли это ядро видеть матрицу C в памяти по строкам или по столбцам? Наши ядра SME вычитывают результат из ZA столбец за столбцом, поэтому объявляют предпочтение столбцового порядка, и BLIS подстраивает C под него. Один логический признак — и фреймворк корректно выбирает порядок хранения для любого вызывающего кода. Такова фактура всего переноса: небольшие, локальные декларации того, чего хочет оборудование, — и фреймворк, который уже продумал последствия.

Вывод. Среднюю ступеньку лестницы SME — упаковку, блокирование под кэш, полный контракт BLAS, многопоточность, обработку краёв — не нужно изобретать: BLIS и есть эта ступенька, причём с тех времён, когда самого оборудования ещё не существовало. Его основополагающее наблюдение: каждый GEMM — одна и та же программа, за вычетом одного маленького ядра и пяти чисел размеров блоков; а контракт его микроядра (накапливать обновления ранга 1 в размещённый в регистрах тайл C) — в точности то, что делает внешнее произведение SME. Перенос сводится к трём действиям: напишите ядро, заполните таблицу, объявите, чего хочет оборудование. Следующие четыре части — эти пропуски, заполненные; и именно в том, чтобы заполнить их хорошо, скрывалось всё удовольствие.

Далее: микроядро sgemm, целиком. В части 4 мы разбираем ядро одинарной точности 2VL×2VL по винтикам: как аккумулятор 32×32 отображается на четыре тайла ZA, почему цикл по K расшит по два, и загружает данные мультивекторными инструкциями SME2, как альфа и бета применяются при чтении ZA столбец за столбцом — и один флаг (beta == 0), с которым нужно обращаться осторожно, иначе значения NaN просочатся в свежую матрицу.

BLIS: недостающую среднюю ступеньку построили тридцать лет назад