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;
}
