==> src/train.c <==1/* train: SGD of y = a + b x + c x^2 + d x^3, one sample a lane, gradients
2 * butterfly-summed across the warp. In a lockstep group of N (manifest
3 * `peer = "lockstep"`, H_GROUP_RANK k, H_GROUP_N N) each step's sums go
4 * through the exchange page and every member applies slot0 + .. + slot(N-1)
5 * in rank order: bitwise-identical params (data parallel, a batch of 32 N a
6 * step). Layout: @ajax/mesh-abi-view-train. OP_INIT generates the samples on
7 * the GPU (sin in f64 arithmetic only: the same bits on every IEEE machine).
8 * The loops mirror mesh-kcheck's CPU reference index by index (bitwise checks).
9 * Ported from alpha-rust-mesh programs/train/src/lib.rs. */
10#include "mesh/rt.h"
11
12static float train_cubic(const float p[4], float x) {
13 float x2 = x * x;
14 float x3 = x2 * x;
15 return p[0] + p[1] * x + p[2] * x2 + p[3] * x3;
16}
17
18/* sin(pi x) for x in [-1, 1]: reflect into [-1/2, 1/2], then the Taylor
19 * series to a^23 in Horner form (error < 1e-20). Only +, -, *, / in f64, so
20 * the CPU copy gives the same bits. Not inlined (one shared copy). */
21__attribute__((noinline)) static float train_sin_pi(float x) {
22 double t = (double)x;
23 if (t > 0.5) t = 1.0 - t;
24 else if (t < -0.5) t = -1.0 - t;
25 double a = 3.141592653589793 * t;
26 double a2 = a * a;
27 double s = 1.0;
28 for (uint32_t k = 11; k >= 1; k--) s = 1.0 - a2 / (double)((2 * k) * (2 * k + 1)) * s;
29 return (float)(a * s);
30}
31
32/* Sample i's x: -1 + 2 i / 255. */
33static float train_sample_x(uint32_t i) { return -1.0f + 2.0f * (float)i / (float)(MESH_TRAIN_SAMPLES - 1); }
34
35/* Lockstep in a group: this member's rank k, the group's size N, the exchange
36 * page, when the manifest grants the page and meshd set H_GROUP_RANK and H_GROUP_N. */
37static bool train_peer(const mesh_ctx *ctx, uint32_t *k, uint32_t *n, mesh_exchange *x) {
38 return mesh_ctx_group(ctx, k, n) && mesh_ctx_exchange(ctx, x);
39}
40
41/* Fresh state: samples, hyperparameters, params and progress zero. The
42 * exchange page is meshd's (zeroed while we are halted): clearing it here could
43 * wipe a peer round mirrored in after the HELLO. */
44static void train_init(const mesh_ctx *ctx) {
45 mesh_view v = mesh_ctx_view(ctx);
46 uint32_t l = mesh_ctx_lane(ctx);
47 for (uint32_t i = l; i < MESH_TRAIN_SAMPLES; i += 32) {
48 float x = train_sample_x(i);
49 mesh_view_set_f32(v, MESH_TRAIN_T_X + 4 * i, x);
50 mesh_view_set_f32(v, MESH_TRAIN_T_Y + 4 * i, train_sin_pi(x));
51 }
52 for (uint32_t j = l; j < 256; j += 32) mesh_view_set_f32(v, MESH_TRAIN_T_LOSS + 4 * j, 0.0f);
53 if (mesh_ctx_leader(ctx)) {
54 for (uint32_t k = 0; k < 4; k++) mesh_view_set_f32(v, MESH_TRAIN_T_PARAMS + 4 * k, 0.0f);
55 mesh_view_set_f32(v, MESH_TRAIN_T_LR, MESH_TRAIN_LR);
56 mesh_view_set_u32(v, MESH_TRAIN_T_STEPS_PER_TASK, MESH_TRAIN_STEPS_PER_TASK);
57 /* A step's batch: 32 samples per system. */
58 uint32_t k, n = 1;
59 mesh_exchange x;
60 if (!train_peer(ctx, &k, &n, &x)) n = 1;
61 mesh_view_set_f32(v, MESH_TRAIN_T_INV_BATCH, 1.0f / (float)(32 * n));
62 const uint32_t zero[5] = {MESH_TRAIN_T_NEXT_STEP, MESH_TRAIN_T_ERR, MESH_TRAIN_T_PENDING, MESH_TRAIN_T_ROUND, MESH_TRAIN_T_STEP_HI};
63 for (int i = 0; i < 5; i++) mesh_view_set_u32(v, zero[i], 0);
64 mesh_ctx_count(ctx, MESH_TRAIN_C_INITS, 1);
65 }
66}
67
68/* Pick xs[j] for a warp-uniform j < 8 without a local-memory array. */
69static inline __attribute__((always_inline)) float train_sel8(const float xs[8], uint32_t j) {
70 float r = xs[0];
71 for (uint32_t k = 1; k < 8; k++)
72 if (j == k) r = xs[k];
73 return r;
74}
75
76/* Round n with the cluster: publish own sums (unless resuming a round whose
77 * slot and counter were already published), wait bounded for every other
78 * system's, return slot0 + slot1 + ... + slot(N-1) folded in system order
79 * (every machine then updates bitwise identically). False on timeout. */
80static bool train_exchange(mesh_exchange ex, uint32_t l, uint32_t k, uint32_t nsys, uint64_t n, const float own[5], bool resuming,
81 float out[5]) {
82 size_t par = MESH_EX_SLOT_BYTES * (size_t)(n % 2);
83 uint8_t *mine = mesh_exchange_slots(ex, k) + par;
84 uint32_t ok = 0;
85 for (int j = 0; j < 5; j++) out[j] = 0.0f;
86 if (l == 0) {
87 if (!resuming) {
88 for (int j = 0; j < 5; j++) mesh_rt_wrf(mine + 4 * j, own[j]);
89 for (int j = 5; j < 8; j++) mesh_rt_wrf(mine + 4 * j, 0.0f);
90 mesh_rt_fence_sc();
91 mesh_rt_st_release64(mesh_exchange_counter(ex, k), n);
92 }
93 /* Bounded by PEER_POLLS polls and by their nominal duration (for all
94 * members together): nanosleep may oversleep, the budget is wall time. */
95 uint64_t deadline = mesh_rt_gtimer() + (uint64_t)MESH_PEER_POLLS * MESH_PEER_POLL_NS;
96 uint32_t polls = 0;
97 for (;;) {
98 bool all = true;
99 for (uint32_t j = 0; j < nsys; j++)
100 if (j != k && mesh_rt_ld_acquire64(mesh_exchange_counter(ex, j)) < n) all = false;
101 if (all) { ok = 1; break; }
102 if (polls >= MESH_PEER_POLLS || mesh_rt_gtimer() >= deadline) break;
103 mesh_rt_nanosleep(MESH_PEER_POLL_NS);
104 polls++;
105 }
106 if (ok) {
107 mesh_rt_fence_acq_rel();
108 uint8_t *s0 = mesh_exchange_slots(ex, 0) + par;
109 for (int j = 0; j < 5; j++) out[j] = mesh_rt_rdf(s0 + 4 * j);
110 for (uint32_t m = 1; m < nsys; m++) {
111 uint8_t *sm = mesh_exchange_slots(ex, m) + par;
112 for (int j = 0; j < 5; j++) out[j] = out[j] + mesh_rt_rdf(sm + 4 * j);
113 }
114 }
115 }
116 if (mesh_rt_shfl_idx(ok, 0) == 0) return false;
117 for (int j = 0; j < 5; j++) out[j] = mesh_rt_shfl_idx_f(out[j], 0);
118 return true;
119}
120
121/* Arg s: the low half of the first step; the steps are base .. base +
122 * T_STEPS_PER_TASK - 1 with base = T_STEP_HI:s (u64, so rounds never wrap).
123 * Returns the steps completed. */
124static uint32_t train_run(const mesh_ctx *ctx, uint32_t s) {
125 mesh_view v = mesh_ctx_view(ctx);
126 uint64_t base = (uint64_t)mesh_view_u32(v, MESH_TRAIN_T_STEP_HI) << 32 | s;
127 uint32_t l = mesh_ctx_lane(ctx);
128 float p[4] = {mesh_view_f32(v, MESH_TRAIN_T_PARAMS), mesh_view_f32(v, MESH_TRAIN_T_PARAMS + 4),
129 mesh_view_f32(v, MESH_TRAIN_T_PARAMS + 8), mesh_view_f32(v, MESH_TRAIN_T_PARAMS + 12)};
130 float lr = mesh_view_f32(v, MESH_TRAIN_T_LR);
131 float inv = mesh_view_f32(v, MESH_TRAIN_T_INV_BATCH);
132 uint32_t steps = mesh_view_u32(v, MESH_TRAIN_T_STEPS_PER_TASK);
133 uint32_t sys = 0, nsys = 1;
134 mesh_exchange ex = {0};
135 bool dual = train_peer(ctx, &sys, &nsys, &ex);
136 if (!dual) { sys = 0; nsys = 1; }
137 /* Lane l holds samples 8l..8l+7; a step's sample i = (8l + off) mod 256 is
138 * then register (off & 7) of lane (l + off / 8) mod 32 -- a shuffle. One
139 * bounds check (a short view faults), then unchecked loads. */
140 mesh_view whole;
141 if (!mesh_view_sub(v, 0, MESH_TRAIN_VIEW_BYTES, &whole)) return 0;
142 uint8_t *vp = mesh_view_ptr(whole);
143 float xs[8], ys[8];
144 for (uint32_t j = 0; j < 8; j++) {
145 xs[j] = mesh_rt_rdf(vp + MESH_TRAIN_T_X + 4 * (8 * l + j));
146 ys[j] = mesh_rt_rdf(vp + MESH_TRAIN_T_Y + 4 * (8 * l + j));
147 }
148 /* A round left pending by a timed-out task is resumed, never recomputed. */
149 bool resume = dual && mesh_view_u32(v, MESH_TRAIN_T_PENDING) != 0 && mesh_view_u32(v, MESH_TRAIN_T_ROUND) == s + 1;
150 if (!resume && mesh_view_u32(v, MESH_TRAIN_T_PENDING) != 0 && l == 0) mesh_view_set_u32(v, MESH_TRAIN_T_PENDING, 0);
151
152 uint32_t done = 0;
153 bool timed_out = false;
154 while (done < steps) {
155 uint64_t t64 = base + done;
156 uint32_t t = (uint32_t)t64; /* sample and loss indices only need t mod 256 */
157 float sums[5] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
158 bool resuming = resume;
159 resume = false;
160 if (!resuming) {
161 /* Step t's batch: offsets N t .. N t + N - 1, system k taking N t + k. */
162 uint32_t off = t * nsys + sys;
163 uint32_t src = (l + (off >> 3)) & 31u;
164 float x = mesh_rt_shfl_idx_f(train_sel8(xs, off & 7u), src);
165 float y = mesh_rt_shfl_idx_f(train_sel8(ys, off & 7u), src);
166 float e = train_cubic(p, x) - y;
167 float g = 2.0f * e;
168 float x2 = x * x;
169 sums[0] = mesh_rt_warp_sum(e * e);
170 sums[1] = mesh_rt_warp_sum(g);
171 sums[2] = mesh_rt_warp_sum(g * x);
172 sums[3] = mesh_rt_warp_sum(g * x2);
173 sums[4] = mesh_rt_warp_sum(g * (x2 * x));
174 }
175 if (dual) {
176 float total[5];
177 if (train_exchange(ex, l, sys, nsys, t64 + 1, sums, resuming, total)) {
178 for (int j = 0; j < 5; j++) sums[j] = total[j];
179 } else {
180 timed_out = true;
181 break;
182 }
183 }
184 for (int k = 0; k < 4; k++) p[k] = p[k] - lr * sums[k + 1] * inv;
185 if (l == 0) mesh_view_set_f32(v, MESH_TRAIN_T_LOSS + 4 * (t % 256), sums[0] * inv);
186 done++;
187 }
188 if (l == 0) {
189 for (uint32_t k = 0; k < 4; k++) mesh_view_set_f32(v, MESH_TRAIN_T_PARAMS + 4 * k, p[k]);
190 if (timed_out) {
191 mesh_view_set_u32(v, MESH_TRAIN_T_ROUND, s + done + 1);
192 mesh_view_set_u32(v, MESH_TRAIN_T_PENDING, 1);
193 mesh_ctx_count(ctx, MESH_TRAIN_C_TIMEOUTS, 1);
194 } else if (dual) {
195 mesh_view_set_u32(v, MESH_TRAIN_T_PENDING, 0);
196 }
197 mesh_view_set_u32(v, MESH_TRAIN_T_ERR, timed_out ? 1u : 0u);
198 uint64_t next = base + done;
199 mesh_view_set_u32(v, MESH_TRAIN_T_STEP_HI, (uint32_t)(next >> 32));
200 mesh_view_set_u32(v, MESH_TRAIN_T_NEXT_STEP, (uint32_t)next);
201 mesh_ctx_count(ctx, MESH_TRAIN_C_STEPS, done);
202 }
203 return done;
204}
205
206static uint32_t train_task(const mesh_ctx *ctx, uint32_t op, uint32_t arg) {
207 if (op == MESH_OP_TASK) return train_run(ctx, arg);
208 if (op == MESH_TRAIN_OP_SET_LR) {
209 if (mesh_ctx_leader(ctx)) mesh_ctx_set_f32(ctx, MESH_TRAIN_T_LR, mesh_f32_from_bits(arg));
210 return MESH_R_NONE;
211 }
212 return MESH_R_BAD_OP;
213}
214
215MESH_PROGRAM(train_init, train_task);