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 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)
ldr q24, [x14]
ldr q25, [x14, #16]
add x14, x14, x2, lsl #2 // move matB pointer by stride N
ldr q26, [x14]
ldr q27, [x14, #16]
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
ldr q28, [x14]
ldr q29, [x14, #16]
add x14, x14, x2, lsl #2
//k=1
fmla16 1, v26, v27
ldr q30, [x14]
ldr q31, [x14, #16]
//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
str q0, [x14]
str q1, [x14, #16]
add x14, x14, x2, lsl #2
str q2, [x14]
str q3, [x14, #16]
add x14, x14, x2, lsl #2
str q4, [x14]
str q5, [x14, #16]
add x14, x14, x2, lsl #2
str q6, [x14]
str q7, [x14, #16]
add x14, x14, x2, lsl #2
str q8, [x14]
str q9, [x14, #16]
add x14, x14, x2, lsl #2
str q10, [x14]
str q11, [x14, #16]
add x14, x14, x2, lsl #2
str q12, [x14]
str q13, [x14, #16]
add x14, x14, x2, lsl #2
str q14, [x14]
str q15, [x14, #16]
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 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.
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
// 8x8 A tile x 8x16 B tile, 128 fmla per K-block.
// matC[m*N+n] += matA[m*K+k] * matB[k*N+n]
madd w14, w11, N, w10 // matB offset
add x14, x4, x14, lsl #1 // matB pointer + offset (bytes)
ldr q8, [x14]
ldr q9, [x14, #16]
add x14, x14, x2, lsl #1 // move matB pointer by stride N
ldr q10, [x14]
ldr q11, [x14, #16]
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
ldr q12, [x14]
ldr q13, [x14, #16]
add x14, x14, x2, lsl #1
//k=1
fmla16 1, v10, v11
ldr q14, [x14]
ldr q15, [x14, #16]
add x14, x14, x2, lsl #1
//k=2
fmla16 2, v12, v13
ldr q8, [x14]
ldr q9, [x14, #16]
add x14, x14, x2, lsl #1
//k=3
fmla16 3, v14, v15
ldr q10, [x14]
ldr q11, [x14, #16]
add x14, x14, x2, lsl #1
//k=4
fmla16 4, v8, v9
ldr q12, [x14]
ldr q13, [x14, #16]
add x14, x14, x2, lsl #1
//k=5
fmla16 5, v10, v11
ldr q14, [x14]
ldr q15, [x14, #16]
//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
str q16, [x14]
str q17, [x14, #16]
add x14, x14, x2, lsl #1
str q18, [x14]
str q19, [x14, #16]
add x14, x14, x2, lsl #1
str q20, [x14]
str q21, [x14, #16]
add x14, x14, x2, lsl #1
str q22, [x14]
str q23, [x14, #16]
add x14, x14, x2, lsl #1
str q24, [x14]
str q25, [x14, #16]
add x14, x14, x2, lsl #1
str q26, [x14]
str q27, [x14, #16]
add x14, x14, x2, lsl #1
str q28, [x14]
str q29, [x14, #16]
add x14, x14, x2, lsl #1
str q30, [x14]
str q31, [x14, #16]
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 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...