Уже немало лет прошло, как я познакомился с инструкциями MMX, SSE, а позже и AVX на процессорах Intel. В своё время они казались какой-то магией на фоне x86 ассемблера, который уже давно стал чем-то обыденным. Они меня настолько зацепили, что пару лет назад у меня появилась идея написать свой собственный софт рендерер для одной известной игры. Сподвигло меня на это то, какую производительность обещали эти инструкции. В какой-то момент я даже думал об этом написать. Но писать текст оказалось куда сложнее кода.
В то время я хотел избежать проблем с поддержкой на разных процессорах. Хотелось иметь возможность проверить мой рендерер на максимально доступном количестве. У меня до сих пор остались знакомые со старыми AMD процессорами, и их потолок был SSE3. Поэтому на тот момент я решил ограничиться максимум SSE3. Так появилась векторная математическая библиотека, чуть менее, чем полностью реализованная на SSE, с редким включением до SSE3. Однако в какой-то момент мне стало интересно, какую максимальную производительность я смогу выжать из процессора для ряда критичных операций векторной математики. Одной из таких операций является умножение матриц float 4 на 4.
Собственно, этим делом решил заняться больше ради развлечения. Давно уже написал и использую умножение матриц для моего софт рендера на SSE и вроде мне хватает. Но тут решил посмотреть, сколько тактов я смогу выжать в принципе из умножения 2-х матриц float4x4. На моём текущем SSE варианте это 16 тактов. Правда недавний переход на IACA 3 стал показывать 19, так как на некоторые инструкции вместо 0* стал писать 1*. Видимо раньше это было просто недоработкой анализатора.
Коротко об используемых утилитах
Для анализа кода использовал известную утилиту Intel Architecture Code Analyzer. Для анализа использую архитектуру Haswell (HSW), как минимальную с поддержкой AVX2. Для написания кода также очень удобно пользоваться: Intel Intrinsics Guide и Intel optimization manual.
Для сборки использую MSVS 2017 Community с консоли. Код пишу в варианте с интринсиками. Пишешь один раз, и обычно он сразу работает на разных платформах. К тому же x64 компилятор на VC++ не поддерживает инлайн ассемблер, а хочется чтобы и под x64 работало.
Поскольку данная статья уже несколько выходит за рамки уровня начинающего в SIMD программировании, я не буду описывать регистры, инструкции, рисовать (или тырить) красивые картинки и пытаться учить программировать с использованием SIMD инструкций. На сайте Intel полно отличной, понятной и подробной документации.
Хотел сделать всё проще.… А получилось как всегда
Вот тут начинается момент, который немало усложняет как реализацию, так и статью. Поэтому немного на нём остановлюсь. Писать умножение матриц со стандартным построчным расположением элементов мне не интересно. Кому было надо, и так изучили в ВУЗах или самостоятельно. Наша же цель — производительность. Во-первых, я давно перешёл на расположение по столбцам. Мой софт рендерер базируется на OpenGL API и поэтому, дабы избежать лишних транспонирований, я начал хранить элементы по столбцам. Также это важно потому, что умножение матриц не так критично. Ну умножили 2-5-10 матриц. И всё. А потом умножаем готовую матрицу на тысячи-миллионы вершин. И вот эта операция уже куда критичнее. Можно, конечно, каждый раз транспонировать. Но зачем, если этого можно избежать.
Но вернёмся исключительно к матрицам. С хранением по столбцам определились. Однако можно ещё усложнить. Мне удобнее хранить старшие элементы векторов и матричных строк в SIMD регистрах так, что x в старшем float (индекс 3), а w в младшем (индекс 0). Тут, видимо, придётся снова сделать отступление на тему почему так.
Дело в том, что в софт рендерере в векторе компонентой w приходится манипулировать чаще (там хранится 1/z), и очень удобно делать это через _ss вариант операции (операции исключительно с компонентой в младшем float регистра xmm), не трогая x, y, z. Поэтому в SSE регистре вектор хранится в понятном порядке x, y, z, w, а в памяти в обратном w, z, y, x.
Далее, все варианты умножения также реализованы отдельными функциями. Так сделано потому, что их я использую для подстановки нужного варианта в зависимости от поддерживаемого типа инструкций. Хорошо описано тут.
Реализуем базовый функционал
Умножение с циклами, row ordered
for (int i = 0; i < 4; ++i) {
for (int j = 0; j < 4; ++j) {
r[i][j] = 0.f;
for (int k = 0; k < 4; ++k) {
r[i][j] += m[i][k] * n[k][j];
}
}
}
Тут всё просто и понятно. На каждый элемент мы делаем 4 умножения и 3 сложения.
В сумме это 256 умножений и 192 сложения. Это без учёта чтения записи элементов.
Всё печально, короче. На этот вариант для внутреннего цикла IACA выдала: 3.65 тактов под x86 сборку и 2.97 под x64 сборку. Не спрашивайте, почему дробные цифры. Не знаю. IACA 2.1 подобным не страдала. В любом случае, эти цифры надо умножить примерно на 4*4*4 = 64. Даже если взять x64, в результате получается около 192 тактов. Понятно, что это примерная оценка. Оценивать производительность точнее для данного варианта не вижу смысла.
Реализация с циклами, column ordered
for (int i = 0; i < 4; ++i) {
for (int j = 0; j < 4; ++j) {
r[j][i] = 0.f;
for (int k = 0; k < 4; ++k) {
r[j][i] += m[k][i] * n[j][k];
}
}
}
Умножение с циклами, SIMD ориентированное хранение
for (int i = 0; i < 4; ++i) {
for (int j = 0; j < 4; ++j) {
r[j][3-i] = 0.f;
for (int k = 0; k < 4; ++k) {
r[j][3-i] += m[k][3-i] * n[j][3-k];
}
}
}
Эта реализация несколько упрощает понимание того, что происходит внутри, но явно недостаточно.
Вспомогательные классы
Для удобства понимания и написания референсного и отладочного кода удобно реализовать пару вспомогательных классов. Ничего лишнего, всё только для понимания. Отмечу, что реализация полноценных классов вектора и матрицы отдельный непростой вопрос, и в тему данной статьи не входит.
struct alignas(sizeof(__m128)) vec4 {
union {
struct { float w, z, y, x; };
__m128 fmm;
float arr[4];
};
vec4() {}
vec4(float a, float b, float c, float d) : w(d), z(c), y(b), x(a) {}
static bool equ(float const a, float const b, float t = .00001f) {
return fabs(a-b) < t;
}
bool operator == (vec4 const& v) const {
return equ(x, v.x) && equ(y, v.y) && equ(z, v.z) && equ(w, v.w);
}
};
struct alignas(sizeof(__m256)) mtx4 {
// тут всё больше для наглядности хранения в памяти и регистрах
union {
struct {
float
_30, _20, _10, _00,
_31, _21, _11, _01,
_32, _22, _12, _02,
_33, _23, _13, _03;
};
__m128 r[4];
__m256 s[2];
vec4 v[4];
};
// добавляем простейшие инициализаторы
mtx4() {}
mtx4(
float i00, float i01, float i02, float i03,
float i10, float i11, float i12, float i13,
float i20, float i21, float i22, float i23,
float i30, float i31, float i32, float i33)
: _00(i00), _01(i01), _02(i02), _03(i03)
, _10(i10), _11(i11), _12(i12), _13(i13)
, _20(i20), _21(i21), _22(i22), _23(i23)
, _30(i30), _31(i31), _32(i32), _33(i33)
{}
// для передачи в реализацию умножения
operator __m128 const* () const { return r; }
operator __m128* () { return r; }
// для тестов
bool operator == (mtx4 const& m) const {
return v[0]==m.v[0] && v[1]==m.v[1] && v[2]==m.v[2] && v[3]==m.v[3];
}
// инициализаторы
static mtx4 identity() {
return mtx4(
1.f, 0.f, 0.f, 0.f,
0.f, 1.f, 0.f, 0.f,
0.f, 0.f, 1.f, 0.f,
0.f, 0.f, 0.f, 1.f);
}
static mtx4 zero() {
return mtx4(
0.f, 0.f, 0.f, 0.f,
0.f, 0.f, 0.f, 0.f,
0.f, 0.f, 0.f, 0.f,
0.f, 0.f, 0.f, 0.f);
}
};
Референсная функция для тестов
Поскольку принятый порядок элементов в матрице немало усложняет понимание, нам также не помешает референсная понятная функция, которая покажет в дальнейших реализациях, что всё работает правильно. Последующие результаты мы будем сравнивать с ней.
void mul_mtx4_mtx4_unroll(__m128* const _r, __m128 const* const _m, __m128 const* const _n) {
mtx4 const& m = **reinterpret_cast<mtx4 const* const*>(&_m);
mtx4 const& n = **reinterpret_cast<mtx4 const* const*>(&_n);
mtx4& r = **reinterpret_cast<mtx4* const*>(&_r);
r._00 = m._00*n._00 + m._01*n._10 + m._02*n._20 + m._03*n._30;
r._01 = m._00*n._01 + m._01*n._11 + m._02*n._21 + m._03*n._31;
r._02 = m._00*n._02 + m._01*n._12 + m._02*n._22 + m._03*n._32;
r._03 = m._00*n._03 + m._01*n._13 + m._02*n._23 + m._03*n._33;
r._10 = m._10*n._00 + m._11*n._10 + m._12*n._20 + m._13*n._30;
r._11 = m._10*n._01 + m._11*n._11 + m._12*n._21 + m._13*n._31;
r._12 = m._10*n._02 + m._11*n._12 + m._12*n._22 + m._13*n._32;
r._13 = m._10*n._03 + m._11*n._13 + m._12*n._23 + m._13*n._33;
r._20 = m._20*n._00 + m._21*n._10 + m._22*n._20 + m._23*n._30;
r._21 = m._20*n._01 + m._21*n._11 + m._22*n._21 + m._23*n._31;
r._22 = m._20*n._02 + m._21*n._12 + m._22*n._22 + m._23*n._32;
r._23 = m._20*n._03 + m._21*n._13 + m._22*n._23 + m._23*n._33;
r._30 = m._30*n._00 + m._31*n._10 + m._32*n._20 + m._33*n._30;
r._31 = m._30*n._01 + m._31*n._11 + m._32*n._21 + m._33*n._31;
r._32 = m._30*n._02 + m._31*n._12 + m._32*n._22 + m._33*n._32;
r._33 = m._30*n._03 + m._31*n._13 + m._32*n._23 + m._33*n._33;
}
Здесь наглядно расписан классический алгоритм, ошибиться сложно (но можно :-) ). На него IACA выдала: x86 — 69.95 такта, x64 — 64 такта. Вот относительно 64 тактов и будем смотреть ускорение данной операции в последующем.
SSE реализация
Классический SSE алгоритм
Почему классический? Потому что он давно уже есть в реализации FVec в составе MSVS. Для начала напишем как у нас представлены элементы матрицы в SSE регистрах. Здесь уже выглядит проще. Просто транспонированная матрица.
// представление матрицы в регистрах
00, 10, 20, 30 // m[0] - в SIMD строках/регистрах храним столбцы
01, 11, 21, 31 // m[1]
02, 12, 22, 32 // m[2]
03, 13, 23, 33 // m[3]
Берём код unroll варианта выше. Какой-то он недружелюбный для SSE. Первая группа строк состоит из результатов по столбцу результирующей матрицы: r._00, r._01, r._02, r._03. У нас это столбец, а нам нужна строка. Да и m, n выглядят неудобно для рассчётов. Поэтому переставляем строчки алгоритма, чтобы результат r был построчным.
// первая группа, это строчка результирующей матрицы r[0]
r00 = m00*n00 + m01*n10 + m02*n20 + m03*n30;
r10 = m10*n00 + m11*n10 + m12*n20 + m13*n30;
r20 = m20*n00 + m21*n10 + m22*n20 + m23*n30;
r30 = m30*n00 + m31*n10 + m32*n20 + m33*n30;
// вторая группа, это строчка результирующей матрицы r[1]
r01 = m00*n01 + m01*n11 + m02*n21 + m03*n31;
r11 = m10*n01 + m11*n11 + m12*n21 + m13*n31;
r21 = m20*n01 + m21*n11 + m22*n21 + m23*n31;
r31 = m30*n01 + m31*n11 + m32*n21 + m33*n31;
// третья группа, это строчка результирующей матрицы r[2]
r02 = m00*n02 + m01*n12 + m02*n22 + m03*n32;
r12 = m10*n02 + m11*n12 + m12*n22 + m13*n32;
r22 = m20*n02 + m21*n12 + m22*n22 + m23*n32;
r32 = m30*n02 + m31*n12 + m32*n22 + m33*n32;
// четвертая группа, это строчка результирующей матрицы r[3]
r03 = m00*n03 + m01*n13 + m02*n23 + m03*n33;
r13 = m10*n03 + m11*n13 + m12*n23 + m13*n33;
r23 = m20*n03 + m21*n13 + m22*n23 + m23*n33;
r33 = m30*n03 + m31*n13 + m32*n23 + m33*n33;
А вот так уже гораздо лучше. Что, собственно, мы видим? По столбцам алгоритма в каждой группе у нас задействованы строчки матрицы m:
m[0]={00,10,20,30}, m[1]={01,11,21,31}, m[2]={02,12,22,32}, m[3]={03,13,23,33},
которые умножаются на один и тот же элемент матрицы n. Например, для первой группы это:n._00,n._10,n._20,n._30. И элементы матрицы n для каждой группы строк алгоритма снова лежат в одной строке матрицы.
Дальше всё просто: строки матрицы m мы просто берём по индексу, а вот что касается элементов n, то мы берём её строку и через инструкцию shuffle раскидываем её всем 4-м элементам регистра, чтобы умножить на строку матрицы m в регистре. Например для элемента n._00 (помним, что его смещение в регистре имеет индекс 3) это будет:
_mm_shuffle_ps(n[0], n[0], _MM_SHUFFLE(3,3,3,3))
В упрощённом виде алгоритм выглядит так:
// задействуется строка n[0]={00,10,20,30}
r[0] = m[0] * n00 + m[1] * n10 + m[2] * n20 + m[3] * n30;
// задействуется строка n[1]={01,11,21,31}
r[1] = m[0] * n01 + m[1] * n11 + m[2] * n21 + m[3] * n31;
// задействуется строка n[2]={02,12,22,32}
r[2] = m[0] * n02 + m[1] * n12 + m[2] * n22 + m[3] * n32;
// задействуется строка n[3]={03,13,23,33}
r[3] = m[0] * n03 + m[1] * n13 + m[2] * n23 + m[3] * n33;
void mul_mtx4_mtx4_sse_v1(__m128* const r, __m128 const* const m, __m128 const* const n) {
r[0] =
_mm_add_ps(
_mm_add_ps(
_mm_mul_ps(m[0], _mm_shuffle_ps(n[0], n[0], _MM_SHUFFLE(3,3,3,3))),
_mm_mul_ps(m[1], _mm_shuffle_ps(n[0], n[0], _MM_SHUFFLE(2,2,2,2)))),
_mm_add_ps(
_mm_mul_ps(m[2], _mm_shuffle_ps(n[0], n[0], _MM_SHUFFLE(1,1,1,1))),
_mm_mul_ps(m[3], _mm_shuffle_ps(n[0], n[0], _MM_SHUFFLE(0,0,0,0)))));
r[1] =
_mm_add_ps(
_mm_add_ps(
_mm_mul_ps(m[0], _mm_shuffle_ps(n[1], n[1], _MM_SHUFFLE(3,3,3,3))),
_mm_mul_ps(m[1], _mm_shuffle_ps(n[1], n[1], _MM_SHUFFLE(2,2,2,2)))),
_mm_add_ps(
_mm_mul_ps(m[2], _mm_shuffle_ps(n[1], n[1], _MM_SHUFFLE(1,1,1,1))),
_mm_mul_ps(m[3], _mm_shuffle_ps(n[1], n[1], _MM_SHUFFLE(0,0,0,0)))));
r[2] =
_mm_add_ps(
_mm_add_ps(
_mm_mul_ps(m[0], _mm_shuffle_ps(n[2], n[2], _MM_SHUFFLE(3,3,3,3))),
_mm_mul_ps(m[1], _mm_shuffle_ps(n[2], n[2], _MM_SHUFFLE(2,2,2,2)))),
_mm_add_ps(
_mm_mul_ps(m[2], _mm_shuffle_ps(n[2], n[2], _MM_SHUFFLE(1,1,1,1))),
_mm_mul_ps(m[3], _mm_shuffle_ps(n[2], n[2], _MM_SHUFFLE(0,0,0,0)))));
r[3] =
_mm_add_ps(
_mm_add_ps(
_mm_mul_ps(m[0], _mm_shuffle_ps(n[3], n[3], _MM_SHUFFLE(3,3,3,3))),
_mm_mul_ps(m[1], _mm_shuffle_ps(n[3], n[3], _MM_SHUFFLE(2,2,2,2)))),
_mm_add_ps(
_mm_mul_ps(m[2], _mm_shuffle_ps(n[3], n[3], _MM_SHUFFLE(1,1,1,1))),
_mm_mul_ps(m[3], _mm_shuffle_ps(n[3], n[3], _MM_SHUFFLE(0,0,0,0)))));
}
Теперь меняем в алгоритме элементы n на соответствующий shuffle, умножение на _mm_mul_ps, сумму на _mm_add_ps, и всё, готово. Оно работает. Код выглядит, правда, куда страшнее, чем выглядел сам алгоритм. На этот код IACA выдала: x86 — 18.89, x64 — 16 тактов. Это в 4 раза быстрее предыдущего. В SSE регистре 4-е float. Почти линейная зависимость.
Украшаем SSE реализацию
И всё-таки в коде это выглядит ужасно. Попытаемся это улучшить, написав немного синтаксического сахара.
// превращаем имена функций в обычные удобочитаемые операции (которые всё-таки лучше прятать в namespace)
__m128 operator + (__m128 const a, __m128 const b) { return _mm_add_ps(a, b); }
__m128 operator - (__m128 const a, __m128 const b) { return _mm_sub_ps(a, b); }
__m128 operator * (__m128 const a, __m128 const b) { return _mm_mul_ps(a, b); }
__m128 operator / (__m128 const a, __m128 const b) { return _mm_div_ps(a, b); }
//_mm_shuffle_ps(u, v, _MM_SHUFFLE(3,2,1,0)) превращается в shuf<3,2,1,0>(u, v)
template <int a, int b, int c, int d> __m128 shuf(__m128 const u, __m128 const v)
{ return _mm_shuffle_ps(u, v, _MM_SHUFFLE(a, b, c, d)); }
template <int a, int b, int c, int d> __m128 shuf(__m128 const v)
{ return _mm_shuffle_ps(v, v, _MM_SHUFFLE(a, b, c, d)); }
// облегчённый одноиндексный вариант
template <int i> __m128 shuf(__m128 const u, __m128 const v)
{ return _mm_shuffle_ps(u, v, _MM_SHUFFLE(i, i, i, i)); }
template <int i> __m128 shuf(__m128 const v)
{ return _mm_shuffle_ps(v, v, _MM_SHUFFLE(i, i, i, i)); }
// для float векторов можно попробовать и такой экзотический вариант,
// который иногда даёт профит, а иногда нет
template <int a, int b, int c, int d> __m128 shufd(__m128 const v)
{ return _mm_castsi128_ps(_mm_shuffle_epi32(_mm_castps_si128(v), _MM_SHUFFLE(a, b, c, d))); }
template <int i> __m128 shufd(__m128 const v)
{ return _mm_castsi128_ps(_mm_shuffle_epi32(_mm_castps_si128(v), _MM_SHUFFLE(i, i, i, i))); }
Данные функции компилятор умеет отлично инлайнить (хотя иногда без __forceinline никак).
void mul_mtx4_mtx4_sse_v2(__m128* const r, __m128 const* const m, __m128 const* const n) {
r[0] = m[0]*shuf<3>(n[0]) + m[1]*shuf<2>(n[0])
+ m[2]*shuf<1>(n[0]) + m[3]*shuf<0>(n[0]);
r[1] = m[0]*shuf<3>(n[1]) + m[1]*shuf<2>(n[1])
+ m[2]*shuf<1>(n[1]) + m[3]*shuf<0>(n[1]);
r[2] = m[0]*shuf<3>(n[2]) + m[1]*shuf<2>(n[2])
+ m[2]*shuf<1>(n[2]) + m[3]*shuf<0>(n[2]);
r[3] = m[0]*shuf<3>(n[3]) + m[1]*shuf<2>(n[3])
+ m[2]*shuf<1>(n[3]) + m[3]*shuf<0>(n[3]);
}
А вот так уже куда лучше и читабельней. На это IACA выдала примерно ожидаемый результат: x86 — 19 (а чего не дробный?), x64 — 16. По сути производительность не изменилась, но код намного красивее и понятнее.
Небольшой вклад в будущую оптимизацию
Введём ещё одно улучшение на уровне функции, не так уж давно появившейся в железном варианте. Операцию multiple-add (fma). fma(a, b, c) = a * b + c.
__m128 mad(__m128 const a, __m128 const b, __m128 const c) {
return _mm_add_ps(_mm_mul_ps(a, b), c);
}
Зачем это надо? Прежде всего, для оптимизации в будущем. Например можно просто в готовом коде заменить mad на fma через те же макросы, кому как удобно. Но основу под оптимизацию мы заложим уже сейчас:
void mul_mtx4_mtx4_sse_v3(__m128* const r, __m128 const* const m, __m128 const* const n) {
r[0] = mad(m[0], shuf<3>(n[0]), m[1]*shuf<2>(n[0]))
+ mad(m[2], shuf<1>(n[0]), m[3]*shuf<0>(n[0]));
r[1] = mad(m[0], shuf<3>(n[1]), m[1]*shuf<2>(n[1])
+ mad(m[2], shuf<1>(n[1]), m[3]*shuf<0>(n[1]));
r[2] = mad(m[0], shuf<3>(n[2]), m[1]*shuf<2>(n[2]))
+ mad(m[2], shuf<1>(n[2]), m[3]*shuf<0>(n[2]));
r[3] = mad(m[0], shuf<3>(n[3]), m[1]*shuf<2>(n[3]))
+ mad(m[2], shuf<1>(n[3]), m[3]*shuf<0>(n[3]));
}
IACA: x86 — 18.89, x64 — 16. Опять дробное. Всё-таки IACA порой выдаёт странные результаты. Код изменился не так сильно. Наверное, даже чуть-чуть похуже стал. Но оптимизация порой требует подобных жертв.
Переходим на сохранение через _mm_stream
Разные руководства по оптимизации рекомендуют лишний раз не дёргать кэш для массовых операций сохранения. Обычно это обоснованно, когда вы занимаетесь обработкой вершин, которых тысячи и больше. Но для матриц это, пожалуй, не так важно. Однако всё равно добавлю.
void mul_mtx4_mtx4_sse_v4(__m128* const r, __m128 const* const m, __m128 const* const n) {
_mm_stream_ps(&r[0].m128_f32[0],
mad(m[0], shuf<3>(n[0]), m[1]*shuf<2>(n[0])) +
mad(m[2], shuf<1>(n[0]), m[3]*shuf<0>(n[0])));
_mm_stream_ps(&r[1].m128_f32[0],
mad(m[0], shuf<3>(n[1]), m[1]*shuf<2>(n[1])) +
mad(m[2], shuf<1>(n[1]), m[3]*shuf<0>(n[1])));
_mm_stream_ps(&r[2].m128_f32[0],
mad(m[0], shuf<3>(n[2]), m[1]*shuf<2>(n[2])) +
mad(m[2], shuf<1>(n[2]), m[3]*shuf<0>(n[2])));
_mm_stream_ps(&r[3].m128_f32[0],
mad(m[0], shuf<3>(n[3]), m[1]*shuf<2>(n[3])) +
mad(m[2], shuf<1>(n[3]), m[3]*shuf<0>(n[3])));
}
По тактам тут ничего не поменялось, от слова совсем. Но, согласно рекомендациям, кэш мы теперь лишний раз не трогаем.
AVX реализация
Базовый AVX вариант
Далее перейдём к следующему этапу оптимизации. В SSE регистр входит 4-е float, а в AVX уже 8. То есть, есть теоретический шанс уменьшить число выполняемых операций, и увеличить производительность если не вдвое, то хотя бы раза в 1.5. Но что-то мне подсказывает, что не всё будет так просто с переходом на AVX. Сможем ли мы получать нужные данные из сдвоенных регистров?
Попробуем разобраться. Снова выпишем наш алгоритм умножения, использованный выше. Можно было бы этого не делать, но с кодом удобнее разбираться, когда всё рядом, и не приходится листать на полстраницы вверх.
//Снова представление матрицы в регистрах:
00, 10, 20, 30,
01, 11, 21, 31,
02, 12, 22, 32,
03, 13, 23, 33
//И алгоритм для SSE:
r0 = m0*n00 + m1*n10 + m2*n20 + m3*n30
r1 = m0*n01 + m1*n11 + m2*n21 + m3*n31
r2 = m0*n02 + m1*n12 + m2*n22 + m3*n32
r3 = m0*n03 + m1*n13 + m2*n23 + m3*n33
На выходе мы ожидаем получить результат в ymm = {r0:r1} и ymm = {r2:r3}. Если в SSE варианте у нас алгоритм обобщался на столбцы, то теперь нам надо его обобщить на строки. Так что действовать как в случае с SSE вариантом уже не выйдет.
Если считать матрицу m в регистры ymm, мы получим ymm = {m0:m1} и ymm = {m2:m3} соответственно. Раньше у нас были только столбцы матрицы в регистре, а теперь и столбцы и строки.
Если попробовать действовать как раньше, то надо ymm={m0:m1} умножать на регистр ymm={n00,n00,n00,n00}:{n10,n10,n10,n10}. Так как n00 и n01 в одной строке матрицы n, то, судя по имеющемуся набору AVX инструкций, раскидывать их по ymm будет дорого. И shuffle, и permute работают отдельно для каждой из двух четверок float (старший и младший xmm) внутри ymm регистров.
Если считать ymm из матрицы n, то получим оба элемента n00 и n10 в старшем из 2-х xmm внутри ymm регистра. {n00,n10,n20,n30}:{n01,n11,n21,n31}. Обычно индекс у имеющихся инструкций от 0 до 3. И адресует float лишь внутри одного xmm регистра из двух внутри ymm регистра. Перекинуть n10 из старшего xmm в младший дёшево не получится. А тут ещё этот фокус надо повторить несколько раз. С такой потерей тактов мы смириться не можем. Надо придумать что-то другое.
Раньше мы обобщали столбцы, а теперь строки. Поэтому попробуем зайти немного с другой стороны. Нам нужно получить результат в {r0:r1}. Значит и улучшать алгоритм надо не по отдельным строчкам алгоритма, а сразу по две. И тут то, что было минусом в работе shuffle и permute, станет для нас плюсом. Смотрим, что у нас будет в регистрах ymm, когда мы считаем матрицу n.
n0n1 = {00, 10, 20, 30} : {01, 11, 21, 31}
n2n3 = {02, 12, 22, 32} : {03, 13, 23, 33}
Ага, замечаем, что в разных xmm частях ymm регистра у нас есть элементы 00 и 01. Их можно размножить по регистру через команду permute в {_00,_00,_00,_00}:{_01,_01,_01,_01}, указав лишь один индекс 3 для обоих xmm частей. Это именно то, что нам надо. Ведь коэффициенты тоже используются в разных строчках. Только теперь в соответствующем ymm регистре для умножения нужно будет держать {m0:m0}, то есть дублированную первую строку матрицы m.
Итак, расписываем алгоритм более подробно. Считываем в ymm регистры сдвоенные строки матрицы m:
mm[0] = {m0:m0}
mm[1] = {m1:m1}
mm[2] = {m2:m2}
mm[3] = {m3:m3}
И тогда вычислять умножение будем как:
r0r1 =
mm[0] * {n00,n00,n00,n00:n01,n01,n01,n01} + // permute<3,3,3,3>(n)
mm[1] * {n10,n10,n10,n10:n11,n11,n11,n11} + // permute<2,2,2,2>(n)
mm[2] * {n20,n20,n20,n20:n21,n21,n21,n21} + // permute<1,1,1,1>(n)
mm[3] * {n30,n30,n30,n30:n31,n31,n31,n31} // permute<0,0,0,0>(n)
r2r3 =
mm[0] * {n02,n02,n02,n02:n03,n03,n03,n03} + // permute<3,3,3,3>(n)
mm[1] * {n12,n12,n12,n12:n13,n13,n13,n13} + // permute<2,2,2,2>(n)
mm[2] * {n22,n22,n22,n22:n23,n23,n23,n23} + // permute<1,1,1,1>(n)
mm[3] * {n32,n32,n32,n32:n33,n33,n33,n33} // permute<0,0,0,0>(n)
Перепишем более наглядно:
r0r1 = mm[0]*n0n1<3,3,3,3>+mm[1]*n0n1<2,2,2,2>+mm[2]*n0n1<1,1,1,1>+mm[3]*n0n1<0,0,0,0>
r2r3 = mm[0]*n2n3<3,3,3,3>+mm[1]*n2n3<2,2,2,2>+mm[2]*n2n3<1,1,1,1>+mm[3]*n2n3<0,0,0,0>
Или в упрощённом виде:
r0r1 = mm[0]*n0n1<3> + mm[1]*n0n1<2> + mm[2]*n0n1<1> + mm[3]*n0n1<0>
r2r3 = mm[0]*n2n3<3> + mm[1]*n2n3<2> + mm[2]*n2n3<1> + mm[3]*n2n3<0>
Вроде всё ясно.
void mul_mtx4_mtx4_avx_v1(__m128* const r, __m128 const* const m, __m128 const* const n) {
__m256 mm0 = _mm256_set_m128(m[0], m[0]);
__m256 mm1 = _mm256_set_m128(m[1], m[1]);
__m256 mm2 = _mm256_set_m128(m[2], m[2]);
__m256 mm3 = _mm256_set_m128(m[3], m[3]);
__m256 n0n1 = _mm256_load_ps(&n[0].m128_f32[0]);
__m256 y1 = _mm256_permute_ps(n0n1, 0xFF);//3,3,3,3
__m256 y2 = _mm256_permute_ps(n0n1, 0xAA);//2,2,2,2
__m256 y3 = _mm256_permute_ps(n0n1, 0x55);//1,1,1,1
__m256 y4 = _mm256_permute_ps(n0n1, 0x00);//0,0,0,0
y1 = _mm256_mul_ps(y1, mm0);
y2 = _mm256_mul_ps(y2, mm1);
y3 = _mm256_mul_ps(y3, mm2);
y4 = _mm256_mul_ps(y4, mm3);
y1 = _mm256_add_ps(y1, y2);
y3 = _mm256_add_ps(y3, y4);
y1 = _mm256_add_ps(y1, y3);
__m256 n2n3 = _mm256_load_ps(&n[2].m128_f32[0]);
__m256 y5 = _mm256_permute_ps(n2n3, 0xFF);
__m256 y6 = _mm256_permute_ps(n2n3, 0xAA);
__m256 y7 = _mm256_permute_ps(n2n3, 0x55);
__m256 y8 = _mm256_permute_ps(n2n3, 0x00);
y5 = _mm256_mul_ps(y5, mm0);
y6 = _mm256_mul_ps(y6, mm1);
y7 = _mm256_mul_ps(y7, mm2);
y8 = _mm256_mul_ps(y8, mm3);
y5 = _mm256_add_ps(y5, y6);
y7 = _mm256_add_ps(y7, y8);
y5 = _mm256_add_ps(y5, y7);
_mm256_stream_ps(&r[0].m128_f32[0], y1);
_mm256_stream_ps(&r[2].m128_f32[0], y5);
}
Вот уже интересные цифры от IACA: x86 — 12.53, x64 — 12. Хотя, конечно, хотелось получше. Что-то упустили.
AVX оптимизация плюс «синтаксический сахар»
Похоже в коде выше AVX использовался не на полную мощь. Находим, что вместо установки двух одинаковых строк в регистр ymm можно задействовать broadcast, который умеет заполнять регистр ymm двумя одинаковыми значениями xmm. Также попутно добавим немного «синтаксического сахара» для AVX функций.
__m256 operator + (__m256 const a, __m256 const b) { return _mm256_add_ps(a, b); }
__m256 operator - (__m256 const a, __m256 const b) { return _mm256_sub_ps(a, b); }
__m256 operator * (__m256 const a, __m256 const b) { return _mm256_mul_ps(a, b); }
__m256 operator / (__m256 const a, __m256 const b) { return _mm256_div_ps(a, b); }
template <int i> __m256 perm(__m256 const v)
{ return _mm256_permute_ps(v, _MM_SHUFFLE(i, i, i, i)); }
template <int a, int b, int c, int d> __m256 perm(__m256 const v)
{ return _mm256_permute_ps(v, _MM_SHUFFLE(a, b, c, d)); }
template <int i, int j> __m256 perm(__m256 const v)
{ return _mm256_permutevar_ps(v, _mm256_set_epi32(i, i, i, i, j, j, j, j)); }
template <int a, int b, int c, int d, int e, int f, int g, int h> __m256 perm(__m256 const v)
{ return _mm256_permutevar_ps(v, _mm256_set_epi32(a, b, c, d, e, f, g, h)); }
__m256 mad(__m256 const a, __m256 const b, __m256 const c) {
return _mm256_add_ps(_mm256_mul_ps(a, b), c);
}
void mul_mtx4_mtx4_avx_v2(__m128* const r, __m128 const* const m, __m128 const* const n) {
__m256 const mm[] {
_mm256_broadcast_ps(m+0),
_mm256_broadcast_ps(m+1),
_mm256_broadcast_ps(m+2),
_mm256_broadcast_ps(m+3)
};
__m256 const n0n1 = _mm256_load_ps(&n[0].m128_f32[0]);
_mm256_stream_ps(&r[0].m128_f32[0],
mad(perm<3>(n0n1), mm[0], perm<2>(n0n1)*mm[1])+
mad(perm<1>(n0n1), mm[2], perm<0>(n0n1)*mm[3]));
__m256 const n2n3 = _mm256_load_ps(&n[2].m128_f32[0]);
_mm256_stream_ps(&r[2].m128_f32[0],
mad(perm<3>(n2n3), mm[0], perm<2>(n2n3)*mm[1])+
mad(perm<1>(n2n3), mm[2], perm<0>(n2n3)*mm[3]));
}
А вот тут результаты уже интереснее. IACA выдаёт цифры: x86 — 10, x64 — 8.58, что выглядит существенно лучше, но всё же не в 2 раза.
AVX+FMA вариант (финальный)
Сделаем ещё одну попытку. Теперь было бы логичным снова вспомнить про набор инструкций FMA, поскольку он был добавлен в процессоры уже после AVX. Просто меняем отдельные mul+add на одну операцию. Хотя инструкцию умножения мы всё равно задействуем, чтобы дать компилятору больше возможностей для оптимизации, а процессору для параллельного выполнения умножений. Обычно я смотрю генерируемый код на ассемблере, чтобы убедиться какой вариант лучше.
В данном случае нам требуется вычислить a*b + c*d + e*f + g*h. Можно сделать это в лоб: fma(a, b, fma(c, d, fma(e, f, g * h))). Но, как мы видим, здесь нельзя выполнить одну операцию, не завершив предыдущую. А это значит, что мы не сможем задействовать возможность делать спаренные умножения, как это позволяет нам SIMD конвейер. Если мы немного преобразуем вычисления fma(a, b, c * d) + fma(e, f, g * h), то увидим, что можно распараллелить вычисления. Сначала сделать два независимых умножения, а потом две независимые fma операции.
__m256 fma(__m256 const a, __m256 const b, __m256 const c) {
return _mm256_fmadd_ps(a, b, c);
}
void mul_mtx4_mtx4_avx_fma(__m128* const r, __m128 const* const m, __m128 const* const n) {
__m256 const mm[]{
_mm256_broadcast_ps(m + 0),
_mm256_broadcast_ps(m + 1),
_mm256_broadcast_ps(m + 2),
_mm256_broadcast_ps(m + 3) };
__m256 const n0n1 = _mm256_load_ps(&n[0].m128_f32[0]);
_mm256_stream_ps(&r[0].m128_f32[0],
fma(perm<3>(n0n1), mm[0], perm<2>(n0n1)*mm[1])+
fma(perm<1>(n0n1), mm[2], perm<0>(n0n1)*mm[3]));
__m256 const n2n3 = _mm256_load_ps(&n[2].m128_f32[0]);
_mm256_stream_ps(&r[2].m128_f32[0],
fma(perm<3>(n2n3), mm[0], perm<2>(n2n3)*mm[1])+
fma(perm<1>(n2n3), mm[2], perm<0>(n2n3)*mm[3]));
}
IACA: x86 — 9.21, x64 — 8. Вот теперь совсем хорошо. Наверное, кто-то скажет, что можно сделать ещё лучше, но я уже не знаю, как.
BONUS из области фантастики
Собственно, из области фантастики он потому, что если я и видел процессоры с поддержкой AVX512, то разве что на картинках. Тем не менее, я попытался реализовать алгоритм. Тут я ничего пояснять не буду, полная аналогия с AVX+FMA. Алгоритм тот же, только операций меньше.
__m512 operator + (__m512 const a, __m512 const b) { return _mm512_add_ps(a, b); }
__m512 operator - (__m512 const a, __m512 const b) { return _mm512_sub_ps(a, b); }
__m512 operator * (__m512 const a, __m512 const b) { return _mm512_mul_ps(a, b); }
__m512 operator / (__m512 const a, __m512 const b) { return _mm512_div_ps(a, b); }
template <int i> __m512 perm(__m512 const v)
{ return _mm512_permute_ps(v, _MM_SHUFFLE(i, i, i, i)); }
template <int a, int b, int c, int d> __m512 perm(__m512 const v)
{ return _mm512_permute_ps(v, _MM_SHUFFLE(a, b, c, d)); }
__m512 fma(__m512 const a, __m512 const b, __m512 const c) {
return _mm512_fmadd_ps(a, b, c);
}
void mul_mtx4_mtx4_avx512(__m128* const r, __m128 const* const m, __m128 const* const _n) {
__m512 const mm[]{
_mm512_broadcast_f32x4(m[0]),
_mm512_broadcast_f32x4(m[1]),
_mm512_broadcast_f32x4(m[2]),
_mm512_broadcast_f32x4(m[3]) };
__m512 const n = _mm512_load_ps(&_n[0].m128_f32[0]);
_mm512_stream_ps(&r[0].m128_f32[0],
fma(perm<3>(n), mm[0], perm<2>(n)*mm[1])+
fma(perm<1>(n), mm[2], perm<0>(n)*mm[3]));
}
Цифры фантастические: x86 — 4.79, x64 — 5.42 (IACA с архитектурой SKX). Это притом, что в алгоритме 256 умножений и 192 сложения.
P.S. Код из статьи
Автор: truthfinder
