Умножение матриц: эффективная реализация шаг за шагом



    Введение


    Умножение матриц — это один из базовых алгоритмов, который широко применяется в различных численных методах, и в частности в алгоритмах машинного обучения. Многие реализации прямого и обратного распространения сигнала в сверточных слоях неронной сети базируются на этой операции. Так порой до 90-95% всего времени, затрачиваемого на машинное обучение, приходится именно на эту операцию. Почему так происходит? Ответ кроется в очень эффективной реализации этого алгоритма для процессоров, графических ускорителей (а в последнее время и специальных ускорителей матричного умножения). Матричное умножение — один из немногих алгоритмов, которые позволяет эффективно задействовать все вычислительные ресурсы современных процессоров и графических ускорителей. Поэтому не удивительно, что многие алгоритмы стараются свести к матричному умножению — дополнительная расходы, связанные с подготовкой данных, как правило с лихвой окупаются общим ускорением алгоритмов.

    Так как реализован алгоритм матричного умножения? Хотя сейчас существуют множество реализаций данного алгоритма, в том числе и в открытых исходных кодах. Но к сожалению, код данных реализаций (большей частью на ассемблере) весьма сложен. Существует хорошая англоязычная статья, подробно описывающая эти алгоритмы. К моему удивлению, я не обнаружил аналогов на Хабре. Как по мне, этого повода вполне достаточно, чтобы написать собственную статью. С целью ограничить объем изложения, я ограничился описанием однопоточного алгоритма для обычных процессоров. Тема многопоточности и алгоритмов для графических ускорителей явно заслуживает отдельной статьи.

    Процесс изложения будет вестись ввиде шагов с примерами по последовательному ускорению алгоритма. Я старался писать максимально упрощая задачу, но не более того. Надеюсь у меня получилось…

    Постановка задачи (0-й шаг)


    В общем случае функция матричного умножения описывается как:

    C[i,j] = a*C[i,j] + b*Sum(A[i,k]*B[k,j]);
    

    Где матрица A имеет размер M х K, матрица B — K х N, и матрица C — M х N.



    Мы без ущерба для изложения, можем считать, что a = 0 и b = 1:

    C[i,j] = Sum(A[i,k]*B[k,j]);
    

    Ее реализация на С++ «в лоб» по формуле будет выглядеть следующим образом:

    void gemm_v0(int M, int N, int K, const float * A, const float * B, float * C)
    {
        for (int i = 0; i < M; ++i)
        {
            for (int j = 0; j < N; ++j)
            {
                C[i*N + j] = 0;
                for (int k = 0; k < K; ++k)
                    C[i*N + j] += A[i*K + k] * B[k*N + j];
            }
        }
    }
    

    Глупо было бы ожидать от нее какой-либо производительности, и действительно тестовые замеры показывают, что при (M=N=K=1152) она выполняется почти 1.8 секунды (тестовая машина — i9-7900X@3.30GHz, ОС — Ubuntu 16.04.6 LTS, компилятор — g++-6.5.0б опции компилятора — "-fPIC -O3 -march=haswell"). Минимальное количество операций для матричного умножения — 2*M*N*K = 2*10^9. Иначего говоря, производительность составляет 1.6 GFLOPS, что очень далеко от теоретического предела однопоточной производительности для данного процессора (~120 GFLOPS (float-32) если ограничится использованием AVX2/FMA и ~200 GFLOPS при использовании AVX-512). Так, что нужно предпринять, чтобы приблизится к теоретическому пределу? Далее мы в ходе ряда последовательных оптимизаций придем к решению, которое во многом воспроизводит то, что используется во многих стандартных библиотеках. В процессе оптимизации, я буду задействовать только AVX2/FMA, AVX-512 я касаться не буду, так как их распостраненность пока невелика.

    Устраняем очевидные недостатки (1-й шаг)


    Сначала устраним самые очевидные недостатки алгоритма:

    1. Вычисление адресов элементов массивов можно упростить — вынести постоянную часть из внутреннего цикла.
    2. В оригинальной версии доступ к элементам массива B производится не последовательно. Его можно упорядочить, если поменять порядок вычисления таким образом, чтобы внутренним циклом был последовательный обход по строчкам для всех трех матриц.

    void gemm_v1(int M, int N, int K, const float * A, const float * B, float * C)
    {
        for (int i = 0; i < M; ++i)
        {
            float * c = C + i * N;
            for (int j = 0; j < N; ++j)
                c[j] = 0;
            for (int k = 0; k < K; ++k)
            {
                const float * b = B + k * N;
                float a = A[i*K + k];
                for (int j = 0; j < N; ++j)
                    c[j] += a * b[j];
            }
        }
    }
    

    Результат тестовых замеров показывает время выполнения в 250 мс, или 11.4 GFLOPS. Т.е. такими небольшими правками мы получили ускорение в 8 раз!

    Векторизуем внутренний цикл (2-й шаг)


    Если внимательно посмотреть на внутренний цикл (по переменной j), то видно, что вычисления можно проводить блоками (векторами). Практически все современные процессоры позволяют проводить вычисления над такими векторами. В частности набор инструкций AVX оперирует с векторами размерностью 256 бит. Что позволяет выполнить 8 операций для вещественных чисел с одинарной точностью за такт. AVX2/FMA делает еще один шаг вперед — он позволяет выполнить слитную операцию умножения и сложения (d = a*b + c) над вектором. Настольные процессоры Интел начиная с 4-го поколения имеют 2 256-bit FMA модуля, что позволяет им теоретически выполнять 2*2*8 = 32 операции (float-32) за такт. К счастью, инструкции AVX2/FMA достаточно легко задействовать напрямую из С/С++ при помощи встроенных функций (intrinsics). Для AVX2/FMA они объявлены в заголовочном файле <immintrin.h>.

    void gemm_v2(int M, int N, int K, const float * A, const float * B, float * C)
    {
        for (int i = 0; i < M; ++i)
        {
            float * c = C + i * N;
            for (int j = 0; j < N; j += 8)
                _mm256_storeu_ps(c + j + 0, _mm256_setzero_ps());
            for (int k = 0; k < K; ++k)
            {
                const float * b = B + k * N;
                __m256 a = _mm256_set1_ps(A[i*K + k]);
                for (int j = 0; j < N; j += 16)
                {
                    _mm256_storeu_ps(c + j + 0, _mm256_fmadd_ps(a, 
                        _mm256_loadu_ps(b + j + 0), _mm256_loadu_ps(c + j + 0)));
                    _mm256_storeu_ps(c + j + 8, _mm256_fmadd_ps(a, 
                        _mm256_loadu_ps(b + j + 8), _mm256_loadu_ps(c + j + 8)));
                }
            }
        }
    }
    

    Запускаем тесты, получаем время 217 мс или 13.1 GFLOPS. Упс! Ускорение всего на 15%. Какже так? Тут нужно учитывать, два фактора:

    1. Компиляторы нынче умные пошли (не все!), и вполне справляются с задачей автовекторизации простых циклов. Уже в 1-м варианте компилятор фактически задействовал инструкции AVX2/FMA, потому ручная оптимизация не дала нам практически никаких преимуществ.
    2. Скорость расчетов в данном случае упирается не в вычислителные возможности процессора, а в скорость загрузки и выгрузки данных. В данном случае процессору для задействования 2 256-bit FMA блоков требуется загрузить 4 и выгрузить 2 256-bit вектора за такт. Это в два раза превышает даже пропускную способность L1 кеша процессора (512/256 bit), не говоря уже о пропускной способности памяти, которая еще на порядок меньше (64-bit на канал)).

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

    Дальнейшие наши шаги по оптимизации алгоритма будут направлены на минимизацию доступа в память.

    Пишем микроядро (3-й шаг)


    В предыдущей версии на 1 FMA операцию приходится 2 загрузки и 1 выгрузка.
    Больше всего загрузок и выгрузок происходит с результирующей матрицей С: данные из нее нужно загрузить, прибавить к ним произведение C[i][j] += A[i][k]*B[k][j], а потом сохранить. И так много раз. Наиболее быстрая память, с которой может работать процессор — это его собственные регистры. Если мы будем хранить результирующее значение матрицы С в регистре процессора, то в процессе расчета нужно будет подгружать только значение матриц A и B. Теперь у нас на 1 FMA операцию приходится только 2 загрузки.

    Если мы будем хранить в регистрах значения двух соседних столбцов матрицы C[i][j] и C[i][j+1], то сможем повторно использовать загруженное значение матрицы A[i][k]. И на 1 FMA операцию потребуется только 1.5 загрузки. Кроме того, сохраняя результат в 2 независимых регистра, мы позволим процессору выполнять 2 FMA операции за такт. Аналогично можно хранить в регистрах значения двух соседних строк — тогда будет осуществляться экономия на загрузке значений матрицы B.



    Всего настольные процессоры Интел начиная с 2-го поколения имеют 16 256-bit векторных регистров (справедливо для 64-bit режима процессора). 12 из них можно использовать для хранения кусочка результирующей матрицы С размером 6x16. В итоге мы сможем выполнить 12*8 = 96 FMA операций загрузив из памяти только 16 + 6 = 22 значений. И того нам удалось сократить доступ к памяти с 2.0 до 0.23 загрузки на 1 FMA операцию — почти в 10 раз!

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

    void micro_6x16(int K, const float * A, int lda, int step, 
        const float * B, int ldb, float * C, int ldc)
    {
        __m256 c00 = _mm256_setzero_ps();
        __m256 c10 = _mm256_setzero_ps();
        __m256 c20 = _mm256_setzero_ps();
        __m256 c30 = _mm256_setzero_ps();
        __m256 c40 = _mm256_setzero_ps();
        __m256 c50 = _mm256_setzero_ps();
        __m256 c01 = _mm256_setzero_ps();
        __m256 c11 = _mm256_setzero_ps();
        __m256 c21 = _mm256_setzero_ps();
        __m256 c31 = _mm256_setzero_ps();
        __m256 c41 = _mm256_setzero_ps();
        __m256 c51 = _mm256_setzero_ps();
        const int offset0 = lda * 0;
        const int offset1 = lda * 1;
        const int offset2 = lda * 2;
        const int offset3 = lda * 3;
        const int offset4 = lda * 4;
        const int offset5 = lda * 5;
        __m256 b0, b1, a0, a1;
        for (int k = 0; k < K; k++)
        {
            b0 = _mm256_loadu_ps(B + 0);
            b1 = _mm256_loadu_ps(B + 8);
            a0 = _mm256_set1_ps(A[offset0]);
            a1 = _mm256_set1_ps(A[offset1]);
            c00 = _mm256_fmadd_ps(a0, b0, c00);
            c01 = _mm256_fmadd_ps(a0, b1, c01);
            c10 = _mm256_fmadd_ps(a1, b0, c10);
            c11 = _mm256_fmadd_ps(a1, b1, c11);
            a0 = _mm256_set1_ps(A[offset2]);
            a1 = _mm256_set1_ps(A[offset3]);
            c20 = _mm256_fmadd_ps(a0, b0, c20);
            c21 = _mm256_fmadd_ps(a0, b1, c21);
            c30 = _mm256_fmadd_ps(a1, b0, c30);
            c31 = _mm256_fmadd_ps(a1, b1, c31);
            a0 = _mm256_set1_ps(A[offset4]);
            a1 = _mm256_set1_ps(A[offset5]);
            c40 = _mm256_fmadd_ps(a0, b0, c40);
            c41 = _mm256_fmadd_ps(a0, b1, c41);
            c50 = _mm256_fmadd_ps(a1, b0, c50);
            c51 = _mm256_fmadd_ps(a1, b1, c51);
            B += ldb; A += step;
        }
        _mm256_storeu_ps(C + 0, _mm256_add_ps(c00, _mm256_loadu_ps(C + 0)));
        _mm256_storeu_ps(C + 8, _mm256_add_ps(c01, _mm256_loadu_ps(C + 8)));
        C += ldc;
        _mm256_storeu_ps(C + 0, _mm256_add_ps(c10, _mm256_loadu_ps(C + 0)));
        _mm256_storeu_ps(C + 8, _mm256_add_ps(c11, _mm256_loadu_ps(C + 8)));
        C += ldc;
        _mm256_storeu_ps(C + 0, _mm256_add_ps(c20, _mm256_loadu_ps(C + 0)));
        _mm256_storeu_ps(C + 8, _mm256_add_ps(c21, _mm256_loadu_ps(C + 8)));
        C += ldc;
        _mm256_storeu_ps(C + 0, _mm256_add_ps(c30, _mm256_loadu_ps(C + 0)));
        _mm256_storeu_ps(C + 8, _mm256_add_ps(c31, _mm256_loadu_ps(C + 8)));
        C += ldc;
        _mm256_storeu_ps(C + 0, _mm256_add_ps(c40, _mm256_loadu_ps(C + 0)));
        _mm256_storeu_ps(C + 8, _mm256_add_ps(c41, _mm256_loadu_ps(C + 8)));
        C += ldc;
        _mm256_storeu_ps(C + 0, _mm256_add_ps(c50, _mm256_loadu_ps(C + 0)));
        _mm256_storeu_ps(C + 8, _mm256_add_ps(c51, _mm256_loadu_ps(C + 8)));
    }
    

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

    void init_c(int M, int N, float * C, int ldc)
    {
        for (int i = 0; i < M; ++i, C += ldc)
            for (int j = 0; j < N; j += 8)
                _mm256_storeu_ps(C + j, _mm256_setzero_ps());
    }
    

    Здесь lda, ldb, ldc — длина строчки (Leading Dimension в общем случае) соответсвующей матрицы.

    Тогда функция умножения примет следующий вид:

    void gemm_v3(int M, int N, int K, const float * A, const float * B, float * C)
    {
        for (int i = 0; i < M; i += 6)
        {
            for (int j = 0; j < N; j += 16)
            {
                init_c(6, 16, C + i*N + j, N);
                micro_6x16(K, A + i*K, K, 1, B + j, N, C + i*N + j, N);
            }
        }
    }
    

    Запускаем ее и получаем время исполнения 78.5 мс или 36.2 GFLOPS. Т.е. использование микроядра позволило ускорить матричное умножение почти в 3 раза. Но полученное быстродействие все еще далеко от максимального. Где теперь узкое место?

    Переупорядочиваем матрицу B (4-й шаг)


    Микроядро за каждую итерацию загружает два 256-bit вектора из матрицы B.



    Причем каждый раз из новой строчки. Это делает невозможным для процессора эффективное кеширование этих данных. Для исправления этой ситуации сделаем два изменения:

    1. Скопируем данные матрицы B во временный буфер таким образом, чтобы данные, необходимые одному микроядру лежали рядом.
    2. Изменим порядок обхода матрицы С: сначала будем ходить по столбцам и только потом по строкам. Это позволит эффективнее использовать переупорядоченные значения матрицы B.

    Для хранения буфера заведем небольшую структуру:

    struct buf_t
    {
        float * p;
        int n;
    
        buf_t(int size) : n(size), p((float*)_mm_malloc(size * 4, 64)) {}
        ~buf_t() { _mm_free(p); }
    };
    

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

    Функция переупорядочивания матрицы B:

    void reorder_b_16(int K, const float * B, int ldb, float * bufB)
    {
        for (int k = 0; k < K; ++k, B += ldb, bufB += 16)
        {
            _mm256_storeu_ps(bufB + 0, _mm256_loadu_ps(B + 0));
            _mm256_storeu_ps(bufB + 8, _mm256_loadu_ps(B + 8));
        }
    }
    

    Ну и собственно 4-я версия функции gemm:

    void gemm_v4(int M, int N, int K, const float * A, const float * B, float * C)
    {
        for (int j = 0; j < N; j += 16)
        {
            buf_t bufB(16*K);
            reorder_b_16(K, B + j, N, bufB.p);
            for (int i = 0; i < M; i += 6)
            {
                init_c(6, 16, C + i*N + j, N);
                micro_6x16(K, A + i*K, K, 1, bufB.p, 16, C + i*N + j, N);
            }
        }
    }
    

    Результаты тестирования (29.5 мс или 96.5 GFLOPS) показывают, что мы на правильном пути. Фактически достигнуто около 80% от теоретически возможного максимума.

    Победа? К сожалению нет. Просто размер матриц, который мы использовали для тестирования (M=N=K=1152) оказался удобным для данной версии алгоритма. Если увеличить К в 100 раз (M=1152, N=1152, K=115200), то эффективность алгоритма упадет до 39.5 GFLOPS — почти в 2.5 раза.

    Локализуем данные в кэше L1 (5-й шаг)


    Так почему же с ростом параметра K, падает эффективность алгоритма? Ответ кроется в величине буфера, который мы использовали для хранения переупорядоченных значений B. При больших значениях K он просто не влазит в кэш процессора. Решением проблемы будет ограничение его величины до размера кэша данных L1. Для процессоров Интел размер кэша данных L1 составляет 32 kb. C ограничением размера буфера, микроядро будет пробегать не по всем значениям K, а только по диапазону, который влазит в L1 кэш. Результаты промежуточных расчетов матрицы С будут храниться в основной памяти.

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

    void macro_v5(int M, int N, int K, const float * A, int lda, 
        const float * B, int ldb, float * bufB, float * C, int ldc)
    {
        for (int j = 0; j < N; j += 16)
        {
            reorder_b_16(K, B + j, ldb, bufB);
            for (int i = 0; i < M; i += 6)
                micro_6x16(K, A + i*lda, lda, 1, bufB, 16, C + i*ldc + j, ldc);
        }
    }
    

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

    void gemm_v5(int M, int N, int K, const float * A, const float * B, float * C)
    {
        const int L1 = 32 * 1024;
        int mK = std::min(L1 / 4 / 16, K);
        buf_t bufB(16 * mK);
        for(int k = 0; k < K; k += mK)
        {
            int dK = std::min(K, k + mK) - k;
            if(k == 0)
                init_c(M, N, C, N);
            macro_v5(M, N, dK, A + k, K, B + k*N, N, bufB.p, C, N);
        }
    }
    

    Результаты замеров показывают, что мы движемся в правильном направлении: для (M=1152, N=1152, K=115200) производительность алгоритма составила 78.1 GFLOPS. Это значительно лучше, чем в прошлой версии, но все еще хуже, чем для матрицы средних размеров.

    Переупорядочиваем матрицу A и локализуем в кэше L2 (6-й шаг)


    Ограничив размер K, который обрабатывается за один проход микроядра, мы сумели локализовать данные матрицы B в кэше L1. Данных, которые подгружаются из матрицы A почти в три раза меньше. Но давайте попробуем локализовать и их, заодно переупорядочив данные, чтобы они лежали последовательно. Напишем для этого специальную функцию:

    void reorder_a_6(const float * A, int lda, int M, int K, float * bufA)
    {
        for (int i = 0; i < M; i += 6)
        {
            for (int k = 0; k < K; k += 4)
            {
                const float * pA = A + k;
                __m128 a0 = _mm_loadu_ps(pA + 0 * lda);
                __m128 a1 = _mm_loadu_ps(pA + 1 * lda);
                __m128 a2 = _mm_loadu_ps(pA + 2 * lda);
                __m128 a3 = _mm_loadu_ps(pA + 3 * lda);
                __m128 a4 = _mm_loadu_ps(pA + 4 * lda);
                __m128 a5 = _mm_loadu_ps(pA + 5 * lda);
                __m128 a00 = _mm_unpacklo_ps(a0, a2);
                __m128 a01 = _mm_unpacklo_ps(a1, a3);
                __m128 a10 = _mm_unpackhi_ps(a0, a2);
                __m128 a11 = _mm_unpackhi_ps(a1, a3);
                __m128 a20 = _mm_unpacklo_ps(a4, a5);
                __m128 a21 = _mm_unpackhi_ps(a4, a5);
                _mm_storeu_ps(bufA + 0, _mm_unpacklo_ps(a00, a01));
                _mm_storel_pi((__m64*)(bufA + 4), a20);
                _mm_storeu_ps(bufA + 6, _mm_unpackhi_ps(a00, a01));
                _mm_storeh_pi((__m64*)(bufA + 10), a20);
                _mm_storeu_ps(bufA + 12, _mm_unpacklo_ps(a10, a11));
                _mm_storel_pi((__m64*)(bufA + 16), a21);
                _mm_storeu_ps(bufA + 18, _mm_unpackhi_ps(a10, a11));
                _mm_storeh_pi((__m64*)(bufA + 22), a21);
                bufA += 24;
            }
            A += 6 * lda;
        }
    }
    

    Так как, данные матрицы A теперь идут последовательно, то параметр lda в макроядре нам больше не нужен. Также поменялись параметры вызова микроядра:

    void macro_v6(int M, int N, int K, const float * A, 
        const float * B, int ldb, float * bufB, float * C, int ldc)
    {
        for (int j = 0; j < N; j += 16)
        {
            reorder_b_16(K, B + j, ldb, bufB);
            for (int i = 0; i < M; i += 6)
                micro_6x16(K, A + i*K, 1, 6, bufB, 16, C + i*ldc + j, ldc);
        }
    }
    

    Размер буфера для переупорядоченной матрицы A ограничиваем размером L2 кэша процессора (он обычно составляет от 256 до 1024 kb для разных типов процессоров). В главной функции добавляется дополнительный цикл по переменной M:

    void gemm_v6(int M, int N, int K, const float * A, const float * B, float * C)
    {
        const int L1 = 32 * 1024, L2 = 256*1024;
        int mK = std::min(L1 / 4 / 16, K) / 4 * 4;
        int mM = std::min(L2 / 4 / mK, M) / 6 * 6;
        buf_t bufB(16 * mK);
        buf_t bufA(mK * mM);
        for(int k = 0; k < K; k += mK)
        {
            int dK = std::min(K, k + mK) - k;
            for (int i = 0; i < M; i += mM)
            {
                int dM = std::min(M, i + mM) - i;
                if (k == 0)
                    init_c(dM, N, C + i * N, N);
                reorder_a_6(A + i * K + k, K, dM, dK, bufA.p);
                macro_v6(dM, N, dK, bufA.p, B + k * N, N, bufB.p, C + i * N, N);
            }
        }
    }
    

    Результаты тестовых замеров для (M=1152, N=1152, K=115200) — 88.9 GFLOPS — приблизились еще на один шаг к результату для матриц среднего размера.

    Задействуем кэш L3 (7-й шаг)


    В процессорах помимо кэша L1 и L2 еще часто бывает кэш L3 (обычно его размер составляет 1-2 MB на ядро). Попробуем задействовать и его, например, для хранения переупорядоченных значений матриц B, чтобы избежать лишних вызовов функции reorder_b_16. В функции макроядра появится дополнительные параметр reorderB, который будет сообщать о том, что данныe матрицы B уже упорядочены:

    void macro_v7(int M, int N, int K, const float * A, 
        const float * B, int ldb, float * bufB, bool reorderB, float * C, int ldc)
    {
        for (int j = 0; j < N; j += 16)
        {
            if(reorderB)
                reorder_b_16(K, B + j, ldb, bufB + K*j);
            for (int i = 0; i < M; i += 6)
                micro_6x16(K, A + i*K, 1, 6, bufB + K*j, 16, C + i*ldc + j, ldc);
        }
    }
    

    В основной функции добавится цикл по N:

    void gemm_v7(int M, int N, int K, const float * A, const float * B, float * C)
    {
        const int L1 = 32 * 1024, L2 = 256*1024, L3 = 2*1024*1024;
        int mK = std::min(L1 / 4 / 16, K) / 4 * 4;
        int mM = std::min(L2 / 4 / mK, M) / 6 * 6;
        int mN = std::min(L3 / 4 / mK, N) / 16 * 16;
        buf_t bufB(mN * mK);
        buf_t bufA(mK * mM);
        for (int j = 0; j < N; j += mN)
        {
            int dN = std::min(N, j + mN) - j;
            for (int k = 0; k < K; k += mK)
            {
                int dK = std::min(K, k + mK) - k;
                for (int i = 0; i < M; i += mM)
                {
                    int dM = std::min(M, i + mM) - i;
                    if (k == 0)
                        init_c(dM, dN, C + i * N + j, N);
                    reorder_a_6(A + i * K + k, K, dM, dK, bufA.p);
                    macro_v7(dM, dN, dK, bufA.p, B + k * N + j, N, 
                        bufB.p, i == 0, C + i * N + j, N);
                }
            }
        }
    }
    

    Результаты замеров для (M=1152, N=1152, K=115200) дают результат в 97.3 GFLOPS. Т.е. мы даже немного превысили результат для матриц среднего размера. Фактически мы получили универсальный алгоритм (на самом деле нет, про ограничения в следующем разделе), который практически одинаково эффективно (порядка 80% от теоретически достижимого макимума) работает для любого размера матриц. На этом предлагаю остановиться и описать, что у нас в итоге получилось.

    Общая схема алгоритма


    На рисунке ниже приведена схема получившегося алгоритма:



    Микро ядро


    • Цикл-1 по переменной k. Переупорядоченные данные из матрицы B лежат в кэше L1, переупорядоченные данные из матрицы A лежат в кэше L2. Сумма аккумулируется в регистрах (кэше L0). Результат записывается в основную память. Размеры микроядра определяются длиной SIMD вектора и количеством векторных регистров. Длина цикла определяется размером кэша L1, где хранится B.

    Макро ядро


    • Цикл-2 по переменной i. Пробегает микроядром по переупорядоченным данным матрицы A, которые лежат в кэше L2.
    • Цикл-3 по переменной j. Пробегает микроядром по переупорядоченным данным матрицы B, которые лежат в кэше L3. Опционально переупорядочивает недостающие данные в B.

    Размеры макроядра определяются величиной кэша.

    Основная функция


    • Цикл-4 по переменной i. Пробегает макроядром по матрице A. На каждой итерации переупорядочивает значения A. Опционально инициализирует значения матрицы С.
    • Цикл-5 по переменной k. Пробегает макроядром по матрицам A и B.
    • Цикл-6 по переменной j. Пробегает макроядром по матрице B.

    Что осталось за кадром?


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

    1. В реальности, к сожалению, размер матриц не всегда кратен размерам микроядра, потому края матриц приходится обрабатывать особым образом. Для чего приходится реализовывать микроядра разных размеров.
    2. Для разных типов процессоров реализуются разные наборы микроядер и функций переупорядочивания. Также свои микроядра будет для чисел с двойной точностью и для комплексных чисел. К счастью, зоопарк микроядер ограничен только ими и на верхнем уровне код достаточно универсальный.
    3. Микроядра часто пишут прямо на ассемблере. Также проводят дополнительное разворачивание циклов. Но это не приводит к существенному ускорению — основные оптимизации заключаются в эффективном использовании кэшевой иерархии памяти процессора.
    4. Для матриц малого размера (по любому измерению) применяют особые алгоритмы — иногда переупорядочивание не эффективно, иногда нужно применять другой порядок обхода матриц. А иногда и реализовывать особые микроядра.
    5. В обобщенном алгоритме матричного умножения все три матрицы могут быть транспонированы. Казалось бы число возможных алгоритмов возрастает в 8 раз! К счастью применение переупорядочивания входных данных, позволяет для всех случаев обойтись унивесальными микроядрами.
    6. Практически все современные процессоры — многоядерны. И библиотеки матричного умножения используют многопоточность для ускорения вычислений. Обычно для этого используется еще 1-3 дополнительных цикла, в которых происходит разбиение задач по разным потокам.

    Заключение


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

    Код проекта с алгоритмами из статьи можно найти на Github.

    Надеюсь вам было интересно!
    Поделиться публикацией

    Комментарии 44

      0

      Интересно было бы сравнить по скорости с библиотекой MKT от того же Intel,
      и например с Eigen еще.

        0
        На больших матрицах думаю будет небольшой проигрыш в пределах 5-10% (для 1 — потока). На малых размерах проигрыш будет в разы (причины этого указаны в статье).
          –1
          А померять? Штрассен должен давать ускорение в 10-100 раз на больших матрицах и в 1000+ на ОЧЕНЬ больших и хорошие библиотеки обязаны его поддерживать в теории
          UPD: вижу ваш ответ ниже. Странно.
            +1
            Штрассен на больших матрицах будет накапливать ошибки округления
        0
        А почему даже после введения правильного выравнивания используются инструкции для чтения/записи по невыровненному адресу? Разве от них не будет оверхеда? И вообще, мне кажется лучшим решением будет просто ввести требование по выравниванию непосредственно для входных данных.
          +2
          Я уже достаточно давно разработываю проект по оптимизации различных алгоритмов (в основном по компьютерному зрению) при помощи SIMD инструкций. В частности я всегда старался выравнивать данные и использовать варианты с выровненными инструкциями (например, _mm_load_ps вместо _mm_loadu_ps), где только можно. Однако в последних поколениях процессоров Интел скорость работы этих инструкций одинаковая (при условии обращения к выровненному адресу конечно). Более того компилятор часто вместо _mm_load_ps подставляет ассемблерную инструкцию соответсвующую _mm_loadu_ps.
            0
            Ясно, спасибо за информацию, и за статью.
              0

              Быстрое гугление говорит, что разница между _mm_load_p и _mm_loadu_p на выравненных данных исчезла начиная с Nehalem (2008).

                –1
                Я уже достаточно давно разработываю проект по оптимизации различных алгоритмов (в основном по компьютерному зрению) при помощи SIMD инструкций.


                Обёртку на python ещё не завезли? :)
                  +1
                  Если подвезете, я возражать не буду. Сам я в этом деле не специалист.
                    –1
                    Это была известная отсылка к принципу: «Для всего есть обёртка на python».

                    Сам я ни разу не писал обёрток, а с графикой работал только в виде лаб в универе, так что ещё больший не специалист, к сожалению.
              0

              Очень информативно, а вы пробовали сравнить результаты для реализации "в лоб" у gcc и clang (-O3 -march=haswell -ffast-math)?
              Я по правде говоря ожидал что оптимизации из первого шага и так должны быть сделаны компилятором (loop unswitch/rotate, const propogation?)

                0
                Вы не пробовали провести что-то подобное для матриц в столбцовом формате?
                  0
                  Основной принцип останется прежним: микроядра, максимальная локализация в кэше. Изменится только функции переупорядочивания.
                  –3
                  GPU просто созданы для быстрого перемножения матриц.
                  Зачем пытаться их обогнать на CPU? Ведь все равно не выйдет.
                    +5
                    Цель статьи — не написать саму быструю реализацию, а показать внутренее устройство алгоритма. Понятно, что GPU будет быстрее. С этим никто не спорит.
                      0
                      А где можно почитать (или может, у вас есть в планах написать статью) об особенностях реализации умножения на картах?

                      Интересно, какие там аппаратные особенности и как они используются.

                      Спасибо.
                    +5
                    Я бы упомянул симпатичную математическую абстракцию часто используемую в подобных в подобных оптимизациях. Для матричного умножения справедливо что если мы поделим матрицу на блоки меньшего размера, то матрицы можно перемножать перемножая блоки

                    Это можно использовать либо явно (tiled версия aka блочное перемножение матриц) либо в рекурсивных схемах (где возникают всякие вкусные субкубические варианты типа алгоритма Штрассена). У Вас как я понимаю по сути tiled версия сильно заточенная под AVX — ширина блоков выбирается под ширину регистров, длина — под размер L1.
                      +1
                      На сколько я видел в разных реализациях, нигде блочное перемножение не используется. Видимо потому, что это приводит к дополнительным копированиям памяти. А как следует из статьи выигрыш от локализации данных в кэше — до 10 раз. Выигрыш от подобных блочных алгоритмов будет проявляться только на матрицах циклопических размеров. А такие уже будет целесообразно считать на GPU.
                        +2
                        На платформе РС может быть и не используется, а мы в обработке сигналов используем. Хотя ваше микроядро по-моему туда же относится.
                          +3
                          Наверное не так выразился. Мысль которую я хотел донести: Не используются субкубические варианты типа алгоритма Штрассена.
                      +1
                      Причём важна форма блоков! Типичная задача — перемножить большие матрицы в медленной памяти, цикл такой: разбиваем на блоки, пересылаем блоки аргументов во внутреннюю память, перемножаем, выгружаем блок результата. Как минимизировать пересылки? Оказывается, чем крупнее и квадратнее блок, тем лучше, поскольку объём работы по перемножению матриц IxK и KxJ пропорционален I*J*K, а объём пересылок — I*K + K*J + I*J.
                        0
                        Вспомнилось, как я году эдак в 2013 матрицы перемножал. Правда, там в ячейках были не float, а int64 (точнее, в A*B=C в A и B — инты, а в C должно было инт64 влезать). Размеры матричек 5000х5000. Подробнее о задаче на fastcomputing.org (вебархив, ибо проект ныне покойный). А вот табличка рекордов. В 40 секунд тогда упихал. Подробности реализации сейчас не вспомню, но вроде там Виноград-Штрассен, который где-то то ли на 200х200, то ли на 1000х1000 переключается уже на перемножение в лоб. Без Штрассена было тормознее, насколько помню. В перемножении в лоб тоже всякие SSE, помню разве что только один хак: мы умножаем не строчку A на строчку B (B транспонирована, конечно), а пару строчек A на пару строчек B, тем самым сокращая число чтений, потом все идет чисто на регистрах процессора, и на выходе получаем 4 значения для этих самых пар строчек.
                          0
                          MKL конечно платный, не всем подходит, но есть же MKL-DNN открытый.
                          Для мелких лучше использовать libxsmm. Он делает JIT компиляцию микрокернела под вашу платформу и бесплатный.

                            0
                            MKL недавно перестал быть платным
                              0
                              Тем более. Берите МКЛ и забудьте про перемножение матриц. По крайней мере кроме краевых примеров 2х2.
                            0
                            Наверное, стоит добавить, что когда матриц много, то очень часто имеет смысл сначала определить порядок умножения матриц.
                              0
                              Очень полезная и хорошая статья. Насколько я знаю, MKL всегда работает на числе потоков, равном числу физических ядер. Насколько результаты статьи зависят от числа использованных виртуальных/физических потоков?
                                0
                                Цитата из введения: С целью ограничить объем изложения, я ограничился описанием однопоточного алгоритма для обычных процессоров. Тема многопоточности и алгоритмов для графических ускорителей явно заслуживает отдельной статьи.
                                  0
                                  Уже увидел. Значит это вопрос к следующей статье. С MKL, как с универсальным средством сравнения для высокопроизводительных вычислений, полезно приводить сравнение всегда.
                                  0
                                  MKL работает:
                                  а) Если выставлен MKL_NUM_THREADS — то количество потоков ограничено этой переменной;
                                  б) Если выставлен OMP_NUM_THREADS и не выставлена а) то ограничивается этой переменной;
                                  По умолчанию считается что количество потоков = количеству ядер. Т.е если ОБЕ переменные не выставлены — будет работать на всем до чего дотянется.
                                    0
                                    Т.е. для использования всех виртуальных ядер с учетом гипертрейдинга достаточно установить MKL_NUM_THREADS или OMP_NUM_THREADS в удвоенное количество физических ядер? Заранее спасибо.
                                      0
                                      С точки зрения ОС — гипертрединг это и так двойное кол-во ядер, так что вообще не ставьте переменную, оно само определит.
                                        0
                                        Может быть и определяет, но использует только физические ядра.
                                      0
                                      Попробовал на dgemm с включенным гипертрейдингом. Уменьшить число потоков через set MKL_NUM_THREADS получается, а увеличить больше, чем число физических ядер — нет. Быстродействие упрямо показывает мах 50% загруженности cpu. Аналогично mkl_set_num_threads не дает загрузить число ядер больше физического. Функция mkl_get_max_threads возвращает значение равное числу физических ядер, а не виртуальных. Помнится, разработчики из Нижнего Новгорода подтверждали этот факт.
                                        0
                                        Тут решение уже за MKL. То, что вы видите загрузку в 50% — к реальности может(и скорее всего так и есть) не иметь никакого отношения, так как каждое из виртуальных ядер может полностью загрузить реальное. Т.е. с точки зрения операционки, одно ядро загружено полностью, на втором вообще ничего не запущено — загрузка 50%, а на деле единственное физическое ядро имеет полностью загруженные конвееры FPU.
                                        MKL библиотека от производителя процессора, я сильно сомневаюсь чтобы они специально занижали производительность своих продуктов.
                                        А в случае когда два потока работают на одном физическом ядре они разделяют как блоки FPU, так и кеш, что скажется на производительности скорее негативно, чем позитивно.
                                          0
                                          Использование гипертридинга в нагруженных кернелах не приведет ни к чему хорошему.
                                          50% загрузка видимо видна в диспетчере процессов Windows, там да. Логическое ядро отображается как нормальное, в связи с чем получается такое недопонимание.
                                            0
                                            Запустив одновременно две задачи, убедился, что время расчета выросло в два раза.
                                      0
                                      Статья читается как какая-то магия. Я, как джавист из энтерпрайза, смог бы самостоятельно придти только к первому варианту решения, и если бы после этого встал вопрос повышения производительности, то только и смог бы придумать что «попробовать в несколько потоков выполнять» или «раскидать вычисления по облаку».
                                        0
                                        Делал подобное лет 10 назад, использовал практически все те же оптимизации, правда оптимизацией под L2/L3 уже не заморачивался, прирост не особо большой от этого получался. Результатами до сих пор пользуюсь — библиотека получилась небольшой, не надо что-то стороннее тянуть большое.
                                        Но все-таки умножение матриц это достаточно простая вещь. Вот сделать то же самое для обращения матрицы (LUP разложения) и распараллелить обращение было посложнее.
                                          0
                                          А что если перед умножением транспонировать матрицу B? Тогда по идее можно получить ускорение за счет того, что элементы матрицы B, которые участвуют в умножении на одну строку матрицы A станут соседями в памяти и будут лучше размещаться в кеше.
                                            +2
                                            К сожалению нет. На первый взгляд действительно все проще: данные A и B лежат одинаково — только считай взаимное скалярное произведение их строчек.

                                            Однако: максимальный размер микроядра получится 3x4, что дает нам (3 + 4)/(3*4) = ~0.58 загрузок на одну fma. Напомню, что при классической схеме с окном 6x16 получается (6 + 16)/(6*16) = ~0.23 загрузок на одну fma. Т.е. предложенная вами схема почти в 2.5 раза более требовательна к пропускной способности памяти. В принципе мои внутренние тесты это подтверждают.
                                            +2
                                            Спасибо, классная статья. Не ожидал, что можно так сильно оптимизировать.

                                            В начале статьи Вы критикуете малопонятность опенсурсных реализаций, ссылаясь в большей степени на ассемблер. Могу покритиковать Ваш код, так как он использует захардкоженные значения :) Думаю, именованные константы мне бы были более понятны, чем всякие 6, 24 и т. д…
                                              0
                                              Как раз имею Core i5 4-го поколения (должно иметь приличные FMA), правда использовал для Matlab и для FFT (вроде как там по такой же логике польза от FMA) на комплексных числах. Вроде как прирост 4670 против Cel. 1005M достаточно большой был.
                                              Ещё раньше использовал Athlon II X3 445, но для немного других задач. Там можно было явно записать решение задачи
                                              dE/dt = f(E)*E + a*(laplas)E
                                              в форме
                                              dE/E = [f(E) — a*k2] dt.
                                              То есть это обычный метод с функцией fft. При этом 3-ядерник грузился процентов на 60, но брат думает, что это тратилось второе ядро только на вывод графики.

                                              Только полноправные пользователи могут оставлять комментарии. Войдите, пожалуйста.

                                              Самое читаемое