DEV Community: james hurley

ARM Matmul Continued: BFMMLA Edition

james hurley — Thu, 01 Oct 2026 03:08:20 +0000

This is a continuation of my previous article on using ARM NEON to do matrix multiplication.

Here's a fun one, we're going to use BFMMLA to do 2x4 @ 4x2 matmul on bf16 inputs. The result is a 2x2 matrix in F32 accumulators in a single float register. We'll break this one up a little more than the last post because it's a little more involved.

#if defined(__APPLE__)
#define MALLOC _malloc
#define FREE   _free
#else
#define MALLOC malloc
#define FREE   free
#endif

.global _mmul_8x8_bf16_mma
.extern MALLOC
.extern FREE

_mmul_8x8_bf16_mma:
    // preserve link register because we'll be calling out to malloc/free
    stp x29, x30, [sp, #-16]!

    // free up some registers
    stp x19, x20, [sp, #-80]!
    stp x21, x22, [sp, #16]
    stp x23, x24, [sp, #32]
    stp x25, x26, [sp, #48]
    stp x27, x28, [sp, #64]

    // free up fp registers since we will use all 32
    stp d8, d9, [sp, #-64]!
    stp d10, d11, [sp, #16]
    stp d12, d13, [sp, #32]
    stp d14, d15, [sp, #48]

    mov x19, x0
    mov x20, x1
    mov x21, x2
    mov x22, x3
    mov x23, x4
    mov x24, x5
    mov x25, x6
    mov x26, x7

    mul w0, w6, w1 // BN * K
    lsl w0, w0, #1 // * 2 (f16)
    uxtw x0, w0
    bl  MALLOC     // alloc scratch memory
    cmp x0, #0
    mov x27, x0
    mov x0, #1
    beq alloc_fail

    mul w0, w26, w20 // BM * K
    lsl w0, w0, #1
    uxtw x0, w0
    bl MALLOC
    cmp x0, #0
    mov x28, x0
    mov x0, #1
    beq alloc_fail

    mov x0, x19
    mov x1, x20
    mov x2, x21
    mov x3, x22
    mov x4, x23
    mov x5, x24
    mov x6, x25
    mov x7, x26

    mov x25, x27 // B scratch pointer in x25
    mov x26, x28 // A scratch pointer in x26

    // q16..q31 -> outputs (16x8=128 / 8 half-precision = 16 registers)
    // q0..q7 -> matA inputs
    // q8..q15 -> matB inputs

    // x0 -> M
    // x1 -> K
    // x2 -> N
    // x3 -> pA
    // x4 -> pB
    // x5 -> pC
    // x6 -> BN
    // x7 -> BM

    // x20..x24 -> counters
    // x25 -> B scratch
    // x26 -> A scratch

mostly housekeeping and prep here. We are allocating two scratch slabs for inputs from the A and B matrices. The reason we are doing this is because we need to swizzle the inputs to be in the right shape for a 2x4 or 4x2 matrix such that we get a sensible output result. We will swizzle as we load data rather than as we perform compute so that we don't have to do it multiple times.

mov x20, #0 // bm
    .M_block: // outer loop blocking M
        cmp w20, M
        bge .M_block_x
        mov x15, x26 // A scratch pointer we can advance
        mov w21, #0 // bm_l
        .M_block_load:
            cmp w21, w7
            bge .M_block_load_x
            mov w22, #0 // bm_k_l
            .M_block_load_k:
                cmp w22, K
                bge .M_block_load_k_x
                // Why are we loading 8 M rows at a time? Because we load 8 N columns at a time due to vector width
                // Why are we reordering the K blocks? Loading 8 K rows in the N block load phase may cause register
                // exhaustion (but mostly I haven't tried it yet).. we make up for this in the compute phase

                // offset = (bm + bm_l) * K + bm_k_l
                add w13, w20, w21
                madd w13, w13, K, w22
                uxtw x13, w13
                add x13, x3, x13, lsl #1 // matA pointer + offset (bytes)
                ldr q0, [x13]
                add x13, x13, x1, lsl #1 // move matA pointer by stride K
                ldr q1, [x13]
                add x13, x13, x1, lsl #1
                ldr q2, [x13]
                add x13, x13, x1, lsl #1
                ldr q3, [x13]
                add x13, x13, x1, lsl #1
                ldr q4, [x13]
                add x13, x13, x1, lsl #1
                ldr q5, [x13]
                add x13, x13, x1, lsl #1
                ldr q6, [x13]
                add x13, x13, x1, lsl #1
                ldr q7, [x13]

                zip1 v8.2d, v0.2d, v1.2d // [m0k0, m0k1, m0k2, ..., m1k3]
                zip2 v9.2d, v0.2d, v1.2d // [m0k4, m0k5, m0k6, ..., m1k7]
                zip1 v10.2d, v2.2d, v3.2d // [m2k0, m2k1, m2k2, ... m3k3]
                zip2 v11.2d, v2.2d, v3.2d // [m2k4, ...]
                zip1 v12.2d, v4.2d, v5.2d // [m4k0, ...]
                zip2 v13.2d, v4.2d, v5.2d // [m4k4, ...]
                zip1 v14.2d, v6.2d, v7.2d // [m6k0, ..., m7k3]
                zip2 v15.2d, v6.2d, v7.2d // [m6k4, ..., m7k7]

                // reorder k blocks such that there are 4 k's for all 8 of the m's
                stp q8, q10, [x15], #32
                stp q12, q14, [x15], #32
                stp q9, q11, [x15], #32
                stp q13, q15, [x15], #32

                add w22, w22, #8 // loading 8 K at a time
                b .M_block_load_k
            .M_block_load_k_x:
            add w21, w21, #8 // loading 8 M rows at a time
            b .M_block_load
        .M_block_load_x:

Here we are loading data from the A matrix input and storing it into A scratch.. Here we want to reorder the input so that we have groups of 2 M rows and 4 K columns arranged contiguously. This is so that we can put these 8 values into a single 128-bit register representing a 2x4 matrix to use with bfmmla. Because of the load width we actually end up loading 8 values in the K dimension from mat A at once. So we put K values 4-7 after we have all 0-3 values because K will be our compute inner loop, so we'll iterate through all of K, then iterate through the N block, then we will move to the next M block. Our loads therefore store all of these values contiguously such that we can just do contiguous 128-bit wide ldr ahead of the compute phase.

Ordering the data in this way is useful for another reason. Because we are not storing all values for a single row contiguously we are able to take advantage of ILP more easily. Our accumulator for a given output value is not going to be used for all input K values at once, it's going to just do a group of 4, then we have to get through the rest of the output values, then we circle back and do the next group of 4. There is no dependency between any of the ops and a single accumulator register.

        mov w21, #0 // bn
        .N_block: // outer loop blocking N
            cmp w21, N
            bge .N_block_x

            // alright we want to load data from mat B into the scratch memory
            // we are going to load each tile contiguously so that they are packed nicely
            mov w22, #0 // bn_l
            mov x15, x25 // B scratch write pointer we can advance
            .N_block_load: // BN
                cmp w22, w6
                bge .N_block_load_x

                mov w23, #0 // k_l
                .N_block_load_k:
                    cmp w23, K
                    bge .N_block_load_k_x
                    // Offset = ((bn + bn_l) + k_l * N) * 2 (bytes)
                    add w12, w21, w22
                    madd w12, w23, N, w12
                    uxtw x12, w12
                    lsl x12, x12, #1
                    add x12, x12, pMatB

                    ldr q0, [x12]
                    add x12, x12, x2, lsl #1
                    ldr q1, [x12]
                    add x12, x12, x2, lsl #1
                    ldr q2, [x12]
                    add x12, x12, x2, lsl #1
                    ldr q3, [x12]
                    // Now we have 4x16 values from B, we need to reshape these into 4x2 vectors for
                    // bfmmla
                    // [N0{k0, k1, k2, k3}, N1{k0, k1, k2, k3}, ...]

                    zip1 v4.8h, v0.8h, v1.8h // [k0n0, k1n0, k0n1, k1n1, ..., k1n3]
                    zip2 v5.8h, v0.8h, v1.8h // [k0n4, k1n4, k0n5, k1n5, ..., k1n7]
                    zip1 v6.8h, v2.8h, v3.8h // [k2n0, k3n0, k2n1, k3n1, ..., k3n3]
                    zip2 v7.8h, v2.8h, v3.8h // [k2n4, k3n4, k2n5, k3n5, ..., k3n7]

                    zip1 v8.4s, v4.4s, v6.4s  // [k0n0, k1n0, k2n0, ..., k3n1]
                    zip2 v9.4s, v4.4s, v6.4s  // [k0n2, k1n2, k2n2, ..., k3n3]
                    zip1 v10.4s, v5.4s, v7.4s // [k0n4, k1n4, k2n4, ..., k3n5]
                    zip2 v11.4s, v5.4s, v7.4s // [k0n6, k1n6, k2n6, ..., k3n7]

                    stp q8, q9, [x15], #32 // wrote 32 bytes, advance x15 pointer
                    stp q10, q11, [x15], #32 // wrote 32 bytes, advance x15 pointer

                    add w23, w23, #4 // advance k_l by 4
                    b .N_block_load_k
                .N_block_load_k_x:
                add w22, w22, #8 // 8-wide
                b .N_block_load
            .N_block_load_x:

For the N block where we load from Matrix B, it's the same idea but our swizzle is a little bit different because the orientation for this matrix is different. Again we want to have 4 values from each column interleaved in pairs so that we can do a 2x4 @ 4x2 bfmmla.

// perform compute on loaded scratch
            mov w22, #0 // m
            .M_inner:
                cmp w22, w7 // m < BM
                bge .M_inner_x

                mov x15, x25 // B scratch read pointer we can advance

                mov w23, #0 // n
                .N_inner:
                    cmp w23, w6 // n < BN?
                    bge .N_inner_x


                    // create an A scratch pointer we can advance in x14
                    mul w13, w22,w1
                    add x14, x26, w13, uxtw #1

                    // clear accumulators for this tile
                    .altmacro
                    .macro clear_reg index
                        movi v\index\().8h, #0
                    .endm
                    .set i, 0
                    .rept 16
                        clear_reg %i + 16
                        .set i, i+1
                    .endr
                    .noaltmacro

                    mov w24, #0 // k
                    .K_inner:
                        cmp w24, K
                        bge .K_inner_x

                        ldp q0, q1, [x14], #32

                        ldp q8, q9, [x15], #32
                        ldp q10, q11, [x15], #32
                        // K0..K3
                        bfmmla v16.4s, v0.8h, v8.8h // M0+1, N0+1
                        bfmmla v17.4s, v0.8h, v9.8h // M0+1, N2+3
                        bfmmla v18.4s, v0.8h, v10.8h // M0+1, N4+5
                        bfmmla v19.4s, v0.8h, v11.8h // M0+1, N6+7

                        ldp q12, q13, [x15], #32
                        ldp q2, q3, [x14], #32

                        bfmmla v20.4s, v1.8h, v8.8h // M2+3, N0+1
                        bfmmla v21.4s, v1.8h, v9.8h // M2+3, N2+3
                        bfmmla v22.4s, v1.8h, v10.8h // M2+3, N4+5
                        bfmmla v23.4s, v1.8h, v11.8h // M2+3, N6+7

                        ldp q14, q15, [x15], #32
                        ldp q4, q5, [x14], #32

                        bfmmla v24.4s, v2.8h, v8.8h // M4+5, N0+1
                        bfmmla v25.4s, v2.8h, v9.8h // M4+5, N2+3
                        bfmmla v26.4s, v2.8h, v10.8h // M4+5, N4+5
                        bfmmla v27.4s, v2.8h, v11.8h // M4+5, N6+7

                        bfmmla v28.4s, v3.8h, v8.8h // M6+7, N0+1
                        bfmmla v29.4s, v3.8h, v9.8h // M6+7, N2+3
                        bfmmla v30.4s, v3.8h, v10.8h // M6+7, N4+5
                        bfmmla v31.4s, v3.8h, v11.8h // M6+7, N6+7

                        ldp q6, q7, [x14], #32
                        // K4..K7
                        bfmmla v16.4s, v4.8h, v12.8h // M0+1, N0+1
                        bfmmla v17.4s, v4.8h, v13.8h // M0+1, N2+3
                        bfmmla v18.4s, v4.8h, v14.8h // M0+1, N4+5
                        bfmmla v19.4s, v4.8h, v15.8h // M0+1, N6+7

                        bfmmla v20.4s, v5.8h, v12.8h // M2+3, N0+1
                        bfmmla v21.4s, v5.8h, v13.8h // M2+3, N2+3
                        bfmmla v22.4s, v5.8h, v14.8h // M2+3, N4+5
                        bfmmla v23.4s, v5.8h, v15.8h // M2+3, N6+7

                        bfmmla v24.4s, v6.8h, v12.8h // M4+5, N0+1
                        bfmmla v25.4s, v6.8h, v13.8h // M4+5, N2+3
                        bfmmla v26.4s, v6.8h, v14.8h // M4+5, N4+5
                        bfmmla v27.4s, v6.8h, v15.8h // M4+5, N6+7

                        bfmmla v28.4s, v7.8h, v12.8h // M6+7, N0+1
                        bfmmla v29.4s, v7.8h, v13.8h // M6+7, N2+3
                        bfmmla v30.4s, v7.8h, v14.8h // M6+7, N4+5
                        bfmmla v31.4s, v7.8h, v15.8h // M6+7, N6+7


                        add w24, w24, #8 // advance k by 8
                        b .K_inner 
                    .K_inner_x:

Here we are doing the actual compute loop. We do 16 bfmmlas at once, then we do another 16 on the same group of accumulators. So we're moving through the K inner dimension 8 elements at a time and as you can see for both scratch pointers we can simply advance the pointer by 32 bytes every load.

                    // we need to swizzle the memory order here because all our registers are holding
                    // 2x2 matrices and we're working on an 8x8 tile

                    zip1 v0.2d, v16.2d, v17.2d // m0n0, m0n1, m0n2, m0n3
                    zip2 v1.2d, v16.2d, v17.2d // m1n0, m1n1, m1n2, m1n3
                    zip1 v2.2d, v18.2d, v19.2d // m0n4, m0n5, m0n6, m0n7
                    zip2 v3.2d, v18.2d, v19.2d // .. etc
                    zip1 v4.2d, v20.2d, v21.2d
                    zip2 v5.2d, v20.2d, v21.2d
                    zip1 v6.2d, v22.2d, v23.2d
                    zip2 v7.2d, v22.2d, v23.2d
                    zip1 v8.2d, v24.2d, v25.2d
                    zip2 v9.2d, v24.2d, v25.2d
                    zip1 v10.2d, v26.2d, v27.2d
                    zip2 v11.2d, v26.2d, v27.2d
                    zip1 v12.2d, v28.2d, v29.2d
                    zip2 v13.2d, v28.2d, v29.2d
                    zip1 v14.2d, v30.2d, v31.2d
                    zip2 v15.2d, v30.2d, v31.2d

                    // matC offset
                    // (m + bn) * N + (bn + n)
                    add w13, w20, w22
                    add w12, w21, w23 // bn + n
                    madd w12, w13, N, w12 
                    uxtw x12, w12
                    // write accumulators to memory
                    add x12, x5, x12, lsl #2
                    stp q0, q2, [x12]
                    add x12, x12, x2, lsl #2
                    stp q1, q3, [x12]
                    add x12, x12, x2, lsl #2
                    stp q4, q6, [x12]
                    add x12, x12, x2, lsl #2
                    stp q5, q7, [x12]
                    add x12, x12, x2, lsl #2
                    stp q8, q10, [x12]
                    add x12, x12, x2, lsl #2
                    stp q9, q11, [x12]
                    add x12, x12, x2, lsl #2
                    stp q12, q14, [x12]
                    add x12, x12, x2, lsl #2
                    stp q13, q15, [x12]

                    add w23, w23, #8 // advance n by 8
                    b .N_inner
                .N_inner_x:
                add w22, w22, #8 // advance m by 8
                b .M_inner
            .M_inner_x:
            add w21, w21, w6 // advance bn by BN
            b .N_block
        .N_block_x:
        add w20, w20, w7 // advance bm by BM
        b .M_block
    .M_block_x:

    mov x0, x25 // free the scratch memory
    bl FREE
    mov x0, x26
    bl FREE
    mov x0, #0

    alloc_fail:
    // pop the registers back from the stack
    ldp d10, d11, [sp, #16]
    ldp d12, d13, [sp, #32]
    ldp d14, d15, [sp, #48]
    ldp d8, d9, [sp], #64

    ldp x21, x22, [sp, #16]
    ldp x23, x24, [sp, #32]
    ldp x25, x26, [sp, #48]
    ldp x27, x28, [sp, #64]
    ldp x19, x20, [sp], #80

    ldp x29, x30, [sp], #16
    ret

Here's the end of it. The most interesting part is writing the accumulators. At this point the accumulator registers are holding 2x2 matrices that need to be swizzled to be in the right format for the 8x8 tile we're working on.

On Graviton 5, this gets us to about 280 Gflops on a single core. The performance is sensitive to the size of the M and N block size. A BN=BM=64 size nets us about 250 Gflops, whereas BN=128, BM=256 we see 280. The performance on my M2 Mac is much worse, I am guessing the instruction is not actually well supported by the M2 processor.

Talk: Real-time Video Upscaling Using CNNs and WebGPU

james hurley — Mon, 21 Sep 2026 19:19:25 +0000

ARM Matrix Multiplication (NEON Edition)

james hurley — Sat, 19 Sep 2026 04:47:51 +0000

it's been a little while since i messed around with writing any sort of neural net kernels. lots of people are doing it these days and the tools are pretty advanced for cuda and whatever else we've got going on in gpu land.

and yet people keep making all these dang ARM computers in my cloud. so how do we make two matrices who love each other very much create a third matrix using arm assembly? and, more importantly, can we make it gooooooo (pretty fast, maybe not gpu fast but you know kinda fast)?

anyway we're gonna try to do the thing using neon assembly to start with. nothing too advanced. i'll make another article about using SME/SVE2 after this.

today's weapon of choice is a macbook air with an M2 processor. The theoretical limit of f32 single-threaded matmul for this processor is 112 GFlops (3.5Ghz * 4 FMA/Cycle * 4 lanes (128-bit register / 32 bits) * 2 flops/fma (multiply, accumulate)).

the performance numbers in this post are based on a sweep of shapes between 64x64 and 1024x1024 (with 20 runs each step), using the average value for each kernel. you can find a chart of sweeps at the bottom of this post.

here's the basic idea:

static void mmul_ref(int M, int K, int N, float* A, float* B, float* C) {
    for (int m = 0; m < M; m++)
        for (int n = 0; n < N; n++) {
            float acc = 0.0f;
            for (int k = 0; k < K; k++)
                acc += A[m * K + k] * B[k * N + n];
            C[m * N + n] = acc;
        }
}

At -O2 this runs at 2.09 Gflops, about 1.9% of our theoretical maximum..

let's try to write this naive implementation in assembly

// function signature is the same as above
// x0 = M
// x1 = K
// x2 = N
// x3 = *A
// x4 = *B
// x5 = *C
_mmul_naive_asm:
    mov w9, #0 // m
    1:
        mov w10, #0 // n
        cmp w9, M
        bge 1f
        2:
            mov w11, #0 // k
            cmp w10, N
            bge 2f
            movi v2.4s, #0
            3:
                cmp w11, K
                bge 3f
                // matC[m * N + n] += matA[m * K + k] * matB[k * N + n]
                madd w13, w9, K, w11 // matA offset
                madd w14, w11, N, w10 // matB offset
                ldr s0, [x3, x13, lsl #2] // load from x3 + offA * 4
                ldr s1, [x4, x14, lsl #2] // load from x4 + offB * 4
                fmadd s2, s0, s1, s2 // FMA acc = acc + (s0 * s1)
                add w11, w11, #1 // increment k by 1
                b 3b
            3:
            str s2, [x5]
            add x5, x5, #4
            add w10, w10, #1 // increment n by 1
            b 2b
        2:
        add w9, w9, #1 // increment m by 1
        b 1b
    1:
    ret

alright we have 1.98 Gflops (or about 1.8% of the theoretical max)... the c compiler is smarter than me.

let's do some simd...
we're going to load 4 values from B into a single 128-bit register and broadcast a single value from A into a 128-bit register and use the fmla instruction to do 4 fma's at the same time across 4 output values. an important caveat here is now N needs to be divisible by 4. we could add some overhang stuff but right now we're in control of the shape of the matrices and i don't feel like doing that.

_mmul_vec_asm:
    mov w9, #0 // m
    1:
        mov w10, #0 // n
        cmp w9, M
        bge 1f
        2:
            mov w11, #0 // k
            cmp w10, N
            bge 2f
            // clear output register
            movi v2.4s, #0
            3:
                cmp w11, K
                bge 3f
                // matC[m * N + n] += matA[m * K + k] * matB[k * N + n]
                madd w13, w9, K, w11 // matA offset
                madd w14, w11, N, w10 // matB offset
                add x15, x3, x13, lsl #2
                lsl w14, w14, #2
                ldr q1, [x4, x14] // load 4 matB values
                ld1r {v0.4s}, [x15] // load a single matA value into all lanes of q0
                fmla v2.4s, v1.4s, v0.4s
                add w11, w11, #1 // increment k by 1
                b 3b
            3:
            str q2, [x5]
            add x5, x5, #16
            add w10, w10, #4 // increment n by 4
            b 2b
        2:
        add w9, w9, #1 // increment m by 1
        b 1b
    1:
    ret

alright this gets us to 7.86 Gflops, about 7% of our target.. but about 4x better than where we were at before - progress!

let's try unrolling K a little more too... get a little more compute happening...


_mmul_vec_asm2:
    mov w9, #0 // m
    1:
        mov w10, #0 // n
        cmp w9, M
        bge 1f
        2:
            mov w11, #0 // k
            cmp w10, N
            bge 2f

            // clear output accumulators
            // We are using 4 here to remove dependencies within K blocks
            movi v0.4s, #0
            movi v1.4s, #0
            movi v2.4s, #0
            movi v3.4s, #0
            // q0..q3 -> outputs
            // q4..q7 -> matB inputs
            // q16..q19 -> matA inputs
            3:
                cmp w11, K
                bge 3f
                // matC[m * N + n] += matA[m * K + k] * matB[k * N + n]
                madd w13, w9, K, w11 // matA offset
                madd w14, w11, N, w10 // matB offset
                add x15, x3, x13, lsl #2
                lsl w14, w14, #2
                ldr q4, [x4, x14] // load 4 matB values
                ld1r {v16.4s}, [x15], #4 // load a single matA value into all lanes
                ld1r {v17.4s}, [x15], #4
                ld1r {v18.4s}, [x15], #4
                ld1r {v19.4s}, [x15], #4
                add w14, w14, N, lsl #2
                ldr q5, [x4, x14] // load 4 matB values
                add w14, w14, N, lsl #2
                ldr q6, [x4, x14] // load 4 matB values
                add w14, w14, N, lsl #2
                ldr q7, [x4, x14] // load 4 matB values
                fmla v0.4s, v4.4s, v16.4s
                fmla v1.4s, v5.4s, v17.4s
                fmla v2.4s, v6.4s, v18.4s
                fmla v3.4s, v7.4s, v19.4s
                add w11, w11, #4
                b 3b
            3:
            // reduce to v0.4s
            // we will sum v0 and v1 into v0, then v2 and v3 into v2, and finally v0 and v2 into v0.
            // this is again to allow multiple floating point instructions to run in parallel
            fadd v0.4s, v0.4s, v1.4s
            fadd v2.4s, v2.4s, v3.4s
            fadd v0.4s, v0.4s, v2.4s
            // write 4 outputs...
            str q0, [x5]
            add x5, x5, #16
            add w10, w10, #4 // do 4 N at once
            b 2b
        2:
        add w9, w9, #1
        b 1b
    1:
    ret

ok now we're at 15.95Gflops (14.2%). So we're starting to move.. something you might have noticed here is we're still writing to a single output but we have 4 accumulators. this is so we can have the cpu pipeline the fmla instructions by eliminating dependencies between fmla executions. if we use a single register here we lose 3-5Gflops.

but the problem here is we're not really getting a lot of compute intensity for each load. one of the classic tricks to improve matrix multiplication performance that you might have seen or used if you've ever done any gpu kernels is tiling. if you're not familiar with it, the idea is to work on patches of the output matrix by sliding windows along the input matrices' K inner dimension. this way we can load a bunch of data and then do a whole bunch of compute to make those loads relatively less expensive via amortization. we are also reducing the overall number of redundant loads by reusing data in registers rather than reloading it.

We have 32 128-bit registers at our disposal here for doing float operations.. so we can have an 8x8 output tile using 16 of those registers and use the other 16 for inputs (8x4 @ 4x8). Now we will do 16 loads and 64 fmla operations on each iteration over the K dimension. We will still unroll K by 4 as we did in the last example, but we'll do 8 values in both the M and N dimensions.

_mmul_r_tile_8x8:
    // free up fp registers since we will use all 32
    stp d8, d9, [sp, #-64]!
    stp d10, d11, [sp, #16]
    stp d12, d13, [sp, #32]
    stp d14, d15, [sp, #48]

    // q0..q15 -> outputs (8x8=64 / 4 single-precision = 16 registers)
    // q16..q23 -> matA inputs
    // q24..q31 -> matB inputs
    mov w9, #0 // m
    1:
        cmp w9, M
        bge 1f
        mov w10, #0 // n
        2:
            cmp w10, N
            bge 2f

            // clear accumulators for this tile
            .altmacro
            .macro clear_reg index
                movi v\index\().4s, #0
            .endm
            .set i, 0
            .rept 16
                clear_reg %i
                .set i, i+1
            .endr
            .noaltmacro

            mov w11, #0 // k
            3:
                cmp w11, K
                bge 3f
                // 8x4 A tile x 4x8 B tile, 64 fmla per K-block.
                // 4 is in the K direction, 8 is in the M or N direction.
                // matC[m*N+n] += matA[m*K+k] * matB[k*N+n]

                madd w14, w11, N, w10 // matB offset
                // load 32 matB values
                add x14, x4, x14, lsl #2 // matB pointer + offset (bytes)
                ldp q24, q25, [x14]
                add x14, x14, x2, lsl #2 // move matB pointer by stride N
                ldp q26, q27, [x14]
                add x14, x14, x2, lsl #2

                madd w13, w9, K, w11 // matA offset
                add x13, x3, x13, lsl #2 // matA pointer + offset (bytes)
                ldr q16, [x13]
                add x13, x13, x1, lsl #2 // move matA pointer by stride K
                ldr q17, [x13]
                add x13, x13, x1, lsl #2
                ldr q18, [x13]
                add x13, x13, x1, lsl #2
                ldr q19, [x13]
                add x13, x13, x1, lsl #2
                ldr q20, [x13]
                add x13, x13, x1, lsl #2
                ldr q21, [x13]
                add x13, x13, x1, lsl #2
                ldr q22, [x13]
                add x13, x13, x1, lsl #2
                ldr q23, [x13]
                add x13, x13, x1, lsl #2

                // template for 16 fmlas... makes life a little easier.
                .macro fmla16 lane, a0, a1
                    fmla v0.4s, \a0\().4s, v16.s[\lane]
                    fmla v1.4s, \a1\().4s, v16.s[\lane]
                    fmla v2.4s, \a0\().4s, v17.s[\lane]
                    fmla v3.4s, \a1\().4s, v17.s[\lane]
                    fmla v4.4s, \a0\().4s, v18.s[\lane]
                    fmla v5.4s, \a1\().4s, v18.s[\lane]
                    fmla v6.4s, \a0\().4s, v19.s[\lane]
                    fmla v7.4s, \a1\().4s, v19.s[\lane]
                    fmla v8.4s, \a0\().4s, v20.s[\lane]
                    fmla v9.4s, \a1\().4s, v20.s[\lane]
                    fmla v10.4s, \a0\().4s, v21.s[\lane]
                    fmla v11.4s, \a1\().4s, v21.s[\lane]
                    fmla v12.4s, \a0\().4s, v22.s[\lane]
                    fmla v13.4s, \a1\().4s, v22.s[\lane]
                    fmla v14.4s, \a0\().4s, v23.s[\lane]
                    fmla v15.4s, \a1\().4s, v23.s[\lane]
                .endm

                // k=0
                fmla16 0, v24, v25

                ldp q28, q29, [x14]
                add x14, x14, x2, lsl #2

                //k=1
                fmla16 1, v26, v27

                ldp q30, q31, [x14]

                //k=2
                fmla16 2, v28, v29

                //k=3
                fmla16 3, v30, v31

                add w11, w11, #4
                b 3b
            3:
            // write accumulators to memory
            madd w14, w9, N, w10 // matC offset
            add x14, x5, x14, lsl #2
            stp q0, q1, [x14]
            add x14, x14, x2, lsl #2
            stp q2, q3, [x14]
            add x14, x14, x2, lsl #2
            stp q4, q5, [x14]
            add x14, x14, x2, lsl #2
            stp q6, q7, [x14]
            add x14, x14, x2, lsl #2
            stp q8, q9, [x14]
            add x14, x14, x2, lsl #2
            stp q10, q11, [x14]
            add x14, x14, x2, lsl #2
            stp q12, q13, [x14]
            add x14, x14, x2, lsl #2
            stp q14, q15, [x14]
            add w10, w10, #8
            b 2b
        2:
        add w9, w9, #8
        b 1b
    1:
    // pop the registers back from the stack
    ldp d10, d11, [sp, #16]
    ldp d12, d13, [sp, #32]
    ldp d14, d15, [sp, #48]
    ldp d8, d9, [sp], #64
    ret

So this gets us to 96.65 Gflops, about 86.3% of our target!

Some details here.. first, hell yea we unrolled that by a lot...

Second in these vector instructions you'll notice we are indexing into the registers loaded in A. This is because we loaded values in the K dimension for A when we did the 128-bit load (whereas for B we were loading values along the N dimension). So as we iterate along K we need to index into the registers holding A data and increment which registers we're using for the B data.

third interesting thing we have some of our B data loads interleaved with flmas. this is to do something like double buffering but maybe in a lazier (as in i had to type less to do this) way. it buys us about 5 Gflops.

Because our fmlas with dependencies are pretty far apart we are also not doing the trick from the previous one where we used multiple accumulators.

if you've done this on a gpu before you might also notice that there's no intermediate load step into workgroup shared memory. first of all there's no work group so there's no workgroup shared memory. secondly with this tile size we end up being compute bound on the M2, so there's no need for us to explicitly shuffle memory around.

but why stop at 97 gflops when we can go for 200 gflops? FP16 time.

this one is going to be an 8x16 tile. we'll use the same number of registers but be able to double the number of values we compute.

_mmul_r_tile_8x8_f16:
    // free up fp registers since we will use all 32
    stp d8, d9, [sp, #-64]!
    stp d10, d11, [sp, #16]
    stp d12, d13, [sp, #32]
    stp d14, d15, [sp, #48]

    // q16..q31 -> outputs (16x8=128 / 8 half-precision = 16 registers)
    // q0..q7 -> matA inputs
    // q8..q15 -> matB inputs
    mov w9, #0 // m
    1:
        cmp w9, M
        bge 1f
        mov w10, #0 // n
        2:
            cmp w10, N
            bge 2f

            // clear accumulators for this tile
            .altmacro
            .macro clear_reg index
                movi v\index\().8h, #0
            .endm
            .set i, 0
            .rept 16
                clear_reg %i + 16
                .set i, i+1
            .endr
            .noaltmacro

            mov w11, #0 // k
            3:
                cmp w11, K
                bge 3f
                // 8x4 A tile x 4x8 B tile, 64 fmla per K-block.
                // 4 is in the K direction, 8 is in the M or N direction.
                // matC[m*N+n] += matA[m*K+k] * matB[k*N+n]

                madd w14, w11, N, w10 // matB offset
                // load 32 matB values
                add x14, x4, x14, lsl #1 // matB pointer + offset (bytes)
                ldp q8, q9, [x14]
                add x14, x14, x2, lsl #1 // move matB pointer by stride N
                ldp q10, q11, [x14]
                add x14, x14, x2, lsl #1

                // load all the needed values from A (8x8)
                madd w13, w9, K, w11 // matA offset
                add x13, x3, x13, lsl #1 // matA pointer + offset (bytes)
                ldr q0, [x13]
                add x13, x13, x1, lsl #1 // move matA pointer by stride K
                ldr q1, [x13]
                add x13, x13, x1, lsl #1
                ldr q2, [x13]
                add x13, x13, x1, lsl #1
                ldr q3, [x13]
                add x13, x13, x1, lsl #1
                ldr q4, [x13]
                add x13, x13, x1, lsl #1
                ldr q5, [x13]
                add x13, x13, x1, lsl #1
                ldr q6, [x13]
                add x13, x13, x1, lsl #1
                ldr q7, [x13]

                .macro fmla16 lane, a0, a1
                    fmla v16.8h, \a0\().8h, v0.h[\lane]
                    fmla v17.8h, \a1\().8h, v0.h[\lane]
                    fmla v18.8h, \a0\().8h, v1.h[\lane]
                    fmla v19.8h, \a1\().8h, v1.h[\lane]
                    fmla v20.8h, \a0\().8h, v2.h[\lane]
                    fmla v21.8h, \a1\().8h, v2.h[\lane]
                    fmla v22.8h, \a0\().8h, v3.h[\lane]
                    fmla v23.8h, \a1\().8h, v3.h[\lane]
                    fmla v24.8h, \a0\().8h, v4.h[\lane]
                    fmla v25.8h, \a1\().8h, v4.h[\lane]
                    fmla v26.8h, \a0\().8h, v5.h[\lane]
                    fmla v27.8h, \a1\().8h, v5.h[\lane]
                    fmla v28.8h, \a0\().8h, v6.h[\lane]
                    fmla v29.8h, \a1\().8h, v6.h[\lane]
                    fmla v30.8h, \a0\().8h, v7.h[\lane]
                    fmla v31.8h, \a1\().8h, v7.h[\lane]
                .endm



                // k=0
                fmla16 0, v8, v9

                ldp q12, q13, [x14]
                add x14, x14, x2, lsl #1

                //k=1
                fmla16 1, v10, v11

                ldp q14, q15, [x14]
                add x14, x14, x2, lsl #1

                //k=2
                fmla16 2, v12, v13

                ldp q8, q9, [x14]
                add x14, x14, x2, lsl #1

                //k=3
                fmla16 3, v14, v15

                ldp q10, q11, [x14]
                add x14, x14, x2, lsl #1

                //k=4
                fmla16 4, v8, v9

                ldp q12, q13, [x14]
                add x14, x14, x2, lsl #1

                //k=5
                fmla16 5, v10, v11


                ldp q14, q15, [x14]

                //k=6
                fmla16 6, v12, v13


                //k=7
                fmla16 7, v14, v15

                add w11, w11, #8
                b 3b
            3:
            // write accumulators to memory
            madd w14, w9, N, w10 // matC offset
            add x14, x5, x14, lsl #1
            stp q16, q17, [x14]
            add x14, x14, x2, lsl #1
            stp q18, q19, [x14]
            add x14, x14, x2, lsl #1
            stp q20, q21, [x14]
            add x14, x14, x2, lsl #1
            stp q22, q23, [x14]
            add x14, x14, x2, lsl #1
            stp q24, q25, [x14]
            add x14, x14, x2, lsl #1
            stp q26, q27, [x14]
            add x14, x14, x2, lsl #1
            stp q28, q29, [x14]
            add x14, x14, x2, lsl #1
            stp q30, q31, [x14]
            add w10, w10, #16 // advance N by 16
            b 2b
        2:
        add w9, w9, #8 // advance M by 8
        b 1b
    1:
    // pop the registers back from the stack
    ldp d10, d11, [sp, #16]
    ldp d12, d13, [sp, #32]
    ldp d14, d15, [sp, #48]
    ldp d8, d9, [sp], #64
    ret

now we're hitting 203.65 gflops, about 91% of the theoretical 224 FP16 GFlops the M2 processor can do on a single core. Something we're doing a little more of in this is interleaving loads with compute to do some double buffering. This also helps recycle registers because we do not have enough registers for a K block size of 8 (although we are sort of forced into that size by loading 128 bits of K from matrix A). You'll notice too that we swapped the input and accumulator registers, this is because the .h[lane] notation does not work above v15.

anyway this was fun... next time we'll do some integer matmul and mess around with sme or sve2 or whatever...

Epilogue: Maybe we do need some cache blocking, after all.

So I started poking at Graviton instances, using C9g which apparently has a 2MB L2 cache (vs the M2's 16MB L2) and saw some pretty steep performance drop off when running the sweep...

the dashed lines are the Graviton instance.. you can see that as the matrix sizes grow the performance decreases. So, we need to load some data into an intermediate scratch buffer first. For this we're going to make a single scratch block to pull data out of B into a block where loads can be contiguous (vs the current stride of N). We can block in M (for mat A) and K as well, but N (mat B) is probably going to be the most fruitful for us to start with because for each M row we are loading the entire matrix B, whereas we only load an 8xK strip from matrix A (which, at the tested sizes, is small enough to stay resident in L1...)

mmul_c_tile_r_tile_8x16_f16:
    // preserve link register because we'll be calling out to malloc/free
    stp x29, x30, [sp, #-16]!

    // free up some registers
    stp x19, x20, [sp, #-64]!
    stp x21, x22, [sp, #16]
    stp x23, x24, [sp, #32]
    stp x25, x26, [sp, #48]

    // free up fp registers since we will use all 32
    stp d8, d9, [sp, #-64]!
    stp d10, d11, [sp, #16]
    stp d12, d13, [sp, #32]
    stp d14, d15, [sp, #48]

    mov x19, x0
    mov x20, x1
    mov x21, x2
    mov x22, x3
    mov x23, x4
    mov x24, x5
    mov x25, x6

    mul w0, w6, w1 // BN * K
    lsl w0, w0, #1 // * 2 (f16)
    uxtw x0, w0
    bl  MALLOC     // alloc scratch memory
    cmp x0, #0
    mov x8, x0
    mov x0, #1
    beq alloc_fail


    mov x0, x19
    mov x1, x20
    mov x2, x21
    mov x3, x22
    mov x4, x23
    mov x5, x24
    mov x6, x25

    // q16..q31 -> outputs (16x8=128 / 8 half-precision = 16 registers)
    // q0..q7 -> matA inputs
    // q8..q15 -> matB inputs

    // x0 -> M
    // x1 -> K
    // x2 -> N
    // x3 -> pA
    // x4 -> pB
    // x5 -> pC
    // x6 -> BN

    // x8 -> pScratch


    mov w17, #0 // bn
    4: // outer loop blocking N
        cmp w17, N
        bge 4f

        // alright we want to load data from mat B into the scratch memory
        // we are going to load each tile contiguously so that they are packed nicely
        mov w10, #0 // bn_l
        mov x15, x8 // scratch write pointer we can advance
        5: // BN
            cmp w10, w6
            bge 5f

            mov w11, #0 // k_l
            6:
                cmp w11, K
                bge 6f
                // Offset = ((bn + bn_l) + k_l * N) * 2 (bytes)
                add w12, w17, w10
                madd w12, w11, N, w12
                uxtw x12, w12
                lsl x12, x12, #1
                add x12, x12, pMatB
                ldp q0, q1, [x12]
                stp q0, q1, [x15], #32 // wrote 32 bytes, advance x15 pointer
                add w11, w11, #1 // advance k_l by 1
                b 6b
            6:

            add w10, w10, #16 // tile is 16-wide
            b 5b
        5:

        // perform compute on loaded scratch
        mov w9, #0 // m
        1:
            cmp w9, M
            bge 1f

            mov x15, x8 // scratch read pointer we can advance
            mov w10, #0 // n
            2:
                cmp w10, w6 // n < BN?
                bge 2f

                // clear accumulators for this tile
                .altmacro
                .macro clear_reg index
                    movi v\index\().8h, #0
                .endm
                .set i, 0
                .rept 16
                    clear_reg %i + 16
                    .set i, i+1
                .endr
                .noaltmacro

                mov w11, #0 // k
                3:
                    cmp w11, K
                    bge 3f

                    // 8x4 A tile x 4x8 B tile, 64 fmla per K-block.
                    // 4 is in the K direction, 8 is in the M or N direction.
                    // matC[m*N+n] += matA[m*K+k] * matB[k*N+n]

                    ldp q8, q9, [x15], #32
                    ldp q10, q11, [x15], #32

                    // load all the needed values from A (8x8)
                    madd w13, w9, K, w11 // matA offset
                    add x13, x3, x13, lsl #1 // matA pointer + offset (bytes)
                    ldr q0, [x13]
                    add x13, x13, x1, lsl #1 // move matA pointer by stride K
                    ldr q1, [x13]
                    add x13, x13, x1, lsl #1
                    ldr q2, [x13]
                    add x13, x13, x1, lsl #1
                    ldr q3, [x13]
                    add x13, x13, x1, lsl #1
                    ldr q4, [x13]
                    add x13, x13, x1, lsl #1
                    ldr q5, [x13]
                    add x13, x13, x1, lsl #1
                    ldr q6, [x13]
                    add x13, x13, x1, lsl #1
                    ldr q7, [x13]

                    .macro fmla16 lane, a0, a1
                        fmla v16.8h, \a0\().8h, v0.h[\lane]
                        fmla v17.8h, \a1\().8h, v0.h[\lane]
                        fmla v18.8h, \a0\().8h, v1.h[\lane]
                        fmla v19.8h, \a1\().8h, v1.h[\lane]
                        fmla v20.8h, \a0\().8h, v2.h[\lane]
                        fmla v21.8h, \a1\().8h, v2.h[\lane]
                        fmla v22.8h, \a0\().8h, v3.h[\lane]
                        fmla v23.8h, \a1\().8h, v3.h[\lane]
                        fmla v24.8h, \a0\().8h, v4.h[\lane]
                        fmla v25.8h, \a1\().8h, v4.h[\lane]
                        fmla v26.8h, \a0\().8h, v5.h[\lane]
                        fmla v27.8h, \a1\().8h, v5.h[\lane]
                        fmla v28.8h, \a0\().8h, v6.h[\lane]
                        fmla v29.8h, \a1\().8h, v6.h[\lane]
                        fmla v30.8h, \a0\().8h, v7.h[\lane]
                        fmla v31.8h, \a1\().8h, v7.h[\lane]
                    .endm



                    // k=0
                    fmla16 0, v8, v9

                    ldp q12, q13, [x15], #32

                    //k=1
                    fmla16 1, v10, v11

                    ldp q14, q15, [x15], #32

                    //k=2
                    fmla16 2, v12, v13

                    ldp q8, q9, [x15], #32

                    //k=3
                    fmla16 3, v14, v15

                    ldp q10, q11, [x15], #32

                    //k=4
                    fmla16 4, v8, v9

                    ldp q12, q13, [x15], #32

                    //k=5
                    fmla16 5, v10, v11


                    ldp q14, q15, [x15], #32

                    //k=6
                    fmla16 6, v12, v13


                    //k=7
                    fmla16 7, v14, v15

                    add w11, w11, #8
                    b 3b // advance k by 8
                3:
                // write accumulators to memory
                // matC offset
                // m * N + (bn + n)
                add w14, w10, w17 // bn + n
                madd w14, w9, N, w14 
                add x14, x5, x14, lsl #1
                stp q16, q17, [x14]
                add x14, x14, x2, lsl #1
                stp q18, q19, [x14]
                add x14, x14, x2, lsl #1
                stp q20, q21, [x14]
                add x14, x14, x2, lsl #1
                stp q22, q23, [x14]
                add x14, x14, x2, lsl #1
                stp q24, q25, [x14]
                add x14, x14, x2, lsl #1
                stp q26, q27, [x14]
                add x14, x14, x2, lsl #1
                stp q28, q29, [x14]
                add x14, x14, x2, lsl #1
                stp q30, q31, [x14]
                add w10, w10, #16 // advance n by 16
                b 2b
            2:
            add w9, w9, #8 // advance m by 8
            b 1b
        1:
        add w17, w17, w6 // advance bn by BN
        b 4b
    4:

    mov x0, x8 // free the scratch memory
    bl FREE
    mov x0, #0

    alloc_fail:
    // pop the registers back from the stack
    ldp d10, d11, [sp, #16]
    ldp d12, d13, [sp, #32]
    ldp d14, d15, [sp, #48]
    ldp d8, d9, [sp], #64

    ldp x21, x22, [sp, #16]
    ldp x23, x24, [sp, #32]
    ldp x25, x26, [sp, #48]
    ldp x19, x20, [sp], #64

    ldp x29, x30, [sp], #16
    ret

This is starting to look a little more like how you would do it on a gpu...

anyway this has given us a nice boost in performance at larger matrix sizes.. For these tests I used BN=128, so the scratch is holding 8 tiles worth of data with their full K dimension.

(the dashed line is non-blocked)