ARM 矩阵乘法（NEON 版）

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 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 Enter fullscreen mode Exit fullscreen mode 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 Enter fullscreen mode Exit fullscreen mode 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. 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 Enter fullscreen mode Exit fullscreen mode 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...