nx_pipeline_test.nx source
↩ module page · 96 lines · 3486 B
1// nx_pipeline_test.nx -- 3-stage pipeline stress + backpressure check.
2//
3// Pipeline: input -> [stage 0: x*2] -> [stage 1: x+1] -> [stage 2: x*10]
4// Each stage runs n_workers parallel workers. Output of input i is
5// ((i*2)+1)*10 = 20i + 10.
6//
7// Push 100 inputs, drain 100 outputs, verify EVERY output's value is
8// in the expected set {20i + 10 | i in 0..99}. Don't require FIFO
9// across stages because parallel workers within a stage may reorder.
10
11import "nx_kernel_v2.nx"
12import "nx_log.nx"
13import "nx_atom.nx"
14import "nx_thread_pool.nx"
15import "nx_pipeline.nx"
16import "nx_hw.nx"
17
18const N_INPUTS: i64 = 64
19
20func stage0_double(x: i64) -> i64 { return x * 2 }
21func stage1_add1(x: i64) -> i64 { return x + 1 }
22func stage2_x10(x: i64) -> i64 { return x * 10 }
23
24func main() -> nx_exit {
25 var workers_per_stage: i64 = nx_hw_worker_count()
26 if workers_per_stage < 2 { workers_per_stage = 2 }
27 let total_pool_workers: i64 = workers_per_stage * 3 + 2 // pipeline + slack
28 let pool: *NxThreadPool = nx_pool_new(total_pool_workers, 64)
29 let pl: *NxPipeline = nx_pipeline_new(pool, 8, workers_per_stage, 16)
30
31 nx_pipeline_add_stage(pl, stage0_double)
32 nx_pipeline_add_stage(pl, stage1_add1)
33 nx_pipeline_add_stage(pl, stage2_x10)
34 nx_pipeline_start(pl)
35
36 println("=== nx_pipeline 3-stage stress ===" as *u8)
37 println("Workers per stage:" as *u8); print_i64(workers_per_stage); println("" as *u8)
38 println("Inputs:" as *u8); print_i64(N_INPUTS); println("" as *u8)
39
40 // Bitmap of which expected values we've seen (set element i when
41 // we receive 20i+10).
42 let seen_raw: *u8 = sys_mmap(N_INPUTS * 8)
43 let seen: *i64 = seen_raw as *i64
44 var z: i64 = 0
45 while z < N_INPUTS { seen[z] = 0; z = z + 1 }
46
47 println("Pushing inputs..." as *u8)
48 // Push N_INPUTS values; channel backpressure makes producer block
49 // when downstream falls behind -- that's the test of bounded queues.
50 var p: i64 = 0
51 while p < N_INPUTS {
52 nx_pipeline_push(pl, p)
53 p = p + 1
54 }
55 println("All pushed; draining..." as *u8)
56
57 // Drain N_INPUTS values, marking the expected slot for each.
58 var got: i64 = 0
59 while got < N_INPUTS {
60 let v: i64 = nx_pipeline_recv(pl)
61 // v should equal (i*2 + 1) * 10 for some i in 0..N_INPUTS.
62 // Solve: i = (v / 10 - 1) / 2. Verify exact reverse:
63 let i_guess: i64 = (v / 10 - 1) / 2
64 let expected: i64 = (i_guess * 2 + 1) * 10
65 if v != expected {
66 println("FAIL: output not in expected set" as *u8)
67 print_i64(v); println("" as *u8); return 1
68 }
69 if i_guess < 0 { println("FAIL: idx <0" as *u8); return 2 }
70 if i_guess >= N_INPUTS { println("FAIL: idx >= N" as *u8); return 3 }
71 if seen[i_guess] != 0 {
72 println("FAIL: duplicate output" as *u8); return 4
73 }
74 seen[i_guess] = 1
75 got = got + 1
76 }
77
78 // Verify all 100 slots were marked.
79 var u: i64 = 0
80 while u < N_INPUTS {
81 if seen[u] != 1 { println("FAIL: missing output" as *u8); return 5 }
82 u = u + 1
83 }
84 println("PASS: all outputs expected, no dups, no missing" as *u8)
85
86 // Clean shutdown: send sentinels, drain residual.
87 nx_pipeline_finish(pl)
88 var c: i64 = 0
89 while c < workers_per_stage {
90 nx_pipeline_recv(pl)
91 c = c + 1
92 }
93 nx_pool_shutdown(pool)
94 println("Pipeline drained + pool shut down cleanly." as *u8)
95 return 0
96}