Заглянем в хрустальный шар: как продвигается разработка стандартных матричных расширений RISC-V

В конце 2023 года в международном консорциуме были сформированы две рабочие группы (Task Group): одна разрабатывает интегрированное матричное расширение (IME), а вторая — независимое (AME). Их роадмапы в целом похожи: месяц на запуск и «раскачку», четыре месяца — на анализ требований, шесть месяцев — на разработку спецификации, шесть месяцев — на proof-of-concept реализации и анализ.

Помимо этого, участники рабочих групп разработают патчи для поддержки стандартных матричных расширений RISC-V в компиляторах, эмуляторах, ядре Linux, oneDNN. А группа разработки независимого расширения также планирует включить оптимизации в TensorFlow и PyTorch.

Теперь давайте рассмотрим каждое из расширений чуть подробнее.

Начнем мы с независимого матричного расширения, которое, судя по всему, во многом будет основано на разработках компании T-Head.

В своей презентации на RISC-V Summit China в августе 2023 года компания T-Head выделила следующие качества, которыми должно обладать стандартное матричное расширение RISC-V:

• Эффективность — повышение производительности AI-приложений за счет ускорения матричных операций.
• Универсальность и масштабируемость — решение для широкого спектра оборудования от маломощных периферийных устройств до высокопроизводительных серверов в дата-центрах.
• Гибкость — поддержка широкого диапазона типов данных и размеров матриц.
• Переносимость двоичного кода (размеры регистров, аппаратная реализация).
• Перспективность — расширяемое до полноценного AI Matrix ++ Extension (за счет добавления большого числа операций над матрицами, встречающихся в AI-приложениях).

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

• На оборудовании разного класса наблюдаются совершенно разные пропорции между скалярными, векторными и матричными нагрузками. Если расширение независимое, подстроиться под них легче.
• Векторное расширение RISC-V (RVV) содержит сотни инструкций, а независимое матричное расширение в его текущей реализации — достаточно легковесная конструкция (порядка 20 инструкций). А это упрощает разработку программного обеспечения.
• Упрощается тепловое проектирование CPU за счет того, что энергозатраты на векторное и матричное расширения можно рассчитывать независимо.
• Более гибкая аппаратная реализация: независимое расширение может быть как юнитом в составе ядра, так и вынесено в сопроцессор и т. д.

Эксперты T-Head отмечали, что независимое расширение не лишено минусов, но они уже выработали подходы для их минимизации. Например, на аппаратном уровне они хотят переиспользовать векторные/матричные регистры, чтобы снизить затраты ресурсов. Кроме того, будут добавлены инструкции, позволяющие избежать или минимизировать пересылки данных между векторными и матричными регистрами, а также специальные флаги используемых регистров. Все это уменьшит накладные расходы на переключение контекста.

Сейчас в реализации T-Head восемь двумерных матричных регистров:

Длина строки матричного регистра равняется RLEN и, соответственно, число столбцов определяется шириной элемента, с которым вы работаете: если ширина элемента составляет N бит, то число столбцов равно RLEN/N. При этом относительно числа строк в регистре действует правило: оно равно RLEN/32. Здесь RLEN = 2^{n+6} = {128, 256, 512, …} бит, N = {4, 8, 16, 32, 64} бит. Таким образом, размер каждого регистра составляет MLEN = RLEN² / 32 = 2^{2n+7} бит.

Выше мы отметили, что T-Head RVM ISА достаточно легковесна и состоит всего из 20 с небольшим инструкций. Их можно условно разделить на несколько основных групп:

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

В презентации разработчики T-Head обозначили следующие майлстоуны дорожной карты разработки независимого матричного расширения:

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

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

Сравнение логов после запуска основных демо-приложений, таких как матричное умножение (GEMM) и инференс нейронной сети ResNet50, для типов данных INT8 и FP16 позволяют сделать вывод, что при использовании матричного расширения число инструкций сокращается от 2 до 14 раз по сравнению с использованием RVV. Насколько дорогими они будут, конечно же, пока неизвестно. Эти цифры стоит рассматривать как предварительные оценки.

В отличие от независимого расширения, второе стандартное матричное расширение RISC-V, интегрированное, в прогрессе отстает. Работы по организации IME TG (Integrated Matrix Extension Task Group) были инициированы RISC-V Vector SIG (Special Interest Group) только в сентябре 2023 года. Катализатором стала активность компании T-Head по созданию группы для разработки стандартного независимого матричного расширения. Пока IME TG формирует лишь общие предположения об использовании векторных регистров для матричных операций.

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

Для оптимального разбиения матриц на блоки применяют описанный в первой статье алгоритм из работы Goto и Geijn. Как и в библиотеке BLIS, оптимизироваться будет последний цикл — micro-kernel.

Архитектурное пространство векторного расширения состоит из 32-х векторных регистров. Руководство RISC-V IME TG предполагает, что половина из них, как и в Power MMA, будет отведена под аккумуляторы. Таким образом, остается еще 16 векторных регистра, и участники группы до сих пор обсуждают, какая их часть будет использоваться для блоков матриц-сомножителей. Это очень важный вопрос, поскольку это число определит, сколько элементов матриц-сомножителей будет загружаться на каждом шаге, — по факту вычислительную интенсивность алгоритма:

Сейчас рабочая группа рассматривает следующие варианты декомпозиции блоков матриц по векторным регистрам:

1. Выделить по четыре регистра на каждый из сомножителей, и в каждом регистре хранить плитку матрицы размера 𝜆 x 𝜆.
2. Брать 16 элементов столбца из матрицы, А и L элементов строки из матрицы B, при этом L равно числу элементов, которые могут быть размещены в одном векторном регистре.
3. Сконфигурировать векторный регистр таким образом, что один его элемент — это матрица 𝜆 x 𝜆. Для элементов матрицы В при этом предлагается выделить два регистра, соответственно из матрицы, А необходимо загрузить 8 таких плиток.
4. Брать по 4𝜆 элементов столбца, А и строки B, при этом в аккумуляторных регистрах так же, как и в первом случае размещаются плитки размера 𝜆 x 𝜆.

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

В предложении о создании IME TG в качестве основной цели указывалось достижение не менее 90% от пиковой производительности для GEMM. А для прочих алгоритмов, в которых внешнее произведение векторов используется в качестве составной части, ожидается ускорение не менее чем в 2 раза по сравнению с векторной реализацией. Поскольку разрабатываемый micro-kernel потенциально может использоваться и в свертках, руководство группы планирует впоследствии добавить поддержку типов данных пониженной и половинной точности, характерных для AI/ML-приложений.

Напоследок — последние новости про разреженные матрицы. Эта тема возникла в декабре 2023 года на заседаниях RISC-V Vector SIG. Председатель группы Хосе Морейра отметил, что операции над разреженными матрицами — сложная проблема, и начинать обсуждать ее нужно как можно раньше. По его мнению, поддержка разреженных матриц обязательно должна быть во второй версии интегрированного матричного расширения, но очень важно учитывать это при проектировании первой версии расширения, чтобы избежать коллизий в будущем.

Кроме того, к настоящему моменту нет библиотек разреженной линейной алгебры, содержащих оптимизации под RISC-V — библиотеки OpenBLAS и BLIS содержат только BLAS/LAPACK-алгоритмы, то есть предназначены для работы с плотными матрицами и векторами. При этом алгоритмы Sparse BLAS/Solvers очень важны для HPC-моделирования (большие разреженные матрицы и системы линейных алгебраических уравнений возникают в результате построения дискретных аналогов дифференциальных уравнений в частных производных, описывающие моделируемые физико-технические процессы) и DL-приложений. Поддержка разреженных матриц в матричных расширениях RISC-V значительно расширила бы возможности для оптимизации алгоритмов такого класса.

Как и в случае плотных матриц, группа разработки интегрированного расширения рассматривает только операцию — умножение разреженной матрицы S на плотную D. Предполагается, что разреженность (процент нулевых элементов, sparsity) матрицы S составляет 50−90%. Такие числа характерны для матриц весов из DL-моделей. Для HPC-приложений разреженность может превышать 99%, поэтому, вероятно, они потребуют иного подхода.

Снова рассматриваются не скалярные, а внешние произведения векторов, но теперь первый вектор — разреженный:

Подчеркнем, что аккумулятор в этом случае остается плотной матрицей.

В чем же заключается отличие от произведения плотных матриц? Допустим, мы не знаем, что левый столбец у нас разреженный и вычисляем внешнее произведение по обычной формуле:

Пусть для определенности у нас выбран следующий способ декомпозиции блоков матриц-операндов по векторным регистрам:

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

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

В рассматриваемом случае мы можем загружать столбец матрицы S не в два векторных регистра, а в один, если уплотнить данные, отбросив нули. Далее внешнее произведение вычисляется как обычно, но его строки нужно прибавлять к корректным строкам матрицы-аккумулятора. Поэтому очевидно, что нужно загружать и маску индексов портрета.

Как выяснилось, в XuanTie Matrix Extension, независимом матричном расширении, разрабатываемом компанией T-Head, уже есть наброски реализации Sparsity Subset. Как и в IME TG, здесь рассматривается пока только умножение с той лишь разницей, что плотная матрица и разреженная поменялись местами, то есть разреженная матрица выступает вторым сомножителем:

Еще одна особенность матричного расширения от компании T-Head заключается в том, что второй сомножитель (в данном случае — разреженную матрицу) мы предварительно транспонируем, а вместо внешнего произведения векторов в алгоритме используется скалярное:

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

Главное отличие от подхода, рассматриваемого в IME TG, в том, что через маску «просеивают» не строки аккумулятора, а элементы строки плотной матрицы-сомножителя: в матричный регистр в итоге загружаются только те элементы плотной матрицы, которые будут умножаться на загруженные ненулевые элементы разреженной матрицы. Далее вычисление скалярного произведения и аккумуляция происходят как обычно.

При таком подходе в независимом матричном расширении появляется:

• два новых дополнительных 2D матричных регистра для маски индексов портрета,
• два control state регистра (CSR) для значения sparsity и числа ненулевых элементов в плитке разреженной матрицы,
• три инструкции группы Sparse Matrix MACC для умножения на разреженную матрицу с накоплением,
• шесть новых инструкций в группе Memory access для доступа к регистрам маски индексов портрета.

Прямо сейчас разрабатываются целых два стандартных матричных расширения открытой архитектуры RISC-V — независимое и интегрированное. Разработчики независимого расширения фокусируются на приложениях AI/ML, а интегрированного — на HPC, но обе группы единогласны в вопросе необходимости поддержки операций над разреженными матрицами. Это значит, что со временем в RISC-V будет много возможностей для оптимизации алгоритмов класса Sparse BLAS и даже Sparse Solvers (решателей разреженных систем уравнений).

Примерно в середине 2025 года мы, скорее всего, уже увидим ратифицированные версии независимого и интегрированного матричных расширений RISC-V. А поскольку это открытые разработки, за всеми изменениями в проектах можно следить на GitHub соответствующих рабочих групп (IME TG и AME TG).