code wiki / _hdl_build / nx_tgrad_core.nx
nx_tgrad_core.nx source
↩ module page · 263 lines · 10461 B
1// nx_tgrad_core.nx -- the T6/T5 tensor-autograd CORE extracted VERBATIM from the GATED
2// _tensor_grad_authored.nx (T6 gate A-F PASS 2026-06-09) so higher rungs (T4 first model)
3// COMPOSE it instead of copy-pasting. The gated artifact is UNTOUCHED and stays the proof;
4// this file is the reusable library form: tape/arena, matvec/addvec/relu/mse/smul forward,
5// one-sweep backward, AdamW ad_step. license_tier: ORIGINAL (extraction, zero new logic).
6import "nx_syscalls.nx"
7import "nx_f32.nx"
8import "nx_f32_div.nx"
9import "nx_f32_cvt.nx"
10import "nx_f32_exp.nx"
11import "nx_f32_log.nx"
12const TG_MAGIC_16777216: i64 = 16777216
13func _tg_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
14func _tg_num(v: i64) -> i64 { let bb: *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;k=1}; while m>0{t[k]=48+(m%10);m=m/10;k=k+1}; var i: i64=0; while i<k{bb[i]=t[k-1-i];i=i+1}; sys_write(1,bb,k); return 0 }
15func tg_q(n: i64, d: i64) -> i64 { return nx_f32_div(nx_i32_to_f32(n), nx_i32_to_f32(d)) }
16func tg_milli(v: i64) -> i64 {
17 var neg: i64 = 0
18 var a: i64 = v
19 if nx_f32_lt(a, 0) == 1 { neg = 1; a = nx_f32_abs(a) }
20 let m: i64 = nx_f32_mul(a, nx_i32_to_f32(1000))
21 var lo: i64 = 0
22 var hi: i64 = TG_MAGIC_16777216
23 while lo < hi {
24 let mid: i64 = (lo + hi + 1) / 2
25 if nx_f32_lt(m, nx_i32_to_f32(mid)) == 1 { hi = mid - 1 } else { lo = mid }
26 }
27 if neg == 1 { return 0 - lo }
28 return lo
29}
30const TG_NCAP: i64 = 256
31const TG_ACAP: i64 = 16384
32func tg_alloc(arena: *i64, ab: *i64, n: i64) -> *i64 {
33 if ab[0] + n > TG_ACAP { _tg_puts("ARENA OVERFLOW (honest fail)\n" as *u8); sys_exit(1) }
34 let base: i64 = arena as i64
35 let q: *i64 = (base + ab[0] * 8) as *i64
36 var i: i64 = 0
37 while i < n { q[i] = 0; i = i + 1 }
38 ab[0] = ab[0] + n
39 return q
40}
41func tg_node(tape: *i64, nb: *i64, op: i64, ai: i64, bi: i64, rows: i64, cols: i64, valp: *i64, gradp: *i64) -> i64 {
42 if nb[0] >= TG_NCAP { _tg_puts("TAPE OVERFLOW (honest fail)\n" as *u8); sys_exit(1) }
43 let r: i64 = nb[0]
44 tape[r*8+0] = op
45 tape[r*8+1] = ai
46 tape[r*8+2] = bi
47 tape[r*8+3] = rows
48 tape[r*8+4] = cols
49 tape[r*8+5] = valp as i64
50 tape[r*8+6] = gradp as i64
51 tape[r*8+7] = 0
52 nb[0] = r + 1
53 return r
54}
55func tg_valp(tape: *i64, x: i64) -> *i64 { return tape[x*8+5] as *i64 }
56func tg_gradp(tape: *i64, x: i64) -> *i64 { return tape[x*8+6] as *i64 }
57func tg_leaf(tape: *i64, nb: *i64, arena: *i64, ab: *i64, p: *i64, rows: i64, cols: i64) -> i64 {
58 return tg_node(tape, nb, 0, 0, 0, rows, cols, p, tg_alloc(arena, ab, rows*cols))
59}
60func tg_matvec(tape: *i64, nb: *i64, arena: *i64, ab: *i64, a: i64, x: i64) -> i64 {
61 let r: i64 = tape[a*8+3]
62 let c: i64 = tape[a*8+4]
63 let av: *i64 = tg_valp(tape, a)
64 let xv: *i64 = tg_valp(tape, x)
65 let y: *i64 = tg_alloc(arena, ab, r)
66 var i: i64 = 0
67 while i < r {
68 var s: i64 = 0
69 var j: i64 = 0
70 while j < c { s = nx_f32_add(s, nx_f32_mul(av[i*c+j], xv[j])); j = j + 1 }
71 y[i] = s
72 i = i + 1
73 }
74 return tg_node(tape, nb, 1, a, x, r, 1, y, tg_alloc(arena, ab, r))
75}
76func tg_addvec(tape: *i64, nb: *i64, arena: *i64, ab: *i64, x: i64, y: i64) -> i64 {
77 let n: i64 = tape[x*8+3]
78 let xv: *i64 = tg_valp(tape, x)
79 let yv: *i64 = tg_valp(tape, y)
80 let o: *i64 = tg_alloc(arena, ab, n)
81 var i: i64 = 0
82 while i < n { o[i] = nx_f32_add(xv[i], yv[i]); i = i + 1 }
83 return tg_node(tape, nb, 2, x, y, n, 1, o, tg_alloc(arena, ab, n))
84}
85func tg_reluvec(tape: *i64, nb: *i64, arena: *i64, ab: *i64, x: i64) -> i64 {
86 let n: i64 = tape[x*8+3]
87 let xv: *i64 = tg_valp(tape, x)
88 let o: *i64 = tg_alloc(arena, ab, n)
89 var i: i64 = 0
90 while i < n {
91 var v: i64 = xv[i]
92 if nx_f32_lt(v, 0) == 1 { v = 0 }
93 o[i] = v
94 i = i + 1
95 }
96 return tg_node(tape, nb, 3, x, 0, n, 1, o, tg_alloc(arena, ab, n))
97}
98func tg_mse(tape: *i64, nb: *i64, arena: *i64, ab: *i64, p: i64, t: i64) -> i64 {
99 let n: i64 = tape[p*8+3]
100 let pv: *i64 = tg_valp(tape, p)
101 let tv: *i64 = tg_valp(tape, t)
102 var acc: i64 = 0
103 var i: i64 = 0
104 while i < n {
105 let d: i64 = nx_f32_sub(pv[i], tv[i])
106 acc = nx_f32_add(acc, nx_f32_mul(d, d))
107 i = i + 1
108 }
109 let o: *i64 = tg_alloc(arena, ab, 1)
110 o[0] = nx_f32_mul(acc, tg_q(1, n))
111 return tg_node(tape, nb, 4, p, t, 1, 1, o, tg_alloc(arena, ab, 1))
112}
113// op 6 = FUSED softmax + cross-entropy loss (T6-remaining rung 2026-06-10; retires the
114// one-hot-MSE workaround). p = logits node (n x 1), t = one-hot target leaf (n x 1).
115// forward: CE = ln(sum exp(l_i - max)) - sum t_i*(l_i - max) (max-shift for stability)
116// backward: dlogit_i += g * (softmax_i - t_i) -- the classic fused identity
117func tg_celoss(tape: *i64, nb: *i64, arena: *i64, ab: *i64, p: i64, t: i64) -> i64 {
118 let n: i64 = tape[p*8+3]
119 let pv: *i64 = tg_valp(tape, p)
120 let tv: *i64 = tg_valp(tape, t)
121 var mx: i64 = pv[0]
122 var i: i64 = 1
123 while i < n {
124 if nx_f32_gt(pv[i], mx) == 1 { mx = pv[i] }
125 i = i + 1
126 }
127 // exps computed ONCE here and the resulting softmax CACHED in node field 7 (was spare):
128 // backward reads it instead of recomputing 2n exps -- measured keystone for vocab-sized
129 // logits (320-wide CE x 1200 positions: minutes -> seconds). Same values, same op order.
130 let sm: *i64 = tg_alloc(arena, ab, n)
131 var s: i64 = 0
132 var dot: i64 = 0
133 i = 0
134 while i < n {
135 let sh: i64 = nx_f32_sub(pv[i], mx)
136 sm[i] = nx_f32_exp(sh)
137 s = nx_f32_add(s, sm[i])
138 dot = nx_f32_add(dot, nx_f32_mul(tv[i], sh))
139 i = i + 1
140 }
141 i = 0
142 while i < n { sm[i] = nx_f32_div(sm[i], s); i = i + 1 }
143 let o: *i64 = tg_alloc(arena, ab, 1)
144 o[0] = nx_f32_sub(nx_f32_log(s), dot)
145 let nd: i64 = tg_node(tape, nb, 6, p, t, 1, 1, o, tg_alloc(arena, ab, 1))
146 tape[nd*8+7] = sm as i64
147 return nd
148}
149func tg_smul(tape: *i64, nb: *i64, arena: *i64, ab: *i64, x: i64, y: i64) -> i64 {
150 let xv: *i64 = tg_valp(tape, x)
151 let yv: *i64 = tg_valp(tape, y)
152 let o: *i64 = tg_alloc(arena, ab, 1)
153 o[0] = nx_f32_mul(xv[0], yv[0])
154 return tg_node(tape, nb, 5, x, y, 1, 1, o, tg_alloc(arena, ab, 1))
155}
156func tg_backward(tape: *i64, nt: i64, root: i64) -> i64 {
157 var i: i64 = 0
158 while i < nt {
159 let gp: *i64 = tg_gradp(tape, i)
160 let n: i64 = tape[i*8+3] * tape[i*8+4]
161 var k: i64 = 0
162 while k < n { gp[k] = 0; k = k + 1 }
163 i = i + 1
164 }
165 let rg: *i64 = tg_gradp(tape, root)
166 rg[0] = nx_i32_to_f32(1)
167 var r: i64 = nt - 1
168 while r >= 0 {
169 let op: i64 = tape[r*8+0]
170 let ai: i64 = tape[r*8+1]
171 let bi: i64 = tape[r*8+2]
172 let g: *i64 = tg_gradp(tape, r)
173 if op == 1 {
174 let rr: i64 = tape[ai*8+3]
175 let cc: i64 = tape[ai*8+4]
176 let av: *i64 = tg_valp(tape, ai)
177 let ag: *i64 = tg_gradp(tape, ai)
178 let xv: *i64 = tg_valp(tape, bi)
179 let xg: *i64 = tg_gradp(tape, bi)
180 var i2: i64 = 0
181 while i2 < rr {
182 var j: i64 = 0
183 while j < cc {
184 ag[i2*cc+j] = nx_f32_add(ag[i2*cc+j], nx_f32_mul(g[i2], xv[j]))
185 xg[j] = nx_f32_add(xg[j], nx_f32_mul(av[i2*cc+j], g[i2]))
186 j = j + 1
187 }
188 i2 = i2 + 1
189 }
190 }
191 if op == 2 {
192 let n2: i64 = tape[r*8+3]
193 let ag2: *i64 = tg_gradp(tape, ai)
194 let bg2: *i64 = tg_gradp(tape, bi)
195 var k2: i64 = 0
196 while k2 < n2 { ag2[k2] = nx_f32_add(ag2[k2], g[k2]); bg2[k2] = nx_f32_add(bg2[k2], g[k2]); k2 = k2 + 1 }
197 }
198 if op == 3 {
199 let n3: i64 = tape[r*8+3]
200 let xv3: *i64 = tg_valp(tape, ai)
201 let ag3: *i64 = tg_gradp(tape, ai)
202 var k3: i64 = 0
203 while k3 < n3 {
204 if nx_f32_gt(xv3[k3], 0) == 1 { ag3[k3] = nx_f32_add(ag3[k3], g[k3]) }
205 k3 = k3 + 1
206 }
207 }
208 if op == 4 {
209 let n4: i64 = tape[ai*8+3]
210 let pv4: *i64 = tg_valp(tape, ai)
211 let tv4: *i64 = tg_valp(tape, bi)
212 let pg4: *i64 = tg_gradp(tape, ai)
213 let c2n: i64 = tg_q(2, n4)
214 var k4: i64 = 0
215 while k4 < n4 {
216 pg4[k4] = nx_f32_add(pg4[k4], nx_f32_mul(g[0], nx_f32_mul(c2n, nx_f32_sub(pv4[k4], tv4[k4]))))
217 k4 = k4 + 1
218 }
219 }
220 if op == 5 {
221 let av5: *i64 = tg_valp(tape, ai)
222 let bv5: *i64 = tg_valp(tape, bi)
223 let ag5: *i64 = tg_gradp(tape, ai)
224 let bg5: *i64 = tg_gradp(tape, bi)
225 ag5[0] = nx_f32_add(ag5[0], nx_f32_mul(g[0], bv5[0]))
226 bg5[0] = nx_f32_add(bg5[0], nx_f32_mul(g[0], av5[0]))
227 }
228 if op == 6 {
229 let n6: i64 = tape[ai*8+3]
230 let tv6: *i64 = tg_valp(tape, bi)
231 let pg6: *i64 = tg_gradp(tape, ai)
232 let smp: *i64 = tape[r*8+7] as *i64
233 var k6: i64 = 0
234 while k6 < n6 {
235 pg6[k6] = nx_f32_add(pg6[k6], nx_f32_mul(g[0], nx_f32_sub(smp[k6], tv6[k6])))
236 k6 = k6 + 1
237 }
238 }
239 r = r - 1
240 }
241 return 0
242}
243func ad_step(w: *i64, g: *i64, m: *i64, v: *i64, n: i64, lr: i64, b1: i64, b2: i64, eps: i64, wd: i64, t: i64) -> i64 {
244 var c1: i64 = nx_i32_to_f32(1)
245 var c2: i64 = nx_i32_to_f32(1)
246 var k: i64 = 0
247 while k < t { c1 = nx_f32_mul(c1, b1); c2 = nx_f32_mul(c2, b2); k = k + 1 }
248 let bc1: i64 = nx_f32_sub(nx_i32_to_f32(1), c1)
249 let bc2: i64 = nx_f32_sub(nx_i32_to_f32(1), c2)
250 let omb1: i64 = nx_f32_sub(nx_i32_to_f32(1), b1)
251 let omb2: i64 = nx_f32_sub(nx_i32_to_f32(1), b2)
252 var i: i64 = 0
253 while i < n {
254 m[i] = nx_f32_add(nx_f32_mul(b1, m[i]), nx_f32_mul(omb1, g[i]))
255 v[i] = nx_f32_add(nx_f32_mul(b2, v[i]), nx_f32_mul(omb2, nx_f32_mul(g[i], g[i])))
256 let mh: i64 = nx_f32_div(m[i], bc1)
257 let vh: i64 = nx_f32_div(v[i], bc2)
258 let upd: i64 = nx_f32_div(mh, nx_f32_add(nx_f32_sqrt(vh), eps))
259 w[i] = nx_f32_sub(w[i], nx_f32_add(nx_f32_mul(lr, upd), nx_f32_mul(nx_f32_mul(lr, wd), w[i])))
260 i = i + 1
261 }
262 return 0
263}