nx_fabric_collective.nx source
↩ module page · 183 lines · 7699 B
1// nx_fabric_collective.nx -- R4 of the SkyHammer fabric exceed ladder: COLLECTIVE COMMUNICATION
2// (the actual AI-training workload the fabric exists to accelerate). Ladder R4, the last MISSING axis F.
3//
4// SkyHammer is "AI-native ... for GPU-era COLLECTIVE workloads" (fab_collective.raw: All-Reduce/All-Gather/
5// Broadcast). All-reduce = element-wise sum of N processes' vectors, delivered to ALL. The naive way routes
6// everything through a root -> the root moves O(N)*data = a bottleneck that grows with N. RING all-reduce
7// (NCCL's bandwidth-optimal algorithm) moves a CONSTANT ~2*data per process regardless of N: a reduce-scatter
8// phase (N-1 steps) then an all-gather phase (N-1 steps), each process touching only data/N per step.
9//
10// This rung performs the REAL reduction (integer, no-float) on N in-memory process vectors and PROVES the
11// result is byte-exact the true sum on every process, while measuring the per-process communication (the
12// bottleneck) -- the honest software expression of "collective acceleration at scale". No external organ to
13// compose (pure algorithm). expect_exit: 0 license_tier: ORIGINAL
14import "nx_syscalls.nx"
15
16const COL_RING: i64 = 0 // bandwidth-optimal ring all-reduce
17const COL_NAIVE: i64 = 1 // root-gather + broadcast (O(N) at the root)
18const COL_RING_NOSUM: i64 = 2 // NC1: ring DATA MOVEMENT but the reduction (add) disabled -> wrong result
19
20struct ColMetrics {
21 max_bytes: i64, // max per-process elements moved (the bottleneck link)
22 correct: i64, // 1 iff EVERY process ended byte-exact == the true element-wise sum
23 n: i64,
24}
25
26func col_mod(a: i64, n: i64) -> i64 { var r: i64 = a % n; if r < 0 { r = r + n } return r }
27
28// Run an all-reduce of N processes x V-element integer vectors. Fills *mo; writes process 0's final
29// vector into out[0..V) (for cross-algorithm agreement checks). Returns 0.
30func col_run(algo: i64, N: i64, V: i64, mo: *ColMetrics, out: *i64) -> i64 {
31 let CS: i64 = V / N // chunk size (V divisible by N)
32 let proc: *i64 = sys_mmap(8 * N * V) as *i64
33 let bytes: *i64 = sys_mmap(8 * N) as *i64 // per-process element-move counter (mmap zero)
34 let truesum: *i64 = sys_mmap(8 * V) as *i64
35
36 // init process vectors + the independent reference sum
37 var i: i64 = 0
38 while i < V {
39 var p: i64 = 0
40 var s: i64 = 0
41 while p < N {
42 let val: i64 = (p * 7 + i * 3 + 5) % 251
43 proc[p * V + i] = val
44 s = s + val
45 p = p + 1
46 }
47 truesum[i] = s
48 i = i + 1
49 }
50
51 if algo == COL_NAIVE {
52 let acc: *i64 = sys_mmap(8 * V) as *i64
53 i = 0
54 while i < V { acc[i] = proc[0 * V + i]; i = i + 1 } // root starts with its own
55 var p: i64 = 1
56 while p < N {
57 i = 0
58 while i < V { acc[i] = acc[i] + proc[p * V + i]; i = i + 1 }
59 bytes[p] = bytes[p] + V // p -> root (send full vector)
60 bytes[0] = bytes[0] + V // root receives
61 p = p + 1
62 }
63 p = 1
64 while p < N { // broadcast the sum
65 i = 0
66 while i < V { proc[p * V + i] = acc[i]; i = i + 1 }
67 bytes[0] = bytes[0] + V // root sends
68 bytes[p] = bytes[p] + V // p receives
69 p = p + 1
70 }
71 i = 0
72 while i < V { proc[0 * V + i] = acc[i]; i = i + 1 }
73 } else {
74 let snap: *i64 = sys_mmap(8 * N * V) as *i64
75 // ---- reduce-scatter: N-1 steps; process p ADDS prev's chunk (p-1-t) into its own ----
76 var t: i64 = 0
77 while t < N - 1 {
78 var x: i64 = 0
79 while x < N * V { snap[x] = proc[x]; x = x + 1 }
80 var p: i64 = 0
81 while p < N {
82 let prev: i64 = col_mod(p - 1, N)
83 let rc: i64 = col_mod(p - 1 - t, N)
84 var j: i64 = 0
85 while j < CS {
86 let idx: i64 = rc * CS + j
87 if algo == COL_RING { proc[p * V + idx] = proc[p * V + idx] + snap[prev * V + idx] }
88 else { proc[p * V + idx] = snap[prev * V + idx] } // NOSUM: move, do NOT reduce
89 j = j + 1
90 }
91 bytes[p] = bytes[p] + CS
92 p = p + 1
93 }
94 t = t + 1
95 }
96 // ---- all-gather: N-1 steps; circulate the summed chunks (overwrite) ----
97 t = 0
98 while t < N - 1 {
99 var x: i64 = 0
100 while x < N * V { snap[x] = proc[x]; x = x + 1 }
101 var p: i64 = 0
102 while p < N {
103 let prev: i64 = col_mod(p - 1, N)
104 let rc: i64 = col_mod(p - t, N)
105 var j: i64 = 0
106 while j < CS {
107 let idx: i64 = rc * CS + j
108 proc[p * V + idx] = snap[prev * V + idx]
109 j = j + 1
110 }
111 bytes[p] = bytes[p] + CS
112 p = p + 1
113 }
114 t = t + 1
115 }
116 }
117
118 // verify EVERY process == the true sum
119 var correct: i64 = 1
120 var p2: i64 = 0
121 while p2 < N {
122 var i2: i64 = 0
123 while i2 < V { if proc[p2 * V + i2] != truesum[i2] { correct = 0 } i2 = i2 + 1 }
124 p2 = p2 + 1
125 }
126 var maxb: i64 = 0
127 p2 = 0
128 while p2 < N { if bytes[p2] > maxb { maxb = bytes[p2] } p2 = p2 + 1 }
129
130 var oi: i64 = 0
131 while oi < V { out[oi] = proc[0 * V + oi]; oi = oi + 1 }
132
133 mo.max_bytes = maxb
134 mo.correct = correct
135 mo.n = N
136 return 0
137}
138
139// ring_allreduce(in, N, V, out) -- R4's ring all-reduce exposed as a LIBRARY taking INPUT vectors.
140// col_run self-generates its data for the benchmark; the DDP trainer (research ยง4) needs the reusable form.
141// Sums N input vectors of V elements (V divisible by N) into out[0..V) via reduce-scatter + all-gather, the
142// SAME order-independent core col_run uses. DRY-debt: a future refactor extracts the shared ring core; kept
143// separate here so the verified col_run is left byte-for-byte untouched.
144func ring_allreduce(src: *i64, N: i64, V: i64, out: *i64) -> i64 {
145 let CS: i64 = V / N
146 let proc: *i64 = sys_mmap(8 * N * V) as *i64
147 var k: i64 = 0
148 while k < N * V { proc[k] = src[k]; k = k + 1 }
149 let snap: *i64 = sys_mmap(8 * N * V) as *i64
150 var t: i64 = 0
151 while t < N - 1 { // reduce-scatter (ADD)
152 var x: i64 = 0
153 while x < N * V { snap[x] = proc[x]; x = x + 1 }
154 var p: i64 = 0
155 while p < N {
156 let prev: i64 = col_mod(p - 1, N)
157 let rc: i64 = col_mod(p - 1 - t, N)
158 var j: i64 = 0
159 while j < CS { let idx: i64 = rc * CS + j; proc[p * V + idx] = proc[p * V + idx] + snap[prev * V + idx]; j = j + 1 }
160 p = p + 1
161 }
162 t = t + 1
163 }
164 t = 0
165 while t < N - 1 { // all-gather (circulate)
166 var x: i64 = 0
167 while x < N * V { snap[x] = proc[x]; x = x + 1 }
168 var p: i64 = 0
169 while p < N {
170 let prev: i64 = col_mod(p - 1, N)
171 let rc: i64 = col_mod(p - t, N)
172 var j: i64 = 0
173 while j < CS { let idx: i64 = rc * CS + j; proc[p * V + idx] = snap[prev * V + idx]; j = j + 1 }
174 p = p + 1
175 }
176 t = t + 1
177 }
178 var oi: i64 = 0
179 while oi < V { out[oi] = proc[0 * V + oi]; oi = oi + 1 }
180 return 0
181}
182
183func main() -> i64 { return 0 }