nx_parallel.nx source
↩ module page · 371 lines · 13103 B
1// nx_parallel.nx -- parallel_for / parallel_map / parallel_reduce.
2//
3// Built on top of nx_thread_pool: each primitive partitions its
4// iteration range into n_workers chunks and submits one task per
5// chunk. Wait until pool.tasks_completed has advanced by n_chunks
6// (snapshot-based, so multiple parallel_* calls compose).
7//
8// Why partition rather than per-iter-submit:
9// * Per-iter submit pays ~20ns chan_send + dispatch per index;
10// for cheap bodies (e.g. element-wise add) that dwarfs the
11// work. Chunked submission amortises chan overhead across
12// ~ITER_COUNT/n_workers items.
13//
14// Static even partition for MVP -- range is divided into roughly
15// equal contiguous chunks. Dynamic / work-stealing partitioning
16// is the L7-evolution follow-up once we measure imbalanced
17// workloads suffering >5% throughput loss.
18//
19// Public API (single namespace; sealed strategy enum):
20// nx_parallel_for(pool, begin, end, fn(i64))
21// nx_parallel_map_i64(pool, in_ptr, out_ptr, n, fn(i64) -> i64)
22// nx_parallel_reduce_i64(pool, ptr, n, init, fn(i64, i64) -> i64)
23//
24// Composes against: [[nx_thread_pool_shared_queue]] (workers),
25// [[vyukov_mpmc_channel]] (task queue), [[atomic_intrinsics_real_amo]]
26// (per-chunk reduction sum), [[nx_hw_dynamic_probes]] (chunk count
27// defaults to nx_hw_worker_count).
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_atom.nx"
37import "nx_thread_pool.nx"
38import "nx_hw.nx"
39
40struct NxPfChunk {
41 start: i64,
42 end: i64,
43 fn_ptr: func(i64) -> i64,
44}
45
46const NX_PF_CHUNK_BYTES: i64 = 24
47
48struct NxMapChunk {
49 start: i64,
50 end: i64,
51 in_ptr: i64, // *i64 (raw)
52 out_ptr: i64, // *i64 (raw)
53 fn_ptr: func(i64) -> i64,
54}
55
56const NX_MAP_CHUNK_BYTES: i64 = 40
57
58struct NxReduceChunk {
59 start: i64,
60 end: i64,
61 arr_ptr: i64, // *i64 (raw)
62 fn_ptr: func(i64, i64) -> i64, // (acc, x) -> acc
63 out_sum: i64, // chunk-local result
64 done_flag: i64, // atomic; 1 when worker finished
65}
66
67const NX_REDUCE_CHUNK_BYTES: i64 = 48
68
69// Worker function for parallel_for: iterate start..end and call fn(i).
70func _nx_pf_worker(ctx: i64) -> i64 {
71 let ch: *NxPfChunk = ctx as *NxPfChunk
72 let fp: func(i64) -> i64 = ch.fn_ptr
73 var i: i64 = ch.start
74 while i < ch.end {
75 fp(i)
76 i = i + 1
77 }
78 return 0
79}
80
81// Worker function for parallel_map: out[i] = fn(in[i]) for i in chunk.
82func _nx_map_worker(ctx: i64) -> i64 {
83 let ch: *NxMapChunk = ctx as *NxMapChunk
84 let fp: func(i64) -> i64 = ch.fn_ptr
85 let in_arr: *i64 = ch.in_ptr as *i64
86 let out_arr: *i64 = ch.out_ptr as *i64
87 var i: i64 = ch.start
88 while i < ch.end {
89 out_arr[i] = fp(in_arr[i])
90 i = i + 1
91 }
92 return 0
93}
94
95// Worker function for parallel_reduce: walk start..end accumulating
96// via fn; store result in chunk.out_sum so the orchestrator can
97// final-combine without atomic contention on a shared accumulator.
98func _nx_reduce_worker(ctx: i64) -> i64 {
99 let ch: *NxReduceChunk = ctx as *NxReduceChunk
100 let fp: func(i64, i64) -> i64 = ch.fn_ptr
101 let arr: *i64 = ch.arr_ptr as *i64
102 var acc: i64 = 0
103 var i: i64 = ch.start
104 while i < ch.end {
105 acc = fp(acc, arr[i])
106 i = i + 1
107 }
108 ch.out_sum = acc
109 let done_addr: *i64 = ((ctx as i64) + 40) as *i64
110 nx_atom_store_i64(done_addr, 1, NX_MO_RELEASE)
111 return 0
112}
113
114// Partition [begin, end) into n_chunks roughly-equal pieces. Returns
115// the start of chunk_idx (0-based). Last chunk takes the remainder.
116func _nx_chunk_start(begin: i64, end: i64, n_chunks: i64, idx: i64) -> i64 {
117 let total: i64 = end - begin
118 let base: i64 = total / n_chunks
119 return begin + idx * base
120}
121
122func _nx_chunk_end(begin: i64, end: i64, n_chunks: i64, idx: i64) -> i64 {
123 if idx == n_chunks - 1 { return end }
124 let total: i64 = end - begin
125 let base: i64 = total / n_chunks
126 return begin + (idx + 1) * base
127}
128
129// parallel_for: apply fn(i) for each i in [begin, end). Returns 0
130// on success, -1 on internal failure (chunk allocation OOM or pool
131// wait timeout).
132func nx_parallel_for(pool: *NxThreadPool, begin: i64, end: i64,
133 fn: func(i64) -> i64) -> i64 {
134 if end <= begin { return 0 }
135
136 let n_workers: i64 = pool.n_workers
137 var n_chunks: i64 = n_workers
138 if n_chunks > end - begin { n_chunks = end - begin }
139 if n_chunks < 1 { n_chunks = 1 }
140
141 let chunks_raw: *u8 = sys_mmap(n_chunks * NX_PF_CHUNK_BYTES)
142 let done_addr: *i64 = ((pool as i64) + NX_POOL_OFF_COMPLETED) as *i64
143 let baseline: i64 = nx_atom_load_i64(done_addr, NX_MO_SEQ_CST)
144
145 var k: i64 = 0
146 while k < n_chunks {
147 let ch_addr: i64 = (chunks_raw as i64) + k * NX_PF_CHUNK_BYTES
148 let ch: *NxPfChunk = ch_addr as *NxPfChunk
149 ch.start = _nx_chunk_start(begin, end, n_chunks, k)
150 ch.end = _nx_chunk_end(begin, end, n_chunks, k)
151 ch.fn_ptr = fn
152 if nx_pool_submit(pool, _nx_pf_worker, ch_addr) != 0 { return -1 }
153 k = k + 1
154 }
155
156 // Wait until pool.tasks_completed has advanced by n_chunks.
157 let target: i64 = baseline + n_chunks
158 var spins: i64 = 0
159 while nx_atom_load_i64(done_addr, NX_MO_SEQ_CST) < target {
160 spins = spins + 1
161 if spins > 2000000000 { return -1 }
162 }
163 return 0
164}
165
166// parallel_map: out[i] = fn(in[i]) for i in [0, n). in and out
167// may NOT alias.
168func nx_parallel_map_i64(pool: *NxThreadPool, in_ptr: *i64, out_ptr: *i64,
169 n: i64, fn: func(i64) -> i64) -> i64 {
170 if n <= 0 { return 0 }
171
172 let n_workers: i64 = pool.n_workers
173 var n_chunks: i64 = n_workers
174 if n_chunks > n { n_chunks = n }
175 if n_chunks < 1 { n_chunks = 1 }
176
177 let chunks_raw: *u8 = sys_mmap(n_chunks * NX_MAP_CHUNK_BYTES)
178 let done_addr: *i64 = ((pool as i64) + NX_POOL_OFF_COMPLETED) as *i64
179 let baseline: i64 = nx_atom_load_i64(done_addr, NX_MO_SEQ_CST)
180
181 var k: i64 = 0
182 while k < n_chunks {
183 let ch_addr: i64 = (chunks_raw as i64) + k * NX_MAP_CHUNK_BYTES
184 let ch: *NxMapChunk = ch_addr as *NxMapChunk
185 ch.start = _nx_chunk_start(0, n, n_chunks, k)
186 ch.end = _nx_chunk_end(0, n, n_chunks, k)
187 ch.in_ptr = in_ptr as i64
188 ch.out_ptr = out_ptr as i64
189 ch.fn_ptr = fn
190 if nx_pool_submit(pool, _nx_map_worker, ch_addr) != 0 { return -1 }
191 k = k + 1
192 }
193
194 let target: i64 = baseline + n_chunks
195 var spins: i64 = 0
196 while nx_atom_load_i64(done_addr, NX_MO_SEQ_CST) < target {
197 spins = spins + 1
198 if spins > 2000000000 { return -1 }
199 }
200 return 0
201}
202
203// parallel_reduce: fold over arr[0..n) with combine fn(acc, x).
204// init is the seed for the orchestrator's final combine of chunk
205// results. Caller's fn must be associative (true sum/product/min/max).
206func nx_parallel_reduce_i64(pool: *NxThreadPool, arr_ptr: *i64, n: i64,
207 init: i64, fn: func(i64, i64) -> i64) -> i64 {
208 if n <= 0 { return init }
209
210 let n_workers: i64 = pool.n_workers
211 var n_chunks: i64 = n_workers
212 if n_chunks > n { n_chunks = n }
213 if n_chunks < 1 { n_chunks = 1 }
214
215 let chunks_raw: *u8 = sys_mmap(n_chunks * NX_REDUCE_CHUNK_BYTES)
216 let done_addr: *i64 = ((pool as i64) + NX_POOL_OFF_COMPLETED) as *i64
217 let baseline: i64 = nx_atom_load_i64(done_addr, NX_MO_SEQ_CST)
218
219 var k: i64 = 0
220 while k < n_chunks {
221 let ch_addr: i64 = (chunks_raw as i64) + k * NX_REDUCE_CHUNK_BYTES
222 let ch: *NxReduceChunk = ch_addr as *NxReduceChunk
223 ch.start = _nx_chunk_start(0, n, n_chunks, k)
224 ch.end = _nx_chunk_end(0, n, n_chunks, k)
225 ch.arr_ptr = arr_ptr as i64
226 ch.fn_ptr = fn
227 ch.out_sum = 0
228 ch.done_flag = 0
229 if nx_pool_submit(pool, _nx_reduce_worker, ch_addr) != 0 { return init }
230 k = k + 1
231 }
232
233 let target: i64 = baseline + n_chunks
234 var spins: i64 = 0
235 while nx_atom_load_i64(done_addr, NX_MO_SEQ_CST) < target {
236 spins = spins + 1
237 if spins > 2000000000 { return init }
238 }
239
240 // Final combine of per-chunk results on the orchestrator thread.
241 var acc: i64 = init
242 var j: i64 = 0
243 while j < n_chunks {
244 let ch_addr: i64 = (chunks_raw as i64) + j * NX_REDUCE_CHUNK_BYTES
245 let ch: *NxReduceChunk = ch_addr as *NxReduceChunk
246 acc = fn(acc, ch.out_sum)
247 j = j + 1
248 }
249 return acc
250}
251
252// ---- L7 x L8 composition: parallel reduce with SIMD inner loop ----
253
254// SIMD-accelerated worker for sum-reduce. Each chunk walks its
255// range in 4-element strides using __simd_vload + vreduce_sum,
256// with a scalar tail handler for any remainder. Result lives in
257// chunk.out_sum (re-using NxReduceChunk to avoid a new struct).
258func _nx_reduce_worker_simd_sum(ctx: i64) -> i64 {
259 let ch: *NxReduceChunk = ctx as *NxReduceChunk
260 let arr: *i64 = ch.arr_ptr as *i64
261 var acc: i64 = 0
262 let span: i64 = ch.end - ch.start
263 let tail: i64 = span - (span / 4) * 4 // 0..3 remainder
264 let end_aligned: i64 = ch.end - tail
265 var i: i64 = ch.start
266 while i < end_aligned {
267 let p: *i64 = ((arr as i64) + i * 8) as *i64
268 let v: i64 = __simd_vload_i64_x4(p)
269 acc = acc + __simd_vreduce_sum_i64_x4(v)
270 i = i + 4
271 }
272 // Scalar tail.
273 while i < ch.end {
274 acc = acc + arr[i]
275 i = i + 1
276 }
277 ch.out_sum = acc
278 let done_addr: *i64 = ((ctx as i64) + 40) as *i64
279 nx_atom_store_i64(done_addr, 1, NX_MO_RELEASE)
280 return 0
281}
282
283// parallel_reduce_sum_simd_i64 -- specialised to associative i64 sum
284// where the per-chunk reducer uses SIMD i64x4 horizontal reduction.
285// Composes the full MIMD x SIMD product: N_cores chunks each running
286// 4 lanes wide. On real silicon expected speedup is ~N_cores *
287// VLEN/8 over scalar single-core (e.g., 16 cores * 4 lanes = 64x).
288func nx_parallel_reduce_sum_simd_i64(pool: *NxThreadPool, arr_ptr: *i64,
289 n: i64, init: i64) -> i64 {
290 if n <= 0 { return init }
291 let n_workers: i64 = pool.n_workers
292 var n_chunks: i64 = n_workers
293 if n_chunks > n { n_chunks = n }
294 if n_chunks < 1 { n_chunks = 1 }
295
296 let chunks_raw: *u8 = sys_mmap(n_chunks * NX_REDUCE_CHUNK_BYTES)
297 let done_addr: *i64 = ((pool as i64) + NX_POOL_OFF_COMPLETED) as *i64
298 let baseline: i64 = nx_atom_load_i64(done_addr, NX_MO_SEQ_CST)
299
300 var k: i64 = 0
301 while k < n_chunks {
302 let ch_addr: i64 = (chunks_raw as i64) + k * NX_REDUCE_CHUNK_BYTES
303 let ch: *NxReduceChunk = ch_addr as *NxReduceChunk
304 ch.start = _nx_chunk_start(0, n, n_chunks, k)
305 ch.end = _nx_chunk_end(0, n, n_chunks, k)
306 ch.arr_ptr = arr_ptr as i64
307 // ch.fn_ptr unused by SIMD worker; leave zero.
308 ch.out_sum = 0
309 ch.done_flag = 0
310 if nx_pool_submit(pool, _nx_reduce_worker_simd_sum, ch_addr) != 0 {
311 return init
312 }
313 k = k + 1
314 }
315
316 let target: i64 = baseline + n_chunks
317 var spins: i64 = 0
318 while nx_atom_load_i64(done_addr, NX_MO_SEQ_CST) < target {
319 nx_thread_yield()
320 spins = spins + 1
321 if spins > 2000000000 { return init }
322 }
323
324 var acc: i64 = init
325 var j: i64 = 0
326 while j < n_chunks {
327 let ch_addr: i64 = (chunks_raw as i64) + j * NX_REDUCE_CHUNK_BYTES
328 let ch: *NxReduceChunk = ch_addr as *NxReduceChunk
329 acc = acc + ch.out_sum
330 j = j + 1
331 }
332 return acc
333}
334
335// ---- self-test ---------------------------------------------------
336
337func _pf_self_test_noop(i: i64) -> i64 { return i + 1 }
338func _map_self_test_double(x: i64) -> i64 { return x * 2 }
339func _reduce_self_test_add(a: i64, b: i64) -> i64 { return a + b }
340
341func main() -> i64 {
342 let pool: *NxThreadPool = nx_pool_new(2, 32)
343
344 // parallel_for empty range -- should be a no-op.
345 if nx_parallel_for(pool, 5, 5, _pf_self_test_noop) != 0 {
346 return __syscall(93, 1, 0, 0, 0, 0, 0)
347 }
348
349 // parallel_map: out[i] = in[i] * 2 over a 32-element array.
350 let in_raw: *u8 = sys_mmap(32 * 8)
351 let out_raw: *u8 = sys_mmap(32 * 8)
352 let in_arr: *i64 = in_raw as *i64
353 let out_arr: *i64 = out_raw as *i64
354 var i: i64 = 0
355 while i < 32 { in_arr[i] = i; i = i + 1 }
356 if nx_parallel_map_i64(pool, in_arr, out_arr, 32, _map_self_test_double) != 0 {
357 return __syscall(93, 2, 0, 0, 0, 0, 0)
358 }
359 var k: i64 = 0
360 while k < 32 {
361 if out_arr[k] != k * 2 { return __syscall(93, 10 + k, 0, 0, 0, 0, 0) }
362 k = k + 1
363 }
364
365 // parallel_reduce: sum 0..31 == 496.
366 let sum: i64 = nx_parallel_reduce_i64(pool, in_arr, 32, 0, _reduce_self_test_add)
367 if sum != 496 { return __syscall(93, 100, 0, 0, 0, 0, 0) }
368
369 nx_pool_shutdown(pool)
370 return 0
371}