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.