fastest matrix multiplication on x86

This challenge requires integration with C, so you can stop reading if you're not interested.

Matrix multiplication is a simple operation, but the performance depends a lot on how efficiently the code is written.

Let's compute m = x * y. The following code is the implementation of the textbook algorithm.

for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
float s = 0;
for (int k = 0; k < n; ++k) {
s += x[i * n + k] * y[k * n + j];
}
m[i * n + j] = s;
}
}


It looks fine, but actually performs quite bad for various reasons.

Simply reordering the operations to let only row access happen is a significant optimization.

//m is zeroed
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
float e = x[i * n + j];
for (int k = 0; k < n; ++k) {
m[i * n + k] += e * y[j * n + k];
}
}
}


This version is about 9.5 times faster than the upper one on my PC. Well, it's quite an optimal environment because in the actual implementation, I gave the compiler a lot of hints to make it auto-vectorize the second version up to the point where I couldn't gain any performance by manually vectorizing. But that's one of the problems of the first version. Column-accessing makes vectorization almost impossible to be done automatically and often not worthy to be done manually.

However, both versions are just the implementation of the most basic algorithm we all might know. I'd like to see what else could be done for a faster matrix multiplication on a restricted x86 environment.

The following restrictions are for simpler and clearer comparison between the algorithms.

• At most one thread can run during the program.
• The allowed x86 extensions are AVX2 and FMA.
• You must provide a C interface with the following prototype.
void (*)(void *m, void *x, void *y)


It does not have to be implemented in C, but you have to provide a way to be called from C.

m, x, and y point to 3 different arrays of 64 * 64 = 4096 floats that are aligned on a 32-byte boundary. These arrays represent a 64x64 matrix in row-major order.

You will be given randomly initialized matrices as input. When all outputs are done, first, the time taken will be measured, and then, a simple checksum will be calculated to check the correctness of the operation.

The table shows how much time the two example implementations shown earlier, each named mat_mul_slow and mat_mul_fast, takes to finish the test.

mat_mul_slow mat_mul_fast
~19.3s ~2s

The measurement will be done on my machine with an Intel Tiger Lake CPU, but I will not run code that seems to be slower than mat_mul_fast.

Built-ins are fine, but it should be open source so that I can dig in their code to see what they've done.

The attached code is the tester and the actual implementation of the examples, compiled by gcc -O3 -mavx2 -mfma.

#include <stdio.h>
#include <stdalign.h>
#include <string.h>
#include <time.h>
#include <math.h>
#include <x86intrin.h>

enum {MAT_N = 64};

typedef struct {
alignas(__m256i) float m[MAT_N * MAT_N];
} t_mat;

static void mat_rand(t_mat *_m, int n) {
unsigned r = 0x55555555;
for (int i = 0; i < n; ++i) {
float *m = _m[i].m;
for (int j = 0; j < MAT_N * MAT_N; ++j) {
r = r * 1103515245 + 12345;
m[j] = (float)(r >> 16 & 0x7fff) / 0x8000p0f;
}
}
}

static unsigned mat_checksum(t_mat *_m, int n) {
unsigned s = 0;
for (int i = 0; i < n; ++i) {
float *m = _m[i].m;
for (int j = 0; j < MAT_N * MAT_N; ++j) {
union {float f; unsigned u;} e = {m[j]};
s += _rotl(e.u, j % 32);
}
}
return s;
}

__attribute__((noinline))
static void mat_mul_slow(void *_m, void *_x, void *_y) {
enum {n = MAT_N};
float *restrict m = __builtin_assume_aligned(((t_mat *)_m)->m, 32);
float *restrict x = __builtin_assume_aligned(((t_mat *)_x)->m, 32);
float *restrict y = __builtin_assume_aligned(((t_mat *)_y)->m, 32);
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
float s = 0.0f;
for (int k = 0; k < n; ++k) {
s = fmaf(x[i * n + k], y[k * n + j], s);
}
m[i * n + j] = s;
}
}
}

__attribute__((noinline))
static void mat_mul_fast(void *_m, void *_x, void *_y) {
enum {n = MAT_N};
float *restrict m = __builtin_assume_aligned(((t_mat *)_m)->m, 32);
float *restrict x = __builtin_assume_aligned(((t_mat *)_x)->m, 32);
float *restrict y = __builtin_assume_aligned(((t_mat *)_y)->m, 32);
memset(m, 0, sizeof(t_mat));
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
float e = x[i * n + j];
for (int k = 0; k < n; ++k) {
m[i * n + k] = fmaf(e, y[j * n + k], m[i * n + k]);
}
}
}
}

/*static void mat_print(t_mat *_m) {
int n = MAT_N;
float *m = _m->m;
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; j += 8) {
for (int k = 0; k < 8; ++k) {
printf("%f ", m[i * n + j + k]);
}
putchar('\n');
}
}
putchar('\n');
}*/

static void test(void (*f)(void *, void *, void *), t_mat *xy, t_mat *z,
int n) {
clock_t c = clock();
for (int i = 0; i < n; ++i) {
f(z + i, xy + i * 2, xy + i * 2 + 1);
}
c = clock() - c;
printf("%u %f\n", mat_checksum(z, n), (double)c / (double)CLOCKS_PER_SEC);
}

#define TEST(f) test(f, xy, z, n)

int main() {
enum {n = 0x20000};
static t_mat xy[n * 2];
static t_mat z[n];
mat_rand(xy, n * 2);
puts("start");
for (int i = 0; i < 2; ++i) {
TEST(mat_mul_slow);
TEST(mat_mul_fast);
}
return 0;
}


https://godbolt.org/z/hE9b9vhEx

• If you are just looking for speed of matrix-matrix multiplications, I suspect LAPACK and BLAS will beat most anything people could write having had 30+ years of hardcore optimisations. Commented May 30, 2022 at 21:45

C (GCC), 2.65× faster than mat_mul_fast

(Compiled with gcc -O3 -mavx2 -mfma, measured on my Intel Comet Lake i7-10710U.)

#include <x86intrin.h>

enum { MAT_N = 64 };

void mat_mul_simd(void *_m, void *_x, void *_y) {
enum { n = MAT_N };
float(*m)[n] = _m, (*x)[n] = _x, (*y)[n] = _y;
for (int i = 0; i < n; ++i) {
__m256 mi[n / 8];
for (int k = 0; k < n / 8; ++k)
mi[k] = _mm256_setzero_ps();
for (int j = 0; j < n; ++j) {
for (int k = 0; k < n / 8; ++k)
}
for (int k = 0; k < n / 8; ++k)
_mm256_store_ps(&m[i][k * 8], mi[k]);
}
}


https://godbolt.org/z/xxvr45b67

C (GCC 11)

This completely auto-vectorized code is very nearly as fast as the above in GCC 11 (gcc -O3 -mavx2 -mfma), but regresses significantly in GCC 12. Auto-vectorization is fragile. 😞

enum { MAT_N = 64 };

void mat_mul_auto(void *_m, void *_x, void *_y) {
enum { n = MAT_N };
float(*restrict m)[n] = __builtin_assume_aligned(_m, 32);
float(*restrict x)[n] = __builtin_assume_aligned(_x, 32);
float(*restrict y)[n] = __builtin_assume_aligned(_y, 32);
for (int i = 0; i < n; ++i) {
float mi[n] = {};
for (int j = 0; j < n; ++j)
for (int k = 0; k < n; ++k)
mi[k] += x[i][j] * y[j][k];
for (int k = 0; k < n; ++k)
m[i][k] = mi[k];
}
}


https://godbolt.org/z/h4GEqGjGW

• Interesting.. If I didn't miss something, this is the exact same algorithm while the difference is that the stores are deferred to the end of the inner loop. Maybe out-of-order execution makes the difference a bit less, but still significant. Commented May 31, 2022 at 3:24
• I get 2.05s vs 0.79s, so very similar to your measurement. Commented May 31, 2022 at 3:25
• Oh, and I also my fast version does memory load for the zeroed vector. That's another problem. Commented May 31, 2022 at 3:30
• Do you get similar results with clang and icc?
– user108721
Commented May 31, 2022 at 7:45