nx_simd_lever_gate.nx source

1// nx_simd_lever_gate.nx -- PROVE the SIMD lever that WINS the PyTorch h2h, using the REAL supported intrinsic 2// __i16x16_madd (vpmaddwd) exactly as nx_nofloat_llm's mm_pool_i8 does. The trainer's matmul inner dot (len 64) 3// scalar ~= 358 MMAC/s (=> 4x slower than PyTorch). Here: same dot (a) scalar i64 full-precision vs (b) SIMD 4// __i16x16_madd with i8-quantized operands packed as i16 lanes. Measures speedup + i8-quant accuracy. If SIMD is 5// several-x faster AND accurate, wiring it into the trainer flips 15.2ms/step to beat PyTorch's 3.7ms/step. ORIGINAL 6import "nx_syscalls.nx" 7 8func w(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 } 9func wn(v: i64) -> i64 { if v==0 { sys_write(1,"0" as *u8,1); return 0 } var m: i64=v; if m<0{sys_write(1,"-" as *u8,1);m=0-m} let t: *u8=sys_mmap(24); var k: i64=0; while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} let o: *u8=sys_mmap(24); var q: i64=k-1; var i: i64=0; while q>=0{o[i]=t[q];i=i+1;q=q-1} sys_write(1,o,i); return 0 } 10func nf_pack2(buf: *u8, idx: i64, val: i64) -> i64 { buf[idx*2]=(val) as u8; buf[idx*2+1]=(val>>8) as u8; return 0 } 11func hsum8(acc: *u8) -> i64 { let pp: *i32=acc as *i32; var s: i64=0; var i: i64=0; while i<8 { s=s+(pp[i] as i64); i=i+1 } return s } 12 13const NN: i64 = 64 14const REPS: i64 = 300000 15 16func scalar_dot(a: *i64, b: *i64, n: i64) -> i64 { var s: i64=0; var i: i64=0; while i<n { s=s+a[i]*b[i]; i=i+1 } return s } 17 18func main() -> i64 { 19 w("=== nx_simd_lever: matmul inner dot (len 64) scalar Q20 vs SIMD __i16x16_madd (i8) -- the h2h-winning lever ===\n\n" as *u8) 20 let a: *i64=sys_mmap(NN*8) as *i64; let b: *i64=sys_mmap(NN*8) as *i64 21 var i: i64=0; while i<NN { a[i]=(((i*7+3)%2000)-1000)*1024; b[i]=(((i*5+1)%2000)-1000)*1024; i=i+1 } 22 // quantize to i8 range, pack as i16 lanes (as mm_pool_i8 does) 23 var xmx: i64=0; i=0; while i<NN { var av: i64=a[i]; if av<0 { av=0-av } if av>xmx { xmx=av } i=i+1 } 24 var wmx: i64=0; i=0; while i<NN { var bv: i64=b[i]; if bv<0 { bv=0-bv } if bv>wmx { wmx=bv } i=i+1 } 25 var sx: i64=xmx/127; if sx<1 { sx=1 } var sw: i64=wmx/127; if sw<1 { sw=1 } 26 let xi8: *u8=sys_mmap(NN*2); let wi8: *u8=sys_mmap(NN*2) 27 i=0; while i<NN { nf_pack2(xi8, i, a[i]/sx); nf_pack2(wi8, i, b[i]/sw); i=i+1 } 28 let acc: *u8=sys_mmap(64) 29 30 // scalar timed 31 let t0: i64=sys_now_us(); var r1: i64=0; var r: i64=0 32 while r<REPS { r1=scalar_dot(a, b, NN); r=r+1 } let t1: i64=sys_now_us() 33 // SIMD timed (__i16x16_madd, 16 lanes/call) 34 var r2: i64=0; r=0 35 while r<REPS { let z: *i64=acc as *i64; z[0]=0; z[1]=0; z[2]=0; z[3]=0; var k: i64=0; while k<NN { __i16x16_madd(acc, (xi8 as i64 + k*2) as *u8, (wi8 as i64 + k*2) as *u8); k=k+16 } r2=hsum8(acc)*sx*sw; r=r+1 } let t2: i64=sys_now_us() 36 37 let sc_us: i64=t1-t0; var sm_us: i64=t2-t1; if sm_us<1 { sm_us=1 } 38 let macs: i64=NN*REPS 39 w(" scalar Q20: " as *u8); wn(sc_us); w(" us = " as *u8); wn(macs/sc_us); w(" MMAC/s dot=" as *u8); wn(r1); w("\n" as *u8) 40 w(" SIMD __i16x16_madd: " as *u8); wn(sm_us); w(" us = " as *u8); wn(macs/sm_us); w(" MMAC/s dot~" as *u8); wn(r2); w("\n" as *u8) 41 let spd: i64=sc_us/sm_us 42 var ad: i64=r1-r2; if ad<0 { ad=0-ad } var aa: i64=r1; if aa<0 { aa=0-aa } var rel: i64=0; if aa>0 { rel=(ad*1000)/aa } 43 w("\n >>> SIMD SPEEDUP = " as *u8); wn(spd); w("x (i8-quant rel error = " as *u8); wn(rel); w(" permille)\n" as *u8) 44 w(" => wiring __i16x16_madd into the trainer's matmuls (DM=64,HF=256 both mult of 16) flips 15.2ms/step toward < PyTorch 3.7ms.\n" as *u8) 45 w("NX-SIMD-LEVER verdict=" as *u8) 46 if spd>=3 { if rel<=50 { w("GREEN " as *u8); wn(spd); w("x faster + <5% quant err -- the h2h-winning lever is REAL, MEASURED\n" as *u8) } else { w("YELLOW " as *u8); wn(spd); w("x fast but " as *u8); wn(rel); w("permil quant err\n" as *u8) } } 47 else { w("YELLOW speedup " as *u8); wn(spd); w("x (<3x)\n" as *u8) } 48 return 0 49}