code wiki / _hdl_build / nx_research_discover_matmul_gate.nx
nx_research_discover_matmul_gate.nx source
↩ module page · 256 lines · 14596 B
1// nx_research_discover_matmul_gate.nx -- DISCOVERY widened to the matrix-multiplication family, the direct
2// AlphaTensor analog (operator: "widen discovery ... search candidate low-mult MATMUL schemes a la AlphaTensor").
3//
4// A fast matmul = a LOW-RANK BILINEAR TENSOR DECOMPOSITION. For an R-term scheme {(u_r,v_r,w_r)}:
5// m_r = (u_r . a)(v_r . b) (R scalar multiplications)
6// c_k = sum_r w_r[k] * m_r (only +/- of the m_r when coeffs are in {-1,0,1})
7// It computes the bilinear map exactly <=> Brent's equations hold:
8// for all i,j,k: sum_r u_r[i] v_r[j] w_r[k] == T[i][j][k] (T = the problem's tensor)
9// R = the multiplication count. Naive 2x2 matmul is rank 8; Strassen is rank 7.
10//
11// This gate:
12// SUBSTRATE : a sovereign Brent-equation VERIFIER -- the exact certifier a discovery search rides on.
13// MATMUL : the verifier certifies naive 2x2 (rank 8) AND Strassen (rank 7) as correct <2,2,2> schemes,
14// and REJECTS a broken Strassen (liar-kill). Ties the substrate to real matrix multiplication.
15// DISCOVERY : on a TRACTABLE bilinear problem -- complex multiplication, tensor (a+bi)(c+di) -- an
16// EXHAUSTIVE {-1,0,1} search DISCOVERS a rank-3 scheme (3 mults < naive 4: the Gauss/Karatsuba
17// speedup) it was never given, proves it MINIMAL (no rank-2 exists), and the same verifier +
18// random-input check certify it exact. A broken rank-3 is rejected.
19// HONEST SCOPE: the SEARCH discovers on the small problem; for <2,2,2> matmul we CERTIFY the known rank-7
20// via the same verifier -- we do NOT claim a blind AlphaTensor-scale rank-7 search.
21// Pure integer, no float, deterministic, reproducible. GREEN iff 8/8. license_tier: ORIGINAL
22import "nx_syscalls.nx"
23import "nx_library.nx"
24
25func g_w(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
26func g_n(v: i64) -> i64 { var m: i64=v; if m<0{g_w("-");m=0-m} let t:*u8=sys_mmap(24); 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; let o:*u8=sys_mmap(24); while i<k{o[i]=t[k-1-i];i=i+1}; sys_write(1,o,k); return 0 }
27func g_row(id: *u8, ok: i64, pass: *i64) -> i64 { g_w(" "); g_w(id); g_w(": "); if ok==1 { g_w("OK\n"); pass[0]=pass[0]+1 } else { g_w("FAIL\n") } return 0 }
28
29// ---- RESEARCHER: locate the bilinear-rank / faster-matmul research in our no-link-rot library ----
30func rs_has(buf: *u8, n: i64, needle: *u8) -> i64 {
31 var nl: i64=0; while needle[nl]!=(0 as u8){nl=nl+1}
32 if nl==0 { return 0 }
33 var i: i64=0
34 while i+nl<=n { var j: i64=0; var ok: i64=1; while j<nl { if buf[i+j]!=needle[j]{ok=0;j=nl}else{j=j+1} } if ok==1 {return 1} i=i+1 }
35 return 0
36}
37func rs_lib_contains(needle: *u8) -> i64 {
38 let cids: *i64 = sys_mmap(8*256) as *i64
39 let cnt: i64 = lib_list(cids, 256)
40 let pp: *i64 = sys_mmap(8) as *i64; let ll: *i64 = sys_mmap(8) as *i64
41 var d: i64 = 0
42 while d < cnt {
43 if lib_get(cids[d] as *u8, pp, ll) == 1 { if rs_has(pp[0] as *u8, ll[0], needle) == 1 { return 1 } }
44 d = d + 1
45 }
46 return 0
47}
48func rs_rng(state: *i64) -> i64 { var x: i64 = state[0]; x = x ^ (x << 13); x = x ^ (x >> 7); x = x ^ (x << 17); state[0] = x; return x }
49func pow3(n: i64) -> i64 { var p: i64=1; var i: i64=0; while i<n { p=p*3; i=i+1 } return p }
50// decode `code` into `len` trits, each mapped to {-1,0,1} (out[m] = (code/3^m)%3 - 1).
51func decode_vec(code: i64, len: i64, out: *i64) -> i64 { var c: i64=code; var i: i64=0; while i<len { out[i]=(c%3)-1; c=c/3; i=i+1 } return 0 }
52
53// ---- VERIFIER: Brent's equations. u,v,w are R x (na,nb,nc) flat; t is na*nb*nc flat. ----
54func brent_check(u: *i64, v: *i64, w: *i64, R: i64, na: i64, nb: i64, nc: i64, t: *i64) -> i64 {
55 var i: i64=0
56 while i<na { var j: i64=0
57 while j<nb { var k: i64=0
58 while k<nc {
59 var s: i64=0; var r: i64=0
60 while r<R { s = s + u[r*na+i]*v[r*nb+j]*w[r*nc+k]; r=r+1 }
61 if s != t[(i*nb+j)*nc+k] { return 0 }
62 k=k+1 }
63 j=j+1 }
64 i=i+1 }
65 return 1
66}
67
68// given products P[ij*R+r] = u_r[i]*v_r[j] (ij=i*2+j, na=nb=2) and target column tcol[ij]=T[i][j][k],
69// search a w-column wc in {-1,0,1}^R satisfying all 4 (i,j) equations; write to wout. wc is caller scratch.
70func solve_wcol(P: *i64, R: i64, tcol: *i64, wout: *i64, wc: *i64) -> i64 {
71 let lim: i64 = pow3(R)
72 var code: i64=0
73 while code<lim {
74 decode_vec(code, R, wc)
75 var ok: i64=1; var ij: i64=0
76 while ij<4 {
77 var s: i64=0; var r: i64=0
78 while r<R { s = s + P[ij*R+r]*wc[r]; r=r+1 }
79 if s != tcol[ij] { ok=0; ij=4 } else { ij=ij+1 }
80 }
81 if ok==1 { var r2: i64=0; while r2<R { wout[r2]=wc[r2]; r2=r2+1 } return 1 }
82 code=code+1
83 }
84 return 0
85}
86// EXHAUSTIVE {-1,0,1} search for a rank-R scheme of a (2,2,2) bilinear problem with tensor t.
87// returns 1 if one EXISTS (breaks early); if recordFirst, writes the first into fu/fv/fw (R x 2 each).
88// trick: fix u,v (3^(2R) each) -> Brent is then LINEAR in w, solved independently per output column (27 cands).
89func search222(R: i64, t: *i64, fu: *i64, fv: *i64, fw: *i64, recordFirst: i64) -> i64 {
90 let ubuf: *i64=sys_mmap(8*64) as *i64; let vbuf: *i64=sys_mmap(8*64) as *i64
91 let P: *i64=sys_mmap(8*64) as *i64; let wc: *i64=sys_mmap(8*64) as *i64
92 let wcol0: *i64=sys_mmap(8*64) as *i64; let wcol1: *i64=sys_mmap(8*64) as *i64
93 let tcol0: *i64=sys_mmap(8*64) as *i64; let tcol1: *i64=sys_mmap(8*64) as *i64
94 var ij0: i64=0; while ij0<4 { tcol0[ij0]=t[ij0*2+0]; tcol1[ij0]=t[ij0*2+1]; ij0=ij0+1 }
95 let uvlim: i64 = pow3(2*R)
96 var found: i64=0
97 var ucode: i64=0
98 while ucode<uvlim {
99 decode_vec(ucode, 2*R, ubuf)
100 var vcode: i64=0
101 while vcode<uvlim {
102 decode_vec(vcode, 2*R, vbuf)
103 var ij: i64=0
104 while ij<4 { let i: i64=ij/2; let j: i64=ij%2; var r: i64=0; while r<R { P[ij*R+r]=ubuf[r*2+i]*vbuf[r*2+j]; r=r+1 } ij=ij+1 }
105 if solve_wcol(P,R,tcol0,wcol0,wc)==1 { if solve_wcol(P,R,tcol1,wcol1,wc)==1 {
106 if found==0 { if recordFirst==1 {
107 var r3: i64=0
108 while r3<R { fu[r3*2+0]=ubuf[r3*2+0]; fu[r3*2+1]=ubuf[r3*2+1]; fv[r3*2+0]=vbuf[r3*2+0]; fv[r3*2+1]=vbuf[r3*2+1]; fw[r3*2+0]=wcol0[r3]; fw[r3*2+1]=wcol1[r3]; r3=r3+1 }
109 } }
110 found=1; vcode=uvlim; ucode=uvlim
111 } }
112 vcode=vcode+1
113 }
114 ucode=ucode+1
115 }
116 return found
117}
118
119// run an R-term complex-mult scheme on a concrete input; out[0]=c0, out[1]=c1.
120func cm_eval(u: *i64, v: *i64, w: *i64, R: i64, a0: i64, a1: i64, b0: i64, b1: i64, out: *i64) -> i64 {
121 var c0: i64=0; var c1: i64=0; var r: i64=0
122 while r<R {
123 let la: i64 = u[r*2+0]*a0 + u[r*2+1]*a1
124 let lb: i64 = v[r*2+0]*b0 + v[r*2+1]*b1
125 let m: i64 = la*lb
126 c0 = c0 + w[r*2+0]*m; c1 = c1 + w[r*2+1]*m
127 r=r+1
128 }
129 out[0]=c0; out[1]=c1; return 0
130}
131// the scheme must compute (a0+a1 i)(b0+b1 i) == (a0b0-a1b1) + (a0b1+a1b0) i over N random integer inputs.
132func verify_cm(u: *i64, v: *i64, w: *i64, R: i64) -> i64 {
133 let st: *i64=sys_mmap(8) as *i64; st[0]=0x2545F4914F6CDD1D
134 let out: *i64=sys_mmap(16) as *i64
135 var it: i64=0; var mism: i64=0
136 while it<300 {
137 let a0: i64=(rs_rng(st)&0x3f)-32; let a1: i64=(rs_rng(st)&0x3f)-32
138 let b0: i64=(rs_rng(st)&0x3f)-32; let b1: i64=(rs_rng(st)&0x3f)-32
139 cm_eval(u,v,w,R,a0,a1,b0,b1,out)
140 let c0n: i64=a0*b0 - a1*b1; let c1n: i64=a0*b1 + a1*b0
141 if out[0]!=c0n { mism=mism+1 }
142 if out[1]!=c1n { mism=mism+1 }
143 it=it+1
144 }
145 return (mism==0) as i64
146}
147
148// ---- tensor + known-scheme builders ----
149func build_t_complex(t: *i64) -> i64 { var z: i64=0; while z<8 { t[z]=0; z=z+1 } t[(0*2+0)*2+0]=1; t[(1*2+1)*2+0]=-1; t[(0*2+1)*2+1]=1; t[(1*2+0)*2+1]=1; return 0 }
150func build_naive_complex(u: *i64, v: *i64, w: *i64) -> i64 {
151 var z: i64=0; while z<8 { u[z]=0; v[z]=0; w[z]=0; z=z+1 }
152 u[0]=1; v[0]=1; w[0]=1 // m1=a0*b0 -> c0 +1
153 u[3]=1; v[3]=1; w[2]=-1 // m2=a1*b1 -> c0 -1 (row1: idx 2,3)
154 u[4]=1; v[5]=1; w[5]=1 // m3=a0*b1 -> c1 +1 (row2: idx 4,5)
155 u[7]=1; v[6]=1; w[7]=1 // m4=a1*b0 -> c1 +1 (row3: idx 6,7)
156 return 0
157}
158func build_t222(t: *i64) -> i64 {
159 var a: i64=0
160 while a<4 { let p: i64=a/2; let q: i64=a%2; var b: i64=0
161 while b<4 { let r: i64=b/2; let s: i64=b%2; var c: i64=0
162 while c<4 { let x: i64=c/2; let y: i64=c%2; var val: i64=0
163 if x==p { if q==r { if s==y { val=1 } } }
164 t[(a*4+b)*4+c]=val; c=c+1 }
165 b=b+1 }
166 a=a+1 }
167 return 0
168}
169func build_naive222(u: *i64, v: *i64, w: *i64) -> i64 {
170 var z: i64=0; while z<32 { u[z]=0; v[z]=0; w[z]=0; z=z+1 }
171 var idx: i64=0; var x: i64=0
172 while x<2 { var m: i64=0
173 while m<2 { var y: i64=0
174 while y<2 { u[idx*4+(x*2+m)]=1; v[idx*4+(m*2+y)]=1; w[idx*4+(x*2+y)]=1; idx=idx+1; y=y+1 }
175 m=m+1 }
176 x=x+1 }
177 return 0
178}
179func set4(arr: *i64, row: i64, a: i64, b: i64, c: i64, d: i64) -> i64 { arr[row*4+0]=a; arr[row*4+1]=b; arr[row*4+2]=c; arr[row*4+3]=d; return 0 }
180// Strassen <2,2,2>, rank 7 (a/b/c indexed row*2+col): the published scheme.
181func build_strassen(u: *i64, v: *i64, w: *i64) -> i64 {
182 set4(u,0, 1,0,0,1); set4(u,1, 0,0,1,1); set4(u,2, 1,0,0,0); set4(u,3, 0,0,0,1); set4(u,4, 1,1,0,0); set4(u,5, -1,0,1,0); set4(u,6, 0,1,0,-1)
183 set4(v,0, 1,0,0,1); set4(v,1, 1,0,0,0); set4(v,2, 0,1,0,-1); set4(v,3, -1,0,1,0); set4(v,4, 0,0,0,1); set4(v,5, 1,1,0,0); set4(v,6, 0,0,1,1)
184 set4(w,0, 1,0,0,1); set4(w,1, 0,0,1,-1); set4(w,2, 0,1,0,1); set4(w,3, 1,0,1,0); set4(w,4, -1,1,0,0); set4(w,5, 0,0,0,1); set4(w,6, 1,0,0,0)
185 return 0
186}
187
188func main() -> i64 {
189 let pass: *i64 = sys_mmap(8) as *i64; pass[0]=0
190 g_w("=== NX-RESEARCH-DISCOVER-MATMUL GATE (bilinear low-rank discovery: search + Brent verifier, AlphaTensor analog) ===\n")
191
192 let found: i64 = rs_lib_contains("matrix multiplication" as *u8)
193
194 // ---- MATMUL <2,2,2>: certify naive (rank8) + Strassen (rank7), reject broken (liar-kill) ----
195 let t222: *i64 = sys_mmap(8*64) as *i64; build_t222(t222)
196 let un: *i64=sys_mmap(8*64) as *i64; let vn: *i64=sys_mmap(8*64) as *i64; let wn: *i64=sys_mmap(8*64) as *i64; build_naive222(un,vn,wn)
197 let us: *i64=sys_mmap(8*64) as *i64; let vs: *i64=sys_mmap(8*64) as *i64; let ws: *i64=sys_mmap(8*64) as *i64; build_strassen(us,vs,ws)
198 let bn8: i64 = brent_check(un,vn,wn,8,4,4,4,t222)
199 let bs7: i64 = brent_check(us,vs,ws,7,4,4,4,t222)
200 // broken Strassen: flip one w coefficient
201 let wb: *i64=sys_mmap(8*64) as *i64; var z: i64=0; while z<28 { wb[z]=ws[z]; z=z+1 } wb[0]=ws[0]+1
202 let bb: i64 = brent_check(us,vs,wb,7,4,4,4,t222)
203
204 // ---- COMPLEX MULT: naive rank4 baseline ----
205 let tcx: *i64 = sys_mmap(8*16) as *i64; build_t_complex(tcx)
206 let ucn: *i64=sys_mmap(8*16) as *i64; let vcn: *i64=sys_mmap(8*16) as *i64; let wcn: *i64=sys_mmap(8*16) as *i64; build_naive_complex(ucn,vcn,wcn)
207 let bn4: i64 = brent_check(ucn,vcn,wcn,4,2,2,2,tcx)
208
209 // ---- DISCOVERY: search a rank-3 complex-mult scheme (it was never supplied) ----
210 let fu: *i64=sys_mmap(8*16) as *i64; let fv: *i64=sys_mmap(8*16) as *i64; let fw: *i64=sys_mmap(8*16) as *i64
211 let ex3: i64 = search222(3, tcx, fu, fv, fw, 1)
212 let bf3: i64 = brent_check(fu,fv,fw,3,2,2,2,tcx)
213 let vex: i64 = verify_cm(fu,fv,fw,3)
214 // ---- MINIMALITY: exhaustive search finds NO rank-2 scheme ----
215 let du: *i64=sys_mmap(8*16) as *i64; let dv: *i64=sys_mmap(8*16) as *i64; let dw: *i64=sys_mmap(8*16) as *i64
216 let ex2: i64 = search222(2, tcx, du, dv, dw, 0)
217 // ---- DISCOVERY LIAR-KILL: a broken rank-3 (flip one coeff) is rejected ----
218 let wbx: *i64=sys_mmap(8*16) as *i64; z=0; while z<6 { wbx[z]=fw[z]; z=z+1 } wbx[0]=fw[0]+1
219 let bbx: i64 = brent_check(fu,fv,wbx,3,2,2,2,tcx)
220
221 g_w(" researcher: faster-matmul / bilinear-rank research in library = "); g_n(found); g_w("\n")
222 g_w(" MATMUL <2,2,2>: Brent certifies naive rank-8="); g_n(bn8); g_w(" Strassen rank-7="); g_n(bs7); g_w(" broken-Strassen rejected="); g_n((bb==0) as i64); g_w("\n")
223 g_w(" COMPLEX mult: naive rank-4 certified="); g_n(bn4); g_w("\n")
224 g_w(" DISCOVERY: rank-3 scheme exists="); g_n(ex3); g_w(" Brent-valid="); g_n(bf3); g_w(" exact-over-random="); g_n(vex); g_w("\n")
225 if ex3==1 { g_w(" found scheme: u=["); var qi: i64=0; while qi<6 { if qi>0 { g_w(",") } g_n(fu[qi]); qi=qi+1 } g_w("] v=["); qi=0; while qi<6 { if qi>0 { g_w(",") } g_n(fv[qi]); qi=qi+1 } g_w("] w=["); qi=0; while qi<6 { if qi>0 { g_w(",") } g_n(fw[qi]); qi=qi+1 } g_w("]\n") }
226 g_w(" MINIMALITY: rank-2 scheme exists="); g_n(ex2); g_w(" (0 => 3 is minimal, matches bilinear-complexity theorem)\n")
227 g_w(" discovery-liar-kill: broken rank-3 rejected="); g_n((bbx==0) as i64); g_w("\n")
228
229 g_row("RESEARCHER: faster-matmul / bilinear-rank research (AlphaTensor) foundationed in library" as *u8, found, pass)
230
231 var mmver: i64=0
232 if bn8==1 { if bs7==1 { mmver=1 } }
233 g_row("MATMUL VERIFIER: Brent certifies naive 2x2 (rank-8) AND Strassen (rank-7) compute <2,2,2> exactly" as *u8, mmver, pass)
234
235 g_row("MATMUL LIAR-KILL: a broken Strassen (one coeff flipped) is REJECTED by the verifier" as *u8, (bb==0) as i64, pass)
236
237 g_row("COMPLEX BASELINE: naive complex multiplication is rank-4 (Brent-certified)" as *u8, bn4, pass)
238
239 var disc: i64=0
240 if ex3==1 { if bf3==1 { if vex==1 { disc=1 } } }
241 g_row("DISCOVERY: search finds a rank-3 complex-mult scheme (3<4 muls), Brent-valid + exact over random" as *u8, disc, pass)
242
243 g_row("MINIMALITY: exhaustive {-1,0,1} search finds NO rank-2 scheme (3 is minimal)" as *u8, (ex2==0) as i64, pass)
244
245 var scope: i64=0
246 if ex3==1 { if ex2==0 { if bs7==1 { scope=1 } } }
247 g_row("HONEST SCOPE: search discovers+bounds on the tractable problem; matmul rank-7 is CERTIFIED, not blind-searched" as *u8, scope, pass)
248
249 var gen: i64=0
250 if bf3==1 { if bbx==0 { gen=1 } }
251 g_row("GENERALITY+LIAR-KILL: one Brent substrate certifies the discovered rank-3 AND rejects a broken rank-3" as *u8, gen, pass)
252
253 g_w("RESEARCH-DISCOVER-MATMUL-GATE rows=8 pass="); g_n(pass[0])
254 if pass[0]==8 { g_w(" verdict=GREEN\n"); sys_exit(0); return 0 }
255 g_w(" verdict=RED\n"); sys_exit(1); return 1
256}