nx_compute_graph_test.nx source
↩ module page · 220 lines · 9151 B
1// nx_compute_graph_test.nx -- typed DAG end-to-end shape verification.
2//
3// Constructs a small computation graph that resembles a tiny ML
4// forward pass: two inputs feed a matmul; the matmul feeds a ReLU;
5// the ReLU is the output. Validates that:
6// * structural validation accepts the well-formed graph
7// * toposort emits a valid linearisation
8// * cycle detection refuses a graph with a back-edge
9// * dangling input detection refuses a graph with unwired slot
10// * bad-edge detection refuses an edge to a non-existent node
11// * content hashes are stable across builds with identical shape
12
13import "nx_syscalls.nx"
14import "nx_tier.nx"
15import "nx_sha256.nx"
16import "nx_compute_node.nx"
17import "nx_compute_graph.nx"
18
19func main() -> nx_int {
20 // ===== Build the happy-path graph =============================
21 //
22 // Nodes:
23 // 0 INPUT -- text embedding tensor (1 output)
24 // 1 CONST -- weight matrix (1 output)
25 // 2 KERNEL MATMUL (2 inputs <- 0 + 1, 1 output)
26 // 3 TENSOR_OP RELU (1 input <- 2, 1 output)
27 // 4 OUTPUT (1 input <- 3, 0 outputs)
28
29 let g: *ComputeGraph = nx_cg_alloc(16)
30
31 let n0: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_INPUT, 0, 0, 1, 0)
32 nx_cg_add_node(g, n0)
33
34 let n1: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_CONST, 0, 0, 1, 0)
35 nx_cg_add_node(g, n1)
36
37 let n2: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_KERNEL,
38 NX_CN_K_MATMUL, 2, 1, 0)
39 nx_cn_set_input(n2, 0, 0, 0) // input 0 <- node 0, port 0
40 nx_cn_set_input(n2, 1, 1, 0) // input 1 <- node 1, port 0
41 nx_cg_add_node(g, n2)
42
43 let n3: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_TENSOR_OP,
44 NX_CN_OP_RELU, 1, 1, 0)
45 nx_cn_set_input(n3, 0, 2, 0) // <- node 2
46 nx_cg_add_node(g, n3)
47
48 let n4: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_OUTPUT, 0, 1, 0, 0)
49 nx_cn_set_input(n4, 0, 3, 0)
50 nx_cg_add_node(g, n4)
51
52 if g.n_nodes != 5 { return 1 }
53
54 // ===== Validate ==============================================
55 let v: nx_int = nx_cg_validate(g)
56 if v != NX_CG_OK { return 2 }
57
58 // ===== Toposort ==============================================
59 let t: nx_int = nx_cg_toposort(g)
60 if t != NX_CG_OK { return 3 }
61 if g.has_topo != 1 { return 4 }
62
63 // Valid topo: 0 and 1 come before 2; 2 before 3; 3 before 4.
64 // Build position-of array so we can check the partial order.
65 let pos: *i64 = (sys_mmap(64)) as *i64
66 var i: nx_int = 0
67 while i < g.n_nodes {
68 pos[nx_cg_topo_at(g, i)] = i
69 i = i + 1
70 }
71 if pos[0] >= pos[2] { return 10 } // 0 must come before 2
72 if pos[1] >= pos[2] { return 11 } // 1 must come before 2
73 if pos[2] >= pos[3] { return 12 } // 2 must come before 3
74 if pos[3] >= pos[4] { return 13 } // 3 must come before 4
75
76 // ===== Compute content hashes ================================
77 nx_cn_compute_hash(n0)
78 nx_cn_compute_hash(n1)
79 nx_cn_compute_hash(n2)
80 nx_cn_compute_hash(n3)
81 nx_cn_compute_hash(n4)
82
83 // Hashes should be non-zero (sanity: SHA256 of any input is essentially
84 // never the zero-byte string)
85 var any_nonzero: nx_int = 0
86 var b: nx_int = 0
87 while b < NX_CN_HASH_BYTES {
88 if n2.content_hash[b] != 0 { any_nonzero = 1 }
89 b = b + 1
90 }
91 if any_nonzero != 1 { return 20 }
92
93 // Two identical kernel nodes (same op, same inputs, same params)
94 // should produce identical content hashes.
95 let m2_dup: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_KERNEL,
96 NX_CN_K_MATMUL, 2, 1, 0)
97 nx_cn_set_input(m2_dup, 0, 0, 0)
98 nx_cn_set_input(m2_dup, 1, 1, 0)
99 nx_cn_compute_hash(m2_dup)
100
101 var bb: nx_int = 0
102 while bb < NX_CN_HASH_BYTES {
103 if m2_dup.content_hash[bb] != n2.content_hash[bb] { return 21 }
104 bb = bb + 1
105 }
106
107 // Different params -> different hash
108 let m2_diff_params: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_KERNEL,
109 NX_CN_K_MATMUL, 2, 1, 1)
110 nx_cn_set_input(m2_diff_params, 0, 0, 0)
111 nx_cn_set_input(m2_diff_params, 1, 1, 0)
112 nx_cn_set_param(m2_diff_params, 0, 42)
113 nx_cn_compute_hash(m2_diff_params)
114
115 var differ: nx_int = 0
116 var bc: nx_int = 0
117 while bc < NX_CN_HASH_BYTES {
118 if m2_diff_params.content_hash[bc] != n2.content_hash[bc] { differ = 1 }
119 bc = bc + 1
120 }
121 if differ != 1 { return 22 }
122
123 // ===== Root hash works after toposort ========================
124 let root_hash: *u8 = sys_mmap(NX_CN_HASH_BYTES)
125 let rh: nx_int = nx_cg_root_hash(g, root_hash)
126 if rh != 0 { return 30 }
127
128 // Build an IDENTICAL second graph and verify root hash matches
129 let g2: *ComputeGraph = nx_cg_alloc(16)
130 let g2n0: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_INPUT, 0, 0, 1, 0)
131 nx_cg_add_node(g2, g2n0)
132 let g2n1: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_CONST, 0, 0, 1, 0)
133 nx_cg_add_node(g2, g2n1)
134 let g2n2: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_KERNEL,
135 NX_CN_K_MATMUL, 2, 1, 0)
136 nx_cn_set_input(g2n2, 0, 0, 0)
137 nx_cn_set_input(g2n2, 1, 1, 0)
138 nx_cg_add_node(g2, g2n2)
139 let g2n3: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_TENSOR_OP,
140 NX_CN_OP_RELU, 1, 1, 0)
141 nx_cn_set_input(g2n3, 0, 2, 0)
142 nx_cg_add_node(g2, g2n3)
143 let g2n4: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_OUTPUT, 0, 1, 0, 0)
144 nx_cn_set_input(g2n4, 0, 3, 0)
145 nx_cg_add_node(g2, g2n4)
146 nx_cg_validate(g2)
147 nx_cg_toposort(g2)
148 nx_cn_compute_hash(g2n0); nx_cn_compute_hash(g2n1); nx_cn_compute_hash(g2n2)
149 nx_cn_compute_hash(g2n3); nx_cn_compute_hash(g2n4)
150 let root_hash_2: *u8 = sys_mmap(NX_CN_HASH_BYTES)
151 nx_cg_root_hash(g2, root_hash_2)
152 var rb: nx_int = 0
153 while rb < NX_CN_HASH_BYTES {
154 if root_hash[rb] != root_hash_2[rb] { return 31 }
155 rb = rb + 1
156 }
157
158 // ===== Bad edge: source node id beyond n_nodes ==============
159 let g3: *ComputeGraph = nx_cg_alloc(8)
160 let g3n0: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_INPUT, 0, 0, 1, 0)
161 nx_cg_add_node(g3, g3n0)
162 let g3n1: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_TENSOR_OP,
163 NX_CN_OP_RELU, 1, 1, 0)
164 nx_cn_set_input(g3n1, 0, 5, 0) // node 5 doesn't exist
165 nx_cg_add_node(g3, g3n1)
166 if nx_cg_validate(g3) != NX_CG_ERR_BAD_EDGE { return 40 }
167
168 // ===== Dangling input (slot unwired) =========================
169 let g4: *ComputeGraph = nx_cg_alloc(8)
170 let g4n0: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_TENSOR_OP,
171 NX_CN_OP_RELU, 1, 1, 0)
172 // intentionally don't set_input -> remains -1
173 nx_cg_add_node(g4, g4n0)
174 if nx_cg_validate(g4) != NX_CG_ERR_DANGLING_INPUT { return 41 }
175
176 // ===== Bad port (src_port beyond src.n_outputs) =============
177 let g5: *ComputeGraph = nx_cg_alloc(8)
178 let g5n0: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_INPUT, 0, 0, 1, 0)
179 nx_cg_add_node(g5, g5n0)
180 let g5n1: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_TENSOR_OP,
181 NX_CN_OP_RELU, 1, 1, 0)
182 nx_cn_set_input(g5n1, 0, 0, 5) // node 0 only has 1 output (port 0)
183 nx_cg_add_node(g5, g5n1)
184 if nx_cg_validate(g5) != NX_CG_ERR_BAD_PORT { return 42 }
185
186 // ===== Cycle detection ======================================
187 //
188 // Build a 2-node cycle: A depends on B, B depends on A.
189 let g6: *ComputeGraph = nx_cg_alloc(8)
190 let g6n0: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_TENSOR_OP,
191 NX_CN_OP_RELU, 1, 1, 0)
192 let g6n1: *ComputeNode = nx_cn_alloc(0, NX_CN_NODE_TENSOR_OP,
193 NX_CN_OP_RELU, 1, 1, 0)
194 nx_cg_add_node(g6, g6n0) // becomes id 0
195 nx_cg_add_node(g6, g6n1) // becomes id 1
196 nx_cn_set_input(g6n0, 0, 1, 0)
197 nx_cn_set_input(g6n1, 0, 0, 0)
198 // Validation passes (every edge points to existing node)
199 if nx_cg_validate(g6) != NX_CG_OK { return 50 }
200 // But toposort refuses
201 if nx_cg_toposort(g6) != NX_CG_ERR_CYCLE { return 51 }
202 if g6.last_verdict != NX_CG_ERR_CYCLE { return 52 }
203
204 // ===== Verdict enum coverage ================================
205 if nx_cg_verdict_is_valid(NX_CG_OK) != 1 { return 60 }
206 if nx_cg_verdict_is_valid(NX_CG_ERR_CYCLE) != 1 { return 61 }
207 if nx_cg_verdict_is_valid(NX_CG_ERR_NODE_LIMIT) != 1 { return 62 }
208 if nx_cg_verdict_is_valid(NX_CG_N_VERDICTS) != 0 { return 63 }
209 if nx_cg_verdict_is_valid(0 - 1) != 0 { return 64 }
210
211 // Sealed node + op + kernel enums
212 if nx_cn_node_kind_is_valid(NX_CN_NODE_KERNEL) != 1 { return 70 }
213 if nx_cn_node_kind_is_valid(NX_CN_NODE_N_KINDS) != 0 { return 71 }
214 if nx_cn_op_is_valid(NX_CN_OP_RELU) != 1 { return 72 }
215 if nx_cn_op_is_valid(NX_CN_OP_N_OPS) != 0 { return 73 }
216 if nx_cn_kernel_is_valid(NX_CN_K_MATMUL) != 1 { return 74 }
217 if nx_cn_kernel_is_valid(NX_CN_K_N_KINDS) != 0 { return 75 }
218
219 return 0
220}