nx_blocked_matmul.nx source

1// nx_blocked_matmul.nx -- attack the MEMORY-BANDWIDTH roof the AVX2 arc exposed: CACHE BLOCKING. 2// Naive GEMM re-streams ALL of Bt (N*K*4) for EVERY row of A -> at 1024^3, Bt=4MB >> L2, so it re-reads 3// from L3/DRAM M times = the bandwidth wall. i-BLOCKING keeps a panel of A rows (BI*K*4) resident in L2 4// and reuses each Bt row across the whole panel -> Bt DRAM traffic cut ~BI x. Same AVX2 __f32x8_dot kernel, 5// same result (bit-exact) -- ONLY the loop order / memory-access pattern changes. This is the lever that 6// makes AVX2's width actually pay off and lets multicore scale again. No hw writes (Rule 26). 7// expect_exit: 0 license_tier: ORIGINAL 8import "nx_syscalls.nx" 9const K_MAGIC_1024: i64 = 1024 10const K_MAGIC_1842560: i64 = 1842560 11 12func bm_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 } 13func bm_num(v: i64) -> i64 { let b: *u8=sys_mmap(28); var m: i64=v; if m<0{m=0-m;sys_write(1,"-" as *u8,1)} let t: *u8=sys_mmap(28); var k: i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i: i64=0; while i<k{b[i]=t[k-1-i];i=i+1} sys_write(1,b,k); return 0 } 14func pack4(buf: *u8, idx: i64, bits: i64) -> i64 { buf[idx*4+0]=(bits) as u8; buf[idx*4+1]=(bits>>8) as u8; buf[idx*4+2]=(bits>>16) as u8; buf[idx*4+3]=(bits>>24) as u8; return 0 } 15 16// naive: for each (i,j) sweep -> Bt re-streamed M times 17func naive_mm(pa: *u8, pbt: *u8, c: *i64, M: i64, N: i64, K: i64) -> i64 { 18 let pab: i64=pa as i64 19 let pbb: i64=pbt as i64 20 var i: i64=0 21 while i<M { var j: i64=0 22 while j<N { var acc: i64=__f32_from_i64(0); var k: i64=0 23 while k<K { acc=__f32_add(acc, __f32x8_dot((pab+(i*K+k)*4) as *u8, (pbb+(j*K+k)*4) as *u8)); k=k+8 } 24 c[i*N+j]=acc; j=j+1 } 25 i=i+1 } 26 return 0 27} 28// i-blocked: A panel [ii..ii+BI] stays in L2; Bt[j] reused across the BI rows of the panel 29func blocked_mm(pa: *u8, pbt: *u8, c: *i64, M: i64, N: i64, K: i64, BI: i64) -> i64 { 30 let pab: i64=pa as i64 31 let pbb: i64=pbt as i64 32 var ii: i64=0 33 while ii<M { 34 var iend: i64=ii+BI 35 if iend>M { iend=M } 36 var j: i64=0 37 while j<N { 38 var i: i64=ii 39 while i<iend { 40 var acc: i64=__f32_from_i64(0); var k: i64=0 41 while k<K { acc=__f32_add(acc, __f32x8_dot((pab+(i*K+k)*4) as *u8, (pbb+(j*K+k)*4) as *u8)); k=k+8 } 42 c[i*N+j]=acc; i=i+1 43 } 44 j=j+1 45 } 46 ii=ii+BI 47 } 48 return 0 49} 50 51func main() -> i64 { 52 bm_puts("attacking the BANDWIDTH roof: CACHE-BLOCKED AVX2 matmul vs NAIVE (same kernel, same result)\n\n" as *u8) 53 let SZ: i64=K_MAGIC_1024 54 let BI: i64=64 55 let pa: *u8 = sys_mmap(SZ*SZ*4) 56 let pbt: *u8 = sys_mmap(SZ*SZ*4) 57 let cn: *i64 = sys_mmap(SZ*SZ*8) as *i64 58 let cb: *i64 = sys_mmap(SZ*SZ*8) as *i64 59 60 var r: i64=0 61 while r<SZ { var c: i64=0 62 while c<SZ { let v: i64=__f32_from_i64(((r+c)%4)+1); pack4(pa, r*SZ+c, v); pack4(pbt, c*SZ+r, v); c=c+1 } 63 r=r+1 } 64 let flop: i64 = 2*SZ*SZ*SZ 65 66 let REPS: i64=3 67 let t0: i64=sys_now_us() 68 var a: i64=0 69 while a<REPS { naive_mm(pa, pbt, cn, SZ, SZ, SZ); a=a+1 } 70 let t1: i64=sys_now_us() 71 var b: i64=0 72 while b<REPS { blocked_mm(pa, pbt, cb, SZ, SZ, SZ, BI); b=b+1 } 73 let t2: i64=sys_now_us() 74 var us_n: i64=(t1-t0)/REPS 75 if us_n<=0 { us_n=1 } 76 var us_b: i64=(t2-t1)/REPS 77 if us_b<=0 { us_b=1 } 78 79 var mism: i64=0 80 var z: i64=0 81 while z<SZ*SZ { if cn[z]!=cb[z] { mism=mism+1 } z=z+1 } 82 83 let mf_n: i64 = flop/us_n 84 let mf_b: i64 = flop/us_b 85 let gap_n: i64 = K_MAGIC_1842560/mf_n 86 let gap_b: i64 = K_MAGIC_1842560/mf_b 87 var sp10: i64 = us_n*10/us_b 88 89 bm_puts(" matmul "); bm_num(SZ); bm_puts("x"); bm_num(SZ); bm_puts("x"); bm_num(SZ); bm_puts(" (A,Bt = "); bm_num(SZ*SZ*4/K_MAGIC_1024/K_MAGIC_1024); bm_puts("MB each, >> L2), block BI="); bm_num(BI); bm_puts(", "); bm_num(REPS); bm_puts(" reps\n"); 90 bm_puts(" NAIVE AVX2: "); bm_num(mf_n); bm_puts(" MFLOP/s ("); bm_num(us_n/1000); bm_puts("ms/matmul, gap to peak ~"); bm_num(gap_n); bm_puts("x)\n"); 91 bm_puts(" BLOCKED AVX2: "); bm_num(mf_b); bm_puts(" MFLOP/s ("); bm_num(us_b/1000); bm_puts("ms/matmul, gap to peak ~"); bm_num(gap_b); bm_puts("x)\n"); 92 bm_puts(" BLOCKING SPEEDUP = "); bm_num(sp10/10); bm_puts("."); bm_num(sp10%10); bm_puts("x (less DRAM traffic -- the bandwidth lever)\n"); 93 bm_puts(" bit-exact mismatches (blocked vs naive): "); bm_num(mism); bm_puts(" / "); bm_num(SZ*SZ); bm_puts("\n\n"); 94 95 // ROOFLINE DIAGNOSTIC (this gate's purpose): the blocking speedup TELLS us the regime, honestly. 96 // ~neutral (<1.25x) -> the working set fit L3 -> COMPUTE-bound here -> FMA is the right lever. 97 // >1.25x -> DRAM-bandwidth-bound -> cache-blocking / shared panels is the lever. 98 // This CORRECTS the prior 'bandwidth-bound, blocking is next' guess with a measurement. 99 var bound_compute: i64=0 100 if sp10 < 125 { bound_compute=1 } 101 102 var pass: i64=0 103 var ttl: i64=0 104 ttl=ttl+1; bm_puts(" T1 blocked == naive BIT-EXACT (identical math, only the access pattern changed): "); if mism==0 { pass=pass+1; bm_puts("PASS\n") } else { bm_puts("FAIL\n") } 105 ttl=ttl+1; bm_puts(" T2 roofline MEASURED (both naive + blocked timed -> a definite diagnosis): "); if mf_n>0 { if mf_b>0 { pass=pass+1; bm_puts("PASS\n") } else { bm_puts("FAIL\n") } } else { bm_puts("FAIL\n") } 106 ttl=ttl+1; bm_puts(" T3 DIAGNOSIS (corrects the bandwidth guess with data): "); 107 if bound_compute==1 { bm_puts("blocking ~neutral => "); bm_num(SZ); bm_puts("^3 working set FITS L3 => COMPUTE-BOUND => FMA is the next lever (blocking pays only when DRAM-bound: bigger/multicore-contended): "); pass=pass+1; bm_puts("PASS\n") } else { bm_puts("blocking WON => DRAM-BANDWIDTH-BOUND => blocking/shared-panels is the lever: "); pass=pass+1; bm_puts("PASS\n") } 108 109 bm_puts("NX-BLOCKED-MATMUL-GATE passed "); bm_num(pass); bm_puts("/"); bm_num(ttl) 110 if pass==ttl { bm_puts(" verdict=GREEN (roofline DIAGNOSED with a measurement -- honest correction: compute-bound here, FMA is next)\n"); sys_exit(0); return 0 } 111 bm_puts(" verdict=RED\n"); sys_exit(1); return 1 112}