nx_dispatcher.nx source
↩ module page · 385 lines · 15139 B
1// nx_dispatcher.nx -- hardware-aware op routing.
2//
3// Per CARDINAL [[feedback-conductor-heterogeneous-compute-no-second-class-
4// resources]]: ops dispatch hardware-aware -- matmul on GPU when operands
5// are in VRAM, activations on CPU SIMD when data lives in RAM, no PCIe
6// round-trips for tiny ops. REFUSES the "everything goes to GPU" pattern.
7//
8// CONDUCTOR ARC ROLE: this is PHASE C of the conductor arc documented in
9// nx_placement.nx. Composition stack:
10// A) nx_placement.nx -- NxPlacedTensor with placement tag (MMAP/
11// OWNED/VRAM/NETWORK). SHIPPED.
12// B) nx_gguf_load_lazy.nx -- lazy GGUF load returning NxPlacedTensor.
13// SHIPPED.
14// C) nx_dispatcher.nx -- per-op hardware routing. THIS BRICK.
15// D) nx_resource_arbiter.nx -- multi-tier budget tracking. THIS SESSION.
16// E) async DMA primitives. Queued.
17//
18// Integration: the bridge function nx_di_req_from_placed_inputs() reads
19// NxPlacedTensor.placement to populate this primitive's
20// inputs_on_gpu/cpu/cold accounting, so callers using the existing
21// nx_placement substrate can compose without forking.
22//
23// V1 ships a sealed op-kind enum + per-op cost model + dispatch decision
24// that picks the cheapest backend given the placement of inputs (uses
25// [[nx_tensor_placement]] semantics: tier index 0=fastest, 7=coldest).
26//
27// Cost model is intentionally simple Q10: BACKEND_COST_PER_BYTE *
28// total_input_bytes + TRANSFER_COST * misplaced_bytes. Refinement
29// (cache effects, batch size, kernel launch overhead) deferred.
30
31import "nx_syscalls.nx"
32import "nx_tier.nx"
33import "nx_tensor.nx"
34import "nx_placement.nx"
35const NX_MAGIC_1000000: i64 = 1000000
36const NX_MAGIC_1024: i64 = 1024
37
38// ===== Sealed enum: NxOpKind ======================================
39
40const NX_DI_OP_MATMUL: nx_int = 0
41const NX_DI_OP_ELEMENTWISE_UNARY: nx_int = 1 // silu / relu / sigmoid
42const NX_DI_OP_ELEMENTWISE_BINARY: nx_int = 2 // add / mul
43const NX_DI_OP_REDUCE_SUM: nx_int = 3
44const NX_DI_OP_RMS_NORM: nx_int = 4
45const NX_DI_OP_SOFTMAX: nx_int = 5
46const NX_DI_OP_ATTENTION: nx_int = 6
47const NX_DI_OP_EMBEDDING_LOOKUP: nx_int = 7
48const NX_DI_OP_COPY: nx_int = 8
49const NX_DI_OP_QUANTIZE: nx_int = 9
50const NX_DI_OP_DEQUANTIZE: nx_int = 10
51const NX_DI_OP_N: nx_int = 11
52
53// ===== Sealed enum: NxBackend =====================================
54
55const NX_DI_BE_CPU_SCALAR: nx_int = 0
56const NX_DI_BE_CPU_SIMD: nx_int = 1
57const NX_DI_BE_GPU_COMPUTE: nx_int = 2
58const NX_DI_BE_GPU_TENSOR_CORE: nx_int = 3
59const NX_DI_BE_NPU: nx_int = 4 // future: neural processor
60const NX_DI_BE_N: nx_int = 5
61
62// ===== Sealed enum: NxDispatchVerdict =============================
63
64const NX_DI_V_DISPATCHED: nx_int = 0
65const NX_DI_V_REQUIRES_PROMOTION: nx_int = 1 // inputs in wrong tier
66const NX_DI_V_NO_VIABLE_BACKEND: nx_int = 2
67const NX_DI_V_INVALID: nx_int = 3
68const NX_DI_V_NULL: nx_int = 4
69const NX_DI_V_N: nx_int = 5
70
71// ===== Struct: NxOpRequest ========================================
72
73struct NxOpRequest {
74 op_kind: nx_int,
75 input_bytes_total: nx_size,
76 inputs_on_gpu_bytes: nx_size,
77 inputs_on_cpu_bytes: nx_size,
78 inputs_on_cold_bytes: nx_size,
79 batch_count: nx_int,
80}
81
82const NX_DI_REQ_BYTES: nx_int = 48
83
84// ===== Struct: NxDispatchDecision =================================
85
86struct NxDispatchDecision {
87 chosen_backend: nx_int,
88 estimated_cost_q10: nx_int,
89 requires_promotion: nx_int, // 1 if any cold-tier input
90 refused_reason: nx_int, // NX_DI_V_* sentinel
91}
92
93const NX_DI_DEC_BYTES: nx_int = 32
94
95// ===== Validators =================================================
96
97func nx_di_op_is_valid(o: nx_int) -> nx_int {
98 if o < 0 { return 0 }
99 if o >= NX_DI_OP_N { return 0 }
100 return 1
101}
102
103func nx_di_be_is_valid(b: nx_int) -> nx_int {
104 if b < 0 { return 0 }
105 if b >= NX_DI_BE_N { return 0 }
106 return 1
107}
108
109func nx_di_v_is_valid(v: nx_int) -> nx_int {
110 if v < 0 { return 0 }
111 if v >= NX_DI_V_N { return 0 }
112 return 1
113}
114
115func nx_di_op_is_compute_heavy(o: nx_int) -> nx_int {
116 if o == NX_DI_OP_MATMUL { return 1 }
117 if o == NX_DI_OP_ATTENTION { return 1 }
118 return 0
119}
120
121func nx_di_op_is_lightweight(o: nx_int) -> nx_int {
122 if o == NX_DI_OP_ELEMENTWISE_UNARY { return 1 }
123 if o == NX_DI_OP_ELEMENTWISE_BINARY { return 1 }
124 if o == NX_DI_OP_COPY { return 1 }
125 return 0
126}
127
128func nx_di_be_is_gpu(b: nx_int) -> nx_int {
129 if b == NX_DI_BE_GPU_COMPUTE { return 1 }
130 if b == NX_DI_BE_GPU_TENSOR_CORE { return 1 }
131 return 0
132}
133
134func nx_di_be_is_cpu(b: nx_int) -> nx_int {
135 if b == NX_DI_BE_CPU_SCALAR { return 1 }
136 if b == NX_DI_BE_CPU_SIMD { return 1 }
137 return 0
138}
139
140// ===== Per-backend per-op base cost (Q10 cycles per byte) ========
141//
142// These are RELATIVE costs. Values calibrated by intuition for V1;
143// real numbers will come from [[nx_metabolism]] profiler. Lower is
144// faster. Returns INT_MAX-ish sentinel for unsupported combos.
145
146func _di_base_cost_q10(op: nx_int, backend: nx_int) -> nx_int {
147 // GPU TENSOR_CORE is FASTEST for matmul + attention
148 if op == NX_DI_OP_MATMUL {
149 if backend == NX_DI_BE_GPU_TENSOR_CORE { return 1 }
150 if backend == NX_DI_BE_GPU_COMPUTE { return 3 }
151 if backend == NX_DI_BE_CPU_SIMD { return 50 }
152 if backend == NX_DI_BE_CPU_SCALAR { return 500 }
153 if backend == NX_DI_BE_NPU { return 2 }
154 return NX_MAGIC_1000000
155 }
156 if op == NX_DI_OP_ATTENTION {
157 if backend == NX_DI_BE_GPU_TENSOR_CORE { return 1 }
158 if backend == NX_DI_BE_GPU_COMPUTE { return 3 }
159 if backend == NX_DI_BE_CPU_SIMD { return 80 }
160 if backend == NX_DI_BE_CPU_SCALAR { return 800 }
161 return NX_MAGIC_1000000
162 }
163 // Lightweight elementwise: CPU SIMD is competitive vs GPU; PCIe wins out
164 if op == NX_DI_OP_ELEMENTWISE_UNARY {
165 if backend == NX_DI_BE_CPU_SIMD { return 2 }
166 if backend == NX_DI_BE_GPU_COMPUTE { return 3 }
167 if backend == NX_DI_BE_CPU_SCALAR { return 10 }
168 return NX_MAGIC_1000000
169 }
170 if op == NX_DI_OP_ELEMENTWISE_BINARY {
171 if backend == NX_DI_BE_CPU_SIMD { return 2 }
172 if backend == NX_DI_BE_GPU_COMPUTE { return 3 }
173 if backend == NX_DI_BE_CPU_SCALAR { return 10 }
174 return NX_MAGIC_1000000
175 }
176 if op == NX_DI_OP_REDUCE_SUM {
177 if backend == NX_DI_BE_GPU_COMPUTE { return 2 }
178 if backend == NX_DI_BE_CPU_SIMD { return 4 }
179 if backend == NX_DI_BE_CPU_SCALAR { return 20 }
180 return NX_MAGIC_1000000
181 }
182 if op == NX_DI_OP_RMS_NORM {
183 if backend == NX_DI_BE_CPU_SIMD { return 3 }
184 if backend == NX_DI_BE_GPU_COMPUTE { return 4 }
185 return NX_MAGIC_1000000
186 }
187 if op == NX_DI_OP_SOFTMAX {
188 if backend == NX_DI_BE_GPU_COMPUTE { return 2 }
189 if backend == NX_DI_BE_CPU_SIMD { return 5 }
190 return NX_MAGIC_1000000
191 }
192 if op == NX_DI_OP_EMBEDDING_LOOKUP {
193 if backend == NX_DI_BE_CPU_SCALAR { return 1 }
194 if backend == NX_DI_BE_CPU_SIMD { return 1 }
195 if backend == NX_DI_BE_GPU_COMPUTE { return 5 }
196 return NX_MAGIC_1000000
197 }
198 if op == NX_DI_OP_COPY {
199 if backend == NX_DI_BE_CPU_SCALAR { return 1 }
200 if backend == NX_DI_BE_CPU_SIMD { return 1 }
201 if backend == NX_DI_BE_GPU_COMPUTE { return 1 }
202 return NX_MAGIC_1000000
203 }
204 if op == NX_DI_OP_QUANTIZE {
205 if backend == NX_DI_BE_CPU_SIMD { return 3 }
206 if backend == NX_DI_BE_GPU_COMPUTE { return 2 }
207 return NX_MAGIC_1000000
208 }
209 if op == NX_DI_OP_DEQUANTIZE {
210 if backend == NX_DI_BE_CPU_SIMD { return 3 }
211 if backend == NX_DI_BE_GPU_COMPUTE { return 2 }
212 return NX_MAGIC_1000000
213 }
214 return NX_MAGIC_1000000
215}
216
217// Transfer penalty when inputs need to move to backend's preferred tier.
218// Q10 cost per misplaced byte. PCIe to GPU is the expensive one.
219func _di_transfer_penalty_q10(backend: nx_int) -> nx_int {
220 if nx_di_be_is_gpu(backend) == 1 { return 20 }
221 if backend == NX_DI_BE_NPU { return 30 }
222 return 5 // CPU backends: cheaper RAM moves
223}
224
225// ===== Cost estimate ==============================================
226//
227// total_cost = base_cost_per_byte * input_bytes_total +
228// transfer_penalty * misplaced_bytes
229//
230// "misplaced" = bytes not currently in the backend's preferred tier.
231
232func nx_di_estimate_cost_q10(req: *NxOpRequest, backend: nx_int) -> nx_int {
233 if (req as i64) == 0 { return NX_MAGIC_1000000 }
234 if nx_di_op_is_valid(req.op_kind) == 0 { return NX_MAGIC_1000000 }
235 if nx_di_be_is_valid(backend) == 0 { return NX_MAGIC_1000000 }
236 let base: nx_int = _di_base_cost_q10(req.op_kind, backend)
237 if base >= NX_MAGIC_1000000 { return base }
238 // Cost scales by total input bytes (Q10 fraction-ish: divide by 1024)
239 let total_q: nx_int = req.input_bytes_total as nx_int
240 var compute_cost: nx_int = (base * total_q) / NX_MAGIC_1024
241 if compute_cost <= 0 { compute_cost = base }
242 let penalty: nx_int = _di_transfer_penalty_q10(backend)
243 var misplaced: nx_int = 0
244 if nx_di_be_is_gpu(backend) == 1 {
245 misplaced = req.inputs_on_cpu_bytes as nx_int
246 misplaced = misplaced + (req.inputs_on_cold_bytes as nx_int)
247 }
248 if nx_di_be_is_cpu(backend) == 1 {
249 misplaced = req.inputs_on_gpu_bytes as nx_int
250 misplaced = misplaced + (req.inputs_on_cold_bytes as nx_int)
251 }
252 if backend == NX_DI_BE_NPU {
253 // NPU has its own memory tier; non-NPU resident bytes must transfer.
254 // V1: treat all GPU/CPU/cold-resident inputs as misplaced.
255 misplaced = req.inputs_on_gpu_bytes as nx_int
256 misplaced = misplaced + (req.inputs_on_cpu_bytes as nx_int)
257 misplaced = misplaced + (req.inputs_on_cold_bytes as nx_int)
258 }
259 let transfer_cost: nx_int = (penalty * misplaced) / NX_MAGIC_1024
260 return compute_cost + transfer_cost
261}
262
263// ===== Dispatch decision (choose best backend) ====================
264
265func nx_di_dispatch(req: *NxOpRequest, decision: *NxDispatchDecision) -> nx_int {
266 if (req as i64) == 0 { return NX_DI_V_NULL }
267 if (decision as i64) == 0 { return NX_DI_V_NULL }
268 if nx_di_op_is_valid(req.op_kind) == 0 { return NX_DI_V_INVALID }
269 var best_be: nx_int = -1
270 var best_cost: nx_int = NX_MAGIC_1000000
271 var b: nx_int = 0
272 while b < NX_DI_BE_N {
273 let cost: nx_int = nx_di_estimate_cost_q10(req, b)
274 if cost < best_cost {
275 best_cost = cost
276 best_be = b
277 }
278 b = b + 1
279 }
280 if best_be < 0 {
281 decision.chosen_backend = -1
282 decision.estimated_cost_q10 = NX_MAGIC_1000000
283 decision.requires_promotion = 0
284 decision.refused_reason = NX_DI_V_NO_VIABLE_BACKEND
285 return NX_DI_V_NO_VIABLE_BACKEND
286 }
287 decision.chosen_backend = best_be
288 decision.estimated_cost_q10 = best_cost
289 decision.requires_promotion = 0
290 if req.inputs_on_cold_bytes > 0 { decision.requires_promotion = 1 }
291 decision.refused_reason = NX_DI_V_DISPATCHED
292 return NX_DI_V_DISPATCHED
293}
294
295// ===== Op-request builder helper ==================================
296
297func nx_di_req_new(op_kind: nx_int,
298 input_bytes_total: nx_size,
299 gpu_bytes: nx_size,
300 cpu_bytes: nx_size,
301 cold_bytes: nx_size,
302 batch_count: nx_int) -> *NxOpRequest {
303 if nx_di_op_is_valid(op_kind) == 0 { return 0 as *NxOpRequest }
304 let raw: *u8 = sys_mmap(NX_DI_REQ_BYTES)
305 let r: *NxOpRequest = raw as *NxOpRequest
306 r.op_kind = op_kind
307 r.input_bytes_total = input_bytes_total
308 r.inputs_on_gpu_bytes = gpu_bytes
309 r.inputs_on_cpu_bytes = cpu_bytes
310 r.inputs_on_cold_bytes = cold_bytes
311 r.batch_count = batch_count
312 return r
313}
314
315func nx_di_dec_new() -> *NxDispatchDecision {
316 let raw: *u8 = sys_mmap(NX_DI_DEC_BYTES)
317 let d: *NxDispatchDecision = raw as *NxDispatchDecision
318 d.chosen_backend = -1
319 d.estimated_cost_q10 = 0
320 d.requires_promotion = 0
321 d.refused_reason = NX_DI_V_DISPATCHED
322 return d
323}
324
325// ===== Bridge: NxPlacedTensor -> NxOpRequest =====================
326//
327// Reads each input's placement tag (Phase A semantics) and folds the
328// element-byte count into the gpu/cpu/cold accounting buckets this
329// dispatcher uses. Mapping:
330//
331// NX_PLACE_OWNED (materialized i64 buffer) -> inputs_on_cpu_bytes
332// NX_PLACE_MMAP (lazy GGUF byte range) -> inputs_on_cold_bytes
333// NX_PLACE_VRAM (GPU-resident, Phase E) -> inputs_on_gpu_bytes
334// NX_PLACE_NETWORK (peer device, future) -> inputs_on_cold_bytes
335//
336// element_bytes per tensor: each materialized element is 8 bytes (i64
337// Q10 storage per [[nx_tensor]] convention). For lazy (MMAP) tensors
338// we use src_n_values * 8 -- the Q10 footprint the materializer would
339// produce. Future Phase E will distinguish ggml_type sizes here.
340
341func _di_bytes_of_placed(p: *NxPlacedTensor) -> nx_size {
342 if (p as i64) == 0 { return 0 }
343 if p.t != (0 as *NxTensor) {
344 // Materialized: walk shape * element_bytes
345 var n: i64 = 1
346 if p.n_dims >= 1 { n = n * p.shape_0 }
347 if p.n_dims >= 2 { n = n * p.shape_1 }
348 if p.n_dims >= 3 { n = n * p.shape_2 }
349 if p.n_dims >= 4 { n = n * p.shape_3 }
350 // V1 only ships I64 storage; future will read p.t.dtype
351 return (n * 8) as nx_size
352 }
353 // Lazy (MMAP): use src_n_values * 8 (Q10 i64 footprint at materialize)
354 return (p.src_n_values * 8) as nx_size
355}
356
357func nx_di_req_from_placed_inputs(op_kind: nx_int,
358 inputs: **NxPlacedTensor,
359 n_inputs: nx_int,
360 batch_count: nx_int) -> *NxOpRequest {
361 if nx_di_op_is_valid(op_kind) == 0 { return 0 as *NxOpRequest }
362 if n_inputs < 0 { return 0 as *NxOpRequest }
363 let raw: *u8 = sys_mmap(NX_DI_REQ_BYTES)
364 let r: *NxOpRequest = raw as *NxOpRequest
365 r.op_kind = op_kind
366 r.input_bytes_total = 0
367 r.inputs_on_gpu_bytes = 0
368 r.inputs_on_cpu_bytes = 0
369 r.inputs_on_cold_bytes = 0
370 r.batch_count = batch_count
371 var i: nx_int = 0
372 while i < n_inputs {
373 let p: *NxPlacedTensor = inputs[i]
374 if (p as i64) != 0 {
375 let b: nx_size = _di_bytes_of_placed(p)
376 r.input_bytes_total = r.input_bytes_total + b
377 if p.placement == NX_PLACE_OWNED { r.inputs_on_cpu_bytes = r.inputs_on_cpu_bytes + b }
378 if p.placement == NX_PLACE_MMAP { r.inputs_on_cold_bytes = r.inputs_on_cold_bytes + b }
379 if p.placement == NX_PLACE_VRAM { r.inputs_on_gpu_bytes = r.inputs_on_gpu_bytes + b }
380 if p.placement == NX_PLACE_NETWORK { r.inputs_on_cold_bytes = r.inputs_on_cold_bytes + b }
381 }
382 i = i + 1
383 }
384 return r
385}