Панорама матричных расширений: от x86 до RISC-V

Эту статью можно поделить на две части:

• Попытаемся понять, зачем вообще нужны матричные расширения и в чем тут соль.
• Соберем воедино все, что можно найти в открытых источниках о ныне существующих матричных расширениях — от Intel AMX до SpacemiT IME.

Матричные операции — это основные хот-споты для приложений из таких областей, как:

• искусственный интеллект и машинное обучение (AI/ML),
• компьютерное зрение (CV),
• дополненная и виртуальная реальность (AR/VR),
• математическое моделирование сложных физико-технических систем (CAE, HPC),
• усовершенствованные системы помощи водителю (ADAS) и др.

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

Очень частое заблуждение, что все математически сложные алгоритмы уже давно выполняются только на специализированных ускорителях: графических (GPU), тензорных (TPU), нейронных (NPU) и так далее.

Безусловно, ускорители очень важны для повышения производительности вычислительно сложных алгоритмов. Но в жизни доля матричных нагрузок велика не только на суперкомпьютерах, кластерах и достаточно мощных пользовательских устройствах. Вы постоянно сталкиваетесь с матричными нагрузками на самых обычных смартфонах, когда обрабатываете фотографии или видео, на планшетах, ноутбуках и не самых мощных персональных компьютерах без видеокарт.

Кроме того, даже если на устройстве пользователя есть ускоритель, нельзя заранее предсказать, сколько у него параллельно идет других задач и насколько этот ускоритель в данный момент свободен. Если ускоритель уже загружен другими задачами, то ваша задача попадает на CPU, и, конечно, хотелось бы, чтобы на CPU она решалась настолько быстро, насколько это возможно.

Итак, для начала представим, что у нас нет не только ускорителей, но и расширений, и рассмотрим самую распространенную матричную операцию — умножение плотных матриц.

Из курса высшей математики вы можете помнить, что произведение матриц вычисляется следующим образом:

То есть, чтобы вычислить матрицу-произведение, вам нужно выполнить M х N скалярных произведений строк на столбцы. В виде алгоритма это можно записать так:

Кажется, что алгоритм предельно прост: мы обнуляем элементы матрицы-аккумулятора и затем в трех вложенных циклах выполняем M х N х K умножений-сложений. Число операций не сократить — значит, алгоритм уже не ускорить?

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

Что же делать?

Решает проблему достаточно простая идея, которая лежит в основе абсолютно всех высокопроизводительных библиотек линейной алгебры. Описана она в знаменитой статье «Анатомия высокопроизводительного матричного умножения» (Anatomy of High-Performance Matrix Multiplication), которую написали Kazushige Goto и Rovert A. Van De Geijn из Техасского университета в Остине.

Идея в том, что вы «нарезаете» матрицу B на горизонтальные полосы, которые по размеру могут уместиться в L3-кэш, матрицу A — на блоки, «влезающие» в L2-кэш, затем полосы матрицы B из L3-кэша разбиваете на столбцы, помещаемые в L1-кэш, как показано на схеме ниже. Аккумулируется произведение уже в регистрах, то есть в самой быстрой памяти.

Называется все это великолепие алгоритмом умножения матриц с многоуровневым блокированием. И по сравнению с тремя циклами обычного алгоритма это уже более монументальная конструкция из шести вложенных циклов:

В плане дальнейшей оптимизации есть некоторый простор для творчества: так, например, в библиотеке OpenBLAS оптимизация начинается с Loop 2 (macro-kernel), а в библиотеке BLIS низкоуровневые оптимизации возникают только в последнем цикле, Loop 0 (micro-kernel).

Представим, что мы реализовали многоуровневое блокирование для матриц произвольного размера, плитки матриц попали в самую быструю память, а число операций, как мы уже выяснили выше, не сократить. Неужели лимит ускорения исчерпан?

Лимит ускорения не исчерпан, если мы сможем распараллелить операции. Так, например, если у нас есть векторное расширение, то мы можем операции векторизовать: загрузить в векторные регистры K элементов из строки матрицы А, K элементов из столбца матрицы В и умножить элементы векторных регистров. Здесь мы предполагаем, что число K после всех описанных выше разбиений матрицы на блоки выбрано с учетом длины векторного регистра.

Таких операций требуется сделать M х N, то есть в K раз меньше, чем в реализации алгоритма матричного умножения без использования расширений. Но, конечно, эти операции более «дорогостоящие»: умножение совершается не над скалярами, а над векторами, плюс неизбежны накладные расходы на загрузку данных в векторные регистры и выгрузку. Поэтому, конечно же, ускорения вычислений в те же самые K раз ожидать не стоит. Однако ускорение умножений матриц даже в 2 раза просто за счет использования векторного расширения — это уже очень неплохо, если такие умножения занимают большую часть времени работы приложения.

И в этот момент у вас может возникнуть вопрос: а нельзя ли не разбивать матрицы на строки и столбцы, а работать прямо с плитками (то есть блоками)? Ответ прост: конечно же, можно, если у вас есть матричное расширение. Вместо одномерных регистров здесь используются двумерные, в которые загружаются плитки матриц-сомножителей, над ними выполняется умножение и плитка матрицы С чудесным образом появляется в еще одном двумерном регистре.

Конечно, такие матричные операции еще более «дорогостоящие» по тактам, чем векторные, но и требуется их меньше. Ведь плиток (блоков) в исходных матрицах явно меньше, чем строк и столбцов.

Но позволит ли матричное расширение получить ускорение по сравнению с векторным? Здесь мы можем вспомнить о такой метрике, как вычислительная интенсивность:

И если мы рассчитаем ее для умножения рассматриваемых блоков матриц, окажется, что для алгоритма с матричным расширением вычислительная интенсивность будет на порядок выше:

Таким образом, с точки зрения теоретических выкладок матричное расширение кажется более привлекательным для повышения производительности по сравнению с векторным. Но, конечно же, интересны показатели ускорения на практике.

Анализ данных из открытых источников показал, что «в среднем по больнице» матричные расширения позволяют получить ускорение от 2 до 12 раз по сравнению с векторными.

Для Arm SME сравнивалось число тактов (на симуляторе). Во всех остальных случаях — время счета.

По степени связанности с векторными расширениями матричные делят на три большие группы: интегрированные, гибридные и независимые.

В случае интегрированных матричных расширений в уже существующую архитектуру не добавляются никакие дополнительные регистры: для хранения как матриц-сомножителей, так и матриц-аккумуляторов используются векторные регистры. Примерами матричных расширений такого класса являются оба кастомных расширения RISC-V: SiFive Intelligence Extension и SpacemiT IME (Integrated Matrix Extension).

Если в архитектуру добавляется независимый блок регистров для матриц-аккумуляторов, то это уже гибридное матричное расширение. Именно к этой группе относятся Power MMA (Matrix Multiply Assist) и Arm SME (Scalable Matrix Extension).

И наконец, если в архитектурное пространство добавлены полноценные двумерные матричные регистры, в которых хранятся как плитки матриц-сомножителей, так и плитки матрицы-аккумулятора, и матричное расширение взаимодействует только с этими регистрами, то перед нами независимое матричное расширение. В этом случае матричное и векторное расширения абсолютно не зависимы друг от друга, и вполне возможна ситуация, когда на CPU есть матричное расширение, но нет векторного. Примеры независимых матричных расширений — Intel AMX (Advanced Matrix Extension) и Apple AMX (Attached Matrix Extension).

Теперь, когда мы получили общее понимание о матричных расширениях и их предназначении, рассмотрим, какие матричные расширения сейчас существуют на рынке, и попробуем понять, какие у них сходства и различия.

Если мы попытаемся построить таймлайн анонсов, появлений и даже обнаружений (было и такое у особо скрытных производителей) матричных расширений, то у нас получится следующий график.

Самое большое расхождение между анонсом и появлением реализации в железе у Intel AMX: в 2020 году оно было анонсировано, но мир увидел его только в 2023 году — на год позже, чем замеры производительности кастомного матричного расширения RISC-V от компании SiFive. То есть в плане матричных расширений RISC-V, несмотря на свою молодость, движется наравне с такой зрелой архитектурой, как х86, в чем-то даже опережая ее. Интересное наблюдение, но давайте посмотрим детальнее, что известно о существующих матричных расширений.

Первое расширение, которое мы рассмотрим, это Intel Advanced Matrix Extension (AMX). Его анонс вышел еще в 2020 году, но в железе оно появилось только в 2023 году, в четвертом поколении серверных процессоров Intel Xeon — Sapphire Rapids. Это независимое от векторного матричное расширение, включающее 12 новых инструкций, которые можно разделить на три группы:

1. конфигурирование матричных регистров («тайлов» или плиток),
2. управление матричными регистрами (load/store, clear, set и другие),
3. операции над плитками матриц в матричных регистрах.

Что же из себя представляют независимые матричные регистры в Intel AMX? Это новый расширяемый двумерный (2D) регистровый файл. Сейчас он состоит из восьми регистров Т0-Т7 по 1 Кб каждый. Из матричных операций пока реализована только одна — TMUL, инструкция матричного умножения (как уже отмечалось выше, это основной хот-спот во многих приложениях). TMUL использует три регистра: T2 += T1 * T0.

Расширение поддерживает типы данных пониженной и половинной точности: INT8, UINT8, BF16, FP16, Complex FP16. Можно сделать вывод, что его разработчики были сфокусированы на ускорении инференса нейронных сетей. При этом аккумуляторы всегда 32-битные: INT32, UINT32 или FP32.

Что мы знаем про производительность этого расширения из открытых источников? Например, есть исследование, в котором рассматривается сверточная нейронная сеть ResNet50 (FP32, batch size = 16). Векторизация на основе Intel VNNI (AVX-512 Vector Neural Network Instructions) позволяет ускорить вычисления в четыре раза по сравнению со скалярной версией, а матричное расширение — почти в 9 раз. Получается, на этом кейсе матричное расширение быстрее векторного более чем в два раза.

Следующее расширение — Apple Attached Matrix Extension (AMX). Это, опять же, независимое матричное расширение. Официальной документации у него нет. А само расширение обнаружили в 2020 году при реверс-инжиниринге процессора Apple M1.

Подробнее об этой детективной истории можно почитать в статье Если кратко, то разработчики Arm долгое время не хотели добавлять кастомные расширения, поскольку это приводит к фрагментации экосистемы. Но, в конце концов, сделали небольшой шаг навстречу в 2019 году. Разрешили производителям добавлять кастомные инструкции и целые расширения при условии, что все взаимодействие с ними будет происходить в библиотеках того же производителя, а пользователь напрямую с инструкциями работать не сможет.

Что мы знаем об этом расширении? Его инструкции были обнаружены при мониторировании вызовов математических библиотек vImage, vDSP, BNNS, BLAS/LAPACK. А значит, оно используется для ускорения обработки изображений и сигналов, операций линейной алгебры и работы с нейронными сетями в приложениях AR/VR, CV и ML. По сравнению с Arm Neon ускорение получается в среднем в два раза.

Архитектурное пространство, выделенное для регистров данного расширения, составляет 5 Кб. Для сомножителей отводится по восемь регистров (64-байтовых векторов) на каждый из двух операндов (x0.7 и y0.7), элементы которых могут иметь следующие типы данных: INT8, UINT8, INT16, UINT16, FP16.64. В Apple M2 также добавлена поддержка BF16. Все остальное пространство расширения (z на схеме ниже) отводится под аккумуляторы и может быть сконфигурировано следующими способами:

• 1 регистр размером 64×64 (32×32) элемента по 8 (32) бита.
• 2 регистра размером 32×32 по 16 бит.
• 4 регистра размером 16×16 по 32 бита.
• 8 регистров размером 8×8 по 64 бита.
• 64 регистра, каждый из которых — 64-байтовый вектор-строка.

Расширение Arm Scalable Matrix Extension (SME) анонсировали в 2022 году для Armv9-A. Профиль «A» (application) соответствует высокопроизводительным ядрам, которые используются в смартфонах, планшетах и других устройствах, вплоть до оборудования центров обработки данных.

Cудя по всему, первым процессором с реализацией Arm SME стал Apple M4, представленный 7 мая 2024 года. Пока еще нет детального сравнения производительности с Arm Neon — в открытом доступе можно найти только первые замеры, поэтому далее мы рассмотрим только результаты сравнения на симуляторе.

Это гибридное расширение, то есть для сомножителей используются векторные регистры — SVE или SVE2, а для матрицы-аккумулятора вводится один дополнительный 2D матричный регистр (по последним данным это дополнительное пространство может конфигурироваться в несколько регистров, как z в описанном выше Apple AMX). Элементы аккумуляторов могут быть 32- или 64-битными, в то время как элементы сомножителей могут иметь ширину от 8 до 64 бит. Известно, что первая версия SME содержит инструкцию умножения матриц, а во второй добавились инструкции для матрично-векторного умножения (GEMV) и нелинейных решателей, а также поддержка разреженных матриц (но процессоров с SME2 пока нет).

На конференции The International Conference for High Performance Computing, Networking, Storage, and Analysis (SC'22), один из докладов был посвящен моделированию производительности Arm SME. Для симуляции авторы использовали модель ядра Fujitsu A64FX (именно оно используется в упомянутом ранее гомогенном суперкомпьютере Fujitsu Fugaku). Тестовая задача состояла из 200 матричных умножений (данные одинарной точности). Сравнивались векторная (на основе SVE) и матричная реализации. В зависимости от размера матриц SME позволило уменьшить число тактов от 2,2 до 6,4 раз по сравнению с SVE. Конечно, это оценка на симуляторе и на реальном железе цифры могут отличаться, однако результат интересный.

Гибридное матричное расширение Power Matrix Multiply Assist (MMA) появилось в 2021 году в Power10 — семействе суперскалярных симметричных мультипроцессоров на базе архитектуры POWER от IBM. Несмотря на то, что дополнительного пространства для аккумуляторных матричных регистров в архитектуре не появилось, расширение Power MMA позиционируется именно как гибридное, поскольку половину векторных регистров выделили под аккумуляторы насовсем. Эти 32 векторных регистра разбили на восемь групп по четыре регистра, и в итоге получили восемь 512-битных аккумуляторных регистров. Остальные 32 векторных регистра используются для операций векторного расширения VSX, а также для матриц-сомножителей в инструкциях матричного расширения. Поддерживаемые типы данных: BF16, FP16.64, INT4.16.

Разработчики Power MMA задались очень правильным вопросом: как с минимальными затратами обеспечить ускорение для широкого спектра приложений? В итоге они добавили реализации матричного умножения с использованием Power MMA для различных типов данных в две библиотеки линейной алгебры: OpenBLAS и Eigen. Эти библиотеки используются во многих популярных фреймворках, таких как NumPy, PyTorch, TensorFlow, OpenCV и других, и влияют на производительность многих приложений.

Рассмотрим результаты, представленные в исследовании производительности производительности Power MMA на нескольких тестовых кейсах. Так, в случае умножения плотных матриц с элементами двойной точности (DGEMM) векторное расширение VSX позволяет ускориться чуть больше чем в два раза по сравнению со скалярным кодом, а матричное — в 5,5 раз.

Также в исследовании рассматривалось влияние типа данных на ускорение инференса нейронной сети ResNet50 по сравнению со скалярной версией: если в случае данных одинарной точности Power MMA дает ускорение в 10 раз, то для 8-битных — уже в 21 раз.

Для LINPACK бенчмарка, представляющего собой решение последовательности плотных систем линейных алгебраических уравнений с матрицами разного размера, Power MMA показывает масштабируемость, превосходящую и скалярную (плато на графике для систем с размерностью, превышающей 1024) и векторную (выходит на плато начиная с 4096 неизвестных) реализации.

Интересующимся будет полезно изучить Best Practices Guide от разработчиков Power MMA. Интересно, что один из авторов этой книги, Хосе Морейра, также является председателем Vector Special Interest Group (SIG) консорциума RISC-V International. Именно он инициировал создание группы разработки интегрированного матричного расширения RISC-V IME TG и стал ее вице-председателем. И, насколько я могу судить по открытым источникам, многие идеи из Power MMA плавно перетекают в это расширение, о котором мы поговорим в следующей «главе».

В 2022 году появилось первое матричное расширение для RISC-V. Это кастомное расширение, интегрированное в векторное, производства компании SiFive. Матричные инструкции оперируют данными, хранящимися в 512-битных векторных регистрах. Расширение позиционируется как подходящее для ускорения инференса, приложений AR/VR, приложений из систем IVI, цифровых и IP-камер, игровых устройств. У него уже есть поддержка в TensorFlow Lite, а на сайте SiFive доступны спецификации для типов данных INT8 и BFloat 16.

В исследовании от компании Google представлены оценки производительности SiFive Intelligence Extension. Так, для умножения матриц с элементами типа INT8 говорится об ускорении в 12 раз по сравнению с векторным расширением RISC-V RVV, а для инференса сверточной нейронной сети MobileNet v1 (batch size = 1, VLEN = 512 бит) — в 6 раз:

Совсем недавно — в конце апреля 2024 года — стало известно о SpacemiT Integrated Matrix Extension (IME). Это еще одно кастомное матричное расширение RISC-V от компании SpacemiT и снова — интегрированное в векторное. Из спецификации расширения следует, что поддерживается только одна матричная операция — умножение плотных матриц, но при помощи параметра Copу можно задавать число операций, выполняемых параллельно. Сомножители и аккумуляторы хранятся в 256-битных векторных регистрах. Поддерживаются следующие типы данных: BF16, FP4.16, INT4.16.

Уже доступна для заказа плата Banana Pi BPI-F3 с 8-ядерным процессором SpacemiT K1 RISC-V, все ядра которого включают векторное расширение RVV 1.0 и четыре из них — матричное расширение SpacemiT IME. Заявлено, что каждое ядро на 30% превосходит по производительности Arm Cortex A55, суммарная производительность 2.0 TOPs AI (подробнее здесь).

Матричные расширения начали появляться относительно недавно — с 2020 года, и в этом направлении x86, Power, Arm и RISC-V движутся практически наравне. Такие расширения подходят для ускорения приложений AI/ML, AR/VR, CV, ADAS, HPC. Вот сводная таблица по основным характеристикам существующих расширений, описанных в тексте:

Очень удачным ходом со стороны разработчиков расширений является добавление оптимизаций в такие широко используемые библиотеки линейной алгебры, как OpenBLAS и Eigen (что было сделано в свое время компанией IBM для Power MMA).