nx_autograd.nx source
↩ module page · 104 lines · 4881 B
1// nx_autograd.nx -- TRAIN-R1: scalar tape-based reverse-mode autograd. THE training keystone: once gradients
2// are correct and an optimizer drives a loss down, every model the team trains (the FNet text model, the
3// pixel engine) is the SAME loop at larger scale. Inference is already sovereign; this is the missing BACKWARD
4// half. Spec: knowledge/specs/2026-06-09-tutoring-training-substrate-rung1.md.
5//
6// TAPE: stride-5 nodes {op, ai, bi, val, grad} in one i64 array. Construction is EAGER -- creating a node
7// computes its f32 `val` via the EXISTING nx_f32 ops (reuse, never reimplement). Backward = ONE reverse sweep
8// (nodes are built in topological order, so reverse order visits parents before children) accumulating into
9// `grad`: add: g->both | mul: ga+=g*vb, gb+=g*va | sub: ga+=g, gb+=-g | relu: pass g iff input val>0.
10// Seed grad[root]=1.0. The whole engine is ~50 lines BECAUSE the f32 numeric tower already exists.
11//
12// genealogy_id: linnainmaa_1970_reverse_mode_ad + rumelhart_1986_backprop
13// lineage_id: sovereign_scalar_tape_autograd_v1
14// license_tier: ORIGINAL verdict: GREEN (nx_train_r1_gate 2026-06-14: gradcheck + model-learns + bit-exact)
15import "nx_f32.nx" // IEEE-754 binary32 add/sub/mul/neg/abs/lt/gt -- the eager forward arithmetic
16import "nx_f32_div.nx" // nx_f32_div -- used by callers building rational constants / MSE
17import "nx_f32_cvt.nx" // nx_i32_to_f32 -- integer -> f32 constant constructor
18import "nx_syscalls.nx" // sys_mmap
19
20const AG_LEAF: i64 = 0
21const AG_ADD: i64 = 1
22const AG_MUL: i64 = 2
23const AG_SUB: i64 = 3
24const AG_RELU: i64 = 4
25const AG_F32_ZERO: i64 = 0 // IEEE-754 binary32 +0.0
26const AG_F32_ONE: i64 = 1065353216 // IEEE-754 binary32 1.0 = 0x3F800000
27
28func ag_val(tape: *i64, k: i64) -> i64 { return tape[5 * k + 3] }
29func ag_grad(tape: *i64, k: i64) -> i64 { return tape[5 * k + 4] }
30
31// append a node; *np is the next free index. returns the new node's index.
32func ag_leaf(tape: *i64, np: *i64, v: i64) -> i64 {
33 let k: i64 = *np
34 tape[5 * k + 0] = AG_LEAF; tape[5 * k + 1] = 0 - 1; tape[5 * k + 2] = 0 - 1
35 tape[5 * k + 3] = v; tape[5 * k + 4] = AG_F32_ZERO
36 *np = k + 1; return k
37}
38func ag_mk(tape: *i64, np: *i64, op: i64, a: i64, b: i64, v: i64) -> i64 {
39 let k: i64 = *np
40 tape[5 * k + 0] = op; tape[5 * k + 1] = a; tape[5 * k + 2] = b
41 tape[5 * k + 3] = v; tape[5 * k + 4] = AG_F32_ZERO
42 *np = k + 1; return k
43}
44func ag_add(tape: *i64, np: *i64, a: i64, b: i64) -> i64 {
45 return ag_mk(tape, np, AG_ADD, a, b, nx_f32_add(tape[5 * a + 3], tape[5 * b + 3]))
46}
47func ag_mul(tape: *i64, np: *i64, a: i64, b: i64) -> i64 {
48 return ag_mk(tape, np, AG_MUL, a, b, nx_f32_mul(tape[5 * a + 3], tape[5 * b + 3]))
49}
50func ag_sub(tape: *i64, np: *i64, a: i64, b: i64) -> i64 {
51 return ag_mk(tape, np, AG_SUB, a, b, nx_f32_sub(tape[5 * a + 3], tape[5 * b + 3]))
52}
53func ag_relu(tape: *i64, np: *i64, a: i64) -> i64 {
54 var v: i64 = tape[5 * a + 3]
55 if nx_f32_gt(v, AG_F32_ZERO) != 1 { v = AG_F32_ZERO }
56 return ag_mk(tape, np, AG_RELU, a, 0 - 1, v)
57}
58
59// reverse-mode sweep: zero grads, seed grad[root]=1, accumulate backward through the tape.
60func ag_backward(tape: *i64, n: i64, root: i64) -> i64 {
61 var i: i64 = 0
62 while i < n { tape[5 * i + 4] = AG_F32_ZERO; i = i + 1 }
63 tape[5 * root + 4] = AG_F32_ONE
64 var k: i64 = n - 1
65 while k >= 0 {
66 let op: i64 = tape[5 * k + 0]
67 let a: i64 = tape[5 * k + 1]
68 let b: i64 = tape[5 * k + 2]
69 let g: i64 = tape[5 * k + 4]
70 if op == AG_ADD {
71 tape[5 * a + 4] = nx_f32_add(tape[5 * a + 4], g)
72 tape[5 * b + 4] = nx_f32_add(tape[5 * b + 4], g)
73 }
74 if op == AG_MUL {
75 tape[5 * a + 4] = nx_f32_add(tape[5 * a + 4], nx_f32_mul(g, tape[5 * b + 3]))
76 tape[5 * b + 4] = nx_f32_add(tape[5 * b + 4], nx_f32_mul(g, tape[5 * a + 3]))
77 }
78 if op == AG_SUB {
79 tape[5 * a + 4] = nx_f32_add(tape[5 * a + 4], g)
80 tape[5 * b + 4] = nx_f32_add(tape[5 * b + 4], nx_f32_neg(g))
81 }
82 if op == AG_RELU {
83 if nx_f32_gt(tape[5 * a + 3], AG_F32_ZERO) == 1 {
84 tape[5 * a + 4] = nx_f32_add(tape[5 * a + 4], g)
85 }
86 }
87 k = k - 1
88 }
89 return 0
90}
91
92// rational f32 constant num/den (e.g. ag_constf(3,2) = 1.5f). Reuses the i32->f32 converter + f32 divide.
93func ag_constf(num: i64, den: i64) -> i64 { return nx_f32_div(nx_i32_to_f32(num), nx_i32_to_f32(den)) }
94
95// compile-only smoke: d/dx(x*x) at x=3 == 6. Real validation in nx_train_r1_gate.nx.
96func main() -> i64 {
97 let tape: *i64 = (sys_mmap(64 * 5 * 8)) as *i64
98 let np: *i64 = (sys_mmap(8)) as *i64
99 *np = 0
100 let x: i64 = ag_leaf(tape, np, nx_i32_to_f32(3))
101 let y: i64 = ag_mul(tape, np, x, x)
102 ag_backward(tape, *np, y)
103 return 0
104}