nx_nofloat_layernorm_gate.nx source
↩ module page · 101 lines · 5751 B
1// nx_nofloat_layernorm_gate.nx -- the THIRD no-float transformer primitive: integer rsqrt -> LayerNorm.
2// Same from-the-hardware-rung-up SYSTEM as the exp/softmax one: range-reduce + Newton iterate in PURE
3// INTEGER -> bit-exact DETERMINISTIC. isqrt(v)=floor(sqrt v) via integer Newton (y=(x+v/x)/2 to a fixed
4// point); fx_rsqrt = 1/sqrt via isqrt; LayerNorm = (x-mean)/sqrt(var+eps), all Q16 integer. Completes the
5// trio every transformer layer needs: matmul (CAP-NOFLOAT-GEMM) + softmax (CAP-NOFLOAT-SOFTMAX) + norm.
6// Float rsqrt/norm is non-deterministic; this is reproducible forever. Verified self-consistently: isqrt
7// on perfect squares is EXACT; LayerNorm output has mean~0 + variance~1; fwd-sum == rev-sum (deterministic).
8// No hw writes (Rule 26). expect_exit: 0 license_tier: ORIGINAL
9import "nx_syscalls.nx"
10import "nx_gate_verdict.nx"
11
12const Q16: i64 = 65536
13
14func ln_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
15func ln_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 }
16func absd(a: i64, b: i64) -> i64 { if a>b { return a-b } return b-a }
17
18// integer floor(sqrt(v)) via Newton -- deterministic (monotone descent to the fixed point)
19func isqrt(v: i64) -> i64 {
20 if v <= 0 { return 0 }
21 if v < 4 { return 1 }
22 var x: i64 = v
23 var y: i64 = (x + 1) >> 1
24 var go: i64 = 1
25 while go==1 { if y < x { x = y; y = (x + v/x) >> 1 } else { go = 0 } }
26 return x
27}
28// 1/sqrt(x) in Q16 for x in Q16 (x>0). sqrt(x_real) in Q16 = isqrt(x<<16); rsqrt = 2^32 / that.
29func fx_rsqrt(x: i64) -> i64 {
30 if x <= 0 { return 0 }
31 let s: i64 = isqrt(x << 16) // sqrt(x) in Q16
32 if s <= 0 { return 0 }
33 return (1 << 32) / s
34}
35// LayerNorm in Q16: out = (x-mean)/sqrt(var+eps). rev toggles sum order (must not change the result).
36func fx_layernorm(x: *i64, out: *i64, n: i64, rev: i64) -> i64 {
37 var s: i64 = 0
38 var i: i64 = 0
39 if rev==0 { i=0; while i<n { s=s+x[i]; i=i+1 } } else { i=n-1; while i>=0 { s=s+x[i]; i=i-1 } }
40 let mean: i64 = s / n
41 var vs: i64 = 0
42 i=0
43 while i<n { let d: i64 = x[i]-mean; vs = vs + ((d*d) >> 16); i=i+1 }
44 let varq: i64 = vs / n // variance in Q16
45 let eps: i64 = 7
46 let sd: i64 = isqrt((varq + eps) << 16) // sqrt(var) in Q16
47 if sd <= 0 { return 0 }
48 i=0
49 while i<n { out[i] = ((x[i]-mean) << 16) / sd; i=i+1 }
50 return sd
51}
52
53func main() -> i64 {
54 ln_puts("THIRD no-float transformer primitive: integer rsqrt -> LayerNorm (deterministic norm)\n\n" as *u8)
55 let TOL: i64 = 1500
56
57 // isqrt KATs (exact on perfect squares)
58 let s144: i64 = isqrt(144)
59 let s1m: i64 = isqrt(1000000)
60 let r4: i64 = fx_rsqrt(4 << 16) // 1/sqrt(4) = 0.5 -> 32768
61 let r025: i64 = fx_rsqrt(Q16 >> 2) // 1/sqrt(0.25) = 2.0 -> 131072
62 ln_puts(" isqrt(144)="); ln_num(s144); ln_puts(" (want 12) isqrt(1000000)="); ln_num(s1m); ln_puts(" (want 1000)\n");
63 ln_puts(" rsqrt(4)="); ln_num(r4); ln_puts(" (want 32768=0.5) rsqrt(0.25)="); ln_num(r025); ln_puts(" (want 131072=2.0)\n\n");
64
65 // LayerNorm on [1,2,3,4,5] (Q16)
66 let n: i64 = 5
67 let x: *i64 = sys_mmap(n*8) as *i64
68 x[0]=65536; x[1]=131072; x[2]=196608; x[3]=262144; x[4]=327680
69 let of: *i64 = sys_mmap(n*8) as *i64
70 let orr: *i64 = sys_mmap(n*8) as *i64
71 fx_layernorm(x, of, n, 0)
72 fx_layernorm(x, orr, n, 1)
73 var omean: i64 = 0
74 var ovar: i64 = 0
75 var detmis: i64 = 0
76 var i: i64 = 0
77 while i<n { omean=omean+of[i]; ovar=ovar+((of[i]*of[i])>>16); if of[i]!=orr[i] { detmis=detmis+1 } i=i+1 }
78 omean = omean / n
79 ovar = ovar / n
80
81 ln_puts(" LayerNorm([1,2,3,4,5]) Q16: ["); i=0; while i<n { ln_num(of[i]); if i<n-1 { ln_puts(", ") } i=i+1 } ln_puts("]\n");
82 ln_puts(" output mean = "); ln_num(omean); ln_puts(" (want ~0) output variance = "); ln_num(ovar); ln_puts(" (want ~65536=1.0) determinism mism = "); ln_num(detmis); ln_puts("\n\n");
83
84 var pass: i64=0
85 var ttl: i64=0
86 ttl=ttl+1; ln_puts(" T1 isqrt EXACT on perfect squares (144->12, 1000000->1000): "); if s144==12 { if s1m==1000 { pass=pass+1; ln_puts("PASS\n") } else { ln_puts("FAIL\n") } } else { ln_puts("FAIL\n") }
87 ttl=ttl+1; ln_puts(" T2 rsqrt accurate (rsqrt(4)~0.5, rsqrt(0.25)~2.0): "); if absd(r4,32768)<=TOL { if absd(r025,131072)<=TOL { pass=pass+1; ln_puts("PASS\n") } else { ln_puts("FAIL\n") } } else { ln_puts("FAIL\n") }
88 ttl=ttl+1; ln_puts(" T3 LayerNorm normalizes: output mean ~0 AND variance ~1.0: "); if absd(omean,0)<=TOL { if absd(ovar,65536)<=TOL { pass=pass+1; ln_puts("PASS\n") } else { ln_puts("FAIL\n") } } else { ln_puts("FAIL\n") }
89 ttl=ttl+1; ln_puts(" T4 DETERMINISTIC: fwd-sum LayerNorm == rev-sum, bit-identical (float cannot): "); if detmis==0 { pass=pass+1; ln_puts("PASS\n") } else { ln_puts("FAIL\n") }
90
91 ln_puts("NX-NOFLOAT-LAYERNORM-GATE passed "); ln_num(pass); ln_puts("/"); ln_num(ttl)
92 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check
93 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled
94 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify.
95 let ctr__dry: *i64 = gv_ctr()
96 ctr__dry[0] = pass
97 ctr__dry[1] = ttl
98 let rc__dry: i64 = gv_verdict("NOFLOAT-LAYERNORM-GATE" as *u8, ctr__dry, "deterministic integer rsqrt+LayerNorm -- the no-float transformer primitive TRIO is complete)" as *u8)
99 sys_exit(rc__dry)
100 return rc__dry
101}