Умножение матриц: пример использования расширения ARM SME2 в Apple M4 Pro / Хабр

Микрофотография скальпированного процессора Apple M4

В конце 2020 года я купил MacBook Pro 13 на процессоре Apple M1, очень хотелось испытать процессоры на архитектуре ARM. Почти сразу на чипе Apple M1 был найден вычислительный блок для матричных операций Apple AMX. Для Apple AMX не было документации, он не использовался в Apple Accelerate, но несколько энтузиастов занимались реверс-инжинирингом и анализом производительности ("https://github.com/corsix/amx ").

В 2024 году вышли компьютеры на базе семейства процессоров Apple M4, у которых блок AMX задействован для выполнения инструкций из Scalable Matrix Extension 2 (сайт ARM недоступен в РФ) (ARM SME2).

В статье рассмотрим использование расширения ARM SME2 на примере умножения заполненных матриц. Увидим, как выжать максимум из процессора и получить прирост производительности в десятки раз.

Умножение заполненных матриц, постановка задачи

На Habr уже была статья про оптимизацию умножения заполненных матриц с использованием Intel AVX2: Умножение матриц: эффективная реализация шаг за шагом и статья об использовании Intel AMX Теоретическая и реальная производительность Intel AMX. Intel AMX с некоторой натяжкой можно рассматривать как аналог расширения ARM SME2, по крайней мере, основная операция в Intel AMX и ARM SME2 одинаковая — внешнее матричное произведение.

Заполненными называют матрицы, у которых подавляющая часть элементов не равна нулю.

Математически алгоритм умножения заполненных матриц записывается как:

$C_{i,j}=\sum_{k=1}^{K}A_{i,k}B_{k,j},$

где $C_{i,j}$ — элемент матрицы в строкеи столбце.

Умножение матриц как в книжке по алгебре

Договоримся, что:

M — число строк в матрицахи;
K — число столбцов в матрицеи число строк в матрице;
N — число столбцов в матрицахи.

Математически для получения одного значения результирующей матрицынадо выполнитьумножений исложение. Поскольку в результирующей матрицевсего $M\times N$ элементов, то для вычисления всей матрицынадо выполнитьопераций:

$Z = \left[K+ \left(K - 1\right)\right]\times N \times M \approx 2\times K\times N \times M$

Число операций будет использоваться при расчёте показателя GFLOPS (Giga FLoating-point Operations Per Second) — миллиарды операций с плавающей точкой в секунду, это наиболее часто используемый показатель для оценки производительности. Например, для размерности матриц $M=1000,\, N=2000,\, K=3000$ и времени умножениясекунды производительность составит $F_{GFLOPS} = \frac{2\times K \times N \times N}{t \times 1.0\cdot e^{9}} = 60.0$ GFLOPS.

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

Fp32 — 32-битное число (4 байта), в C/C++ обозначается float,
Fp64 — 64-битное число (8 байт), в C/C++ обозначается double.

Это важно, так как скорость вычислений существенно зависит от битности чисел (чем больше битность числа, тем медленнее вычисления). Плюс для хранения одной и той же матрицы с типом Fp64 требуется в 2 раза больше памяти, чем для типа Fp32, что может вносить коррективы в скорость вычислений за счёт ограничений пропускной способности шины памяти. В общем случае, сравнивать можно только скорости для одинаковых типов данных.

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

// размерность матрицы
const int M = 100;
const int N = 200;

// выделяем линейный массив
double *C = new double[N * M];

// циклы по строкам и столбцам
for (int i = 0; i < M; ++i){
	for (int j = 0; j < N; ++j){
		// обращение к элементу матрицы в строке i и столбце j
		C[i * N + j] = i + j;
	}
}
...
// не забываем вернуть память системе
delete[] C;

Из личного опыта: во многих книгах по С/С++, в институтских лекциях и других источниках для хранения матриц используют двойные указатели double **C. Это порождает проблему со скоростью расчетов, так как двойные указатели приводят к двойному обращению к памяти: сначала получаем адрес i-й строки C[i] — первое обращение к памяти, потом само значение C[i][j] — второе обращение. Каждое обращение к памяти — это долгая операция, занимающая от 20 тактов процессора если данные нашлись в кеше и до 100–150 тактов процессора, если в кэше данных нет (Infographics: Operation Costs in CPU Clock Cycles).
При использовании линейного массива чтение одно, а целочисленная арифметика вычисления смещения i * N + jпрактически бесплатна по сравнению с лишним обращением к памяти, поэтому во всех высокопроизводительных матричных пакетах используются линейные массивы.

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

Базовый вариант умножения плотных матриц

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

A, B и C — указатели на память, где по строчкам лежат значения матриц;
lda, ldb, ldc — соответствующие смещения следующих элементов в столбце для матриц A, B и C. То есть Aij = A[i * lda + j] и т.д.

Кому интересно может прочитать про индексирование с использованием strides в NumPy.

Самая простая реализация алгоритма (прямо по математической формуле):

void gemm_v00(uint64_t M, uint64_t N, uint64_t K,         
			  uint64_t lda, const float *A,       
			  uint64_t ldb, const float *B,         
			  uint64_t ldc, float *C) {

    for (uint64_t i = 0; i < M; ++i) {        
	    for (uint64_t j = 0; j < N; ++j) {            
		    float s = 0.0;            
			for (uint64_t k = 0; k < K; ++k) {                
				s += A[i * lda + k] * B[k * ldb + i];  
            }            
            C[i * ldc + j] = s;  
        }  
    }  
}

Для тестирования будем использовать следующие размерности матриц,,.

Получаем:

Версия	Время (сек.)	FP32, операций в секунду (GFLOPS)
V00	22.86	2.93

Достаточно медленно, всего в районе 3 GFLOPS, учитывая, что пиковая производительность быстрых ядер Apple M4 Pro близка к 100 GFLOPS.

Основная проблема заключается в доступе к матрице: матрица хранится по строкам, а в алгоритме доступ по столбцам. Процессор загружает данные из оперативной памяти кеш-линиями, длина кеш-линии — от 32 байт и более. В случае Apple M4 — 128 байт. Получается, что для чисел Fp32 одинарной точности (4 байта), загружая первое значение в первом столбце матрицы, мы загружаем сразучисла из первой строки матрицы. Если число столбцов в матрицебольше 32, то для следующего значения в первом столбце матрицы нам опять придется грузить данные из памяти. Загрузка из оперативной памяти — достаточно медленная операция, требующая от 100 до 150 тактов процессора. И, дополнительно, из загруженных 32 чисел мы используем только одно, что составляет примерно $\frac{1}{32}\cdot 100\%=3\%$ . То есть числа мы грузим, может и быстро, но зря! Мы не используем $97\%$ пропускной способности шины памяти.

Расположение в памяти матрицы B, жёлтым выделены элементы первого столбца

Оптимизированный по доступу к памяти алгоритм

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

void gemm_v01(uint64_t M, uint64_t N, uint64_t K,         
			  uint64_t lda, const float *A,         
			  uint64_t ldb, const float *B,         
			  uint64_t ldc, float *C) {  
	// цикл по строкам матриц А и С  
	for (uint64_t row = 0; row < M; ++row) {        
		const float *a_row = A + row * lda;        
		float *c_row = C + row * ldc;  
        memset(c_row, 0, sizeof(float) * K);
        // Считаем добавку к строке из матрицы С используя
        // одно число из матрицы A и строку из матрицы B        
        for (uint64_t k = 0; k < K; ++k) {            
	        float a_row_k = a_row[k];            
	        const float *b_k = B + k * ldb;            
	        for (uint64_t col = 0; col < N; ++col) {                
		        c_row[col] += b_k[col] * a_row_k;  
            }  
        }  
    }  
}

Алгоритм умножения матриц, оптимизированный по доступу к памяти

Версия	Время (сек.)	FP32, операций в секунду (GFLOPS)
V00	22.86	2.93
V01 (optimized memory pattern)	2.23	30.13

Скорость выросла в 10 раз! Теперь у нас последовательный доступ ко всем матрицам.

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

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

Инструкции матричного расширения ARM SME2

ARM SME/SME2 — это архитектурное расширение Armv9, созданное для ускорения матричных и плотных линейно-алгебраических вычислений (GEMM, свёрточные ядра и т. п.).

Основные отличия ARM SME2 от ARM NEON, Intel AVX512 и подобных:

ZA Tile — аппаратная матрица-аккумулятор ZA, большое регистровое 2D-хранилище, в котором накапливаются результаты матричных операций. Основная операция — внешнее произведение векторов. Для Apple M4 это 64 линии по 64 байта каждая, итого $64\times 64 = 4\text{KB}$ каждая.
Streaming State — специальный режим выполнения SME-кода, уменьшающий накладные расходы на сохранение/восстановление состояния и оптимизированный под «долгие» вычисления. Это означает, что при переходе в Streaming State инструкции отправляются в вычислительный блок ARM SME2, общий для целого кластера ядер.
Vector Length-agnostic — инструкции не зависят от ширины регистра, по стандарту размерность регистра может меняться от 128 до 2048 бит в зависимости от физической реализации. За счет предикатов (маски) один и тот же «бинарь» будет работать на разном «железе».

ARM SME/SME2 — новая технология. Первые доступные на рынке компьютеры с процессорами семейства Apple M4 появились в конце 2024 года. Но в 2025 году анонсированы несколько процессоров с поддержкой ARM SME2: мобильные MediaTek Dimensity 9500, Qualcomm Snapdragon 8 Elite Gen 5 и ноутбучные Qualcomm X2 Elite. Массовые поставки телефонов и ноутбуков с этими процессорами будут в 2026 году.

На сайте ARM (недоступен в РФ) есть статья, в которой показано, что задержки при использовании ARM SME2 от 2 до 4 раз меньше, чем при использовании NPU процессора, так как NPU — это отдельный ускоритель, а блок вычислений ARM SME2 — часть ядерного кластера. Соответственно, авторами статьи рекомендуется использовать ARM SME2 для интерактивных задач: сегментация изображений, распознавание речи, генерация аудио и работа с LLM.

Краткое описание ARM SME2

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

Одна из ключевых особенностей ARM SME2 — это отсутствие привязки к размерности регистра. По стандарту размерность регистров может быть от 128 до 2048 бит. На этапе компиляции размерность регистров неизвестна, соответственно, исполняемый код универсален и может быть запущен на любом «железе», поддерживающем расширение ARM SME2. Конкретная размерность регистров доступна только в процессе исполнения инструкций.

Пример функции, считающей скалярное произведение 2-х векторов:

__arm_new("za") __arm_locally_streaming		 
float sve_dot(const float* restrict a, const float* restrict b, uint64_t n) {
    // Инициализируем аккумулятор нулями
    svfloat32_t acc = svdup_f32(0.0f);
    // Размерность регистров в 32 битных элементах 
    // Для Apple M4 размерность 16 элементов (512 бит)
    uint64_t SVL = svcntw()

    for (uint64_t i = 0; i < n; i += SVL) {
        // Создаем предикат для оставшихся элементов
        svbool_t pg = svwhilelt_b32(i, n);
        
        // Загружаем данные из a и b в регистры согласно предикатам
        svfloat32_t va = svld1_f32(pg, &a[i]);
        svfloat32_t vb = svld1_f32(pg, &b[i]);
        
        // Смешанное умножение и сложение acc[k] = acc[k] + va[k] * vb[k]
        acc = svmla_f32_m(pg, acc, va, vb);
    }
    
    // Возвращаем горизонтальную сумму sum(acc[k])
    return svaddv_f32(svptrue_b32(), acc);
}

Разберем основные особенности кода:

uint64_t SVL = svcntw() — тут при выполнении получаем длину регистров ARM SME2 в элементах Fp32.
for (uint64_t i = 0; i < n; i += SVL) — основное тут, что цикл у нас один, нам не надо отдельно обрабатывать остаток, когда длина массива не кратна размерности регистров.
svbool_t pg = svwhilelt_b32(i, n); — функция расчета предиката. Предикат можно рассматривать как массив-маску длиной SVL, где для каждого элемента 1 — элемент участвует в вычислениях, 0 — нет. Псевдокод для данной функции мог бы выглядеть так (это только псевдокод, приведённый для понимания, не скомпилируется и в реальности реализован в железе как инструкция процессора):

svbool_t svwhilelt_b32(uint64_t i, uint64_t n){
	uint64_t SVL = svcntw();
	svbool_t p;
	for (uint64_t k = 0; k < SVL; ++k){
		if (i + k < n) {
			p[k] = true;
		} else {
			p[k] = false;
		}
	}
	return p;
}

svld1_f32, svmla_f32_m и svaddv_f32 — функции загрузки, смешанного умножения и сложения, горизонтального сложения соответственно, получают на вход маску-предикат.

Всего в ARM SME2 предусмотрено:

32 регистра для данных;
16 регистров для предикатов.

Следующее важное отличие ARM SME2 — это наличие специального двумерного регистра ZA, через который реализуются инструкции внешнего матричного произведения. В примере ниже: pMDim — предикат для zL, pNDim — предикат для zR, синим и серым показано состояние плитки 0 регистра ZA после вызова инструкции внешнего матричного умножения svmopa, и учитывая, что изначально в ZA были нули.

Математически это $ZA += zL \cdot zR^T$ , где — вектор-столбец, а — вектор-строка.

Размерность регистра ZA для Apple M4 — $64\times 64$ байта. Если используются 2-байтовые типы, например Fp16, то регистр ZA делится на 2 плитки (Tile 0 и Tile 1) по $32\times 32$ элемента каждая. Соответственно для типа Fp32 будет 4 плитки по $16\times 16$ элементов, а для Fp64 — 8 плиток по $8\times 8$ элементов. Пример ниже показывает плитки для типа Fp32 и размерности регистра ZA $16\times16$ байт (ширина — 128 бит).

Также для регистра ZA предусмотрены инструкции записи/чтения как по строкам, так и по столбцам. Это очень удобно использовать, например, для транспонирования данных. По сравнению с AVX2 это очень просто, не нужны инструкции типа shuffle и permute.

Перемножение матриц с использованием ARM SME2, пример от ARM

Данный алгоритм приведён на сайте ARM (недоступен в РФ). Давайте посмотрим и разберём его. Будем рассматривать алгоритм для типа Fp32, и обозначим SVL — длину регистров ARM SME2 для типа Fp32. Для Apple M4 SVL=16.

На рисунке приведён самый внутренний цикл. Идем по блоку строк матрицы высотой и блоку столбцов матрицы шириной. Элементарная операция: внешнее умножение столбца из длиной и строки издлиной и добавление результата к блоку матрицы размером $SVL \times SVL$ .

Таким образом, за выполнение инструкций внешнего умножения и сложения мы получаем готовый блок из матрицы размером $SVL \times SVL$ .

Без векторных вычислений нам бы потребовалось $2\times K \times SVL \times SVL$ инструкций, то есть для, выигрыш в количестве инструкций $\approx 500$ раз. Для сравнения, при использовании AVX512 потребуется приблизительно $K \times SVL$ инструкций, то есть в 16 раз больше.

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

Как видно из алгоритма выше матрица должна быть подготовлена и переведена из хранения по строкам в хранение по блокам строк высотой, а внутри блока — с хранением по столбцам. Звучит сложно, но можно посмотреть на рисунок ниже. Блоки размерности $SVL \times SVL$ отмечены разными цветами.

Как видно, матрица справа требует больше памяти для хранения (серые блоки внизу).

Код, реализующий данную трансформацию

// Преобразование матрицы из хранения по строкам в хранение по блокам строк
// высотой SVL внутри хранение по столбцам
// M -- количество строк в матрице
// K -- количество столбцов в матрице
// SVL -- размерность регистров ARM SME 2 в элементах float
// a -- указатель на элементы матрицы записанные по строкам
// a_mod -- указатель на память, выделенную под модифицированную 
// матрицу A_mod
// __arm_streaming -- означает, что функция может быть вызвана
// только в потоковом режиме
// __arm_inout("za") -- означает, что функция делит состояние регистра ZA
// с вызывающим. т.е. ZA не будет очищаться перед вызовом функции
// и состояние ZA у вызывающего после вызова функции будет такое же как в
// конце функции
void preprocess_l_intr(uint64_t M, uint64_t K, uint64_t SVL, 
	const float *restrict a, 
	float *restrict a_mod) __arm_streaming __arm_inout("za") {
	
	// считаем количество строк в модифицированной матрице A_mod
	// M_mod >= M и M_mod % SVL == 0
    const uint64_t M_mod = SVL * (M / SVL + (M % SVL != 0 ? 1 : 0));

    // Внешний цикл по строкам матрицы A, идем блоками по SVL строк
    for (uint64_t row = 0; row < M; row += SVL) {
		// предикат для строк
        svbool_t pMDim = svwhilelt_b32(row, M);

        // Внутренний цикл по столбцам матрицы A, идем по 2 блока SVL
        // Грузим по 2*SVL так как у Apple M4 регистры 64 байта, а 
        // длина КЕШ линии 128 байт. т.е. 2*SVL элементов будут загружены 
        // за раз одной инструкцией
        for (uint64_t col = 0; col < K; col += 2 * SVL) {
			// предикат для столбцов длины 2*SVL
            svcount_t pKDim = svwhilelt_c32(col, K, 2);

            // Грузим по строкам, по 4 строки длиной 2*SVL за шаг цикла
            for (uint64_t trow = 0; trow < SVL; trow += 4) {
		        // делаем предикаты для каждой из 4-х строк
		        // если в pMDim[trow + 0] == true то возвращает pNDim
		        // иначе возвращает нулевую маску
                svcount_t p0 = svpsel_lane_c32(pKDim, pMDim, trow + 0);
                svcount_t p1 = svpsel_lane_c32(pKDim, pMDim, trow + 1);
                svcount_t p2 = svpsel_lane_c32(pKDim, pMDim, trow + 2);
                svcount_t p3 = svpsel_lane_c32(pKDim, pMDim, trow + 3);
				// рассчитываем смещение угла блока SVL строк и 2*SVL столбцов
                const uint64_t tile_UL_corner = (row + trow) * K + col;
                // грузим 4 строки длиной 2*SVL
                svfloat32x2_t zp0 = svld1_x2(p0, &a[tile_UL_corner + 0 * K]);
                svfloat32x2_t zp1 = svld1_x2(p1, &a[tile_UL_corner + 1 * K]);
                svfloat32x2_t zp2 = svld1_x2(p2, &a[tile_UL_corner + 2 * K]);
                svfloat32x2_t zp3 = svld1_x2(p3, &a[tile_UL_corner + 3 * K]);
				// собираем из каждой строки только первые SVL элементов 
                svfloat32x4_t zq0 = svcreate4(svget2(zp0, 0), svget2(zp1, 0),
                                              svget2(zp2, 0), svget2(zp3, 0));
                // собираем из каждой строки вторые SVL элементов
                svfloat32x4_t zq1 = svcreate4(svget2(zp0, 1), svget2(zp1, 1),
                                              svget2(zp2, 1), svget2(zp3, 1));
                // Записываем первые 4*SVL элементов в tile 0 регистра ZA
                svwrite_hor_za32_f32_vg4(
                    /* tile: */ 0, /* slice: */ trow, zq0);
                // Записываем вторые 4*SVL элементов в tile 0 регистра ZA
                svwrite_hor_za32_f32_vg4(
                    /* tile: */ 1, /* slice: */ trow, zq1);
            }

            // Читаем записанные 2 блока SVL*SVL из регистра ZA 
            // и сохраняем в A_mod
            // смещение первого блока в A_mod
            const uint64_t dest_0 = row * K + col * SVL;
            // смещение второго блока в A_mod
            const uint64_t dest_1 = dest_0 + SVL * SVL;
            // читаем и пишем по 4 столбца в 2-х блоках
            for (uint64_t tcol = 0; tcol < SVL; tcol += 4) {
		        // предикаты на проверку, что не выходим за границы выделенной
		        // для A_mod памяти для блока 1 и 2
                svcount_t p0 = svwhilelt_c32(dest_0 + tcol * SVL, K * M_mod, 4);
                svcount_t p1 = svwhilelt_c32(dest_1 + tcol * SVL, K * M_mod, 4);
                // читаем 4 столбца из tile 0 регистра ZA
                svfloat32x4_t zq0 =
                    svread_ver_za32_f32_vg4(/* tile: */ 0, /* slice: */ tcol);
                // читаем 4 столбца из tile 1 регистра ZA
                svfloat32x4_t zq1 =
                    svread_ver_za32_f32_vg4(/* tile: */ 1, /* slice: */ tcol);
                // сохраняем в память
                svst1(p0, &a_mod[dest_0 + tcol * SVL], zq0);
                svst1(p1, &a_mod[dest_1 + tcol * SVL], zq1);
            }
        }
    }
}

Код матричного произведения

// матричное умножение
// M -- число строк матриц A и C
// N -- число столбцов матриц B и С
// K -- число столбцов матрицы A и число строк матрицы B
// SVL -- размерность регистров ARM SME 2 в элементах float
// matLeft_mod -- указатель на память модифицированной матрицы A_mod
// matRight -- указатель на память матрицы B
// matResult -- указатель на память матрицы С
// __arm_streaming -- означает, что функция может быть вызвана
// только в потоковом режиме
// __arm_inout("za") -- означает, что функция делит состояние регистра ZA
// с вызывающим. т.е. ZA не будет очищаться перед вызовом функции
// и состояние ZA у вызывающего после вызова функции будет такое же как в
// конце функции
void matmul_intr_impl(
    uint64_t M, uint64_t K, uint64_t N, uint64_t SVL,
    const float *restrict matLeft_mod, const float *restrict matRight,
    float *restrict matResult) __arm_streaming __arm_inout("za") {

    // Строим матрицу С плитка за плиткой 
    for (uint64_t row = 0; row < M; row += SVL) {
		// предикат для строк
        svbool_t pMDim = svwhilelt_b32(row, M);

        for (uint64_t col = 0; col < N; col += SVL) {
			// предикат для столбцов
            svbool_t pNDim = svwhilelt_b32(col, N);

            // Внешнее произведение + накопление
            // обнуляем регистр ZA
            svzero_za();
            // Угол в модифицированной матрице A
            const uint64_t matLeft_pos = row * K;
            // Угол в матрице B
            const uint64_t matRight_UL_corner = col;
            // Внутренний цикл по K, берем столбец длиной SVL из A_mod
            // берем строку длиной SVL из B
            // Результат внешнего произведения добавляем в tile 0 регистра ZA
            for (uint64_t k = 0; k < K; k++) {
                svfloat32_t zL =
                    svld1(pMDim, &matLeft_mod[matLeft_pos + k * SVL]);
                svfloat32_t zR =
                    svld1(pNDim, &matRight[matRight_UL_corner + k * N]);
                svmopa_za32_m(0, pMDim, pNDim, zL, zR);
            }

            // Сохраняем результат накопленный в регистре ZA в matResult.
            // смещение угла результирующей матрицы
            const uint64_t result_tile_UL_corner = row * N + col;
            // Сохраняем за раз по 4 строки
            for (uint64_t trow = 0; trow < SVL && row + trow < M; trow += 4) {
	            // предикаты для 4-х строк
                svbool_t p0 = svpsel_lane_b32(pNDim, pMDim, row + trow + 0);
                svbool_t p1 = svpsel_lane_b32(pNDim, pMDim, row + trow + 1);
                svbool_t p2 = svpsel_lane_b32(pNDim, pMDim, row + trow + 2);
                svbool_t p3 = svpsel_lane_b32(pNDim, pMDim, row + trow + 3);
				// сохранение данных из регистра ZA в результирующую матрицу
                svst1_hor_za32(/* tile: */ 0, /* slice: */ trow + 0, p0,
                    &matResult[result_tile_UL_corner + (trow + 0) * N]);
                svst1_hor_za32(/* tile: */ 0, /* slice: */ trow + 1, p1,
                    &matResult[result_tile_UL_corner + (trow + 1) * N]);
                svst1_hor_za32(/* tile: */ 0, /* slice: */ trow + 2, p2,
                    &matResult[result_tile_UL_corner + (trow + 2) * N]);
                svst1_hor_za32(/* tile: */ 0, /* slice: */ trow + 3, p3,
                    &matResult[result_tile_UL_corner + (trow + 3) * N]);
            }
        }
    }
}

Результаты работы данного алгоритма:

Версия	Время (сек.)	FP32, операций в секунду (GFLOPS)
V00	22.86	2.93
V01 (optimized memory pattern)	2.23	30.13
V02 SME2 (ARM Example)	0.29	228.54

Использование инструкций ARM SME2 ускорило умножение матриц еще в 7.7 раз относительно версии V01 и почти в 79 раз относительно версии V00.

Использование ARM SME2 + L1, L2 кеш

На следующем шаге оптимизации используем L1 и L2 кеш процессора.

Идеи:

Для матрицы выделим блок высотой и такой длиной $K_{L1}$ , чтобы он помещался в L1 кеш ядра процессора.
Для матрицы выделим блок высотой $K_{L1}$ и шириной $N_{L2}$ такой, что $N_{L2}$ кратно и блок матрицы помещается в кеш L2 процессора.

На рисунке показаны циклы умножения матриц: красный — самый внутренний цикл, серый — цикл следующего уровня.

Код с использованием ARM SME2 и кеш L1, L2

/// Для матрицы A[M, K] записанной по строкам, строим матрицу A_MOD[SVL, K]
/// записанную по столбцам
/// ## Параметры
/// - M -- Число строк в матрице, M <= SVL
/// - K -- Число столбцов в матрице
/// - lda -- сдвиг следующей строки в матрице A
/// - a -- указатель на матрицу
/// - a_mod -- Результат матрица с хранением по столбцам, 
/// память должна быть выделена вне, минимум K * SVL элементов
/// ## Атрибуты
/// __arm_inout("za") -- функция имеет общее состояние регистра ZA 
/// с вызывающей функцией
/// __arm_streaming -- функция работает в потоковом режиме
inline void prepare_a_mat_f32(uint64_t M, uint64_t K, 
		uint64_t lda, const float* restrict a,
        float* restrict a_mod) __arm_inout("za") __arm_streaming {
    
    // длина вектора SVL в элементах типа float (4)
    uint64_t SVL = svcntw();

    // получаем предикат (маску) т.е. в pM для всех индексов < M будет
    // выставлено True, далее False
    // например если SVL = 16, M = 10 тогда
    // pM = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0], внутри это
    // сохранено в битовом регистре
    svbool_t pM = svwhilelt_b32((uint64_t)0, M);

    // ходим по столбцам рассматриваем по 2 блока длины SVL
    for(uint64_t col = 0; col < K; col += 2 * SVL) {

        // маска из 2 * SVL значений 1 если индекс меньше K, иначе 0
        // псевдокод:
        // for (int i = 0; i < 2 * SVL; ++i)
        //     pK[i] = (col + i < K) ? 1 : 0
        svcount_t pK = svwhilelt_c32(col, K, 2);

        // рассматриваем по 4 строки за раз всего строк в тайле SVL
        for(uint64_t trow = 0; trow < M; trow += 4) {
            // p0 первая строка c индексом row + trow + 0
            // если в маске pM на индексе trow + 0 стоит 1, то p0 =
            // pK, иначе p0 все забито нулями
            svcount_t p0 = svpsel_lane_c32(pK, pM, trow + 0);
            svcount_t p1 = svpsel_lane_c32(pK, pM, trow + 1);
            svcount_t p2 = svpsel_lane_c32(pK, pM, trow + 2);
            svcount_t p3 = svpsel_lane_c32(pK, pM, trow + 3);

            // Угол тайла, (матрица записана по строкам, соответсвенно
            // индексы [i, j] и [i+1, j] отличаются на lda)
            const float* restrict a_off = a + trow * lda + col;
            // zp[n] -- содержит строку c a[trow + n, col] 
            // до a[trow + n, col +  2 * SVL]
            svfloat32x2_t zp0 = svld1_x2(p0, a_off + 0 * lda);
            svfloat32x2_t zp1 = svld1_x2(p1, a_off + 1 * lda);
            svfloat32x2_t zp2 = svld1_x2(p2, a_off + 2 * lda);
            svfloat32x2_t zp3 = svld1_x2(p3, a_off + 3 * lda);

            // Собираем в строку длиной 4 SVL первые подстроки из zp[0:4]
            svfloat32x4_t zq0 = svcreate4(svget2(zp0, 0), svget2(zp1, 0), 
						                  svget2(zp2, 0), svget2(zp3, 0));
            // Собираем в строку длиной 4 SVL вторые подстроки из zp[0:4]
            svfloat32x4_t zq1 = svcreate4(svget2(zp0, 1), svget2(zp1, 1), 
							              svget2(zp2, 1), svget2(zp3, 1));

            // записываем в таил 0 регистра ZA четыре строки 
            // начиная со строки trow
            svwrite_hor_za32_f32_vg4(/*t*/ 0, /*s*/ trow, zq0);
            // записываем в таил 1 регистра ZA четыре строки 
            // начиная со строки trow
            svwrite_hor_za32_f32_vg4(/*t*/ 1, /*s*/ trow, zq1);
        }

        // К текущему моменту в тайлах 0 и 1 собраны M строк
        // Необходимо прочитать по столбцам и сохранить в матрицу a_mod
        // a_mod имеет размерность SVL строк (SVL >= M) 
        // и K столбцов сохранена должна быть в виде
        // непрерывного массива нумерацией A_MOD[i, j] = a_mod[j * SVL + i]

        // Смещение в массиве a_mod двух блоков SVL * SVL
        const uint64_t dest_0 = col * SVL;
        const uint64_t dest_1 = dest_0 + SVL * SVL;
        // За раз работаем с 4 столбцами в тайлах 0, 1 регистра ZA
        for(uint64_t tcol = 0; tcol < SVL; tcol += 4) {
            // Проверяем, что при записи в блоки со смещениями dest_0,
            // dest_1
            // не выйдем за граница массива a_mod
            // При записи пишем куски с мусором, потом их перезаписываем,
            // но по результату выходов за границы массива нет
            svcount_t p0 = svwhilelt_c32(dest_0 + tcol * SVL, K * SVL, 4);
            svcount_t p1 = svwhilelt_c32(dest_1 + tcol * SVL, K * SVL, 4);
            // Загружаем 4 столбца из тайла 0
            svfloat32x4_t zq0 = svread_ver_za32_f32_vg4(/* tile: */ 0, 
											            /* slice: */ tcol);
            // Загружаем 4 столбца из тайла 1
            svfloat32x4_t zq1 = svread_ver_za32_f32_vg4(/* tile: */ 1, 
											            /* slice: */ tcol);
            // записываем 4 столбца в блок dest_0
            svst1(p0, &a_mod[dest_0 + tcol * SVL], zq0);
            // записываем 4 столбца в блок dest_1
            svst1(p1, &a_mod[dest_1 + tcol * SVL], zq1);
        }
    }
}

// Загрузка блока M*SVL из С в регистр ZA
inline void 
load2za_1x1_f32(svbool_t pM, svbool_t pN, uint64_t M, 
				uint64_t ldc, const float* restrict c) 
				__arm_inout("za") __arm_streaming {
	// загружаем по 4 строки за раз в tile 0
    for(uint64_t trow = 0; trow < M; trow += 4) {

        svbool_t p0 = svpsel_lane_b32(pN, pM, trow + 0);
        svbool_t p1 = svpsel_lane_b32(pN, pM, trow + 1);
        svbool_t p2 = svpsel_lane_b32(pN, pM, trow + 2);
        svbool_t p3 = svpsel_lane_b32(pN, pM, trow + 3);

        const float* restrict c_off = c + trow * ldc;
        svld1_hor_za32(/*t*/ 0, /*s*/ trow + 0, p0, c_off + 0 * ldc);
        svld1_hor_za32(/*t*/ 0, /*s*/ trow + 1, p1, c_off + 1 * ldc);
        svld1_hor_za32(/*t*/ 0, /*s*/ trow + 2, p2, c_off + 2 * ldc);
        svld1_hor_za32(/*t*/ 0, /*s*/ trow + 3, p3, c_off + 3 * ldc);
    }
}

// Сохранение блока M*SVL из регистра ZA в C 
inline void 
store2za_1x1_f32(svbool_t pM, svbool_t pN, uint64_t M, 
				 uint64_t ldc, float* restrict c) 
				 __arm_inout("za") __arm_streaming {
	// Сохраняем по 4 строки из tile 0
    for(uint64_t trow = 0; trow < M; trow += 4) {

        svbool_t p0 = svpsel_lane_b32(pN, pM, trow + 0);
        svbool_t p1 = svpsel_lane_b32(pN, pM, trow + 1);
        svbool_t p2 = svpsel_lane_b32(pN, pM, trow + 2);
        svbool_t p3 = svpsel_lane_b32(pN, pM, trow + 3);

        float* restrict c_off = c + trow * ldc;
        svst1_hor_za32(/*t*/ 0, /*s*/ trow + 0, p0, c_off + 0 * ldc);
        svst1_hor_za32(/*t*/ 0, /*s*/ trow + 1, p1, c_off + 1 * ldc);
        svst1_hor_za32(/*t*/ 0, /*s*/ trow + 2, p2, c_off + 2 * ldc);
        svst1_hor_za32(/*t*/ 0, /*s*/ trow + 3, p3, c_off + 3 * ldc);
    }
}

// микро ядро реализующее умножение Am размера [SVL, K] 
// на B размера [K, SVL]
inline void 
micro_gemm_f32(uint64_t M, uint64_t K, uint64_t N, 
			   uint64_t lda, const float* restrict a, 
			   uint64_t ldb, const float* restrict b,
               uint64_t ldc, float* restrict c) 
               __arm_inout("za") __arm_streaming {

    const uint64_t SVL = svcntw();

    svbool_t pM = svwhilelt_b32((uint64_t)0, M);
    svbool_t pN = svwhilelt_b32((uint64_t)0, N);
    // Загружаем C в ZA.
    load2za_1x1_f32(pM, pN, M, ldc, c);
    for(uint64_t k = 0; k < K; ++k) {
        // Читаем k столбец матрицы A
        svfloat32_t zL = svld1(pM, a + k * SVL);
        // Читаем k строку матрицы B
        svfloat32_t zR = svld1(pN, b + k * ldb);
        // внешнее произведение столбца zL и строки zR прибавляется к регистру
        // ZA
        svmopa_za32_m(0, pM, pN, zL, zR);
    }
    // Сохраняем ZA в C
    store2za_1x1_f32(pM, pN, M, ldc, c);
}


// Версия с использованием ARM SME2 и L1, L2 КЕШ процессора
// __arm_new("za") -- говорит о том, что ZA должен быть обнулен 
// перед вызовом функции
// __arm_locally_streaming -- функция должна выполняться в потоковом режиме
__arm_new("za") __arm_locally_streaming void 
internal_gemm_f32_v03(uint64_t M, uint64_t K, uint64_t N, 
					  uint64_t lda, const float* restrict a, 
					  uint64_t ldb, const float* restrict b, 
					  uint64_t ldc, float* restrict c) {
	// размерность регистров SME в типах float
    const uint64_t SVL = svcntw();
    // длина L1 КЕШа (64KB)
    const uint64_t L1  = 64 * 1024;
    const uint64_t L1d = L1 / sizeof(float);
    // длина L2 КЕШа (4MB)
    const uint64_t L2 = 4 * 1024 * 1024;

    // хотим, чтобы блок матрицы A длиной K_block столбцов 
    // и высотой SVL * sizeof(float) строк
    // помещался в L1 КЕШ
    uint64_t K_block = L1 / SVL / sizeof(float);
    K_block = MIN(K_block, K);

    // хотим, чтобы блок матрицы B длиной N_block столбцов 
    // и высотой K_block строк помещался в L2 КЕШ
    uint64_t N_block = (L2 / (K_block * SVL * sizeof(float))) * SVL;
    N_block = MIN(N_block, N);

	// память на стеке для хранения блока матрицы A, 
	// записанного по столбцам
    float a_t[L1d];
	
	// цикл по столбцам матриц B и С рассматриваем по N_block столбцов за раз
    for(uint64_t c_j = 0; c_j < N; c_j += N_block) {
	    // количество столбцов в блоках матриц B и C
        const uint64_t N_micro = MIN(N_block, N - c_j);
        // цикл по столбцам матрицы A и строкам матрицы B
        for(uint64_t k = 0; k < K; k += K_block) {
		    // количество столбцов в блоке батрицы A 
		    // и строк в блоке матрицы B
            const uint64_t K_micro = MIN(K_block, K - k);
            // указатель на блок матрицы B[k, cj]
            const float* restrict b_tile = b + c_j + k * ldb;
            // цикл по строкам матриц A и С, рассматриваем по SVL строк за раз
            // переиспользуем блок матрицы B слезающий в КЕШ L2
            for(uint64_t row = 0; row < M; row += SVL) {
	            // указатель на блок матрицы A[row, k]
                const float* restrict a_tile = a + k + row * lda;
                // указатель на блок матрицы C[row, k]
                float* restrict c_tile = c + c_j + row * ldc;
                // Число строк в блоке матрицы A и столбцов в блоке матрицы B
                const uint64_t M_micro = MIN(SVL, M - row);
                // Готовим блок матрицы A
                // транспонируя его хранение и сохраняя по столбцам в a_t
                prepare_a_mat_f32(M_micro, K_micro, lda, a_tile, a_t);
                // цикл по столбцам внутри блока матрицы B, 
                // рассматриваем по SVL столбцов за раз 
                // Переиспользуем блок матрицы A (a_t) влезающий в КЕШ L1
                for(uint64_t col = 0; col < N_micro; col += SVL) {
		            // укахатель на микроблок матрицы B размер [K_micro, SVL]
                    const float* restrict b_micro = b_tile + col;
                    // укахатель на микроблок матрицы C размер [SVL, SVL]
                    float* restrict c_micto = c_tile + col;
                    micro_gemm_f32(M_micro, K_micro, 
				                   N_micro - col, SVL, a_t, 
				                   ldb, b_micro, ldc, c_micto);
                }
            }
        }
    }
}

Основные изменения в коде:

Нет выделения динамической памяти под копию матрицы A, память зарезервирована на стеке только под блок матрицы A, помещающийся в КЕШ L1 ядра процессора.
В матрицах A и B выделены блоки, помещающиеся в КЕШ L1 и L2 соответственно.
Циклы по блокам реализованы так, чтобы в самом внутреннем цикле, не считая micro_gemm_f32, переиспользуется блок матрицы A из КЕШ L1. На следующем внешнем цикле переиспользуется блок матрицы B из КЕШ L2.

Результаты:

Основные изменения в коде:

Нет выделения динамической памяти под копию матрицы. Память зарезервирована на стеке только под блок матрицы, помещающийся в кеш L1 ядра процессора.
В матрицах A и B выделены блоки, помещающиеся в кеш L1 и L2 соответственно.
Циклы по блокам реализованы так, чтобы в самом внутреннем цикле (не считая micro_gemm_f32) переиспользовался блок матрицы A из кеша L1. На следующем внешнем цикле переиспользуется блок матрицы B из кеша L2.

Результаты:

Версия	Время (сек.)	FP32, операций в секунду (GFLOPS)
V00	22.86	2.93
V01 (optimized memory pattern)	2.23	30.13
V02 SME2 (ARM Example)	0.29	228.54
V03 SME2 (L1, L2 cache)	0.17	398.32

Использование кешей L1 и L2 ускорило перемножение матриц относительно версии V02 SME2 еще в 1.74 раза. Относительно версии V01 — в 13 раз, относительно версии V00 — в почти 136 раз.

Конвейер

Для типа Fp32 регистр ZA делится на 4 плитки. Можно ли задействовать их одновременно?

Идея: загрузить в ZA одновременно 2 или 4 блока из матрицы C.

Будем называть конфигурацией пару чисел $X \times Y$ , где — число блоков по элементов, одновременно читаемых из матрицы, а — число блоков по элементов, одновременно читаемых из матрицы. Таким образом в матрице одновременно рассчитывается $X \cdot Y$ блоков размером $SVL \times SVL$ .

Например, в конфигурации $1\times 2$ будем одновременно считать 2 блока $SVL \times SVL$ из матрицы. Первый блок (красный) в tile 0 регистра ZA, второй блок (зеленый) в tile 1 регистра ZA. В данном случае на 2 внешних умножения требуется 1 чтение элементов из и 2 чтения элементов из.

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

Код, реализующий данный алгоритм:

// Загружаем [SVL, 2*SVL]
inline void 
load2za_1x2_f32(svbool_t pM, svcount_t pN, uint64_t M, 
			    uint64_t ldc, const float* restrict c) 
			    __arm_inout("za") __arm_streaming {
    for(uint64_t trow = 0; trow < M; trow += 4) {

        svcount_t p0 = svpsel_lane_c32(pN, pM, trow + 0);
        svcount_t p1 = svpsel_lane_c32(pN, pM, trow + 1);
        svcount_t p2 = svpsel_lane_c32(pN, pM, trow + 2);
        svcount_t p3 = svpsel_lane_c32(pN, pM, trow + 3);

        const float* restrict c_off = c + trow * ldc;
        svfloat32x2_t zp0           = svld1_x2(p0, c_off + 0 * ldc);
        svfloat32x2_t zp1           = svld1_x2(p1, c_off + 1 * ldc);
        svfloat32x2_t zp2           = svld1_x2(p2, c_off + 2 * ldc);
        svfloat32x2_t zp3           = svld1_x2(p3, c_off + 3 * ldc);

        svfloat32x4_t zq0 = svcreate4(svget2(zp0, 0), svget2(zp1, 0), 
							          svget2(zp2, 0), svget2(zp3, 0));
        svfloat32x4_t zq1 = svcreate4(svget2(zp0, 1), svget2(zp1, 1), 
							          svget2(zp2, 1), svget2(zp3, 1));
        svwrite_hor_za32_f32_vg4(0, trow, zq0);
        svwrite_hor_za32_f32_vg4(1, trow, zq1);
    }
}

inline void 
store2za_1x2_f32(svbool_t pM, svcount_t pN, uint64_t M, 
				 uint64_t ldc, float* restrict c) 
				 __arm_inout("za") __arm_streaming {
    for(uint64_t trow = 0; trow < M; trow += 4) {

        svcount_t p0 = svpsel_lane_c32(pN, pM, trow + 0);
        svcount_t p1 = svpsel_lane_c32(pN, pM, trow + 1);
        svcount_t p2 = svpsel_lane_c32(pN, pM, trow + 2);
        svcount_t p3 = svpsel_lane_c32(pN, pM, trow + 3);

        float* restrict c_off = c + trow * ldc;
        svfloat32x4_t zq0 = svread_hor_za32_f32_vg4(/*tile*/ 0, 
											        /*slice*/ trow);
        svfloat32x4_t zq1 = svread_hor_za32_f32_vg4(/*tile*/ 1, 
											        /*slice*/ trow);
        svst1(p0, c_off + 0 * ldc, svcreate2(svget4(zq0, 0), svget4(zq1, 0)));
        svst1(p1, c_off + 1 * ldc, svcreate2(svget4(zq0, 1), svget4(zq1, 1)));
        svst1(p2, c_off + 2 * ldc, svcreate2(svget4(zq0, 2), svget4(zq1, 2)));
        svst1(p3, c_off + 3 * ldc, svcreate2(svget4(zq0, 3), svget4(zq1, 3)));
    }
}

inline void 
micro_gemm_1x2_f32(uint64_t M, uint64_t K, uint64_t N, 
				   uint64_t lda, const float* restrict a, 
				   uint64_t ldb, const float* restrict b,
				   uint64_t ldc, float* restrict c) 
				   __arm_inout("za") __arm_streaming {

    const uint64_t SVL = svcntw();

    svbool_t  pM = svwhilelt_b32((uint64_t)0, M);
    svcount_t pN = svwhilelt_c32((uint64_t)0, N, 2);
    // Load C to ZA.
    load2za_1x2_f32(pM, pN, M, ldc, c);
    for(uint64_t k = 0; k < K; ++k) {
        // Читаем k столбец блок строки матрицы L
        svfloat32_t zL = svld1(pM, a + k * SVL);

        // Читаем k строку блок столбца матрицы R
        svfloat32x2_t zRx2 = svld1_x2(pN, b + k * ldb);
        // внешнее произведение столбца zL и строки zR прибавляется к регистру
        // ZA
        svmopa_za32_m(0, pM, svpext_lane_c32(pN, 0), zL, svget2(zRx2, 0));
        svmopa_za32_m(1, pM, svpext_lane_c32(pN, 1), zL, svget2(zRx2, 1));
    }
    // Store ZA to C
    store2za_1x2_f32(pM, pN, M, ldc, c);
}

__arm_locally_streaming void 
internal_gemm_1x2_f32_v04(uint64_t M, uint64_t K, uint64_t N, 
						  uint64_t lda, const float* restrict a, 
						  uint64_t ldb, const float* restrict b, 
						  uint64_t ldc, float* restrict c) {

    const uint64_t SVL  = svcntw();
    const uint64_t SVL2 = SVL * 2;
    // length  of L1 cache (64KB)
    const uint64_t L1  = 96 * 1024;
    const uint64_t L1d = L1 / sizeof(float);
    // length of L2 cache (4MB)
    const uint64_t L2 = 4 * 1024 * 1024;

    // хотим, чтобы блок матрицы A длиной K_block столбцов и высотой SVL * sizeof(float) строк
    // помещался в L1 cache
    uint64_t K_block = L1 / SVL / sizeof(float);
    K_block          = MIN(K_block, K);

    // хотим, чтобы блок матрицы B длиной N_block столбцов и высотой K_block строк помещался в L2
    // cache
    uint64_t N_block = (L2 / (K_block * SVL * sizeof(float))) * SVL;
    N_block          = MIN(N_block, N);

    float a_t[L1d];

    for(uint64_t c_j = 0; c_j < N; c_j += N_block) {
        const uint64_t N_micro = MIN(N_block, N - c_j);
        for(uint64_t k = 0; k < K; k += K_block) {
            const uint64_t K_micro = MIN(K_block, K - k);
            // float* restrict a_tile = a + k;
            const float* restrict b_tile = b + c_j + k * ldb;
            for(uint64_t row = 0; row < M; row += SVL) {
                const float* restrict a_tile = a + k + row * lda;
                float* restrict c_tile       = c + c_j + row * ldc;
                const uint64_t M_micro       = MIN(SVL, M - row);
                // transpose a -> a_t
                prepare_a_mat_f32(M_micro, K_micro, lda, a_tile, a_t);
                for(uint64_t col = 0; col < N_micro; col += SVL2) {
                    const float* restrict b_micro = b_tile + col;
                    float* restrict c_micto       = c_tile + col;
                    micro_gemm_1x2_f32(M_micro, K_micro, MIN(SVL2, N_micro - col), SVL, a_t, ldb,
                                       b_micro, ldc, c_micto);
                }
            }
        }
    }
}

Для конфигурации $1\times 4$ будем одновременно считать 4 блока $SVL \times SVL$ из матрицы. Первый блок (красный) в tile 0 регистра ZA, второй блок (зеленый) в tile 1 регистра ZA и так далее. В данном случае на 4 внешних умножения требуется 1 чтение элементов из и 4 чтения элементов из.

Вычисление 4-х блоков запускаются подряд для использования конвейера

Код для 4-х блоков

inline void 
load2za_1x4_f32(svbool_t pM, svcount_t pN, uint64_t M, 
				uint64_t ldc, const float* restrict c) 
				__arm_inout("za") __arm_streaming {
    for(uint64_t trow = 0; trow < M; trow += 2) {
        svcount_t p0 = svpsel_lane_c32(pN, pM, trow + 0);
        svcount_t p1 = svpsel_lane_c32(pN, pM, trow + 1);

        const float* restrict c_off = c + trow * ldc;
        svfloat32x4_t zp0           = svld1_x4(p0, c_off + 0 * ldc);
        svfloat32x4_t zp1           = svld1_x4(p1, c_off + 1 * ldc);

        svwrite_hor_za32_f32_vg2(0, trow, svcreate2(svget4(zp0, 0), 
										            svget4(zp1, 0)));
        svwrite_hor_za32_f32_vg2(1, trow, svcreate2(svget4(zp0, 1), 
											        svget4(zp1, 1)));
        svwrite_hor_za32_f32_vg2(2, trow, svcreate2(svget4(zp0, 2), 
											        svget4(zp1, 2)));
        svwrite_hor_za32_f32_vg2(3, trow, svcreate2(svget4(zp0, 3), 
											        svget4(zp1, 3)));
    }
}

inline void 
store2za_1x4_f32(svbool_t pM, svcount_t pN, uint64_t M, 
				 uint64_t ldc, float* restrict c) 
				 __arm_inout("za") __arm_streaming {
    for(uint64_t trow = 0; trow < M; trow += 2) {
        svcount_t p0 = svpsel_lane_c32(pN, pM, trow + 0);
        svcount_t p1 = svpsel_lane_c32(pN, pM, trow + 1);

        float* restrict c_off = c + trow * ldc;
        svfloat32x2_t zq0     = svread_hor_za32_f32_vg2(/*t*/ 0, /*s*/ trow);
        svfloat32x2_t zq1     = svread_hor_za32_f32_vg2(/*t*/ 1, /*s*/ trow);
        svfloat32x2_t zq2     = svread_hor_za32_f32_vg2(/*t*/ 2, /*s*/ trow);
        svfloat32x2_t zq3     = svread_hor_za32_f32_vg2(/*t*/ 3, /*s*/ trow);

        svst1(p0, c_off + 0 * ldc, svcreate4(svget2(zq0, 0), svget2(zq1, 0), 
								             svget2(zq2, 0), svget2(zq3, 0)));
        svst1(p1, c_off + 1 * ldc, svcreate4(svget2(zq0, 1), svget2(zq1, 1), 
								             svget2(zq2, 1), svget2(zq3, 1)));
    }
}

inline void 
micro_gemm_1x4_f32(uint64_t M, uint64_t K, uint64_t N, 
				   uint64_t lda, const float* restrict a, 
				   uint64_t ldb, const float* restrict b,
				   uint64_t ldc, float* restrict c) 
				   __arm_inout("za") __arm_streaming {

    const uint64_t SVL = svcntw();

    svbool_t  pM = svwhilelt_b32((uint64_t)0, M);
    svcount_t pN = svwhilelt_c32((uint64_t)0, N, 4);
    // Load C to ZA.
    load2za_1x4_f32(pM, pN, M, ldc, c);
    for(uint64_t k = 0; k < K; ++k) {
        // Читаем k столбец блок строки матрицы L
        svfloat32_t zL = svld1(pM, a + k * SVL);

        // Читаем k строку блок столбца матрицы R
        svfloat32x4_t zRx4 = svld1_x4(pN, b + k * ldb);
        // внешнее произведение столбца zL и строки zR прибавляется к регистру
        // ZA
        svmopa_za32_m(0, pM, svpext_lane_c32(pN, 0), zL, svget4(zRx4, 0));
        svmopa_za32_m(1, pM, svpext_lane_c32(pN, 1), zL, svget4(zRx4, 1));
        svmopa_za32_m(2, pM, svpext_lane_c32(pN, 2), zL, svget4(zRx4, 2));
        svmopa_za32_m(3, pM, svpext_lane_c32(pN, 3), zL, svget4(zRx4, 3));
    }
    // Store ZA to C
    store2za_1x4_f32(pM, pN, M, ldc, c);
}

__arm_new("za") __arm_locally_streaming void 
internal_gemm_1x4_f32_v05(uint64_t M, uint64_t K, uint64_t N, 
						  uint64_t lda, const float* restrict a, 
						  uint64_t ldb, const float* restrict b, 
						  uint64_t ldc, float* restrict c) {

    const uint64_t SVL  = svcntw();
    const uint64_t SVL4 = SVL * 4;

    // length  of L1 cache (64KB)
    const uint64_t L1  = 64 * 1024;
    const uint64_t L1d = L1 / sizeof(float);
    // length of L2 cache (4MB)
    const uint64_t L2 = 4 * 1024 * 1024;

    // хотим, чтобы блок матрицы A длиной K_block столбцов и высотой SVL * sizeof(float) строк
    // помещался в L1 cache
    uint64_t K_block = L1 / SVL / sizeof(float);
    K_block          = MIN(K_block, K);

    // хотим, чтобы блок матрицы B длиной N_block столбцов и высотой K_block строк помещался в L2
    // cache
    uint64_t N_block = (L2 / (K_block * SVL * sizeof(float))) * SVL;
    N_block          = MIN(N_block, N);

    float a_t[L1d];

    for(uint64_t c_j = 0; c_j < N; c_j += N_block) {
        const uint64_t N_micro = MIN(N_block, N - c_j);
        for(uint64_t k = 0; k < K; k += K_block) {
            const uint64_t K_micro = MIN(K_block, K - k);
            // float* restrict a_tile = a + k;
            const float* restrict b_tile = b + c_j + k * ldb;
            for(uint64_t row = 0; row < M; row += SVL) {
                const float* restrict a_tile = a + k + row * lda;
                float* restrict c_tile       = c + c_j + row * ldc;
                const uint64_t M_micro       = MIN(SVL, M - row);
                // transpose a -> a_t
                prepare_a_mat_f32(M_micro, K_micro, lda, a_tile, a_t);
                for(uint64_t col = 0; col < N_micro; col += SVL4) {
                    const float* restrict b_micro = b_tile + col;
                    float* restrict c_micto       = c_tile + col;
                    micro_gemm_1x4_f32(M_micro, K_micro, 
				                       MIN(SVL4, N_micro - col), 
				                       SVL, a_t, ldb, b_micro, 
				                       ldc, c_micto);
                }
            }
        }
    }
}

Результаты:

Версия	Время (сек.)	FP32, операций в секунду (GFLOPS)
V00	22.86	2.93
V01 (optimized memory pattern)	2.23	30.13
V02 SME2 (ARM Example)	0.29	228.54
V03 SME2 (L1, L2 cache)	0.17	398.32
V04 SME2 (1x2 pipeline)	0.09	746.77
V05 SME2 (1x4 pipeline)	0.05	1315.44

Вычисление 2-х и 4-х блоков в матрице одновременно, позволяет получить ускорение в 1.87 и в 3.3 раза соответственно относительно версии V03, рассчитывающей один блок. Это указывает на то, что: либо сопроцессор умеет выполнять несколько инструкций внешнего умножения одновременно, либо у сопроцессора каждая инструкция раскладывается на несколько (минимум 4) подинструкции, которые выполняются на конвейере. Относительно версии V01 лучший результат почти в 44 раза быстрее, относительно версии V00 — в почти 450 раз.

Есть еще конфигурация $2 \times 2$ , когда из матрицы и считываются блоки длиной $2\times SVL$ . В матрице одновременно рассчитывается 4 блока $SVL \times SVL$ . При такой конфигурации требуется 4 чтения длиной элементов на расчет 4-х блоков, что должно быть выгоднее чем в конфигурации $1 \times 4$ , где на расчет 4-х блоков надо 5 чтений длиной элементов. Но при такой конфигурации ширина блока матрицы , помещающегося в L1 кеш в 2 раза меньше и в результате такая конфигурация для матриц одинарной точности Fp32 работает медленнее чем $1 \times 4$ .

Два чтения по 2SVL элементов из матриц A и B, 4 одновременно вычисляемых блока матрицы C

Результаты:

Версия	Время (сек.)	FP32, операций в секунду (GFLOPS)
V00	22.86	2.93
V01 (optimized memory pattern)	2.23	30.13
V02 SME2 (ARM Example)	0.29	228.54
V03 SME2 (L1, L2 cache)	0.17	398.32
V04 SME2 (1x2 pipeline)	0.09	746.77
V05 SME2 (1x4 pipeline)	0.05	1315.44
V06 SME2 (2x2 pipeline)	0.053	1270.40
Apple BLAS	0.052	1284.35

Дополнительно приведены результаты для функции cblas_sgemm из набора BLAS фреймворка Apple Accelerate. Как видно для одного потока наш код по скорости находится на уровне оптимизированных функций от самой Apple.

Прежде чем двигаться дальше, приведу результаты, для чисел с двойной точностью Fp64:

Версия	Время (сек.)	FP64, операций в секунду (GFLOPS)
V00	68.31	0.98
V01 (optimized memory pattern)	4.43	15.13
V02 SME2 (ARM Example)	1.99	33.63
V03 SME2 (L1, L2 cache)	0.67	100.70
V04 SME2 (1x2 pipeline)	0.36	185.93
V05 SME2 (1x4 pipeline)	0.20	334.48
V06 SME2 (2x2 pipeline)	0.18	380.27
Apple BLAS	0.18	365.56

Как видно при переходе от одинарной точности Fp32 к двойной Fp64 скорость лучшего метода падает в 3.46 раза. Отмечу, что для двойной точности в регистре ZA есть целых 8 независимых плиток, тем не менее конфигурации $2 \times 4$ не дают прироста производительности относительно $2 \times 2$ . Это связано с тем, что инструкция внешнего умножения выполняется за 4 шага в рамках конвейера, поэтому использование более чем 4 плиток (tile) регистра ZA не приносит ускорения, аналогичные результаты получены в Йенском университете.

Параллельность

До сих пор все наши результаты были получены при запуске в однопоточном режиме. Пришло время воспользоваться многоядерностью. На моей тестовой машине с процессором Apple M4 Pro имеется 10 быстрых ядер и 4 медленных. Быстрые ядра организованы в 2 кластера по 5 ядер. В каждом быстром кластере есть 1 сопроцессор, выполняющий инструкции ARM SME2. Медленные ядра организованы в 1 кластер с 1 сопроцессором для ARM SME2.

Вывод lstopo для процессора Apple M4 Pro — Вывод `lstopo` для процессора Apple M4 Pro

Получается, что для выполнения ARM SME2 у нас 2 быстрых сопроцессора и 1 медленный. Соответственно, от параллельной версии с ARM SME2 ожидаемый результат — ускорение чуть больше чем в 2 раза относительно однопоточной версии.

Для простоты реализуем параллельность через деление матриц и на блоки строк, где Nt — число потоков выполнения:

Код, реализующий параллельность, представлен ниже. Так получилось, что эксперименты с ARM SME2 выпали на время, когда я разбирался с языком ZIG, поэтому параллельность реализована через него. Вообще, заставить ZIG компилировать файлы с ARM SME2 инструкциями — не очень простая задача, поэтому в build.zig прописаны "танцы с бубном" и явно указано, где находятся нужные файлы из набора инструментов LLVM. Так что тем, кто скачает исходники в конце статьи, придется несколько повозиться, чтобы их собрать. За это заранее приношу свои извинения. Тем не менее, код для параллельной реализации прост, понятен и может быть реализован на любом языке.

const std = @import("std");
const Allocator = std.mem.Allocator;

/// C = beta * C + alpha * A@B
pub fn mm_par_sme_max(comptime T: type, alloc: Allocator, NTREADS: usize,
				      M: usize, N: usize, K: usize, alpha: T, 
				      lda: usize, A: []const T,
					  ldb: usize, B: []const T, 
					  beta: T, ldc: usize, C: []T) !void {
    // Считаем количество строк в матрице, первые m_first потоков получат 
    // +1 строку 
    const M_thread = M / NTREADS;
    const m_first = M % NTREADS;

    var pool: std.Thread.Pool = undefined;
    try pool.init(std.Thread.Pool.Options{
        .allocator = alloc,
        .n_jobs = NTREADS, // количество потоков
    });
    defer pool.deinit();

    var wg: std.Thread.WaitGroup = .{};
    for (0..NTREADS) |id| {
	    // Количество строк в блоке для потока выполнения
        const t_m = if (id < m_first) M_thread + 1 else M_thread;
        // смещение блока матрицы A
        const a_disp = lda * M_thread * id + @min(id, m_first) * lda;
        // блок А матрицы
        const a_slice = A[a_disp .. a_disp + t_m * lda];
        // смещение блока матрицы C
        const c_disp = ldc * M_thread * id + @min(m_first, id) * ldc;
        // блок матрицы C
        const c_slice = C[c_disp .. c_disp + t_m * lda];
        // Запускаем на обработку
        pool.spawnWg(&wg, sme_gemm_max, 
			         .{ T, t_m, N, K, alpha, lda, a_slice, 
			            ldb, B, beta, ldc, c_slice });
    }
    wg.wait();
}

Результаты для типа Fp32 (10 потоков выполнения):

Версия	Время (сек.)	Fp32, операций в секунду (GFLOPS)
V01 (optimized memory pattern)	2.23	30.13
Parallel V01	0.27	244.02
V06 SME2 (2x2 pipeline)	0.054	1245.47
Parallel V06 SME2	0.023	2887.21
Parallel Apple BLAS	0.024	2824.90

Ускорение для параллельной версии V01 без использования векторных инструкций составляет 8.1 раза. Ускорение для параллельной версии V06 с использованием ARM SME2 в конфигурации $2\times 2$ составляет 2.32 раза. Это как раз укладывается в наше предположение, что ускорение должно быть чуть больше, чем в 2 раза, за счёт 2-х быстрых сопроцессоров и одного медленного. Время параллельного расчёта нашего кода хорошо коррелирует со временем работы функции cblas_sgemm из набора BLAS фреймворка Apple Accelerate.

Результаты для типа Fp64 (10 потоков выполнения):

Версия	Время (сек.)	Fp64, операций в секунду (GFLOPS)
V01 (optimized memory pattern)	4.53	14.81
Parallel V01	0.64	104.12
V06 SME2 (2x2 pipeline)	0.17	384.24
Parallel V06 SME2	0.082	813.37
Parallel Apple BLAS	0.092	725.60

Ускорение для параллельной версии V01 без использования векторных инструкций составляет 7 раз. Ускорение для параллельной версии V06 с использованием ARM SME2 в конфигурации $2\times 2$ составляет 2.17 раза.

По данным университета Friedrich Schiller University Jena, пиковая производительность блока исполнения ARM SME2 для типа Fp32 в процессоре Apple M4 Pro составляет $2 \times 2008\text{ GFLOPS} + 357\text{ GFLOPS} = 4373\text{ GFLOPS}$ (2 быстрых блока и 1 медленный). Для типа Fp64 — $2 \times 503\text{ GFLOPS} + 89\text{ GFLOPS}=1095\text{ GFLOPS}$ . Соответственно, для типа Fp32 наш код Parallel V06 SME2 достигает 66% от пиковой производительности, а для типа Fp64 — 74%. Тяжело сказать, насколько это хорошо или плохо, но если сравниваться с эффективностью в $\approx 38\%$ для Intel AMX (данные из статьи), то реализация блока матричных вычислений в Apple M4 Pro является эффективной и сбалансированной по производительности и пропускной способности шины памяти.

Странности

Наблюдаются странности со временем выполнения при размерностях матриц, равных степени двойки. Наиболее яркий пример — при. Сравните количество операций в секунду (GFLOPS) при $M = 4096, \, K=4096$ и разных для типа Fp64:

Версия	N=4096	N=4095	N=4097
V00	0.63	1.00	1.00
V01 (optimized memory pattern)	15.02	12.25	12.21
V02 SME2 (ARM Example)	33.74	38.17	38.23
V03 SME2 (L1, L2 cache)	40.48	93.98	92.36
V04 SME2 (1x2 pipeline)	65.58	165.55	165.31
V05 SME2 (1x4 pipeline)	123.37	291.38	288.30
V06 SME2 (2x2 pipeline)	264.21	324.88	318.61
Apple BLAS	359.75	358.85	349.88

При падение производительности относительно и составляет 2.5 раза, при этом функции от Apple работают стабильно, то есть это известная проблема, и у неё есть решение.

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

Исходный код и компиляция

Для сборки вам понадобится:

Компьютер с процессором семейства Apple M4 или Apple M5.
Исходный код, можно найти тут: Исходники.
Пакет Xcode Command Line Tools, можно установить в терминале командой xcode-select --install
Пакетный менеджер Homebrew, инструкцию по установке можно найти на их сайте.
Компилятор языка ZIG, у меня установлена версия 0.15.12, установленная через brew install zig
Набор инструментов LLVM, у меня версия 22.1.1, установленная через brew install llvm

Команда в терминале для сборки и запуска:

zig build --release=fast run -- --m=3000 --n=3000 --k=4000 --nt=1

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

Параметры --m=, --n=, --k= отвечают за размеры матриц A, B, C. --nt= — количество потоков исполнения.

Приведу флаги компиляции, которые надо указать компилятору clang++, если вы захотите использовать инструкции ARM SME2 отдельно от ZIG проекта:

clang++ -O3 -march=native+sme2+sme+sme-f64f64 ...

Выводы

Использовать расширение ARM SME2 довольно приятно: есть очень широкий набор инструкций, а реализация, основанная на предикатах, избавляет от написания лишнего кода, связанного с обработкой концевых остатков в массивах.
Производительность сопроцессора ARM SME2 в Apple M4 Pro достигает 1300 GFLOPS для Fp32 в однопоточном режиме и 2900 GFLOPS в многопоточном режиме. Для сравнения: результаты из статьи для 24-ядерного процессора Intel Xeon(R) Gold 5412U с — 1400 GFLOPS для Fp16 с использованием Intel AMX и 240 GFLOPS для Fp32 при использовании AVX512.
Для чисел с двойной точностью Fp64 сопроцессор Apple M4 Pro выдаёт 380+ GFLOPS в однопоточном режиме и 800+ GFLOPS в многопоточном, что тоже очень и очень хорошо для потребительского сегмента.
Сопроцессор, реализующий ARM SME2, достаточно энергоэффективный. Так, в однопоточном режиме Apple M4 Pro без использования ARM SME2 потребляет 8 ватт, при использовании ARM SME2 — 3 ватта. В многопоточном режиме Apple M4 Pro без использования ARM SME2 потребляет 42 ватта, при использовании ARM SME2 — 8 ватт. При этом производительность отличается в 8 раз в пользу ARM SME2 (тут надо оговориться, что у меня нет замеров производительности с использованием инструкций ARM NEON, с ними ситуация может быть немного иная). Итоговая производительность на ватт при использовании ARM SME2 выше, как минимум на порядок.

P.S. Если у вас в доступе техника Apple с процессорами M4, M4 Max, M5, M5 Pro, M5 Max и есть желание, буду очень рад результатам замеров производительности на этих машинах.