nx_intfp_mlp_gradcheck_gate.nx source

1// nx_intfp_mlp_gradcheck_gate.nx -- SECOND brick of integer training: prove the CHAIN RULE in fixed-point. 2// A 2-layer MLP h=W1@x ; a=ReLU(h) ; y=W2@a ; L=sum(y^2) done ENTIRELY in Q16 integer, with full backprop 3// (gradients flow y<-a<-h<-W1 through the ReLU kink) and a gold-standard INTEGER finite-difference gradcheck on 4// BOTH W1 (the DEEP matrix -- its gradient traverses the whole chain) and W2. If the deep gradients match numeric, 5// the fixed-point chain rule is sound and a multi-op integer autograd tape is buildable. No float anywhere. 6// Q16 (S=2^16): h_q=(W1_q@x_q)>>16 ; a=max(0,h) ; y_q=(W2_q@a_q)>>16 ; L_q32=sum(y_q^2). 7// gy=2y ; gW2=(gy*a)>>16 ; ga=(sum_o gy*W2)>>16 ; gh=ga*[h>0] ; gW1=(gh*x)>>16 (all Q16 of the real grad) 8// license_tier: ORIGINAL 9import "nx_syscalls.nx" 10 11func w(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 } 12func 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 } 13func iabs(v: i64) -> i64 { if v<0 { return 0-v } return v } 14 15const S: i64 = 65536 16const K: i64 = 3 17const H: i64 = 4 18const O: i64 = 2 19 20// forward, fills h,a,y buffers; returns L_q32 = sum(y_q^2) 21func fwd(W1: *i64, x: *i64, W2: *i64, hb: *i64, ab: *i64, yb: *i64) -> i64 { 22 var i: i64=0 23 while i<H { 24 var acc: i64=0; var kk: i64=0 25 while kk<K { acc=acc + W1[i*K+kk]*x[kk]; kk=kk+1 } 26 let hv: i64=acc/S; hb[i]=hv; if hv<0 { ab[i]=0 } else { ab[i]=hv } // ReLU 27 i=i+1 28 } 29 var Lq: i64=0; var o: i64=0 30 while o<O { 31 var acc2: i64=0; var j: i64=0 32 while j<H { acc2=acc2 + W2[o*H+j]*ab[j]; j=j+1 } 33 let y: i64=acc2/S; yb[o]=y; Lq=Lq + y*y; o=o+1 34 } 35 return Lq 36} 37 38func main() -> i64 { 39 w("=== nx_intfp_mlp_gradcheck: Q16 2-layer MLP (W1->ReLU->W2->L=sum y^2) -- CHAIN-RULE gradcheck, no float ===\n\n" as *u8) 40 let W1: *i64=sys_mmap(H*K*8) as *i64; let W2: *i64=sys_mmap(O*H*8) as *i64 41 let x: *i64=sys_mmap(K*8) as *i64 42 let hb: *i64=sys_mmap(H*8) as *i64; let ab: *i64=sys_mmap(H*8) as *i64; let yb: *i64=sys_mmap(O*8) as *i64 43 // fill: mixed signs so some ReLU units are DEAD (exercises the kink/branch) 44 var i: i64=0; while i<H { var kk: i64=0; while kk<K { W1[i*K+kk]=((i*2-kk*3+1)*S)/10; kk=kk+1 } i=i+1 } 45 var o: i64=0; while o<O { var j: i64=0; while j<H { W2[o*H+j]=((o-j+2)*S)/10; j=j+1 } o=o+1 } 46 var kk2: i64=0; while kk2<K { x[kk2]=((kk2+2)*S)/10; kk2=kk2+1 } // 0.2,0.3,0.4 47 48 let L0: i64=fwd(W1, x, W2, hb, ab, yb) 49 w(" forward L_q32=" as *u8); wn(L0); w(" h_q=[" as *u8); i=0; while i<H { wn(hb[i]); if i<H-1 { w("," as *u8) } i=i+1 } w("] a_q=[" as *u8); i=0; while i<H { wn(ab[i]); if i<H-1 { w("," as *u8) } i=i+1 } w("] (dead ReLU units test the chain)\n\n" as *u8) 50 51 // analytic grads at base point 52 let gy: *i64=sys_mmap(O*8) as *i64; o=0; while o<O { gy[o]=2*yb[o]; o=o+1 } 53 let gW2: *i64=sys_mmap(O*H*8) as *i64; o=0; while o<O { var j: i64=0; while j<H { gW2[o*H+j]=(gy[o]*ab[j])/S; j=j+1 } o=o+1 } 54 let ga: *i64=sys_mmap(H*8) as *i64; var j2: i64=0; while j2<H { var acc: i64=0; o=0; while o<O { acc=acc + gy[o]*W2[o*H+j2]; o=o+1 } ga[j2]=acc/S; j2=j2+1 } 55 let gh: *i64=sys_mmap(H*8) as *i64; i=0; while i<H { if hb[i]<0 { gh[i]=0 } else { gh[i]=ga[i] } i=i+1 } 56 let gW1: *i64=sys_mmap(H*K*8) as *i64; i=0; while i<H { var kk: i64=0; while kk<K { gW1[i*K+kk]=(gh[i]*x[kk])/S; kk=kk+1 } i=i+1 } 57 58 let DELTA: i64=66; let TOLP: i64=80 59 var npass: i64=0; var ncell: i64=0; var worst: i64=0 60 61 // gradcheck W2 (shallow) 62 w(" --- W2 (shallow) --- o j analytic numeric rel(permille)\n" as *u8) 63 o=0 64 while o<O { 65 var j: i64=0 66 while j<H { 67 let ana: i64=gW2[o*H+j]; let save: i64=W2[o*H+j] 68 W2[o*H+j]=save+DELTA; let Lp: i64=fwd(W1,x,W2,hb,ab,yb) 69 W2[o*H+j]=save-DELTA; let Lm: i64=fwd(W1,x,W2,hb,ab,yb) 70 W2[o*H+j]=save; let d: i64=fwd(W1,x,W2,hb,ab,yb) 71 let num: i64=(Lp-Lm)/(2*DELTA); let rel: i64=(iabs(num-ana)*1000)/(iabs(ana)+100) 72 w(" " as *u8); wn(o); w(" " as *u8); wn(j); w(" " as *u8); wn(ana); w(" " as *u8); wn(num); w(" " as *u8); wn(rel) 73 ncell=ncell+1; if rel<=TOLP { npass=npass+1; w(" ok\n" as *u8) } else { w(" FAIL\n" as *u8) } 74 if rel>worst { worst=rel } 75 j=j+1 76 } 77 o=o+1 78 } 79 // gradcheck W1 (DEEP -- gradient flows through ReLU and W2) 80 w(" --- W1 (DEEP, thru ReLU+W2) --- i k analytic numeric rel(permille)\n" as *u8) 81 i=0 82 while i<H { 83 var kk: i64=0 84 while kk<K { 85 let ana: i64=gW1[i*K+kk]; let save: i64=W1[i*K+kk] 86 W1[i*K+kk]=save+DELTA; let Lp: i64=fwd(W1,x,W2,hb,ab,yb) 87 W1[i*K+kk]=save-DELTA; let Lm: i64=fwd(W1,x,W2,hb,ab,yb) 88 W1[i*K+kk]=save; let d: i64=fwd(W1,x,W2,hb,ab,yb) 89 let num: i64=(Lp-Lm)/(2*DELTA); let rel: i64=(iabs(num-ana)*1000)/(iabs(ana)+100) 90 w(" " as *u8); wn(i); w(" " as *u8); wn(kk); w(" " as *u8); wn(ana); w(" " as *u8); wn(num); w(" " as *u8); wn(rel) 91 ncell=ncell+1; if rel<=TOLP { npass=npass+1; w(" ok\n" as *u8) } else { w(" FAIL\n" as *u8) } 92 if rel>worst { worst=rel } 93 kk=kk+1 94 } 95 i=i+1 96 } 97 w("\n CHAIN-RULE gradcheck: " as *u8); wn(npass); w("/" as *u8); wn(ncell); w(" cells within " as *u8); wn(TOLP); w(" permille; worst=" as *u8); wn(worst); w("\n" as *u8) 98 w("NX-INTFP-MLP-GRADCHECK verdict=" as *u8) 99 if npass==ncell { w("GREEN passes=" as *u8); wn(npass); w("/" as *u8); wn(ncell); w(" -- fixed-point CHAIN RULE proven (deep grads thru ReLU); integer tape is buildable\n" as *u8) } 100 else { w("RED passes=" as *u8); wn(npass); w("/" as *u8); wn(ncell); w(" -- chain-rule Q-format bug\n" as *u8) } 101 return 0 102}