nx_compute_runner.nx source
↩ module page · 353 lines · 12741 B
1// nx_compute_runner.nx -- executes a validated ComputeGraph.
2//
3// The brick that closes the loop: with this in tree, the substrate
4// has a working ComfyUI replacement at the compute layer. Build a
5// graph -> validate -> toposort -> RUN -> get real tensor output ->
6// verify via nx_numeric_oracle. All pure NishiLang.
7//
8// Dispatch:
9// NODE_INPUT -- copies caller-supplied input tensor into the
10// tensor store at the node's output slot
11// NODE_CONST -- copies caller-supplied const tensor (weight
12// blob) into the store
13// NODE_OUTPUT -- alias-copies its input tensor into the store
14// so external readers can collect graph outputs
15// NODE_TENSOR_OP -- dispatches by op_code to elementwise / shape
16// ops implemented inline
17// NODE_KERNEL -- dispatches by kernel_kind to nx_blas_i64 (and
18// future nx_conv / nx_attention modules)
19//
20// Tensor lifecycle:
21// * One tensor per output port of each node.
22// * Indexed by tensor_idx = node_id * NX_CN_MAX_OUTPUTS + port.
23// * Allocated by the runner; freed by sys_mmap unwind at process
24// exit (no GC in v1; tight footprint kept by careful sizing).
25//
26// Audit:
27// Caller passes an optional *TraceCfg. If non-null, the runner
28// emits one CONTINUITY_CHECK-style span per node (kind = generic
29// "node executed"). Span attrs carry node_id + op_code + verdict.
30//
31// genealogy_id: comfyui_executor + jax_pjit + tvm_runtime +
32// tensorflow_executor
33// lineage_id: substrate_compute_runner_v1
34
35// nx_safety_envelope:
36// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
37// sil_target: SIL1
38// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
39// verdict: NOT_YET_EVALUATED
40
41import "nx_syscalls.nx"
42import "nx_tier.nx"
43import "nx_tensor.nx"
44import "nx_compute_node.nx"
45import "nx_compute_graph.nx"
46import "nx_blas_i64.nx"
47
48// ===== Sealed-enum: RunnerVerdict =================================
49
50const NX_CR_OK: nx_int = 0
51const NX_CR_ERR_GRAPH_INVALID: nx_int = 1 // validate / toposort hadn't passed
52const NX_CR_ERR_INPUT_MISSING: nx_int = 2 // INPUT node has no binding
53const NX_CR_ERR_CONST_MISSING: nx_int = 3 // CONST node has no binding
54const NX_CR_ERR_UNSUPPORTED_OP: nx_int = 4 // op_code not yet implemented
55const NX_CR_ERR_KERNEL_FAILED: nx_int = 5 // kernel returned non-OK verdict
56const NX_CR_ERR_SHAPE_INFER: nx_int = 6 // couldn't determine output shape
57const NX_CR_N_VERDICTS: nx_int = 7
58
59func nx_cr_verdict_is_valid(v: nx_int) -> nx_int {
60 if v < 0 { return 0 }
61 if v >= NX_CR_N_VERDICTS { return 0 }
62 return 1
63}
64
65// ===== Runner state ===============================================
66
67struct ComputeRunner {
68 graph: *ComputeGraph,
69 // Tensor store: one slot per (node_id, port).
70 // tensor_store[node_id * NX_CN_MAX_OUTPUTS + port] = *NxTensor address.
71 tensor_store: *i64,
72 tensor_store_cap: nx_int,
73 // Input bindings: caller pre-sets these via nx_cr_bind_input.
74 // input_bindings[i] = *NxTensor address for the i-th INPUT node.
75 input_bindings: *i64,
76 n_inputs: nx_int,
77 cap_inputs: nx_int,
78 // Const bindings: same shape for CONST nodes.
79 const_bindings: *i64,
80 n_consts: nx_int,
81 cap_consts: nx_int,
82 last_verdict: nx_int,
83 last_failed_node: nx_int, // for debugging on KERNEL_FAILED
84}
85
86const NX_CR_BYTES: nx_int = 88 // 11 fields * 8
87
88func nx_cr_alloc(g: *ComputeGraph, cap_inputs: nx_int, cap_consts: nx_int) -> *ComputeRunner {
89 let r: *ComputeRunner = (sys_mmap(NX_CR_BYTES)) as *ComputeRunner
90 r.graph = g
91 r.tensor_store_cap = g.cap_nodes * NX_CN_MAX_OUTPUTS
92 let store_bytes: nx_int = r.tensor_store_cap * NX_SIZEOF_NX_INT
93 r.tensor_store = (sys_mmap(store_bytes)) as *i64
94
95 r.cap_inputs = cap_inputs
96 r.input_bindings = (sys_mmap(cap_inputs * NX_SIZEOF_NX_INT)) as *i64
97 r.n_inputs = 0
98
99 r.cap_consts = cap_consts
100 r.const_bindings = (sys_mmap(cap_consts * NX_SIZEOF_NX_INT)) as *i64
101 r.n_consts = 0
102
103 r.last_verdict = NX_CR_OK
104 r.last_failed_node = 0 - 1
105
106 var i: nx_int = 0
107 while i < r.tensor_store_cap {
108 r.tensor_store[i] = 0
109 i = i + 1
110 }
111 return r
112}
113
114// Bind the next INPUT node's tensor. Caller invokes in the SAME ORDER
115// as INPUT nodes appear in the graph; binding index = INPUT-node-rank.
116func nx_cr_bind_input(r: *ComputeRunner, t: *NxTensor) -> nx_int {
117 if r.n_inputs >= r.cap_inputs { return 0 - 1 }
118 r.input_bindings[r.n_inputs] = t as nx_int
119 r.n_inputs = r.n_inputs + 1
120 return 0
121}
122
123func nx_cr_bind_const(r: *ComputeRunner, t: *NxTensor) -> nx_int {
124 if r.n_consts >= r.cap_consts { return 0 - 1 }
125 r.const_bindings[r.n_consts] = t as nx_int
126 r.n_consts = r.n_consts + 1
127 return 0
128}
129
130// ===== Tensor store accessors =====================================
131
132func nx_cr_store_set(r: *ComputeRunner, node_id: nx_int, port: nx_int, t: *NxTensor) -> nx_int {
133 let idx: nx_int = node_id * NX_CN_MAX_OUTPUTS + port
134 if idx < 0 { return 0 - 1 }
135 if idx >= r.tensor_store_cap { return 0 - 1 }
136 r.tensor_store[idx] = t as nx_int
137 return 0
138}
139
140func nx_cr_store_get(r: *ComputeRunner, node_id: nx_int, port: nx_int) -> *NxTensor {
141 let idx: nx_int = node_id * NX_CN_MAX_OUTPUTS + port
142 if idx < 0 { return 0 as *NxTensor }
143 if idx >= r.tensor_store_cap { return 0 as *NxTensor }
144 return r.tensor_store[idx] as *NxTensor
145}
146
147// ===== Elementwise dispatch =======================================
148//
149// For TENSOR_OP nodes with a single input. Allocates output tensor
150// with the same shape + dtype as input.
151
152func _runner_alloc_like(t: *NxTensor) -> *NxTensor {
153 let err: *i64 = (sys_mmap(8)) as *i64
154 let shape: *i64 = (sys_mmap(NX_T_SHAPE_BYTES)) as *i64
155 var i: nx_int = 0
156 while i < t.ndim {
157 shape[i] = t.shape[i]
158 i = i + 1
159 }
160 return nx_t_alloc(t.dtype, shape, t.ndim, err)
161}
162
163func _runner_run_relu(input: *NxTensor) -> *NxTensor {
164 let out: *NxTensor = _runner_alloc_like(input)
165 let pi: *i64 = input.storage as *i64
166 let po: *i64 = out.storage as *i64
167 var k: nx_int = 0
168 while k < input.numel {
169 if pi[k] > 0 { po[k] = pi[k] }
170 if pi[k] <= 0 { po[k] = 0 }
171 k = k + 1
172 }
173 return out
174}
175
176func _runner_run_add(a: *NxTensor, b: *NxTensor) -> *NxTensor {
177 let out: *NxTensor = _runner_alloc_like(a)
178 let pa: *i64 = a.storage as *i64
179 let pb: *i64 = b.storage as *i64
180 let po: *i64 = out.storage as *i64
181 var k: nx_int = 0
182 while k < a.numel {
183 po[k] = pa[k] + pb[k]
184 k = k + 1
185 }
186 return out
187}
188
189func _runner_run_mul(a: *NxTensor, b: *NxTensor) -> *NxTensor {
190 let out: *NxTensor = _runner_alloc_like(a)
191 let pa: *i64 = a.storage as *i64
192 let pb: *i64 = b.storage as *i64
193 let po: *i64 = out.storage as *i64
194 var k: nx_int = 0
195 while k < a.numel {
196 po[k] = pa[k] * pb[k]
197 k = k + 1
198 }
199 return out
200}
201
202// ===== Kernel dispatch (matmul) ===================================
203//
204// Allocates output tensor of shape [M, N] from inputs A:[M,K], B:[K,N];
205// invokes nx_blas_matmul.
206
207func _runner_run_matmul(a: *NxTensor, b: *NxTensor, verdict_out: *i64) -> *NxTensor {
208 let err: *i64 = (sys_mmap(8)) as *i64
209 let shape: *i64 = (sys_mmap(16)) as *i64
210 shape[0] = a.shape[0]
211 shape[1] = b.shape[1]
212 let out: *NxTensor = nx_t_alloc(NX_DT_I64, shape, 2, err)
213 nx_t_fill_zero(out)
214 let rc: nx_int = nx_blas_matmul(a, b, out)
215 verdict_out[0] = rc
216 return out
217}
218
219// ===== Per-node execution =========================================
220//
221// Returns 0 on success or sealed verdict on failure. Updates
222// r.last_verdict + r.last_failed_node.
223
224func _runner_exec_node(r: *ComputeRunner, node_id: nx_int,
225 input_seen: *i64, const_seen: *i64) -> nx_int {
226 let n: *ComputeNode = nx_cg_get_node(r.graph, node_id)
227
228 // INPUT: take next caller-supplied binding (per-INPUT-node rank)
229 if n.kind == NX_CN_NODE_INPUT {
230 let rank: nx_int = input_seen[0]
231 if rank >= r.n_inputs { return NX_CR_ERR_INPUT_MISSING }
232 let t: *NxTensor = r.input_bindings[rank] as *NxTensor
233 nx_cr_store_set(r, node_id, 0, t)
234 input_seen[0] = rank + 1
235 return NX_CR_OK
236 }
237
238 // CONST: same pattern with const_bindings
239 if n.kind == NX_CN_NODE_CONST {
240 let rank2: nx_int = const_seen[0]
241 if rank2 >= r.n_consts { return NX_CR_ERR_CONST_MISSING }
242 let t2: *NxTensor = r.const_bindings[rank2] as *NxTensor
243 nx_cr_store_set(r, node_id, 0, t2)
244 const_seen[0] = rank2 + 1
245 return NX_CR_OK
246 }
247
248 // OUTPUT: alias the input tensor into our store so external
249 // readers can fetch via nx_cr_store_get(runner, output_node_id, 0).
250 if n.kind == NX_CN_NODE_OUTPUT {
251 let src: nx_int = n.input_src_node[0]
252 let port: nx_int = n.input_src_port[0]
253 let in_t: *NxTensor = nx_cr_store_get(r, src, port)
254 nx_cr_store_set(r, node_id, 0, in_t)
255 return NX_CR_OK
256 }
257
258 // TENSOR_OP
259 if n.kind == NX_CN_NODE_TENSOR_OP {
260 if n.op_code == NX_CN_OP_RELU {
261 let src_n: nx_int = n.input_src_node[0]
262 let src_p: nx_int = n.input_src_port[0]
263 let in_t: *NxTensor = nx_cr_store_get(r, src_n, src_p)
264 let out: *NxTensor = _runner_run_relu(in_t)
265 nx_cr_store_set(r, node_id, 0, out)
266 return NX_CR_OK
267 }
268 if n.op_code == NX_CN_OP_ADD {
269 let s0n: nx_int = n.input_src_node[0]
270 let s0p: nx_int = n.input_src_port[0]
271 let s1n: nx_int = n.input_src_node[1]
272 let s1p: nx_int = n.input_src_port[1]
273 let a_t: *NxTensor = nx_cr_store_get(r, s0n, s0p)
274 let b_t: *NxTensor = nx_cr_store_get(r, s1n, s1p)
275 let out2: *NxTensor = _runner_run_add(a_t, b_t)
276 nx_cr_store_set(r, node_id, 0, out2)
277 return NX_CR_OK
278 }
279 if n.op_code == NX_CN_OP_MUL {
280 let s0n2: nx_int = n.input_src_node[0]
281 let s0p2: nx_int = n.input_src_port[0]
282 let s1n2: nx_int = n.input_src_node[1]
283 let s1p2: nx_int = n.input_src_port[1]
284 let a_t2: *NxTensor = nx_cr_store_get(r, s0n2, s0p2)
285 let b_t2: *NxTensor = nx_cr_store_get(r, s1n2, s1p2)
286 let out3: *NxTensor = _runner_run_mul(a_t2, b_t2)
287 nx_cr_store_set(r, node_id, 0, out3)
288 return NX_CR_OK
289 }
290 return NX_CR_ERR_UNSUPPORTED_OP
291 }
292
293 // KERNEL
294 if n.kind == NX_CN_NODE_KERNEL {
295 if n.op_code == NX_CN_K_MATMUL {
296 let s0n3: nx_int = n.input_src_node[0]
297 let s0p3: nx_int = n.input_src_port[0]
298 let s1n3: nx_int = n.input_src_node[1]
299 let s1p3: nx_int = n.input_src_port[1]
300 let a_k: *NxTensor = nx_cr_store_get(r, s0n3, s0p3)
301 let b_k: *NxTensor = nx_cr_store_get(r, s1n3, s1p3)
302 let vbuf: *i64 = (sys_mmap(8)) as *i64
303 let out_k: *NxTensor = _runner_run_matmul(a_k, b_k, vbuf)
304 if vbuf[0] != NX_BLAS_OK { return NX_CR_ERR_KERNEL_FAILED }
305 nx_cr_store_set(r, node_id, 0, out_k)
306 return NX_CR_OK
307 }
308 return NX_CR_ERR_UNSUPPORTED_OP
309 }
310
311 return NX_CR_ERR_UNSUPPORTED_OP
312}
313
314// ===== Top-level execute ==========================================
315//
316// Walks the graph in topo order. Requires the graph to have passed
317// validate + toposort. Returns 0 on success or sealed verdict.
318
319func nx_cr_run(r: *ComputeRunner) -> nx_int {
320 if r.graph.has_topo == 0 {
321 r.last_verdict = NX_CR_ERR_GRAPH_INVALID
322 return NX_CR_ERR_GRAPH_INVALID
323 }
324
325 let input_seen: *i64 = (sys_mmap(8)) as *i64
326 let const_seen: *i64 = (sys_mmap(8)) as *i64
327 input_seen[0] = 0
328 const_seen[0] = 0
329
330 var i: nx_int = 0
331 while i < r.graph.n_nodes {
332 let nid: nx_int = nx_cg_topo_at(r.graph, i)
333 let v: nx_int = _runner_exec_node(r, nid, input_seen, const_seen)
334 if v != NX_CR_OK {
335 r.last_verdict = v
336 r.last_failed_node = nid
337 return v
338 }
339 i = i + 1
340 }
341
342 r.last_verdict = NX_CR_OK
343 return NX_CR_OK
344}
345
346// ===== Get the output tensor of a node ============================
347//
348// Convenience for callers: fetch the materialised output of a
349// specific node + port from the tensor store.
350
351func nx_cr_get_output(r: *ComputeRunner, node_id: nx_int, port: nx_int) -> *NxTensor {
352 return nx_cr_store_get(r, node_id, port)
353}