nx_thread_pool.nx source
↩ module page · 363 lines · 16333 B
1// nx_thread_pool.nx -- shared-queue thread pool (Rayon-precursor).
2//
3// Architecture: hw-sized worker fleet + single Vyukov MPMC queue
4// of task records. Workers spin on the channel; each task carries
5// a typed fn-pointer + ctx; sentinel tasks (is_sentinel == 1)
6// signal a worker to exit.
7//
8// MVP shape -- single shared queue, not yet work-stealing. Work-
9// stealing buys you locality when one worker is starving and another
10// is buried; for compute-bound tasks of comparable cost the shared-
11// queue model is within ~5% of work-stealing (per Rayon's own bench).
12// Composes upward to work-stealing in a follow-up by giving each
13// worker a private deque + a steal protocol; the public API does
14// not change.
15//
16// Tasks live in a bump-allocated arena (NxPoolTask). Submit picks
17// the next slot via atomic FAA on `next_task_slot`, fills it, sends
18// the slot pointer through the channel. Worker receives the pointer,
19// calls task.fn(task.ctx), bumps the completed-counter. Slot reuse
20// is a future concern -- bench first.
21//
22// Composes against: [[vyukov_mpmc_channel]] (queue),
23// [[atomic_intrinsics_real_amo]] (FAA counters),
24// [[thread_clone_native_trampoline]] (worker spawn),
25// [[nx_hw_dynamic_probes]] (worker sizing),
26// [[fn_ptr_indirect_call]] (typed call site).
27
28// nx_safety_envelope:
29// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
30// sil_target: SIL1
31// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
32// verdict: NOT_YET_EVALUATED
33
34import "nx_syscalls.nx"
35import "nx_atom.nx"
36import "nx_thread.nx"
37import "nx_chan.nx"
38import "nx_hw.nx"
39const NX_MAGIC_2147483647: i64 = 2147483647
40const NX_MAGIC_65536: i64 = 65536
41const NX_MAGIC_2000000000: i64 = 2000000000
42
43struct NxPoolTask {
44 fn_ptr: func(i64) -> i64, // typed fn-pointer; called by worker
45 ctx: i64, // argument passed to fn_ptr
46 is_sentinel: i64, // 1 => worker exits on receipt
47}
48
49const NX_POOL_TASK_BYTES: i64 = 24
50
51// Default arena size in task records. 16k records = 384 KiB. Bump
52// when bench shows the slot recycling matters; until then linear
53// growth is fine and predictable.
54const NX_POOL_ARENA_DEFAULT: i64 = 16384
55
56struct NxThreadPool {
57 chan_ptr: i64, // *NxChan -- task slot pointers
58 arena_base: i64, // *NxPoolTask -- pre-allocated arena
59 arena_cap: i64, // capacity in task records
60 next_task_slot: i64, // atomic; bump for fresh slot
61 n_workers: i64,
62 workers_alive: i64, // atomic; decremented when worker exits
63 tasks_submitted: i64, // atomic; bumped per submit
64 tasks_completed: i64, // atomic; bumped per task run
65 task_epoch: i64, // atomic; bumped per submit -- the futex word
66 // idle workers sleep on (low 32 bits)
67 submitter_waiting: i64, // atomic; 1 while nx_pool_wait is FUTEX-blocked,
68 // 0 while it spins or runs. Workers wake the
69 // submitter ONLY when this is 1 -> skips the
70 // futile per-task completion wake during the
71 // submitter's spin phase (2026-07-10).
72}
73
74const NX_POOL_OFF_NEXT_SLOT: i64 = 24
75const NX_POOL_OFF_ALIVE: i64 = 40
76const NX_POOL_OFF_SUBMITTED: i64 = 48
77const NX_POOL_OFF_COMPLETED: i64 = 56
78const NX_POOL_OFF_TASK_EPOCH: i64 = 64
79const NX_POOL_OFF_WAITING: i64 = 72
80
81// ===== spin-then-BLOCK (futex eventcount, 2026-07-10) ==============
82// Pre-fix, idle workers yield-spun FOREVER in nx_chan_recv and the
83// submitter yield-spun in nx_pool_wait: a 14-worker pool burned 14
84// cores while idle (two LLM seats = 28 spinning threads) and every
85// dispatch paid ~1.5ms of scheduler round-trips through the yield
86// storm (MEASURED nx_q8_matmul_micro: a 128-wide matmul took LONGER
87// than a 4864-wide one). Now: bounded spin (catches back-to-back
88// dispatch), then futex-sleep. Idle pool = ZERO CPU. Protocol is
89// the standard eventcount: sleeper loads the epoch, re-checks the
90// condition, then FUTEX_WAITs on the epoch's low 32 bits; the waker
91// makes work visible FIRST, then bumps the epoch and FUTEX_WAKEs.
92// A bump between load and wait makes the wait return immediately
93// (EAGAIN) -- no lost wakeups. rv64 futex=98 -> x86 202 (table row
94// blessed 2026-07-10).
95
96const NX_FUTEX: i64 = 98 // rv64 futex (table -> x86_64 202)
97const NX_FUTEX_WAIT_PRIV: i64 = 128 // FUTEX_WAIT | FUTEX_PRIVATE_FLAG
98const NX_FUTEX_WAKE_PRIV: i64 = 129 // FUTEX_WAKE | FUTEX_PRIVATE_FLAG
99const NX_POOL_SPIN: i64 = 64 // yields before sleeping. (TESTED 2026-07-10: busy-spin@50k helped pteam +54% but pool only +3% = noise, at the cost of idle-CPU burn for ALL shared-pool users -> reverted to idle-friendly yield-spin.)
100
101// Sleep until *addr's low 32 bits differ from seen's low 32 bits.
102// NATIVE-THREADED mode: workers spawn via sys_thread_create (->CreateThread on a
103// native PE) and futex is UNAVAILABLE (no WaitOnAddress bridge) -> the idle-wait
104// becomes a BUSY-SPIN (futex_wait/wake become no-ops; the channel is already
105// lock-free-spin). Additive, default 0 = normal futex pool (WSL2/Linux).
106static g_pool_native: i64
107func nx_pool_set_native(v: i64) -> i64 { g_pool_native = v; return 0 }
108func nx_pool_is_native() -> i64 { return g_pool_native }
109
110// NAMED futex shims (patchable): Linux = the futex syscall; on a native PE the
111// emitter OVERWRITES these with jmp -> WaitOnAddress / WakeByAddressAll thunks,
112// so the pool's idle-wait SLEEPS (not busy-spin) on native too.
113func sys_futex_wait(addr: i64, seen: i64) -> i64 {
114 return __syscall(NX_FUTEX, addr, NX_FUTEX_WAIT_PRIV, seen & 0xFFFFFFFF, 0, 0, 0)
115}
116func sys_futex_wake(addr: i64) -> i64 {
117 return __syscall(NX_FUTEX, addr, NX_FUTEX_WAKE_PRIV, NX_MAGIC_2147483647, 0, 0, 0)
118}
119func _pool_futex_wait(addr: *i64, seen: i64) -> i64 {
120 return sys_futex_wait(addr as i64, seen)
121}
122// Wake every sleeper on *addr.
123func _pool_futex_wake_all(addr: *i64) -> i64 {
124 return sys_futex_wake(addr as i64)
125}
126
127// Worker thread main. Spins on the channel; runs each task; bumps
128// the completed counter; exits when is_sentinel == 1.
129func _nx_pool_worker_main(arg: *u8) -> i64 {
130 let pool: *NxThreadPool = arg as *NxThreadPool
131 let c: *NxChan = pool.chan_ptr as *NxChan
132 let done_addr: *i64 = ((arg as i64) + NX_POOL_OFF_COMPLETED) as *i64
133 let alive_addr: *i64 = ((arg as i64) + NX_POOL_OFF_ALIVE) as *i64
134 let epoch_addr: *i64 = ((arg as i64) + NX_POOL_OFF_TASK_EPOCH) as *i64
135 let box: *i64 = sys_mmap(16) as *i64 // try_recv out box, ONCE per worker
136
137 var keep_going: i64 = 1
138 while keep_going == 1 {
139 // spin-then-block receive (see eventcount comment above).
140 var got: i64 = 0
141 while got == 0 {
142 var spin: i64 = 0
143 while spin < NX_POOL_SPIN {
144 got = nx_chan_try_recv(c, box)
145 if got == 1 { spin = NX_POOL_SPIN } else {
146 nx_thread_yield()
147 spin = spin + 1
148 }
149 }
150 if got == 0 {
151 let e: i64 = nx_atom_load_i64(epoch_addr, NX_MO_SEQ_CST)
152 got = nx_chan_try_recv(c, box) // re-check after epoch load
153 if got == 0 { _pool_futex_wait(epoch_addr, e) }
154 }
155 }
156 let task: *NxPoolTask = box[0] as *NxPoolTask
157 if task.is_sentinel == 1 {
158 keep_going = 0
159 } else {
160 let fp: func(i64) -> i64 = task.fn_ptr
161 fp(task.ctx)
162 nx_atom_faa_i64(done_addr, 1, NX_MO_SEQ_CST)
163 // wake the submitter ONLY if it is actually FUTEX-blocked (flag=1);
164 // during its spin phase the wake is a futile syscall. Safe: the
165 // submitter sets waiting=1 BEFORE re-checking completed + waiting,
166 // and futex_wait's compare-value handles any residual race.
167 let wa_addr: *i64 = ((arg as i64) + NX_POOL_OFF_WAITING) as *i64
168 if nx_atom_load_i64(wa_addr, NX_MO_SEQ_CST) == 1 {
169 _pool_futex_wake_all(done_addr)
170 }
171 }
172 }
173 nx_atom_faa_i64(alive_addr, -1, NX_MO_SEQ_CST)
174 return 0
175}
176
177// Create a new pool with n_workers worker threads. If n_workers <= 0,
178// uses nx_hw_worker_count(). queue_cap is the MPMC channel depth;
179// submit blocks when full.
180func nx_pool_new(n_workers: i64, queue_cap: i64) -> *NxThreadPool {
181 if n_workers < 1 { n_workers = nx_hw_worker_count() }
182 if queue_cap < n_workers { queue_cap = n_workers * 4 }
183
184 let raw: *u8 = sys_mmap(128)
185 let pool: *NxThreadPool = raw as *NxThreadPool
186
187 let chan: *NxChan = nx_chan_new(queue_cap)
188 pool.chan_ptr = chan as i64
189
190 let arena_bytes: i64 = NX_POOL_ARENA_DEFAULT * NX_POOL_TASK_BYTES
191 let arena: *u8 = sys_mmap(arena_bytes)
192 pool.arena_base = arena as i64
193 pool.arena_cap = NX_POOL_ARENA_DEFAULT
194 pool.next_task_slot = 0
195 pool.n_workers = n_workers
196 pool.workers_alive = n_workers
197 pool.tasks_submitted = 0
198 pool.tasks_completed = 0
199 pool.task_epoch = 0
200 pool.submitter_waiting = 0
201
202 // Spawn workers.
203 var i: i64 = 0
204 while i < n_workers {
205 var tid: i64 = 0
206 if g_pool_native != 0 {
207 // native PE: spawn via the patchable sys_thread_create shim (->CreateThread)
208 let argp: *i64 = sys_mmap(16) as *i64
209 argp[0] = _nx_pool_worker_main as i64
210 argp[1] = raw as i64
211 tid = sys_thread_create(argp as i64)
212 } else {
213 tid = nx_thread_spawn_fn(_nx_pool_worker_main, raw, NX_MAGIC_65536)
214 }
215 if tid <= 0 {
216 // Couldn't spawn -- workers_alive will under-count. Caller
217 // can probe via nx_pool_n_alive.
218 return pool
219 }
220 i = i + 1
221 }
222 return pool
223}
224
225// Submit a task. Returns 0 (always succeeds now). Blocks (via
226// nx_chan_send) when the channel is full.
227//
228// RING ARENA (2026-07-08): the slot index wraps modulo arena_cap so a
229// long-lived shared pool (e.g. the LLM forward's process pool doing
230// thousands of submits per token) never "exhausts" the arena. The
231// old bump allocator returned -1 past arena_cap, and callers that
232// then delta-waited on tasks that were never queued spun to the
233// 2-billion timeout (found 2026-07-08: nx_llm_forward_profile wedged
234// 428s/token once F32 block matmuls started submitting 168x/token).
235// SAFE because queue_cap (channel depth) << arena_cap: a slot written
236// at submit N is recv'd + fully read by a worker within queue_cap
237// submits, long before submit N+arena_cap could overwrite it. The
238// completed-counter (tasks_completed) is a separate monotonic FAA,
239// unaffected by slot reuse, so delta-waits stay correct.
240// SERIAL fallback: on a target with no working threads (e.g. native Windows PE,
241// where the inline-clone builtin the pool spawns workers with is a no-op), set
242// this to run every submitted task INLINE so nx_pool_wait passes immediately.
243// Default 0 = normal threaded pool (WSL2/Linux); additive, no behavior change.
244static g_pool_serial: i64
245func nx_pool_set_serial(v: i64) -> i64 { g_pool_serial = v; return 0 }
246
247func nx_pool_submit(pool: *NxThreadPool, fn: func(i64) -> i64, ctx: i64) -> i64 {
248 if g_pool_serial != 0 {
249 fn(ctx)
250 let dc: *i64 = ((pool as i64) + NX_POOL_OFF_COMPLETED) as *i64
251 nx_atom_faa_i64(dc, 1, NX_MO_SEQ_CST)
252 return 0
253 }
254 let slot_idx_addr: *i64 = ((pool as i64) + NX_POOL_OFF_NEXT_SLOT) as *i64
255 let idx: i64 = nx_atom_faa_i64(slot_idx_addr, 1, NX_MO_SEQ_CST)
256 let slot_i: i64 = idx % pool.arena_cap
257 let slot: *NxPoolTask = ((pool.arena_base) + slot_i * NX_POOL_TASK_BYTES) as *NxPoolTask
258 slot.fn_ptr = fn
259 slot.ctx = ctx
260 slot.is_sentinel = 0
261 let c: *NxChan = pool.chan_ptr as *NxChan
262 nx_chan_send(c, slot as i64)
263 let sub_addr: *i64 = ((pool as i64) + NX_POOL_OFF_SUBMITTED) as *i64
264 nx_atom_faa_i64(sub_addr, 1, NX_MO_SEQ_CST)
265 // work is VISIBLE (sent) -> bump the epoch, wake sleepers (eventcount).
266 let ep_addr: *i64 = ((pool as i64) + NX_POOL_OFF_TASK_EPOCH) as *i64
267 nx_atom_faa_i64(ep_addr, 1, NX_MO_SEQ_CST)
268 _pool_futex_wake_all(ep_addr)
269 return 0
270}
271
272// Wait until tasks_completed >= expected: bounded yield-spin (catches
273// fast completions), then futex-sleep on the completed counter itself
274// (workers wake it after every task). Returns 0 on success, -1 on the
275// iteration safety guard (would need ~2B wakeups -- effectively a hang
276// detector, same guard meaning as the old spin).
277func nx_pool_wait(pool: *NxThreadPool, expected: i64) -> i64 {
278 let done_addr: *i64 = ((pool as i64) + NX_POOL_OFF_COMPLETED) as *i64
279 let wa_addr: *i64 = ((pool as i64) + NX_POOL_OFF_WAITING) as *i64
280 var iters: i64 = 0
281 while nx_atom_load_i64(done_addr, NX_MO_SEQ_CST) < expected {
282 if iters < NX_POOL_SPIN {
283 nx_thread_yield()
284 } else {
285 // announce we are about to block, THEN re-check (eventcount): a
286 // worker completing after this store sees waiting=1 and wakes us;
287 // one completing before is caught by the re-checked `seen`.
288 nx_atom_store_i64(wa_addr, 1, NX_MO_SEQ_CST)
289 let seen: i64 = nx_atom_load_i64(done_addr, NX_MO_SEQ_CST)
290 if seen < expected { _pool_futex_wait(done_addr, seen) }
291 nx_atom_store_i64(wa_addr, 0, NX_MO_SEQ_CST)
292 }
293 iters = iters + 1
294 if iters > NX_MAGIC_2000000000 { nx_atom_store_i64(wa_addr, 0, NX_MO_SEQ_CST); return -1 }
295 }
296 nx_atom_store_i64(wa_addr, 0, NX_MO_SEQ_CST)
297 return 0
298}
299
300// Worker exit signal: send n_workers sentinel tasks then spin until
301// workers_alive reaches 0. Idempotent only in the sense that calling
302// shutdown twice would over-send sentinels (with no live workers to
303// consume them); don't.
304func nx_pool_shutdown(pool: *NxThreadPool) -> i64 {
305 let slot_idx_addr: *i64 = ((pool as i64) + NX_POOL_OFF_NEXT_SLOT) as *i64
306 let c: *NxChan = pool.chan_ptr as *NxChan
307 var i: i64 = 0
308 while i < pool.n_workers {
309 let idx: i64 = nx_atom_faa_i64(slot_idx_addr, 1, NX_MO_SEQ_CST)
310 let slot_i: i64 = idx % pool.arena_cap
311 let slot: *NxPoolTask = ((pool.arena_base) + slot_i * NX_POOL_TASK_BYTES) as *NxPoolTask
312 slot.is_sentinel = 1
313 slot.ctx = 0
314 nx_chan_send(c, slot as i64)
315 i = i + 1
316 }
317 // sentinels are visible -> wake every sleeping worker to consume them
318 // (submit's epoch bump does this for tasks; shutdown must do its own).
319 let ep_addr2: *i64 = ((pool as i64) + NX_POOL_OFF_TASK_EPOCH) as *i64
320 nx_atom_faa_i64(ep_addr2, 1, NX_MO_SEQ_CST)
321 _pool_futex_wake_all(ep_addr2)
322 let alive_addr: *i64 = ((pool as i64) + NX_POOL_OFF_ALIVE) as *i64
323 var spins: i64 = 0
324 while nx_atom_load_i64(alive_addr, NX_MO_SEQ_CST) > 0 {
325 nx_thread_yield()
326 spins = spins + 1
327 if spins > NX_MAGIC_2000000000 { return -1 }
328 }
329 return 0
330}
331
332func nx_pool_n_alive(pool: *NxThreadPool) -> i64 {
333 let alive_addr: *i64 = ((pool as i64) + NX_POOL_OFF_ALIVE) as *i64
334 return nx_atom_load_i64(alive_addr, NX_MO_SEQ_CST)
335}
336
337func nx_pool_n_completed(pool: *NxThreadPool) -> i64 {
338 let done_addr: *i64 = ((pool as i64) + NX_POOL_OFF_COMPLETED) as *i64
339 return nx_atom_load_i64(done_addr, NX_MO_SEQ_CST)
340}
341
342// ---- self-test ---------------------------------------------------
343
344// Trivial task that returns 0 (used only by the smoke).
345func _nx_pool_self_test_task(ctx: i64) -> i64 {
346 return ctx + 1
347}
348
349func main() -> i64 {
350 let pool: *NxThreadPool = nx_pool_new(2, 16)
351 if pool.n_workers != 2 { return __syscall(93, 1, 0, 0, 0, 0, 0) }
352 if pool.arena_cap != NX_POOL_ARENA_DEFAULT { return __syscall(93, 2, 0, 0, 0, 0, 0) }
353 if pool.tasks_submitted != 0 { return __syscall(93, 3, 0, 0, 0, 0, 0) }
354
355 // Submit one task; wait; check tasks_completed.
356 nx_pool_submit(pool, _nx_pool_self_test_task, 42)
357 if nx_pool_wait(pool, 1) != 0 { return __syscall(93, 4, 0, 0, 0, 0, 0) }
358 if nx_pool_n_completed(pool) != 1 { return __syscall(93, 5, 0, 0, 0, 0, 0) }
359
360 nx_pool_shutdown(pool)
361 if nx_pool_n_alive(pool) != 0 { return __syscall(93, 6, 0, 0, 0, 0, 0) }
362 return 0
363}