nx_companion_compose_test.nx source
↩ module page · 199 lines · 9897 B
1// nx_companion_compose_test.nx -- end-to-end integration smoke proving the
2// conductor + parallel-companion substrate composes with the existing
3// Nishi AI LLM stack.
4//
5// Wires together (all bits-up NishiLang -- no external imports):
6//
7// Other-agent stack (already shipped):
8// nx_tensor.nx -- NxTensor (Q10 i64 storage)
9// nx_placement.nx -- NxPlacedTensor (Phase A: OWNED/MMAP/VRAM/NETWORK)
10//
11// This-agent stack (this session):
12// nx_tensor_placement.nx -- 8-tier placement registry (refinement on Phase A)
13// nx_dispatcher.nx -- Phase C op routing
14// nx_resource_arbiter.nx -- Phase D budget tracking
15// nx_actor.nx -- cooperative actor scheduler
16// nx_message.nx -- typed inter-actor passing
17// nx_stream_multiplexer.nx -- output stream interleaving
18// nx_interrupt_broker.nx -- sub-50ms preemption
19//
20// Scenario: a 3-actor DND-companion mini-session.
21// Actor 1 (LLM): generates tokens, emits LLM_TOKEN messages
22// Actor 2 (TTS): subscribes to LLM_TOKEN, pushes AUDIO_FRAME chunks to mux
23// Actor 3 (GAME): pushes GAME_ACTION chunks to mux
24// Operator clicks -> interrupt broker yields LLM checkpoint within 50ms
25
26import "nx_syscalls.nx"
27import "nx_tier.nx"
28import "nx_tensor.nx"
29import "nx_placement.nx"
30import "nx_tensor_placement.nx"
31import "nx_dispatcher.nx"
32import "nx_resource_arbiter.nx"
33import "nx_actor.nx"
34import "nx_message.nx"
35import "nx_stream_multiplexer.nx"
36import "nx_interrupt_broker.nx"
37import "nx_meristem.nx"
38
39func main() -> i64 {
40 let now: nx_size = 1000000
41
42 // ===== Stage 1: build placed tensors from the other agent's stack =====
43 //
44 // Build a [256, 4096] embedding matrix in OWNED placement (synthetic;
45 // simulates a materialized weight tensor in CPU RAM).
46 let shape_buf: *u8 = sys_mmap(32)
47 let shape: *i64 = shape_buf as *i64
48 shape[0] = 256
49 shape[1] = 4096
50 let err_buf: *u8 = sys_mmap(8)
51 let err_p: *i64 = err_buf as *i64
52 let t_embed: *NxTensor = nx_t_alloc(NX_DT_I64, shape, 2, err_p)
53 if err_p[0] != NX_T_OK { return 1 }
54 if (t_embed as i64) == 0 { return 2 }
55 let p_embed: *NxPlacedTensor = nx_placed_tensor_new_owned(t_embed)
56 if p_embed.placement != NX_PLACE_OWNED { return 3 }
57
58 // Build a lazy NxPlacedTensor (simulates a GGUF-resident weight in MMAP).
59 // src_buf is a dummy byte buffer; we won't materialize here, just
60 // verify the placement bridge reads it correctly.
61 let dummy_bytes: *u8 = sys_mmap(1024)
62 let shape2_buf: *u8 = sys_mmap(32)
63 let shape2: *i64 = shape2_buf as *i64
64 shape2[0] = 1024
65 shape2[1] = 1024
66 let p_lazy: *NxPlacedTensor = nx_placed_tensor_new_lazy(dummy_bytes, 0, NX_GGML_TYPE_F32, 1048576, shape2, 2)
67 if p_lazy.placement != NX_PLACE_MMAP { return 4 }
68 if nx_placed_tensor_is_live(p_lazy) != 0 { return 5 }
69
70 // ===== Stage 2: dispatcher reads placed-tensor placement =====
71 //
72 // Build a matmul OpRequest from [p_embed (OWNED -> cpu_bytes),
73 // p_lazy (MMAP -> cold_bytes)] inputs. Dispatcher should still
74 // choose GPU_TENSOR_CORE despite cold inputs because the compute
75 // savings dominate transfer cost for big matmuls.
76 let placed_inputs_buf: *u8 = sys_mmap(16)
77 let placed_inputs: **NxPlacedTensor = placed_inputs_buf as **NxPlacedTensor
78 placed_inputs[0] = p_embed
79 placed_inputs[1] = p_lazy
80 let req: *NxOpRequest = nx_di_req_from_placed_inputs(NX_DI_OP_MATMUL, placed_inputs, 2, 1)
81 if req.op_kind != NX_DI_OP_MATMUL { return 6 }
82 // p_embed OWNED -> cpu (256*4096*8 = 8388608 bytes)
83 if req.inputs_on_cpu_bytes != 8388608 { return 7 }
84 // p_lazy MMAP -> cold (1024*1024*8 = 8388608 bytes)
85 if req.inputs_on_cold_bytes != 8388608 { return 8 }
86 if req.input_bytes_total != 16777216 { return 9 }
87
88 let dec: *NxDispatchDecision = nx_di_dec_new()
89 if nx_di_dispatch(req, dec) != NX_DI_V_DISPATCHED { return 10 }
90 // GPU tensor cores still win for big matmul + we flag promotion required
91 if nx_di_be_is_gpu(dec.chosen_backend) != 1 { return 11 }
92 if dec.requires_promotion != 1 { return 12 } // cold inputs flag this
93
94 // ===== Stage 3: resource arbiter reserves VRAM under contention =====
95 let arb: *NxResourceArbiter = nx_ra_new(16)
96 nx_ra_set_budget(arb, NX_RA_KIND_VRAM_BYTES, 16777216) // 16MB toy budget
97 nx_ra_set_max_share(arb, NX_RA_KIND_VRAM_BYTES, 717) // 70% cap per cell
98 let out_buf: *u8 = sys_mmap(8)
99 let out_g: *i64 = out_buf as *i64
100 // LLM actor requests 8MB VRAM (50% of budget; under 70% cap)
101 if nx_ra_request(arb, 1001, NX_RA_KIND_VRAM_BYTES, 8388608, 80, 1, out_g, now) != NX_RA_V_GRANTED_FULL { return 13 }
102 if out_g[0] != 8388608 { return 14 }
103 // TTS actor requests 4MB; budget has 8MB free, 4MB < 70% cap -> OK
104 if nx_ra_request(arb, 1002, NX_RA_KIND_VRAM_BYTES, 4194304, 60, 1, out_g, now) != NX_RA_V_GRANTED_FULL { return 15 }
105 if nx_ra_total_held(arb, NX_RA_KIND_VRAM_BYTES) != 12582912 { return 16 }
106
107 // ===== Stage 4: actor scheduler with 3 cells =====
108 let sched: *NxActorScheduler = nx_ac_sched_new(8, 1000000, 4, now)
109 nx_ac_spawn(sched, 1001, NX_ME_ROLE_LLM_INFERENCE, 80, 0, now) // LLM
110 nx_ac_spawn(sched, 1002, NX_ME_ROLE_AUDIO_TTS, 70, 0, now) // TTS
111 nx_ac_spawn(sched, 1003, NX_ME_ROLE_GAME_LOGIC, 90, 0, now) // GAME
112 if nx_ac_actor_count(sched) != 3 { return 17 }
113
114 // pick_next should prefer GAME (priority 90)
115 let first: *NxActor = nx_ac_pick_next(sched)
116 if first.actor_id != 1003 { return 18 }
117
118 // ===== Stage 5: message bus wiring =====
119 let bus: *NxMessageBus = nx_ms_bus_new(8, 16, 4)
120 nx_ms_register(bus, 1001)
121 nx_ms_register(bus, 1002)
122 nx_ms_register(bus, 1003)
123 // TTS subscribes to LLM_TOKEN
124 nx_ms_subscribe(bus, 1002, NX_MS_KIND_LLM_TOKEN)
125 // GAME subscribes to OPERATOR_INPUT
126 nx_ms_subscribe(bus, 1003, NX_MS_KIND_OPERATOR_INPUT)
127 if nx_ms_subscription_count(bus, NX_MS_KIND_LLM_TOKEN) != 1 { return 19 }
128
129 // ===== Stage 6: LLM actor steps + emits LLM_TOKEN =====
130 if nx_ac_step(sched, 1001, 5000, now + 100) != NX_AC_V_STEPPED { return 20 }
131 // LLM emits a token via fanout (TTS catches it; GAME does not, since
132 // GAME isn't subscribed to LLM_TOKEN)
133 if nx_ms_send_fanout(bus, 1001, NX_MS_KIND_LLM_TOKEN, 0xCAFE, 8, now + 100) != NX_MS_V_DELIVERED { return 21 }
134 if nx_ms_pending(bus, 1002) != 1 { return 22 }
135 if nx_ms_pending(bus, 1003) != 0 { return 23 }
136
137 // ===== Stage 7: TTS receives token, steps, pushes audio chunk =====
138 let llm_msg: *NxMessage = nx_ms_receive(bus, 1002)
139 if llm_msg.kind != NX_MS_KIND_LLM_TOKEN { return 24 }
140 if llm_msg.sender_actor_id != 1001 { return 25 }
141 nx_ac_step(sched, 1002, 3000, now + 200)
142
143 // ===== Stage 8: stream multiplexer interleaves chunks =====
144 let mux: *NxStreamMultiplexer = nx_sm_new(16, 50000)
145 // TTS emits AUDIO_FRAME, GAME emits GAME_ACTION
146 nx_sm_push(mux, 1002, NX_SM_CLASS_AUDIO_TTS, 100, 32, 0xCA02, now + 200)
147 nx_sm_push(mux, 1003, NX_SM_CLASS_GAME_ACTION, 200, 16, 0xCA03, now + 250)
148 nx_sm_push(mux, 1001, NX_SM_CLASS_LLM_TEXT, 300, 64, 0xCA01, now + 100)
149 if nx_sm_pending_count(mux) != 3 { return 26 }
150
151 // mux delivers GAME_ACTION first (priority 90), then LLM_TEXT (80), then AUDIO_TTS (70)
152 let out_class_buf: *u8 = sys_mmap(24)
153 let oc: *i64 = out_class_buf as *i64
154 let os: *i64 = (out_class_buf + 8) as *i64
155 let oid: *i64 = (out_class_buf + 16) as *i64
156 nx_sm_next(mux, oc, os, oid, now + 300)
157 if oc[0] != (NX_SM_CLASS_GAME_ACTION as i64) { return 27 }
158 nx_sm_next(mux, oc, os, oid, now + 301)
159 if oc[0] != (NX_SM_CLASS_LLM_TEXT as i64) { return 28 }
160 nx_sm_next(mux, oc, os, oid, now + 302)
161 if oc[0] != (NX_SM_CLASS_AUDIO_TTS as i64) { return 29 }
162
163 // ===== Stage 9: operator interrupt preempts LLM =====
164 let broker: *NxInterruptBroker = nx_ib_new(8)
165 // Operator clicks at now + 1000 with REALTIME deadline (50ms)
166 nx_ib_file(broker, 5001, NX_IB_KIND_OPERATOR_INPUT, NX_IB_DL_REALTIME, 95, 0xABCD, now + 1000)
167 // LLM op currently running, yieldable at +1500us (well within 50ms window)
168 let llm_op: *NxRunningOp = nx_ib_op_new(1001, 80, now, 200000, now + 1500, 0)
169 if nx_ib_arbitrate(broker, 5001, llm_op, now + 5000) != NX_IB_V_YIELD { return 30 }
170 if nx_ib_grants(broker) != 1 { return 31 }
171 nx_ib_resolve(broker, 5001)
172
173 // Now publish OPERATOR_INPUT to actors; GAME picks it up (subscribed)
174 nx_ms_send_fanout(bus, 0, NX_MS_KIND_OPERATOR_INPUT, 0xABCD, 8, now + 1500)
175 if nx_ms_pending(bus, 1003) != 1 { return 32 }
176
177 // ===== Stage 10: tensor placement registry tracks weights across tiers =====
178 let reg: *NxPlacementRegistry = nx_tp_new(32, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF)
179 nx_tp_register(reg, 5001, NX_TP_TIER_RAM_HOT, 8388608, 0xDEAD, now) // embedding weights
180 nx_tp_register(reg, 5002, NX_TP_TIER_NVME_MMAP, 8388608, 0xBEEF, now) // lazy GGUF weights
181 if nx_tp_count_in_tier(reg, NX_TP_TIER_RAM_HOT) != 1 { return 33 }
182 if nx_tp_count_in_tier(reg, NX_TP_TIER_NVME_MMAP) != 1 { return 34 }
183
184 // Promote lazy weight to VRAM
185 nx_tp_promote(reg, 5002, NX_TP_TIER_VRAM_HOT, 0xCAFE, now + 2000)
186 if nx_tp_count_in_tier(reg, NX_TP_TIER_VRAM_HOT) != 1 { return 35 }
187 if nx_tp_count_in_tier(reg, NX_TP_TIER_NVME_MMAP) != 0 { return 36 }
188
189 // ===== Stage 11: release VRAM + verify arbiter unwound =====
190 if nx_ra_release(arb, 1001, NX_RA_KIND_VRAM_BYTES) != 8388608 { return 37 }
191 if nx_ra_total_held(arb, NX_RA_KIND_VRAM_BYTES) != 4194304 { return 38 }
192
193 // ===== Stage 12: epoch tick resets multiplexer + scheduler budgets =====
194 nx_sm_epoch_tick(mux)
195 nx_ac_epoch_tick(sched, now + 5000)
196 if sched.epoch_count != 1 { return 39 }
197
198 return 0
199}