nx_pteam.nx source
↩ module page · 236 lines · 10718 B
1// nx_pteam.nx -- PERSISTENT BARRIER TEAM (fork-join parallel-for), the pool-
2// scaling fix. MEASURED (2026-07-10): the channel pool scales only ~4.5x on
3// 14 threads (gcc OpenMP 8.8x) because workers pick tasks off the MPMC
4// channel via CONTENDED CAS -> serialized start. A fork-join team fixes it:
5// each worker owns a FIXED band index and reads it from a shared descriptor,
6// so all workers start SIMULTANEOUSLY -- no channel, no per-task queue, one
7// generation bump instead of 14 sends+wakes. Idle = futex-sleep on the
8// generation word (0 CPU), same eventcount discipline as the futex pool.
9//
10// API: nx_pteam_new(n) ; nx_pteam_run(t, fn, ctx) calls fn(band, ctx) once
11// per band in [0,n) across the team and returns when all bands finish.
12// SINGLE-DRIVER (one thread calls run at a time) -- matches the LLM forward.
13//
14// genealogy_id: openmp_static_forkjoin + eventcount_futex
15// lineage_id: substrate_pteam_v1
16import "nx_syscalls.nx"
17import "nx_tier.nx"
18import "nx_atom.nx"
19import "nx_thread.nx"
20import "nx_thread_pool.nx"
21import "nx_hw.nx"
22
23const PT_FUTEX: i64 = 98 // rv64 futex -> x86 202 (blessed table row)
24const PT_WAIT_PRIV: i64 = 128 // FUTEX_WAIT | PRIVATE
25const PT_WAKE_PRIV: i64 = 129 // FUTEX_WAKE | PRIVATE
26const PT_SPIN: i64 = 50000 // BUSY-spin iters before sleeping (no sched_yield syscall;
27 // ~500us hot window through the matmul burst). Isolated to
28 // the pteam (matmul-only) so no shared-pool idle-CPU cost.
29const PT_SETAFF: i64 = 122 // rv64 sched_setaffinity -> x86 203 (blessed table row)
30
31// Pin the calling thread to a single core (gcc/OpenMP default -- no thread
32// migration -> cache/memory locality; the barrier-quality lever).
33func _pt_setaffinity(core: i64) -> i64 {
34 let mask: *i64 = sys_mmap(8) as *i64
35 mask[0] = 1 << core
36 return __syscall(PT_SETAFF, 0, 8, mask as i64, 0, 0, 0)
37}
38
39struct NxPTeam {
40 n_workers: i64, // off 0
41 fn_ptr: i64, // off 8 func(band:i64, ctx:i64)->i64 as i64
42 ctx: i64, // off 16
43 generation: i64, // off 24 atomic; bumped per run (the release word)
44 done_count: i64, // off 32 atomic; (legacy path -- unused by flag barrier)
45 alive: i64, // off 40 atomic; workers dec on exit
46 shutdown: i64, // off 48 atomic; 1 => workers exit
47 driver_wait: i64, // off 56 atomic; 1 while the driver is futex-blocked
48 flags: i64 // off 64 ptr to n_workers*64B: per-worker DONE flag on
49 // its OWN cache line (gcc-class barrier; each
50 // worker writes its own line -> no counter bounce)
51}
52const NX_PTEAM_BYTES: i64 = 128 // >64 for the added field + head-line pad
53const PT_OFF_GEN: i64 = 24
54const PT_OFF_DONE: i64 = 32
55const PT_OFF_ALIVE: i64 = 40
56const PT_OFF_SD: i64 = 48
57const PT_OFF_DWAIT: i64 = 56
58const PT_OFF_FLAGS: i64 = 64
59const PT_FLAG_STRIDE: i64 = 64 // one cache line per worker flag (no false sharing)
60
61// route through the patchable futex shims (Linux futex; native WaitOnAddress).
62func _pt_wait(addr: *i64, seen: i64) -> i64 {
63 return sys_futex_wait(addr as i64, seen)
64}
65func _pt_wake(addr: *i64) -> i64 {
66 return sys_futex_wake(addr as i64)
67}
68
69// per-worker bootstrap arg: [team_ptr, band_index]
70func _pteam_worker(arg: *u8) -> i64 {
71 let wa: *i64 = arg as *i64
72 let team: *NxPTeam = wa[0] as *NxPTeam
73 let band: i64 = wa[1]
74 // NOTE: _pt_setaffinity(band) MEASURED-HARMFUL on WSL2 (2026-07-10):
75 // pinning to VM vCPUs fights the hypervisor's float scheduling -> pteam
76 // 3.43x->1.57x. gcc's 8.8x uses NO affinity (OMP_PROC_BIND off) -> it was
77 // never the lever. Helper kept (may help on a NATIVE target host).
78 let gen_addr: *i64 = ((team as i64) + PT_OFF_GEN) as *i64
79 let alive_addr: *i64 = ((team as i64) + PT_OFF_ALIVE) as *i64
80 let sd_addr: *i64 = ((team as i64) + PT_OFF_SD) as *i64
81 let dw_addr: *i64 = ((team as i64) + PT_OFF_DWAIT) as *i64
82 let flags_base: i64 = nx_atom_load_i64(((team as i64) + PT_OFF_FLAGS) as *i64, NX_MO_SEQ_CST)
83 // this worker's OWN done-flag, on its OWN cache line (no false sharing);
84 // also the address this worker WAKES (driver sleeps per-pending-flag)
85 let my_flag: *i64 = (flags_base + band * PT_FLAG_STRIDE) as *i64
86
87 var last_gen: i64 = 0
88 var run: i64 = 1
89 while run == 1 {
90 // wait for generation != last_gen (spin then futex).
91 var got: i64 = 0
92 while got == 0 {
93 var spin: i64 = 0
94 while spin < PT_SPIN {
95 let g: i64 = nx_atom_load_i64(gen_addr, NX_MO_SEQ_CST)
96 if g != last_gen { got = 1; spin = PT_SPIN } else {
97 if nx_atom_load_i64(sd_addr, NX_MO_SEQ_CST) == 1 { got = 2; spin = PT_SPIN }
98 else { spin = spin + 1 } // BUSY-spin (no sched_yield)
99 }
100 }
101 if got == 0 {
102 if nx_atom_load_i64(sd_addr, NX_MO_SEQ_CST) == 1 { got = 2 } else {
103 _pt_wait(gen_addr, last_gen) // re-checks gen inside the loop
104 }
105 }
106 }
107 if got == 2 { run = 0 } else {
108 let g2: i64 = nx_atom_load_i64(gen_addr, NX_MO_SEQ_CST)
109 let fp: func(i64, i64) -> i64 = team.fn_ptr as func(i64, i64) -> i64
110 fp(band, team.ctx)
111 // publish DONE by writing THIS generation into MY OWN cache line.
112 // driver polls all N flags == generation -> no shared-counter bounce.
113 // RACE FIX: wake MY OWN flag line (the driver now sleeps on the
114 // first PENDING worker's line, not flags[0]). seq_cst order makes
115 // the skip safe: dw==0 here => this store precedes the driver's
116 // dw=1, so the driver's later snapshot of my flag sees gen.
117 nx_atom_store_i64(my_flag, g2, NX_MO_SEQ_CST)
118 if nx_atom_load_i64(dw_addr, NX_MO_SEQ_CST) == 1 { _pt_wake(my_flag) }
119 last_gen = g2
120 }
121 }
122 nx_atom_faa_i64(alive_addr, -1, NX_MO_SEQ_CST)
123 return 0
124}
125
126func nx_pteam_new(n_workers: i64) -> *NxPTeam {
127 if n_workers < 1 { n_workers = nx_hw_worker_count() }
128 let raw: *u8 = sys_mmap(NX_PTEAM_BYTES)
129 let t: *NxPTeam = raw as *NxPTeam
130 t.n_workers = n_workers
131 t.fn_ptr = 0
132 t.ctx = 0
133 t.generation = 0
134 t.done_count = 0
135 t.alive = n_workers
136 t.shutdown = 0
137 t.driver_wait = 0
138 // per-worker DONE flags, one cache line each (init 0 == "gen 0 not started")
139 let flags: *i64 = sys_mmap(n_workers * PT_FLAG_STRIDE) as *i64
140 t.flags = flags as i64
141 var f: i64 = 0
142 while f < n_workers {
143 let fa: *i64 = ((flags as i64) + f * PT_FLAG_STRIDE) as *i64
144 fa[0] = 0
145 f = f + 1
146 }
147 var i: i64 = 0
148 while i < n_workers {
149 let wa: *i64 = sys_mmap(16) as *i64 // [team, band] -- one per worker
150 wa[0] = raw as i64
151 wa[1] = i
152 // spawn via the patchable shim: [fn, ctx] (Linux=clone, native=CreateThread)
153 let sargp: *i64 = sys_mmap(16) as *i64
154 sargp[0] = _pteam_worker as i64
155 sargp[1] = wa as i64
156 let tid: i64 = sys_thread_create(sargp as i64)
157 if tid <= 0 { return t }
158 i = i + 1
159 }
160 return t
161}
162
163// Run fn(band, ctx) for band in [0, n_workers) across the team; block until
164// all bands complete. SINGLE-DRIVER.
165func _pt_all_done(flags: i64, nw: i64, gen: i64) -> i64 {
166 var i: i64 = 0
167 while i < nw {
168 if nx_atom_load_i64((flags + i * PT_FLAG_STRIDE) as *i64, NX_MO_SEQ_CST) != gen { return 0 }
169 i = i + 1
170 }
171 return 1
172}
173
174func nx_pteam_run(t: *NxPTeam, fn: func(i64, i64) -> i64, ctx: i64) -> i64 {
175 let gen_addr: *i64 = ((t as i64) + PT_OFF_GEN) as *i64
176 let dw_addr: *i64 = ((t as i64) + PT_OFF_DWAIT) as *i64
177 let flags: i64 = t.flags
178 let nw: i64 = t.n_workers
179
180 t.fn_ptr = fn as i64
181 t.ctx = ctx
182 // bump generation (the release word); workers run then write gen into their
183 // OWN flag line. seq_cst: a worker seeing the new gen runs fn from a clean
184 // state; the driver seeing flag==gen sees fn's results.
185 let gen: i64 = nx_atom_faa_i64(gen_addr, 1, NX_MO_SEQ_CST) + 1
186 _pt_wake(gen_addr) // wake any sleeping workers
187
188 // wait until every worker's flag == gen (busy-spin then futex on flag[0]).
189 var iters: i64 = 0
190 while _pt_all_done(flags, nw, gen) == 0 {
191 if iters < PT_SPIN {
192 iters = iters // BUSY-spin (no sched_yield)
193 } else {
194 nx_atom_store_i64(dw_addr, 1, NX_MO_SEQ_CST)
195 // RACE FIX (2026-07-15, found as a LIVE nondeterministic native
196 // deadlock): sleeping on flags[0] with its snapshot violates the
197 // eventcount invariant -- if worker 0 finished EARLY the snapshot
198 // already equals gen, the remaining workers' wakes land BEFORE the
199 // driver blocks, and the kernel compare (flags[0]==snapshot) then
200 // sleeps FOREVER with all work done. Sleep instead on the FIRST
201 // PENDING worker's flag with ITS snapshot: that word unchanged
202 // <=> that worker genuinely pending; its finish-store flips the
203 // compare or its wake lands after we block. Workers wake their
204 // OWN line (see _pteam_worker).
205 var pi: i64 = 0 - 1
206 var si: i64 = 0
207 var j: i64 = 0
208 while j < nw {
209 let vj: i64 = nx_atom_load_i64((flags + j * PT_FLAG_STRIDE) as *i64, NX_MO_SEQ_CST)
210 if vj != gen { pi = j; si = vj; j = nw } else { j = j + 1 }
211 }
212 if pi >= 0 { _pt_wait((flags + pi * PT_FLAG_STRIDE) as *i64, si) }
213 nx_atom_store_i64(dw_addr, 0, NX_MO_SEQ_CST)
214 }
215 iters = iters + 1
216 if iters > 2000000000 { nx_atom_store_i64(dw_addr, 0, NX_MO_SEQ_CST); return 0 - 1 }
217 }
218 nx_atom_store_i64(dw_addr, 0, NX_MO_SEQ_CST)
219 return 0
220}
221
222func nx_pteam_shutdown(t: *NxPTeam) -> i64 {
223 let gen_addr: *i64 = ((t as i64) + PT_OFF_GEN) as *i64
224 let alive_addr: *i64 = ((t as i64) + PT_OFF_ALIVE) as *i64
225 let sd_addr: *i64 = ((t as i64) + PT_OFF_SD) as *i64
226 nx_atom_store_i64(sd_addr, 1, NX_MO_SEQ_CST)
227 nx_atom_faa_i64(gen_addr, 1, NX_MO_SEQ_CST) // nudge waiters to re-check shutdown
228 _pt_wake(gen_addr)
229 var spins: i64 = 0
230 while nx_atom_load_i64(alive_addr, NX_MO_SEQ_CST) > 0 {
231 nx_thread_yield()
232 spins = spins + 1
233 if spins > 2000000000 { return 0 - 1 }
234 }
235 return 0
236}