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

МЕНЮ


Искусственный интеллект. Новости
Поиск
Регистрация на сайте
Сбор средств на аренду сервера для ai-news

ТЕМЫ


Новости ИИРазработка ИИВнедрение ИИРабота разума и сознаниеМодель мозгаРобототехника, БПЛАТрансгуманизмОбработка текстаТеория эволюцииДополненная реальностьЖелезоКиберугрозыНаучный мирИТ индустрияРазработка ПОТеория информацииМатематикаЦифровая экономика

Авторизация



RSS


RSS новости

Новостная лента форума ailab.ru


Введение

Умножение матриц — это один из базовых алгоритмов, который широко применяется в различных численных методах, и в частности в алгоритмах машинного обучения. Многие реализации прямого и обратного распространения сигнала в сверточных слоях неронной сети базируются на этой операции. Так порой до 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. Надеюсь вам было интересно!


Источник: habr.com

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