Погружение в матрицу: расширение RISC-V от T-Head

17 min

879

YADRO corporate blogHigh performance*Algorithms*Mathematics*Machine learning*

Продолжим нашу «антологию матричных расширений» текстом про независимое матричное расширение RISC-V от компании T-Head.

Почему мы рассматриваем именно его? Интересно понять, что из себя представляет будущее стандартное матричное расширение RISC-V, попробовать реализовать алгоритм с его использованием, соотнести это со своим предыдущим опытом низкоуровневых оптимизаций. Кроме того, это интересная возможность попробовать написать программу для расширения, которого еще нет ни в одном процессоре, и запустить ее на эмуляторе.

А еще ISA этого расширения весьма минималистична и, на мой взгляд, идеально подходит для тех, кто никогда не использовал матричные расширения в своем коде, но хочет попробовать (или узнать, как это выглядит изнутри). Не переживайте, текст не требует опыта низкоуровневых оптимизаций математических библиотек: погружение в матрицу будет постепенным.

Этот текст — третий в «цикле» про матричные расширения. Сначала я рассказала про матричные расширения в целом: зачем нужны, как работают и какие существуют. А затем сфокусировала внимание на разрабатываемых расширениях для RISC-V.

Интро про T-Head RVM

Вспомним, что мы узнали из предыдущего текста. T-Head RVM — это независимое матричное расширение RISC-V. В базовой комплектации (без Sparsity Subset, который сейчас активно разрабатывается) содержит восемь двумерных матричных регистров, которые используются и для сомножителей, и для аккумуляторов. Длина строки регистра равна RLEN бит, а число строк определяется как RLEN/32. Число столбцов зависит от ширины элементов матрицы и равно RLEN/N бит, где N — ширина элемента.

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

ISA расширения достаточно легковесна: содержит порядка 20 инструкций. Это инструкции для:

умножения блоков матриц и аккумуляции,
прочих операций над матрицами (список таких операций T-Head сейчас расширяет),
загрузки и выгрузки данных из матричных регистров,
копирования данных между регистрами разного типа,
конфигурирования матричных регистров, их освобождения и т.д.

В статье мы разберем эти инструкции, а также соответствующие им интринсики, и рассмотрим их особенности.

Репозиторий

Спецификацию расширения, демо-приложения и другие полезные вещи вы можете найти в репозитории. Вы можете склонировать его себе для дальнейшей работы. У репозитория следующая структура:

|--spec                    ## The RISC-V Matrix Extension specification
|--doc/                     ## The user guide for tools  
    |--shl                   ## The SHL 2.0 user guide
    |--abi                   ## The Matrix Extension ABI Manual
    |--intrinsic           ## The Matrix Extension intrinsic API Reference Manual
|--shl/                      ## A neural networks library using RISC-V Matrix Extension
|--hhb/                     ## A toolkit used for deploying neural network models
|--qemu/                  ## Emulator
|--xuantie-gnu-toolchain/   ## GNU toolchain
    |--riscv-gcc/                   ## Compiler
    |--riscv-binutils-gdb/      ## Assembler
|--demos/               
    |--resnet50            ## A resnet50 evaluation demo using nn library
    |--GEMM              ## A GEMM evaluation demo using intrinsic

Для быстрого запуска демо-приложений необходимо выполнить два скрипта из папки demos:

cd demos/
./env.sh
./run.sh

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

Хочу обратить внимание, что репозиторий содержит библиотеку shl. Это высокопроизводительная библиотека гетерогенных вычислений от T-Head. Среди ее исходников вы можете найти ассемблерные реализации алгоритмов с использованием матричного расширения и C-реализации алгоритмов с использованием интринсиков векторного расширения. Если хотите углубиться в тему низкоуровневых оптимизаций математических библиотек, советую их изучить. Это может быть полезно для обучения и вдохновения при написании своих реализаций алгоритмов.

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

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

Спецификация

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

с целочисленными типами, как со знаком, так и без него, шириной от 8 до 64 бит,
с числами с плавающей точкой шириной от 16 до 64 бит.

В таблице ниже приведены используемые в репозитории обозначения. Типу данных <dtype>_t в API интринсиков соответствует аббревиатура <atype>, а в названиях инструкций ширина элемента обозначается суффиксом <sbit>.

*Поддерживаемые типы данных и используемые обозначения*

Префикс “m” в названии типа говорит о том, что это не скаляр, а матрица: так, тип m<dtype>_t соответствует одному матричному регистру, элементы которого имеют тип <dtype>_t. Для чисел с плавающей точкой в названии матричных типов вам может встретиться постфикс “x2”: m<dtype>x2_t — это объединение двух матричных регистров, элементы которых имеют тип <dtype>_t.

Зачем такие конфигурации существуют и почему они встречаются только для float16..64, мы обсудим дальше, когда будем разбирать инструкции для умножения матриц. Пока этот момент надо просто запомнить.

Кроме того, в названиях матричных типов вы можете встретить постфикс “v”: m<dtype>v_t — это обозначение вектора в матричном регистре. В текущей реализации такому типу данных соответствует непосредственно матричный регистр и номер строки внутри него.

Теперь перейдем к операциям, которые мы можем выполнять.

Конфигурирование матричных регистров

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

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

Первые 8 бит этого числа соответствуют значению (это число строк в блоках первого сомножителя, матрицы , и аккумулятора — матрицы ),
Следующие 8 бит — это (число столбцов в блоках второго сомножителя, матрицы , и аккумулятора — матрицы ).
Последние 16 бит — это $k = K \cdot \mbox{element width}$ , где — число столбцов в блоке первого сомножителя, матрицы , и строк в блоке второго сомножителя, матрицы .

Последний момент необходимо запомнить и учитывать при конфигурировании регистров.

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

// set M = mnk[0:7], N = mnk[8:15], k = mnk[16:31]
uint32_t mcfg(uint32_t mnk);

или соответствующую ему ассемблерную инструкцию:

# rd - new matrix configuration
mcfg rd, mnk

Кроме того, вы можете конфигурировать по одной размерности. В случае конфигурирования размерностей / вы передаете число строк/столбцов в следующие интринсики:

// set <M/N> = size 
uint32_t mcfg<m/n> (uint8_t size);

B случае конфигурирования размерности не забывайте, что задавать надо не , а $k = K \cdot \mbox{element width}$ :

// set k = size ( ! K = k / sizeof(element) ! )
uint32_t mcfgk (uint16_t size);

Соответствуют этим интринсикам конфигурирования отдельных размерностей следующие инструкции:

mcfg<m/n/k> rd, size

Например, если мы хотим сконфигурировать матричные регистры для случая и dtype = float32, то мы можем выполнить следующее:

mcfgm(4); 
mcfgn(4); 
uint32_t mnk = mcfgk(4 * sizeof(float));

После этого вы можете проверить значение mnk, чтобы узнать, удалось ли сконфигурировать регистры желаемым образом.

Загрузка данных

Вы сконфигурировали матричные регистры. Теперь нам необходимо загрузить в них данные. Здесь возможны два варианта.

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

// [stream] matrix load (row length - stride byte) 
m<dtype>_t m[s]ld_<atype>(<dtype>_t *base, long stride);

Этому интринсику соответствует следующая инструкция:

# md - matrix register (m0..7)
m[s]ld<sbit> md, (*base), stride

Например, для корректной загрузки блока по указателю block_b с размером, соответствующим конфигурации матричного регистра, из матрицы с элементами типа float и числом столбцов ROW_SIZE, необходимо выполнить:

mfloat32_t mb = mld_f32(block_b, ROW_SIZE * sizeof(float));

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

// whole matrix register load
m<dtype>t mld1m<atype>(<dtype>t base);

Данному интринсику соответствует следующая инструкция:

mld1m md, (base)

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

// whole 2 matrix registers load - only for float!
m<dtype>x2_t mld2m_<atype>(<dtype>_t *base);

Последнему интринсику соответствует инструкция:

mld2m md, (base)

Выгрузка данных

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

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

// [stream] matrix store (row length - stride byte) 
void m[s]st_<atype>_m<atype>(<dtype>_t *base, long stride, m<dtype>_t ms);

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

mst_f32_mf32(block_b, ROW_SIZE * sizeof(float), mb);

Соответствующая инструкция выглядит так:

# ms – matrix register
m[s]st<sbit> ms, stride, (*base)

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

// whole matrix register store
void mst1m_<atype>_m<atype>(<dtype>_t *base, m<dtype>_t ms);

При реализации на ассемблере для этих целей используйте команду:

mst1m ms, (*base)

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

// whole 2 matrix registers store - only for float!
void mst2m_<atype>_m<atype>x2(<dtype>_t *base, m<dtype>x2_t ms);

Соответствует этому следующая инструкция:

mst2m ms, (*base)

Матричное умножение

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

# md, ms2, ms1 - matrix registers
fmmacc.<sbit> md, ms2, ms1

Здесь происходит умножение двух блоков матриц из двух матричных регистров ms1 и ms2, а также добавление их произведения к содержимому третьего матричного регистра, md, с сохранением в него же. У элементов матриц — тип float<sbit>.

Какие здесь есть подводные камни? Во-первых, матрицу , согласно документации, необходимо предварительно транспонировать, то есть инструкция выполняет следующее:

$C_{M\times N}+=A_{M\times K}\cdot B^{T}_{N\times K}$

Таким образом, в регистре ms2 находится блок размером не $K \times N$ элементов, а $N \times K.$ Кроме того, выше уже упоминалось о том, что есть ограничения на максимальное число строк и столбцов в матричном регистре. А в случае матричного умножения у нас еще и взаимосвязаны размеры блоков операндов и аккумулятора, что накладывает дополнительные ограничения. С учетом всего перечисленного получаем следующие максимальные значения размеров блоков при умножении матриц с элементами типа float16..64:

*Конфигурирование матричных регистров для операции fmmacc*

И вот здесь есть интересный момент.

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

А вот в случае элементов с половинной или двойной точностью у нас возникает более интересная конфигурация: появляются блоки разных размеров. Например, для float16 у блока матрицы размер $M \times K = 4 \times 8$ . Блок матрицы имеет размер $M \times N$ , что в случае float16, как видно из таблицы выше, также соответствует размеру $4 \times 8$ . Однако размер блока матрицы B от них отличается: $N \times K = 8 \times 8$ .

Возникает естественный вопрос: как это соотнести с тем, что остальные регистры вмещают блоки $4 \times 8$ ? На самом деле, достаточно просто. Это означает, что для блока матрицы нам необходимо использовать два регистра $4 \times 8$ , и в сумме они дадут $8 \times 8$ . Собственно, для этого и требуются упоминаемые выше сдвоенные регистры.

По этой причине API интринсиков для умножения матриц выглядит не так просто, как соответствующая инструкция: в интринсиках необходимо указывать, для каких операндов требуются сдвоенные регистры. Напоминаю, что в таких типах данных у нас появляется постфикс “x2” после <dtype>, а у интринсиков следующий вид:

// dest_{MxN} += src2_{MxK} * src1_{NxK}
mfloat16_t fmmacc_mf16x2_mf16(mfloat16_t dest,  mfloat16x2_t src1, mfloat16_t src2);
mfloat32_t fmmacc_mf32(mfloat32_t dest, mfloat32_t src1, mfloat32_t src2);
mfloat64x2_t fmmacc_mf64(mfloat64x2_t dest, mfloat64_t src1, mfloat64_t src2);

Матричное умножение с увеличением ширины аккумулятора

Казалось бы, с умножением матриц с элементами типа float16..64 мы разобрались. Но в рассмотренном случае элементы аккумулятора и сомножителей были одного и того же типа, то есть имели одинаковую ширину. Однако во многих задачах для получения результатов с приемлемой погрешностью при использовании половинной или одинарной точности для элементов матриц-сомножителей требуется использовать более широкие аккумуляторы. Поэтому спецификация матричного расширения T-Head содержит инструкцию для умножения матриц с элементами типа float16 или float32 и аккумуляцией результата в матрицу с элементами типа float32 или float64 соответственно, то есть с увеличением ширины аккумулятора в два раза:

# md, ms2, ms1 - matrix registers
fwmmacc.<sbit> md, ms2, ms1

Отличие от названия предыдущей инструкции — в букве “w” (widen).

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

*Конфигурирование матричных регистров для операции fwmmacc*

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

// widen dest_{MxN} += src2_{MxK} * src1_{NxK}
mfloat32_t fwmmacc_mf16(mfloat32_t dest, mfloat16_t src1, mfloat16_t src2);
mfloat64x2_t fwmmacc_mf32(mfloat64x2_t dest, mfloat32_t src1, mfloat32_t src2);

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

Матричное умножение целочисленных матриц

Матричное расширение T-Head позволяет выполнять такую операцию для сомножителей с 8-битными или 16-битными элементами. При этом для целочисленного умножения в T-Head RVM аккумулятор всегда шире в четыре раза по сравнению с сомножителями, в названии соответствующих команд это подчеркивается буквой “q” — quarter. Поэтому ограничения на максимальные значения , и в зависимости от типа элементов матриц-сомножителей отличаются от рассмотренных выше случаев:

*Конфигурирование матричных регистров для операций mmaqa[u], mmaqaus, mmaqasu*

Проанализировав таблицу, можно сделать вывод, что в случае целочисленного умножения матриц с 16-битными элементами для аккумулятора нам потребуется сдвоенный матричный регистр. Если тип элементов обеих матриц-сомножителей знаковый либо, наоборот, беззнаковый, API интринсиков следующее:

// 4x widen dest_{MxN} += src2_{MxK} * src1_{NxK}
mint32_t mmaqa[u]_m[u]i8(mint32_t dest, m[u]int8_t src1, m[u]int8_t src2);
mint64x2_t mmaqa[u]_m[u]i16(mint64x2_t dest, m[u]int16_t src1, m[u]int16_t src2);

Соответствует этим функциям инструкция:

# md, ms2, ms1 - matrix registers
mmaqa[u].<sbit> md, ms2, ms1

Если же тип элементов матрицы беззнаковый, а матрицы — нет, то необходимо использовать функции:

// 4x widen dest_{MxN} += src2_{MxK} * src1_{NxK}
mint32_t mmaqaus_mui8_mi8(mint32_t dest, muint8_t src1, mint8_t src2);
mint64x2_t mmaqaus_mui16_mi16(mint64x2_t dest, muint16_t src1, mint16_t src2);

На ассемблере им соответствует следующая инструкция:

# md, ms2, ms1 - matrix registers
mmaqaus.<sbit> md, ms2, ms1

Если же наоборот, элементы матрицы беззнаковые, а — нет, то для их умножения используем функции:

// 4x widen dest_{MxN} += src2_{MxK} * src1_{NxK}
mint32_t mmaqasu_mi8_mui8(mint32_t dest, mint8_t src1, muint8_t src2);
mint64x2_t mmaqasu_mi16_mui16(mint64x2_t dest, mint16_t src1, muint16_t src2);

Инструкция в этом случае следующая:

# md, ms2, ms1 - matrix registers
mmaqasu.<sbit> md, ms2, ms1

Поэлементные операции

Теперь рассмотрим более простые в плане конфигурирования регистров операции — поэлементные. Пока их можно производить только для матриц с элементами типа <dtype> = int{32,64}.

Поэлементное сложение. Вы можете поэлементно сложить элементы из двух матричных регистров и результат сохранить в третий матричный регистр:

// matrix-matrix add: src1 + src2 
m<dtype>_t madd_m<atype>(m<dtype>_t src1, m<dtype>_t src2);

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

// matrix-vector add (each row of src2 + src1[row index])
m<dtype>_t madd_m<atype>_m<atype>v(m<dtype>_t src1, m<dtype>_t src2, uint8_t index);

Кроме того, можно к каждому элементу матричного регистра прибавить одно и то же число:

// matrix-scalar add (each element of src1 + src2)
m<dtype>_t madd_m<atype>_<atype>(m<dtype>_t src1, <dtype>_t src2);

Перечисленным операциям поэлементного сложения соответствуют следующие инструкции:

# md, ms, ms1, ms2 -  matrix registers
madd.<sbit>.mm md, ms2, ms1
madd.<sbit>.mv.i md, ms2, ms1[index]
madd.<sbit>.mx md, ms, rs

Поэлементное вычитание. Операция происходит аналогично, с той лишь разницей, что “add” в названиях интринсиков меняется на “sub”:

// matrix-matrix sub: src1 - src2 
m<dtype>_t msub_m<atype>(m<dtype>_t src1, m<dtype>_t src2);
// matrix-vector sub (each row of src2 - src1[row index])
m<dtype>_t msub_m<atype>_m<atype>v(m<dtype>_t src1, m<dtype>_t src2, uint8_t index);
// matrix-scalar sub (each element of src1 - src2)
m<dtype>_t msub_m<atype>_<atype>(m<dtype>_t src1, <dtype>_t src2);

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

# md, ms, ms1, ms2 - matrix registers
msub.<sbit>.mm md, ms2, ms1
msub.<sbit>.mv.i md, ms2, ms1[index]
msub.<sbit>.mx md, ms, rs

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

// matrix-matrix pointwise mul: src1 .* src2 
m<dtype>_t mmul_m<atype>(m<dtype>_t src1, m<dtype>_t src2);
// matrix-vector pointwise mul (each row of src2 .* src1[row index])
m<dtype>_t mmul_m<atype>_m<atype>v(m<dtype>_t src1, m<dtype>_t src2, uint8_t index);
// matrix-scalar pointwise mul (each element of src1 * src2)
m<dtype>_t mmul_m<atype>_<atype>(m<dtype>_t src1, <dtype>_t src2);

Точно так же как и с соответствующими им инструкциями:

# md, ms, ms1, ms2 -  matrix registers
mmul.<sbit>.mm md, ms2, ms1
mmul.<sbit>.mv.i md, ms2, ms1[index]
mmul.<sbit>.mx md, ms, rs

Тем не менее, по сравнению с поэлементным сложением и вычитанием у поэлементного умножения есть одна особенность. Для матриц с элементами типа int32 вышеперечисленные функции и инструкции для поэлементного умножения сохраняют младшую половину 64-битного результата. Если необходимо сохранить старшую половину, следует заменить “mul” в названии интринсиков/инструкций на “mulh”.

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

// matrix-matrix pointwise sra (shift right arithm.) src1 by src2 bit 
m<dtype>_t msra_m<atype>_mu<atype>(m<dtype>_t src1, mu<dtype>_t src2);

Можно задать строку сдвигов, и этот вектор будет применяться к каждой строке матричного регистра:

// matrix-vector pointwise sra (src1 each row by src2[row index] bit)
m<dtype>_t msra_m<atype>_mu<atype>v(m<dtype>_t src1,  mu<dtype>_t src2, uint8_t index);

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

// matrix-scalar pointwise shift (each element of src1 by src2 bit)
m<dtype>_t msra_m<atype>_u<atype>(m<dtype>_t src1, u<dtype>_t src2);

Вышеперечисленным интринсикам для поэлементного арифметического сдвига вправо соответствуют следующие инструкции:

# md, ms, ms1, ms2 -  matrix registers
msra.<sbit>.mm md, ms2, ms1
msra.<sbit>.mv.i md, ms2, ms1[index]
msra.<sbit>.mx md, ms, rs

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

Копирование данных

С поддерживаемыми матричными операциями мы разобрались. Теперь — о копировании данных между регистрами.

Для копирования данных из одного матричного регистра в другой (тоже матричный) предназначена функция:

// matrix-matrix mov 
m<dtype>_t mmov_m<atype>(m<dtype>_t src);

Можно скопировать строку с номером index из матричного регистра src в каждую строку возвращаемого матричного регистра:

// matrix-vector mov (it duplicates a vector to each row of src)
m<dtype>_t mmov_m<atype>v(m<dtype>_t src, uint8_t index);

Можно инициализировать каждый элемент матричного регистра одним и тем же скаляром:

// matrix-scalar mov with duplicates scalar to each element of src
m<dtype>_t mdup_<atype>(<dtype>_t src);

Если указать индекс элемента матричного регистра, то можно скопировать скаляр только в этот элемент, не изменяя другие:

// matrix-scalar mov
m<dtype>_t mmov_mx_<atype>(m<dtype>_t dest, <dtype>_t src, uint8_t index);

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

// scalar-matrix mov 
<dtype>_t mmov_xm_<atype>(m<dtype>_t src, uint8_t index);

Перечисленным интринсикам копирования соответствуют следующие инструкции:

# md, ms - matrix registers
mmov.mm md, ms
mmov.mv.i md, ms[index]
mdup<sbit>.m.x md, src
mmov<sbit>.m.x md, src, index
mmov<sbit>.x.m rd, ms, index

Вспомогательные операции

Нам осталось рассмотреть небольшую группу вспомогательных операций.

Первая — для того, чтобы обнулить все элементы матричного регистра:

// matrix register with zero element (may be x2 for float)
m<dtype>_t mzero_m<atype>();

Ему соответствует следующая инструкция:

# md - matrix register
mzero md

Фактически это частный случай рассмотренного в предыдущем подразделе интринсика копирования скаляра в каждый элемент матричного регистра, если в качестве скаляра передавать в него ноль:

m<dtype>_t mdup_<atype>(0.0);

Вторая — возвращение всего матричного пространства в исходное состояние:

// it sets all matrix space to Initial State 
void mrelease();

На ассемблере это выполняется следующим образом:

mrelease

Небольшой пример

Реализуем при помощи матричного расширения от T-Head простейший пример — поэлементное произведение блоков матриц.

Начинаем мы с подключения соответствующего заголовочного файла:

#include <riscv_matrix.h>

Как уже было сказано, поэлементные операции определены только для целочисленных матриц. Инициализируем 3 матрицы размера $4\times4$ с элементами одинарной точности:

#define N 16
...
  int32_t x[N] = {16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1};
  int32_t y[N] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};
  int32_t z[N] = {0};

В традиционном двумерном представлении заданные матрицы имеют следующий вид:

Поскольку мы хотим выполнять умножение не для полных матриц, а только для отдельных блоков, сразу рассчитаем stride. Он необходим для корректной загрузки блока в матричный регистр и выгрузки из него:

uint8_t row_size = 4;
  long stride = row_size * sizeof(int32_t);Дальше необходимо загрузить в матричные регистры те блоки, над которыми мы собираемся производить операцию. Пусть для примера мы хотим перемножить только блоки размера 2 \times 2, выделенные на рисунке:

Дальше необходимо загрузить в матричные регистры те блоки, над которыми мы собираемся производить операцию. Пусть для примера мы хотим перемножить только блоки размера $2 \times 2$ , выделенные на рисунке:

Задаем размер блока:

uint8_t block_size = 2;

Вспоминаем, что третья размерность у нас будет конфигурироваться как число элементов в строке, умноженное на их ширину:

uint16_t m_k = block_size * sizeof(int32_t);

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

/* Configuration matrix size */
  mcfgm(block_size);
  mcfgn(block_size);
  mcfgk(m_k);

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

Теперь мы можем загрузить блоки-сомножители в матричные регистры, не забывая указывать stride:

 /* load matrix blocks to registers */
  mint32_t ma = mld_i32(x, stride);
  mint32_t mb = mld_i32(y, stride);

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

mint32_t ans = mmul_mi32(ma, mb); // pointwise mul

Нам осталось только лишь выгрузить результат из матричного регистра ans в соответствующий блок матрицы z. Здесь, опять же, необходимо помнить, что мы работаем не со всей матрицей, а только лишь с блоком внутри нее. Поэтому необходимо указывать stride для выгрузки, чтобы элементы из матричного регистра попали на корректные позиции внутри матрицы:

mst_i32_mi32(z, stride, ans); // store

В итоге получаем следующий результат:

Сохраним представленный выше код в файл pointwise_matmul.c. Его можно скомпилировать и запустить с помощью тулчейна и эмулятора из репозитория. При компиляции важно указывать архитектуру, поддерживающую данное матричное расширение:

riscv64-unknown-linux-gnu-gcc   
    -static -O2 
    -mtune=c908v 
    -march=rv64g_xtheadmatrix  
    pointwise_matmul.c 
    -o pointwise_matmul

При запуске на эмуляторе также необходимо указывать соответствующую архитектуру и флаг матричного расширения:

qemu-riscv64 
       -cpu c908v,
       x-matrix=on 
       ./pointwise_matmul

Заключение

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

Теперь вы полностью готовы к тому, чтобы реализовать при помощи T-Head RVM один из наиболее критичных для производительности многих приложений алгоритмов — умножение плотных матриц. Основные идеи для высокопроизводительной реализации этого алгоритма мы рассмотрели в первом тексте нашего «цикла» про матричные расширения.

В следующем тексте мой коллега, Андрей Соколов (@Andy31), подробно объяснит, как применять эти идеи на практике, а также покажет, как использовать T-Head RVM для низкоуровневой оптимизации матричного умножения.

Tags:

Hubs: