nx_compute_node.nx source
↩ module page · 326 lines · 12133 B
1// nx_compute_node.nx -- typed compute node, foundation of the
2// ComfyUI-replacement DAG.
3//
4// ComfyUI's node is a Python class with INPUT_TYPES / RETURN_TYPES +
5// a `def execute(self, **kwargs)`. Stringly-typed everywhere; the
6// type system only catches mismatches at runtime; the workflow JSON
7// is opaque; failures have no audit trail.
8//
9// Our compute node is:
10//
11// * **typed** -- NodeKind sealed enum + OpCode sealed enum + typed
12// input/output ports. Type mismatches caught at graph-build time.
13// * **content-addressed** -- each node carries a 32-byte SHA-256
14// hash of (op_code, params, input hashes). Identical sub-graphs
15// produce identical hashes -> deterministic caching.
16// * **audit-able** -- every node carries a stable id; the runner
17// emits a span per execution (compose with nx_trace_emit).
18// * **flat-array layout** -- inputs/params live in i64 buffers per
19// the substrate convention; no Python dict-of-anything.
20//
21// This module ships the NODE primitive only. nx_compute_graph adds
22// the DAG container + topo sort. The runner (executes a graph
23// against tensor inputs) lands once we have real kernels.
24//
25// genealogy_id: comfyui_node_pattern + mlir_op + tvm_relay +
26// jax_jaxpr + onnx_node + tensorflow_xla_hlo
27// lineage_id: substrate_compute_node_v1
28
29// nx_safety_envelope:
30// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
31// sil_target: SIL1
32// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
33// verdict: NOT_YET_EVALUATED
34
35import "nx_syscalls.nx"
36import "nx_tier.nx"
37import "nx_tensor.nx"
38import "nx_sha256.nx"
39
40// ===== Sealed-enum: NodeKind =======================================
41//
42// What KIND of node this is. Runner dispatches on this.
43
44const NX_CN_NODE_CONST: nx_int = 0 // pre-loaded tensor (weight blob)
45const NX_CN_NODE_INPUT: nx_int = 1 // graph-level input port
46const NX_CN_NODE_OUTPUT: nx_int = 2 // graph-level output port
47const NX_CN_NODE_TENSOR_OP: nx_int = 3 // elementwise / reduction / shape
48const NX_CN_NODE_KERNEL: nx_int = 4 // heavy op dispatched to a kernel
49const NX_CN_NODE_CONTROL: nx_int = 5 // sub-graph / scheduler hook
50const NX_CN_NODE_N_KINDS: nx_int = 6
51
52func nx_cn_node_kind_is_valid(k: nx_int) -> nx_int {
53 if k < 0 { return 0 }
54 if k >= NX_CN_NODE_N_KINDS { return 0 }
55 return 1
56}
57
58// ===== Sealed-enum: OpCode (for TENSOR_OP) ========================
59//
60// Elementwise, reduction, shape transforms. Light-weight ops that
61// don't justify a full kernel module. Heavy ops use NODE_KERNEL
62// with a KernelKind below.
63
64const NX_CN_OP_ADD: nx_int = 0
65const NX_CN_OP_SUB: nx_int = 1
66const NX_CN_OP_MUL: nx_int = 2
67const NX_CN_OP_DIV: nx_int = 3
68const NX_CN_OP_NEG: nx_int = 4
69const NX_CN_OP_RELU: nx_int = 5
70const NX_CN_OP_GELU: nx_int = 6
71const NX_CN_OP_SILU: nx_int = 7
72const NX_CN_OP_SWISH: nx_int = 8
73const NX_CN_OP_SOFTMAX: nx_int = 9
74const NX_CN_OP_SIGMOID: nx_int = 10
75const NX_CN_OP_TANH: nx_int = 11
76const NX_CN_OP_RESHAPE: nx_int = 12
77const NX_CN_OP_PERMUTE: nx_int = 13
78const NX_CN_OP_SLICE: nx_int = 14
79const NX_CN_OP_CONCAT: nx_int = 15
80const NX_CN_OP_SPLIT: nx_int = 16
81const NX_CN_OP_SUM_REDUCE: nx_int = 17
82const NX_CN_OP_MEAN_REDUCE: nx_int = 18
83const NX_CN_OP_MAX_REDUCE: nx_int = 19
84const NX_CN_OP_MIN_REDUCE: nx_int = 20
85const NX_CN_OP_TRANSPOSE: nx_int = 21
86const NX_CN_OP_VIEW: nx_int = 22
87const NX_CN_OP_N_OPS: nx_int = 23
88
89func nx_cn_op_is_valid(o: nx_int) -> nx_int {
90 if o < 0 { return 0 }
91 if o >= NX_CN_OP_N_OPS { return 0 }
92 return 1
93}
94
95// ===== Sealed-enum: KernelKind (for NODE_KERNEL) ==================
96//
97// Heavy ops backed by dedicated kernel modules (BLAS, conv, attention,
98// etc.). Stay open for extension; new kernels add entries here.
99
100const NX_CN_K_MATMUL: nx_int = 0 // standard matmul (gemm)
101const NX_CN_K_GEMM: nx_int = 1 // alpha*A@B + beta*C
102const NX_CN_K_DOT: nx_int = 2 // 1-D dot product
103const NX_CN_K_CONV2D: nx_int = 3
104const NX_CN_K_CONV2D_DEPTHWISE: nx_int = 4
105const NX_CN_K_CONV2D_TRANSPOSE: nx_int = 5
106const NX_CN_K_LAYERNORM: nx_int = 6
107const NX_CN_K_RMSNORM: nx_int = 7
108const NX_CN_K_GROUPNORM: nx_int = 8
109const NX_CN_K_BATCHNORM: nx_int = 9
110const NX_CN_K_ATTENTION_QKV: nx_int = 10 // self-attention
111const NX_CN_K_ATTENTION_CROSS: nx_int = 11 // cross-attention
112const NX_CN_K_EMBEDDING_LOOKUP: nx_int = 12
113const NX_CN_K_ROTARY_EMBED: nx_int = 13 // RoPE
114const NX_CN_K_FFT: nx_int = 14
115const NX_CN_K_IFFT: nx_int = 15
116const NX_CN_K_N_KINDS: nx_int = 16
117
118func nx_cn_kernel_is_valid(k: nx_int) -> nx_int {
119 if k < 0 { return 0 }
120 if k >= NX_CN_K_N_KINDS { return 0 }
121 return 1
122}
123
124// ===== ComputeNode struct =========================================
125//
126// Each node owns its own metadata. Inputs reference OTHER nodes by
127// (node_id, output_port_id) -- the graph stores edges; the node
128// itself just stores the local view.
129//
130// Output tensors live in a separate per-graph tensor store that the
131// runner allocates; this struct holds the indexed REFERENCE so we
132// don't tie node lifecycle to tensor lifecycle.
133
134const NX_CN_MAX_INPUTS: nx_int = 8 // covers attention's Q+K+V+mask+...
135const NX_CN_MAX_OUTPUTS: nx_int = 4 // most ops are 1; split takes a few
136const NX_CN_MAX_PARAMS: nx_int = 16 // axis, scale, dim, etc.
137const NX_CN_HASH_BYTES: nx_int = 32 // SHA-256
138
139struct ComputeNode {
140 node_id: nx_int,
141 kind: nx_int, // NX_CN_NODE_*
142 op_code: nx_int, // NX_CN_OP_* if kind==TENSOR_OP, NX_CN_K_* if kind==KERNEL
143 n_inputs: nx_int,
144 n_outputs: nx_int,
145 n_params: nx_int,
146 // Input edges: input_src_node[i] is the source node id; input_src_port[i]
147 // is which output port of that source node.
148 input_src_node: *i64, // [n_inputs]
149 input_src_port: *i64, // [n_inputs]
150 // Op-specific scalar parameters (axis, dim, scale_q10, etc.)
151 params: *i64, // [n_params]
152 // Output tensor refs (slots in the graph's tensor store; -1 = unallocated)
153 output_tensor_id: *i64, // [n_outputs]
154 // Content-addressed hash; computed by nx_cn_compute_hash.
155 content_hash: *u8 // [NX_CN_HASH_BYTES]
156}
157
158const NX_CN_BYTES: nx_int = 88 // 11 fields * 8
159
160// ===== Builder ====================================================
161//
162// Allocates a node with the given kind + op_code + slot capacities.
163// Slots beyond n_inputs/n_outputs/n_params are valid storage but
164// uninitialised; caller fills via the _set helpers before adding the
165// node to a graph.
166
167func nx_cn_alloc(node_id: nx_int, kind: nx_int, op_code: nx_int,
168 n_inputs: nx_int, n_outputs: nx_int,
169 n_params: nx_int) -> *ComputeNode {
170 let n: *ComputeNode = (sys_mmap(NX_CN_BYTES)) as *ComputeNode
171 n.node_id = node_id
172 n.kind = kind
173 n.op_code = op_code
174 n.n_inputs = n_inputs
175 n.n_outputs = n_outputs
176 n.n_params = n_params
177
178 n.input_src_node = (sys_mmap(NX_CN_MAX_INPUTS * NX_SIZEOF_NX_INT)) as *i64
179 n.input_src_port = (sys_mmap(NX_CN_MAX_INPUTS * NX_SIZEOF_NX_INT)) as *i64
180 n.params = (sys_mmap(NX_CN_MAX_PARAMS * NX_SIZEOF_NX_INT)) as *i64
181 n.output_tensor_id = (sys_mmap(NX_CN_MAX_OUTPUTS * NX_SIZEOF_NX_INT)) as *i64
182 n.content_hash = (sys_mmap(NX_CN_HASH_BYTES)) as *u8
183
184 var i: nx_int = 0
185 while i < NX_CN_MAX_INPUTS {
186 n.input_src_node[i] = 0 - 1
187 n.input_src_port[i] = 0 - 1
188 i = i + 1
189 }
190 var j: nx_int = 0
191 while j < NX_CN_MAX_OUTPUTS {
192 n.output_tensor_id[j] = 0 - 1
193 j = j + 1
194 }
195 return n
196}
197
198// ===== Input edge setter ==========================================
199//
200// Caller wires inputs after alloc: "input i of this node comes from
201// node N's output port P". Returns 0 on success, -1 if slot index
202// out of range.
203
204func nx_cn_set_input(n: *ComputeNode, input_slot: nx_int,
205 src_node_id: nx_int, src_port: nx_int) -> nx_int {
206 if input_slot < 0 { return 0 - 1 }
207 if input_slot >= n.n_inputs { return 0 - 1 }
208 if input_slot >= NX_CN_MAX_INPUTS { return 0 - 1 }
209 n.input_src_node[input_slot] = src_node_id
210 n.input_src_port[input_slot] = src_port
211 return 0
212}
213
214// ===== Parameter setter ===========================================
215
216func nx_cn_set_param(n: *ComputeNode, slot: nx_int, value: nx_int) -> nx_int {
217 if slot < 0 { return 0 - 1 }
218 if slot >= n.n_params { return 0 - 1 }
219 if slot >= NX_CN_MAX_PARAMS { return 0 - 1 }
220 n.params[slot] = value
221 return 0
222}
223
224// ===== Output tensor slot ID assignment ===========================
225//
226// Graph runner assigns tensor IDs as it walks the graph; this is the
227// setter the runner uses.
228
229func nx_cn_set_output_tensor_id(n: *ComputeNode, port: nx_int,
230 tensor_id: nx_int) -> nx_int {
231 if port < 0 { return 0 - 1 }
232 if port >= n.n_outputs { return 0 - 1 }
233 n.output_tensor_id[port] = tensor_id
234 return 0
235}
236
237// ===== Read accessors ============================================
238
239func nx_cn_get_input_src_node(n: *ComputeNode, slot: nx_int) -> nx_int {
240 if slot < 0 { return 0 - 1 }
241 if slot >= n.n_inputs { return 0 - 1 }
242 return n.input_src_node[slot]
243}
244
245func nx_cn_get_input_src_port(n: *ComputeNode, slot: nx_int) -> nx_int {
246 if slot < 0 { return 0 - 1 }
247 if slot >= n.n_inputs { return 0 - 1 }
248 return n.input_src_port[slot]
249}
250
251func nx_cn_get_param(n: *ComputeNode, slot: nx_int) -> nx_int {
252 if slot < 0 { return 0 - 1 }
253 if slot >= n.n_params { return 0 - 1 }
254 return n.params[slot]
255}
256
257func nx_cn_get_output_tensor_id(n: *ComputeNode, port: nx_int) -> nx_int {
258 if port < 0 { return 0 - 1 }
259 if port >= n.n_outputs { return 0 - 1 }
260 return n.output_tensor_id[port]
261}
262
263// ===== Content hash (caching key) =================================
264//
265// Per QMDB 2025 + Merkle DAG + the existing nx_module_cas pattern:
266// hash = SHA-256 of (kind, op_code, params, input_src_node[i],
267// input_src_port[i]). Two nodes with identical (op, params, inputs)
268// produce identical hashes; the runner can dedupe.
269//
270// Hashes are computed AFTER inputs are wired up. Per the substrate
271// content-addressed cardinal: cache by hash, not by name or address.
272
273const NX_CN_HASH_HDR_BYTES: nx_int = 24 // kind + op + n_inputs + n_params (each 8 bytes? or 4?)
274
275// LE i64 write into a byte buffer (hoisted above the hasher because
276// nxc2's resolver is single-pass).
277func _cn_write_i64_le(buf: *u8, off: nx_int, value: nx_int) -> nx_int {
278 var v: nx_int = value
279 var i: nx_int = 0
280 while i < 8 {
281 let b: nx_int = v - (v / 256) * 256
282 var b_pos: nx_int = b
283 if b_pos < 0 { b_pos = b_pos + 256 }
284 buf[off + i] = b_pos
285 v = v / 256
286 i = i + 1
287 }
288 return 0
289}
290
291func nx_cn_compute_hash(n: *ComputeNode) -> nx_int {
292 // Build the canonical-bytes buffer in a fixed layout:
293 // 8 bytes kind (i64 LE)
294 // 8 bytes op_code
295 // 8 bytes n_inputs
296 // 8 bytes n_params
297 // for each input: 8 bytes src_node, 8 bytes src_port
298 // for each param: 8 bytes value
299 let total: nx_int = 32 + n.n_inputs * 16 + n.n_params * 8
300 let buf: *u8 = sys_mmap(total)
301
302 _cn_write_i64_le(buf, 0, n.kind)
303 _cn_write_i64_le(buf, 8, n.op_code)
304 _cn_write_i64_le(buf, 16, n.n_inputs)
305 _cn_write_i64_le(buf, 24, n.n_params)
306
307 var off: nx_int = 32
308 var i: nx_int = 0
309 while i < n.n_inputs {
310 _cn_write_i64_le(buf, off, n.input_src_node[i])
311 _cn_write_i64_le(buf, off + 8, n.input_src_port[i])
312 off = off + 16
313 i = i + 1
314 }
315 var j: nx_int = 0
316 while j < n.n_params {
317 _cn_write_i64_le(buf, off, n.params[j])
318 off = off + 8
319 j = j + 1
320 }
321
322 // Reuse the existing nx_sha256 primitive
323 sha256_digest(buf, total, n.content_hash)
324 return 0
325}
326