Matrix Multiplications
There are not that many problems in computer science I care about, but matrix multiplication is one that's been seriously bothering me for a while now. So I did some experiments with my Apple silicon computer, in the hopes of gaining some new insights.
First, I did three tests: a naïve approach with C++20+ features, another approach with my attempt at ARM hardware intrinsics, and a third one with Strassen's algorithm, all single-threaded.
Sample:
void base_multiply_vec(StridedView A, StridedView B, StridedView C, bool accumulate = true) {
size_t n = A.extent(0);
for (size_t i = 0; i < n; ++i) {
for (size_t k = 0; k < n; ++k) {
float64x2_t a_vec = vdupq_n_f64(A[i, k]);
const double* b_ptr = &B[k, 0];
double* c_ptr = &C[i, 0];
for (size_t j = 0; j < n; j += 2) {
float64x2_t b_vec = vld1q_f64(b_ptr + j);
if (accumulate) {
float64x2_t c_vec = vld1q_f64(c_ptr + j);
c_vec = vfmaq_f64(c_vec, a_vec, b_vec);
vst1q_f64(c_ptr + j, c_vec);
} else {
if (k == 0) {
float64x2_t c_vec = vmulq_f64(a_vec, b_vec); // c = a * b
vst1q_f64(c_ptr + j, c_vec);
} else {
float64x2_t c_vec = vld1q_f64(c_ptr + j);
c_vec = vfmaq_f64(c_vec, a_vec, b_vec);
vst1q_f64(c_ptr + j, c_vec);
}
}
}
}
}
}
A single-threaded benchmark dump:
Benchmarking Matrix Dimension: 256 x 256
Standard O(N^3) Loop Time: 6.99287 ms
C++23 Strassen Time: 5.8185 ms
Validation Max Divergence: 2.41585e-13
----------------------------------------
Benchmarking Matrix Dimension: 512 x 512
Standard O(N^3) Loop Time: 30.4854 ms
C++23 Strassen Time: 17.4701 ms
Validation Max Divergence: 8.81073e-13
----------------------------------------
Benchmarking Matrix Dimension: 1024 x 1024
Standard O(N^3) Loop Time: 139.558 ms
C++23 Strassen Time: 102.393 ms
Validation Max Divergence: 3.80851e-12
Then, I added in some multithreading using OpenMP, and tested both the intrinsics version and the Strassen's version with both block-contiguous tiling and without.
Without:
Benchmarking Matrix Dimension: 256 x 256 (INTEGER VALIDATION)
Standard O(N^3) Loop Time: 6.24008 ms
Parallel Strassen Time: 1.90571 ms
Validation Max Divergence: 0
----------------------------------------
Benchmarking Matrix Dimension: 512 x 512 (INTEGER VALIDATION)
Standard O(N^3) Loop Time: 24.523 ms
Parallel Strassen Time: 4.37921 ms
Validation Max Divergence: 0
----------------------------------------
Benchmarking Matrix Dimension: 1024 x 1024 (INTEGER VALIDATION)
Standard O(N^3) Loop Time: 128.556 ms
Parallel Strassen Time: 30.2994 ms
Validation Max Divergence: 0
With tiling:
Benchmarking Matrix Dimension: 256 x 256 (BLOCK CONCONTIGUOUS TILED - INTEGER VALIDATION)
Standard O(N^3) Loop Time: 6.67008 ms
Parallel Tiled Strassen: 7.17817 ms
Validation Max Divergence: 0
----------------------------------------
Benchmarking Matrix Dimension: 512 x 512 (BLOCK CONCONTIGUOUS TILED - INTEGER VALIDATION)
Standard O(N^3) Loop Time: 23.609 ms
Parallel Tiled Strassen: 4.00054 ms
Validation Max Divergence: 0
----------------------------------------
Benchmarking Matrix Dimension: 1024 x 1024 (BLOCK CONCONTIGUOUS TILED - INTEGER VALIDATION)
Standard O(N^3) Loop Time: 126.672 ms
Parallel Tiled Strassen: 29.4473 ms
Validation Max Divergence: 0
Usually, Strassen's introduces some drift between it and the naïve version because of floating point math. Some of the benchmarks here show integer validation. That's because of debugging, and I'm too lazy to print everything out.
Anyway next, I tested tiled GEMM, just to compare it to my relatively successful Strassen attempt.
Sample:
void gemm_rectangular_tiled_safe(
std::span A, std::span B, std::span C,
size_t M, size_t K, size_t N)
{
size_t tiles_M = M / THRESHOLD;
size_t tiles_N = N / THRESHOLD;
size_t tiles_K = K / THRESHOLD;
size_t tile_size = THRESHOLD * THRESHOLD;
size_t total_tasks = tiles_M * tiles_N;
std::atomic next_task{0};
unsigned int num_threads = std::thread::hardware_concurrency();
if (num_threads == 0) num_threads = 4;
std::vector workers;
workers.reserve(num_threads);
for (unsigned int t = 0; t < num_threads; ++t) {
workers.emplace_back([&]() {
size_t task_id;
while ((task_id = next_task.fetch_add(1, std::memory_order_relaxed)) < total_tasks)
{
size_t mi = task_id / tiles_N;
size_t ni = task_id % tiles_N;
size_t c_offset = (mi * tiles_N + ni) * tile_size;
std::span c_tile = C.subspan(c_offset, tile_size);
for (size_t ki = 0; ki < tiles_K; ++ki) {
size_t a_offset = (mi * tiles_K + ki) * tile_size;
size_t b_offset = (ki * tiles_N + ni) * tile_size;
std::span a_tile = A.subspan(a_offset, tile_size);
std::span b_tile = B.subspan(b_offset, tile_size);
base_multiply_tiled_safe(a_tile, b_tile, c_tile, (ki > 0));
}
}
});
}
}
Here's a random Strassen's sample for good measure:
#pragma omp parallel sections
{
// M1 = (A11 + A22) * (B11 + B22)
#pragma omp section
{
double* t_scratch_ptr = base_next_scratch + (branch_stride * 0);
double* t1_ptr = t_scratch_ptr + (sub_size * 7);
double* t2_ptr = t_scratch_ptr + (sub_size * 14);
double* child_scratch = t_scratch_ptr + (sub_size * 21);
add_matrices_tiled(a11, a22, t1_ptr, sub_size);
add_matrices_tiled(b11, b22, t2_ptr, sub_size);
strassen_serial_tiled(t1_ptr, t2_ptr, m1_ptr, n / 2, child_scratch, false);
}
// M2 = (A21 + A22) * B11
#pragma omp section
{
double* t_scratch_ptr = base_next_scratch + (branch_stride * 1);
double* t1_ptr = t_scratch_ptr + (sub_size * 7);
double* child_scratch = t_scratch_ptr + (sub_size * 21);
add_matrices_tiled(a21, a22, t1_ptr, sub_size);
strassen_serial_tiled(t1_ptr, b11, m2_ptr, n / 2, child_scratch, false);
}
// M3 = A11 * (B12 - B22)
#pragma omp section
{
double* t_scratch_ptr = base_next_scratch + (branch_stride * 2);
double* t2_ptr = t_scratch_ptr + (sub_size * 7);
double* child_scratch = t_scratch_ptr + (sub_size * 21);
sub_matrices_tiled(b12, b22, t2_ptr, sub_size);
strassen_serial_tiled(a11, t2_ptr, m3_ptr, n / 2, child_scratch, false);
}
// M4 = A22 * (B21 - B11)
#pragma omp section
{
double* t_scratch_ptr = base_next_scratch + (branch_stride * 3);
double* t2_ptr = t_scratch_ptr + (sub_size * 7);
double* child_scratch = t_scratch_ptr + (sub_size * 21);
sub_matrices_tiled(b21, b11, t2_ptr, sub_size);
strassen_serial_tiled(a22, t2_ptr, m4_ptr, n / 2, child_scratch, false);
}
// M5 = (A11 + A12) * B22
#pragma omp section
{
double* t_scratch_ptr = base_next_scratch + (branch_stride * 4);
double* t1_ptr = t_scratch_ptr + (sub_size * 7);
double* child_scratch = t_scratch_ptr + (sub_size * 21);
add_matrices_tiled(a11, a12, t1_ptr, sub_size);
strassen_serial_tiled(t1_ptr, b22, m5_ptr, n / 2, child_scratch, false);
}
// M6 = (A21 - A11) * (B11 + B12)
#pragma omp section
{
double* t_scratch_ptr = base_next_scratch + (branch_stride * 5);
double* t1_ptr = t_scratch_ptr + (sub_size * 7);
double* t2_ptr = t_scratch_ptr + (sub_size * 14);
double* child_scratch = t_scratch_ptr + (sub_size * 21);
sub_matrices_tiled(a21, a11, t1_ptr, sub_size);
add_matrices_tiled(b11, b12, t2_ptr, sub_size);
strassen_serial_tiled(t1_ptr, t2_ptr, m6_ptr, n / 2, child_scratch, false);
}
// M7 = (A12 - A22) * (B21 + B22)
#pragma omp section
{
double* t_scratch_ptr = base_next_scratch + (branch_stride * 6);
double* t1_ptr = t_scratch_ptr + (sub_size * 7);
double* t2_ptr = t_scratch_ptr + (sub_size * 14);
double* child_scratch = t_scratch_ptr + (sub_size * 21);
sub_matrices_tiled(a12, a22, t1_ptr, sub_size);
add_matrices_tiled(b21, b22, t2_ptr, sub_size);
strassen_serial_tiled(t1_ptr, t2_ptr, m7_ptr, n / 2, child_scratch, false);
}
}
Just to get you in the mood. Here's some more benchmark dumps:
Tiled C++23 Engine: A(256x256) * B(256x256)
Standard O(N^3) Rectangular Loop Time: 6.99425 ms
Parallel Tiled C++23 GEMM Time: 2.09221 ms
Validation Max Divergence: 0
----------------------------------------
Tiled C++23 Engine: A(512x512) * B(512x512)
Standard O(N^3) Rectangular Loop Time: 29.3373 ms
Parallel Tiled C++23 GEMM Time: 3.3995 ms
Validation Max Divergence: 0
----------------------------------------
Tiled C++23 Engine: A(1024x1024) * B(1024x1024)
Standard O(N^3) Rectangular Loop Time: 130.619 ms
Parallel Tiled C++23 GEMM Time: 20.7029 ms
Validation Max Divergence: 0
----------------------------------------
Tiled C++23 Engine: A(640x384) * B(384x896)
Standard O(N^3) Rectangular Loop Time: 27.2447 ms
Parallel Tiled C++23 GEMM Time: 4.42667 ms
Validation Max Divergence: 0
----------------------------------------
Tiled C++23 Engine: A(2176x2176) * B(2176x3840)
Standard O(N^3) Rectangular Loop Time: 2225.44 ms
Parallel Tiled C++23 GEMM Time: 332.537 ms
Validation Max Divergence: 0Finally, I also masochistically wanted to take a big dump on my efforts and run Apple's own hardware-accelerated matrix multiplication (included in the Accelerate header). Benchmark follows:
Apple Accelerate Engine: A(4096x4096) * B(4096x4096)
Apple Accelerate cblas_dgemm Time: 343.114 ms
Validation Max Divergence: 0
So, what we used here in that last one is a highly secretive, specialized piece of hardware integrated directly inside the CPU cluster: the Apple Matrix Coprocessor (AMX) alongside ARM's Scalable Matrix Extension (SME). Not the GPU.
The experiment wouldn't be complete without taking the Metal framework for a spin:
Apple MLX GPU Engine: A(4096x4096) * B(4096x4096)
Apple MLX GPU mx::matmul Time: 44.4126 ms
Validation Matrix Shape: 4096x4096
A marked speed improvement.
Since I know this is just my writing and you're going to have to take my word for everything unless you test this out yourself, I might as well just tell you what I found, and it was rather interesting, actually. Some findings were unintuitive to me.
Without further ado: first, on my Mac at least, writing regular old naïve C++23 actually resulted in more performant code than writing out the intrinsics manually! I'm not sure if this is due to my crappy understanding of intrinsics or whether the compiler's just that smart. But regardless, this was actually a positive surprise, because it means that there's no particular reason to steer away from platform-independent code. At least in my case here.
Second, Strassen's algorithm is just horrible. It's genuinely difficult to get right, and it loses to tiled GEMM, which is infinitely simpler and more maintainable. I suppose Strassen's might actually catch up to GEMM with huge inputs, but that would require matrix sizes of 10,000^2 at least. I haven't exactly encountered many cases where those sizes are necessary.
And third, OpenMP is a bad way to parallelize things. It seems like a good idea at first, but C++23 comes with really nice thread abstractions that don't require all that much boilerplate anymore. The thing that you get for free with OpenMP is the thread scheduler, but if you manage to avoid it, you get better performance.
Finally, I learned that what you should seriously do in any real-life calculations is cheat: Apple's hardware acceleration is both relatively easy to use and practically blows everything else out of the water. I know NVIDIA has similar stuff for their GPUs and I think Intel might have some too, like Advanced Matrix Extensions (AMX) and Xe Matrix Extensions (XMX). Yes, it seems Intel and Apple have technologies with the same abbreviations.
Of course, the GPU is still faster than Apple's AMX. However, AMX has one good thing going for it, and that's double precision. GPU's traditionally use 32-bit floats all around, although I suppose you can come up with some trickery as well to get around that limitation.
So, all in all, what I took home from this little experiment is that the dumbest code is the best code: plain C++ with either tiled scans in O(n^3) or hardware-accelerated shenanigans beats the theoretical time complexity of O(n^{log_{2} 7}) on relatively recent hardware using "commodity" matrix sizes.
Hardware creeps me out, because who knows what horrible things you could do with it.