5
\$\begingroup\$

I would appreciate a review of the following Rust implementation of high performance matrix multiplication.

After reviewing available literature, including Anatomy of High Performance Matrix Multiplication paper and How to Optimize GEMM, I came up with an implementation that performs very poorly compared to publically available libraries. I would like to see if anyone can identify issues and/or opportunities for improvement.

My benchmark is against numpy, multiplying two 1024x1024 matcies. On my laptop, numpy does 220+ gflops (measured by 1024^3/time).

My implementation tiles the input matrices and then does AVX2 SIMD multiplication on 8x8 blocks. Tile size is 128x128. I backed this size out of the L1d cache size, 192K. 128x128x4 = 64K per tile. Larger tiles would exceed the AVX2 cache size.

I pack A, B and C, into 64-byte aligned arrays so they load cleanly into cache lines and so that I can use the aligned-load/store AVX2 instructions.

The best I've gotten from this code is 40 gflops, which is not even 20% of the performance that numpy has. I'm clearly leaving a lot of performance on the table but I cannot identify where.

People familiar with high performance matrix multiplication, I would very much appreciate your eyeballs and opinion on the following code. Executable project is available on github.

#![allow(clippy::too_many_arguments)]

const S: usize = 128; // process 64x64 tiles

#[repr(align(64))]
struct PackedA([f32; 8 * 8]);

#[repr(align(64))]
struct PackedB([f32; S * S]);

#[repr(align(64))]
struct PackedC([f32; S * S]);

pub fn matmul(a: &[f32], b: &[f32], c: &mut [f32], m: usize, p: usize, n: usize) {
    gemm(c, a, b, n, p, n, m, p, n);
}

pub fn gemm(
    c: &mut [f32],
    a: &[f32],
    b: &[f32],
    ldc: usize,
    lda: usize,
    ldb: usize,
    m: usize,
    p: usize,
    n: usize,
) {
    for k in (0..p).step_by(S) {
        for i in (0..n).step_by(S) {
            let ib = (k * ldb) + i;
            let mut bp = PackedB([0.0; S * S]);
            pack(&mut bp.0, &b[ib..], S, ldb); // pack b
            for j in (0..m).step_by(S) {
                let ia = (j * lda) + k;
                let ic = (j * ldc) + i;
                let mut cp = PackedC([0.0; S * S]);
                gemm_block(&mut cp.0, &a[ia..], &bp.0, S, lda, S, S, S, S); // use packed b
                for j2 in 0..S {
                    for i2 in 0..S {
                        c[ic + (j2 * ldc) + i2] += cp.0[(j2 * S) + i2]; // unpack c
                    }
                }
            }
        }
    }
}

pub fn gemm_block(
    c: &mut [f32],
    a: &[f32],
    b: &[f32],
    ldc: usize,
    lda: usize,
    ldb: usize,
    m: usize,
    p: usize,
    n: usize,
) {
    for j in (0..m).step_by(8) {
        for k in (0..p).step_by(8) {
            let ia = (j * lda) + k;
            let mut ap = PackedA([0.0; 8 * 8]);
            pack(&mut ap.0, &a[ia..], 8, lda); // pack a
            for i in (0..n).step_by(8) {
                let ib = (k * ldb) + i;
                let ic = (j * ldc) + i;
                mm8_simd(&mut c[ic..], &ap.0, &b[ib..], ldc, 8, ldb); // use packed a
            }
        }
    }
}

// ------------------------------------------------------------------------------------------------

pub fn pack(r: &mut [f32], x: &[f32], ldr: usize, ldx: usize) {
    for j in 0..ldr {
        for i in 0..ldr {
            r[(j * ldr) + i] = x[(j * ldx) + i];
        }
    }
}

// ------------------------------------------------------------------------------------------------

#[allow(clippy::erasing_op)]
#[allow(clippy::identity_op)]
#[allow(clippy::missing_safety_doc)]
pub fn mm8_simd(c: &mut [f32], a: &[f32], b: &[f32], ldc: usize, lda: usize, ldb: usize) {
    unsafe {
        use std::arch::x86_64::*;
        let mut cv0 = _mm256_load_ps(c.get_unchecked_mut(0 * ldc));
        let mut cv1 = _mm256_load_ps(c.get_unchecked_mut(1 * ldc));
        let mut cv2 = _mm256_load_ps(c.get_unchecked_mut(2 * ldc));
        let mut cv3 = _mm256_load_ps(c.get_unchecked_mut(3 * ldc));
        let mut cv4 = _mm256_load_ps(c.get_unchecked_mut(4 * ldc));
        let mut cv5 = _mm256_load_ps(c.get_unchecked_mut(5 * ldc));
        let mut cv6 = _mm256_load_ps(c.get_unchecked_mut(6 * ldc));
        let mut cv7 = _mm256_load_ps(c.get_unchecked_mut(7 * ldc));
        for k in 0..8 {
            let bv = _mm256_load_ps(b.get_unchecked(k * ldb));
            cv0 = _mm256_fmadd_ps(_mm256_broadcast_ss(a.get_unchecked(0 * lda + k)), bv, cv0);
            cv1 = _mm256_fmadd_ps(_mm256_broadcast_ss(a.get_unchecked(1 * lda + k)), bv, cv1);
            cv2 = _mm256_fmadd_ps(_mm256_broadcast_ss(a.get_unchecked(2 * lda + k)), bv, cv2);
            cv3 = _mm256_fmadd_ps(_mm256_broadcast_ss(a.get_unchecked(3 * lda + k)), bv, cv3);
            cv4 = _mm256_fmadd_ps(_mm256_broadcast_ss(a.get_unchecked(4 * lda + k)), bv, cv4);
            cv5 = _mm256_fmadd_ps(_mm256_broadcast_ss(a.get_unchecked(5 * lda + k)), bv, cv5);
            cv6 = _mm256_fmadd_ps(_mm256_broadcast_ss(a.get_unchecked(6 * lda + k)), bv, cv6);
            cv7 = _mm256_fmadd_ps(_mm256_broadcast_ss(a.get_unchecked(7 * lda + k)), bv, cv7);
        }
        _mm256_store_ps(c.get_unchecked_mut(0 * ldc), cv0);
        _mm256_store_ps(c.get_unchecked_mut(1 * ldc), cv1);
        _mm256_store_ps(c.get_unchecked_mut(2 * ldc), cv2);
        _mm256_store_ps(c.get_unchecked_mut(3 * ldc), cv3);
        _mm256_store_ps(c.get_unchecked_mut(4 * ldc), cv4);
        _mm256_store_ps(c.get_unchecked_mut(5 * ldc), cv5);
        _mm256_store_ps(c.get_unchecked_mut(6 * ldc), cv6);
        _mm256_store_ps(c.get_unchecked_mut(7 * ldc), cv7);
    }
}
\$\endgroup\$
2
  • 1
    \$\begingroup\$ I really like your narrative analysis, for example the L1 cache sizing. I feel it is helpful. But what's missing is observations of performance counters. You have a hypothesis about cache behavior, cool. But now we need to gather evidence about it, right? Plus maybe other things, like fraction of bus cycles that were idle. We want to contrast system behavior under this code versus behavior under the numpy code. (And BTW, in fairness, a great many people have scrutinized the numpy BLAS performance, over many more years than you have yet devoted to this subroutine.) \$\endgroup\$
    – J_H
    Commented Jan 24 at 5:50
  • 2
    \$\begingroup\$ This question is about low level performance details, where the generated object code really matters. Could we have a godbolt link, please? \$\endgroup\$
    – J_H
    Commented Jan 24 at 5:55

1 Answer 1

3
\$\begingroup\$

Without doing a lot of deep analysis, a couple of things stand out to me. In the previous question with very similar code, these things were not concerns, because for a 8x8 there was no alternative. There are alternatives now, so now those things are concerns.

I backed this size out of the L1d cache size, 192K.

You don't have 192K of L1d, unless you're counting multiple cores, but sizing the block for cache that's fragmented across multiple cores doesn't make sense since there's only one thread working on it. However, Goto recommended sizing the blocks to L2 anyway. But that changes something else too, which I'll mentioned later.

The inner-most loop only makes 8 iterations, so the bulk loads and stores to/from c happen relatively often. I don't have an exact recommendation for how many iterations that loop should make, but probably more than 8.

The basic kernel here is square in terms of scalars, but it's acting like an 8x8 by 8x1 multiplication if the SIMD is factored in. The thing with that is that every FMA is paired with a load, while if you had a wider matrix then you could reuse the result of a load multiple times. CPUs these days can perform a ton of loads per cycle, so it may seem like this shouldn't matter, and maybe it wouldn't if the data was all in L1, but loading that much data cannot be sustained from L2 (for example on Skylake, L2 can supply something like 30 bytes per cycle sustained, a decent amount but much less than 64). The trick to get around that is working with a slightly wider b matrix, so that in each iteration of the inner loop there are (for example) 2 loads from b, and 4 corresponding loads from a, but still 8 independent FMAs. Maybe 4x2 is not quite the right shape either, but that's the general idea: using w+h loads and doing w*h FMAs, such that w+l < w*h.

\$\endgroup\$
3
  • \$\begingroup\$ That's Goto's equation 2. The paper indicated it's best if the the kernel was square, so I tried 8x8, but doing more fmas with a wider B makes sense. I will try 8x16. \$\endgroup\$
    – Ana
    Commented Jan 24 at 15:10
  • \$\begingroup\$ After trying 8x16, I got to 44 gflops, still a long way from 220 via numpy. \$\endgroup\$
    – Ana
    Commented Jan 24 at 15:30
  • \$\begingroup\$ @Ana the part where a square kernel is recommended does so on the basis of minimizing loads, that minimum is obtained when the kernel is square with SIMD factored out, so it ends up being 8x wider than tall with AVX. But a 4x4 kernel (4x32 scalars) is too big to fit in the registers. \$\endgroup\$
    – user555045
    Commented Jan 24 at 15:39

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.