SIMD для умножения матриц

[vak]

Во всех современных процессорах имеются арифметические инструкции для векторных вычислений (SIMD). Однако средний программист не имеет представления о них, и не использует. А зря: малыми усилиями можно ускорить код в три-четыре раза. Покажу, как это делается, на примере классического умножения матриц.

Напишем на Си умножение целочисленных матриц размера 4x4:

void matmul(int32_t result[4][4], const int32_t foo[4][4], const int32_t bar[4][4])
{
    for (int r = 0; r < 4; r++) {
        for (int c = 0; c < 4; c++) {
            int sum = 0;
            for (int k = 0; k < 4; k++) {
                sum += foo[r][k] * bar[k][c];
            }
            result[r][c] = sum;
        }
    }
}

Intel и AMD предлагают набор инструкций, известный в узких кругах как SSE или AVX. В Си компиляторе имеются соответствующие встроенные типы и функции (intrinsics). Перепишем наше умножение матриц с этими интринсиками.

#if __x86_64__
#include <immintrin.h>

void matmul_simd(int32_t result[4][4], const int32_t foo[4][4], const int32_t bar[4][4])
{
    for (int r = 0; r < 4; ++r) {

        // Сумма для этой итерации: строка значений int32, изначально нули.
        __m128i sum = _mm_setzero_si128();

        for (unsigned k = 0; k < 4; k++) {

            // Загружаем одно число первой матрицы, размножаем четыре копии.
            const __m128i a = _mm_set1_epi32(foo[r][k]);
            
            // Загружаем строку второй матрицы.
            const __m128i b = _mm_loadu_si128((const __m128i *)&bar[k][0]);

            // Умножаем поэлементно.
            const __m128i product = _mm_mullo_epi32(a, b);
    
            // Суммируем поэлементно.
            sum = _mm_add_epi32(sum, product);
        }
                
        // Сохраняем строку результата.
        _mm_storeu_si128((__m128i *)&result[r][0], sum);
    }
}
#endif

В архитектуре ARM64 аналогичный набор инструкций называется Neon. Вот то же самое матричное умножение с интринсиками ARM Neon:

#if __aarch64__
#include <arm_neon.h>
    
void matmul_simd(int32_t result[4][4], const int32_t foo[4][4], const int32_t bar[4][4])
{
    for (int r = 0; r < 4; ++r) {

        // Сумма для этой итерации: строка значений int32, изначально нули.
        int32x4_t sum = vmovq_n_s32(0);

        for (unsigned k = 0; k < 4; k++) {

            // Загружаем одно число первой матрицы, размножаем четыре копии.
            const int32x4_t a = vdupq_n_s32(foo[r][k]);

            // Загружаем строку второй матрицы.
            const int32x4_t b = vld1q_s32(&bar[k][0]);

            // Умножаем поэлементно.
            const int32x4_t product = vmulq_s32(a, b);

            // Суммируем поэлементно.
            sum = vaddq_s32(sum, product);
        }

        // Сохраняем строку результата.
        vst1q_s32(&result[r][0], sum);
    }
}
#endif

Законченный исходный файл, включая тестовые данные, можно взять здесь: matmul4.c. Компилировать надо с флагом -march=native.

Такую картинку нарисовала Gemini, когда я попросил её изобразить умножение матриц 4x4.

Comments

[juan_gandhi]

Очень хорошо! Это ж можно и в джаву вставлять.

[vak]

Наверное как-то можно и в Джаву запендюрить.
Вот тут народ такое обсуждает: https://dzone.com/articles/vectorised-algorithms-in-java

[juan_gandhi]

Hmm. it does look pretty revolutionary. And for some reasom they don't want JNI.

[sobriquet9]

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

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

[vak]

Уже больше 20 лет как SSE2 присутствует во всех процессорах как Intel, так и AMD. Neon тоже обязательная часть всякого 64-битного процессора ARM. Нет причин не задействовать.

Конечно, у Интела имеется определённый бардак с вариациями архитектур. SSE3 к примеру.

[beldmit]

OpenSSL, например, в run-time это детектит, но всё равно ассемблерный код для всех наборов инструкций.

С другой стороны, современные версии Linux требования к минимальному набору команд потихоньку поднимают, и добавляют возможность держать несколько сборок и брать оптимальную run-time.

[ikg]

Для упомянутого в примере интринсика _mm_mullo_epi32() не хватит уровня SSE2, требуется SSE4.1. Т.е. полный набор всех SSE.

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

[sassa_nf]

JNI is too slow for this.

Project Panama for the win.