Когда профилируешь вычисления или инференс небольших моделей в Python, часто натыкаешься на странную вещь: вроде бы под капотом у NumPy крутится мощный OpenBLAS или Intel MKL, но на небольших матрицах (16x16, 32x32) умножение A @ B работает подозрительно медленно — по 3–5 микросекунд на одну операцию.

Если ты перемножаешь матрицы 2048x2048 раз в минуту — эти микросекунды незаметны. Но когда у тебя токенная генерация с batch=1, фильтр Калмана, физический движок или обработка сенсоров в реальном времени с миллионом итераций, профайлер начинает показывать, что программа большую часть времени просто простаивает на оверхеде библиотек.

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

Почему гиганты пасуют на мелочи?

OpenBLAS и MKL великолепны на гигантских массивах. Но плата за их универсальность — огромный диспетчеризационный стек:

1. Проверка типов, флагов выравнивания и шагов (strides) массивов.

2. Логика эвристик: решение о том, стоит ли будить пул потоков OpenMP. Будить потоки ради матрицы 16x16 — это чистый убыток по тактам из-за синхронизации мьютексов.

3. Лишние аллокации памяти под промежуточные буферы.

Для матрицы 16x16 из float32 весь объём данных — это всего 16 × 16 × 4 байта = 1 КБ. Она целиком со свистом помещается даже в L1-кэш ядра (32 КБ), но NumPy тратит на неё больше времени, чем уходит на сам счёт.

Архитектура микроядра: регистровый тайлинг 6x16

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

В x86-64 с AVX2 у нас есть 16 векторных регистров ymm0–ymm15 по 256 бит (по 8 чисел float каждый). Я выбрал раскладку микроядра размером 6 строк на 16 столбцов (блок 6x16).

Почему именно 6x16:

• Чтобы хранить 16 столбцов для 6 строк, нужно ровно 12 регистров-аккумуляторов ().

• 2 регистра нужны под загрузку строки из матрицы (также 16 float).

• 1–2 регистра остаются свободными для бродкаста элементов матрицы через mm256set1_ps.

Итого заняты все 16 регистров, ни одного лишнего сброса на стек (spilling) нет.

Фрагмент рабочего цикла микроядра выглядит так:

// Загружаем 16 float из матрицы B (два YMM-регистра)
__m256 b0 = _mm256_loadu_ps(&B[k * ldb + 0]);
__m256 b1 = _mm256_loadu_ps(&B[k * ldb + 8]);

// Для каждой из 6 строк матрицы A размножаем элемент a_ik на весь вектор
// и сразу делаем FMA (Fused Multiply-Add)
__m256 a0 = _mm256_set1_ps(A[0 * lda + k]);
c00 = _mm256_fmadd_ps(a0, b0, c00);
c01 = _mm256_fmadd_ps(a0, b1, c01);

__m256 a1 = _mm256_set1_ps(A[1 * lda + k]);
c10 = _mm256_fmadd_ps(a1, b0, c10);
c11 = _mm256_fmadd_ps(a1, b1, c11);
// ... развёрнуто до 6-й строки

Инструкция mm256fmadd_ps делает умножение и сложение за 1 такт, не теряя точности на промежуточном округлении.

Для матриц большего размера микроядро оборачивается в L1/L2 кэш-блокировку с тайлами Mc=64,Nc=128,Kc=128Mc​=64,Nc​=128,Kc​=128.

Грабли интеграции с Python: Buffer Protocol без копирования

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

Пришлось написать нативный C-extension для CPython через стандартный Buffer Protocol (Py_buffer):

Py_buffer viewA, viewB;
if (PyObject_GetBuffer(objA, &viewA, PyBUF_SIMPLE) != 0) return NULL;
if (PyObject_GetBuffer(objB, &viewB, PyBUF_SIMPLE) != 0) {
    PyBuffer_Release(&viewA);
    return NULL;
}

const float* ptrA = (const float*)viewA.buf;
const float* ptrB = (const float*)viewB.buf;

Это даёт честный zero-copy: мы просто забираем сырой указатель на непрерывный кусок памяти прямо из объекта NumPy, не копируя ни единого байта, и отдаём ядру.

Что получилось по замерам

Тестировал на стандартном x86-64 процессоре (AVX2/FMA). Сравнивал чистый вызов NumPy A @ B (OpenBLAS) и своё ядро NanoGEMM. Замер — медиана из 10 000 итераций.

Начиная со 128x128 и выше, OpenBLAS уже подключает многопоточность и уходит вперёд, но в диапазоне 16x16 – 64x64 отсутствие диспетчеризационного шума даёт ощутимый выигрыш.

Исходники и как пощупать

Библиотеку собрал в пакет nanogemm, весит всего около 100 КБ, ставится через pip:

pip install nanogemm

В коде всё просто:

import numpy as np
import nanogemm

a = np.random.randn(16, 16).astype(np.float32)
b = np.random.randn(16, 16).astype(np.float32)

c = nanogemm.matmul(a, b)

Репозиторий проекта (MIT): github.com/eminsk/nanogemm

Будет интересно, если кто-то прогонит бенчмарк на своих процессорах (особенно интересно поведение на AMD Zen и серверах) — делитесь результатами в комментариях.