GEMM performance on M2

I modified the test code to [print latency](https://github.com/philipturner/metal-flash-attention/blob/main/Tests/FlashAttentionTests/GEMM/LaplacianTest.swift#L232) and for multiplying two 4096x4096 matrices the latency is about a second on my M2. For that size, MPSMatrixMultiplication takes about 0.06 seconds on my machine. @liuliu, @philipturner mentioned over email you might be a good person to talk to about this. In my experiments, I've written a double-buffered tiled matrix multiply for the M2 which takes 0.1 seconds (below), but even then it is limited by how quickly it can feed data to the ALU, getting ~80% utilization. I tried both the morton order loading and the simd_async_copy instructions in this repo, but both slowed down loading. @philipturner mentioned this may be because they have disabled the simd_async_copy instructions, and my benchmarks are also suggestive of this, as using simd_async_copy causes a lot of integer math to be done by the ALU, suggesting it may be being emulated. The mystery remains though how MPS achieves such good performance, so I thought I'd make this issue to try and figure it out. <details> <summary>The fastest kernel I've been able to come up with</summary> ```c++ // assumes n,m,k are divisible by 8. // assumes thread width is 32. // Let (SW*8)x(SW*8) be the dimensions of a threadgroup tile in C. // SW*SW is the number of 8x8 simdgroup tiles in a threadgroup tile. // Let TW be the number of threads in a simdgroup and TMAX the max // threads in a threadgroup. Pick SW s.t. TW*SW*SW <= T_MAX. // // Launch with threadgroup shape (TW,SW,SW) and grid shape // (TW,SW*ceil(ceil(m/8)/SW),SW*ceil(ceil(n/8)/SW). // Grid shape is (cols,rows) as simdgroups use (col,row) indicies. // // Grid shape needs to be a multiple of SW so we always have a full // threadgroup. To see this, consider the case where a full // threadgroup would hang off the edge of C to the right. In this // case, we want to load tiles of A as normal, but don't nesecarially // need to load all the columns of B. The requirement that the A tile // is loaded as normal means we need an entire threadgroup working. #include <metal_simdgroup_matrix> #include <metal_compute> using namespace metal; inline void load_tile( const device float* src, const ushort2 src_tg_pos, const ushort2 src_shape, // (cols, rows) threadgroup float* dst, const ushort start ) { for (ushort i = 0; i < 2; i++) { ushort idx = start+i; ushort2 pos = src_tg_pos + ushort2(idx%(SW*8), idx/(SW*8)); dst[idx] = select(0., src[pos.x+pos.y*src_shape.x], all(pos<src_shape)); } } template<ushort DIM> inline void simdgroup_multiply( threadgroup float* A, threadgroup float* B, ushort2 c_pos, thread simdgroup_float8x8 &acc ) { simdgroup_float8x8 A_simd[DIM]; simdgroup_float8x8 B_simd[DIM]; for (ushort i = 0; i < DIM; ++i) simdgroup_load(A_simd[i], A, DIM*8, ulong2(i*8, c_pos.y*8)); for (ushort i = 0; i < DIM; ++i) simdgroup_load(B_simd[i], B, DIM*8, ulong2(c_pos.x*8, i*8)); for (ushort i = 0; i < DIM; ++i) simdgroup_multiply_accumulate(acc, A_simd[i], B_simd[i], acc); } kernel void matmul( constant ushort& n, constant ushort& k, constant ushort& m, constant float& alpha, constant float& beta, const device float* A, const device float* B, device float* C, ushort3 t_pos [[thread_position_in_grid]], ushort3 t_tg_pos [[thread_position_in_threadgroup]], ushort3 tg_pos [[threadgroup_position_in_grid]] ) { ulong2 c_origin = ulong2(t_pos.yz*8); ushort2 c_tg_origin = tg_pos.yz*SW*8; ushort2 a_tg_origin = ushort2(0,c_tg_origin.y); ushort2 b_tg_origin = ushort2(c_tg_origin.x,0); alignas(16) threadgroup float A_tg[2][SW*8*SW*8]; alignas(16) threadgroup float B_tg[2][SW*8*SW*8]; simdgroup_float8x8 acc(0.); ushort buffer = 0; ushort start = (t_tg_pos.y+t_tg_pos.z*SW)*8*8 + t_tg_pos.x*2; ushort k_tiles = ((k/8)+SW-1)/SW; // preload load_tile(A, a_tg_origin, ushort2(k,n), A_tg[buffer], start); load_tile(B, b_tg_origin, ushort2(m,k), B_tg[buffer], start); for (ushort l = 1; l < k_tiles; l++) { buffer ^= 1; ushort2 a_tg_pos = a_tg_origin+ushort2(l*SW*8,0); ushort2 b_tg_pos = b_tg_origin+ushort2(0,l*SW*8); threadgroup_barrier(metal::mem_flags::mem_threadgroup); load_tile(A, a_tg_pos, ushort2(k,n), A_tg[buffer], start); load_tile(B, b_tg_pos, ushort2(m,k), B_tg[buffer], start); simdgroup_multiply<SW>(A_tg[buffer^1],B_tg[buffer^1],t_tg_pos.yz,acc); } // post threadgroup_barrier(metal::mem_flags::mem_threadgroup); simdgroup_multiply<SW>(A_tg[buffer],B_tg[buffer],t_tg_pos.yz,acc); if (c_origin.x<m&&c_origin.y<n) { simdgroup_float8x8 c_simd; simdgroup_load(c_simd,C,m,c_origin); simdgroup_multiply(c_simd,c_simd, simdgroup_float8x8(beta)); simdgroup_multiply_accumulate(c_simd,acc,simdgroup_float8x8(alpha),c_simd); simdgroup_store(c_simd,C,m,c_origin); } } ``` </details>