==> src/train.c <==1/* train for the wasm executor (alpha-rust-mesh programs/train/wasm): SGD of
2 * y = a + b x + c x^2 + d x^3 on 256 samples of sin(pi x), lockstep with the cluster through the
3 * exchange page, in the cluster float mode. The warp of the GPU image -- 32 lanes, one sample each,
4 * butterfly-summed -- is mesh-ref's step_sums, written operation for operation as the GPUs compute
5 * (Z: every f32 op flushed), so the params, losses and exchanged sums are the GPUs' bits (kcheck X:
6 * the cross-arch digests). Layout: mesh-abi views::train. */
7#include "mesh/abi/view_train.h"
8#include "mesh/ref/cubic.h"
9#include "mesh/rt_wasm.h"
10
11/* Lockstep in a cluster: (this system k, the size N, the exchange page). */
12static bool peer_system(const mesh_ctx *ctx, uint32_t *k, uint32_t *n, mesh_exchange *x) {
13 return mesh_ctx_group(ctx, k, n) && mesh_ctx_exchange(ctx, x);
14}
15
16static float xs[MESH_TRAIN_SAMPLES], ys[MESH_TRAIN_SAMPLES];
17
18/* Fresh state: samples (sin_pi in f64, the f32 grid in Z: the GPUs' bits), hyperparameters,
19 * params and progress zero. The exchange page is meshd's. */
20static void init(const mesh_ctx *ctx) {
21 mesh_view v = mesh_ctx_view(ctx);
22 mesh_cubic_samples(xs, ys);
23 mesh_view_write_f32s(v, MESH_TRAIN_T_X, xs, MESH_TRAIN_SAMPLES);
24 mesh_view_write_f32s(v, MESH_TRAIN_T_Y, ys, MESH_TRAIN_SAMPLES);
25 static const float zeros[256];
26 mesh_view_write_f32s(v, MESH_TRAIN_T_LOSS, zeros, 256);
27 mesh_view_write_f32s(v, MESH_TRAIN_T_PARAMS, zeros, 4);
28 mesh_view_set_f32(v, MESH_TRAIN_T_LR, MESH_TRAIN_LR);
29 mesh_view_set_u32(v, MESH_TRAIN_T_STEPS_PER_TASK, MESH_TRAIN_STEPS_PER_TASK);
30 /* A step's batch: 32 samples per system. */
31 uint32_t k, n = 1;
32 mesh_exchange x;
33 if (!peer_system(ctx, &k, &n, &x)) n = 1;
34 mesh_view_set_f32(v, MESH_TRAIN_T_INV_BATCH, mesh_ftz_div(1.0f, mesh_ftz_from_u32(32 * n)));
35 const uint32_t zero[5] = {MESH_TRAIN_T_NEXT_STEP, MESH_TRAIN_T_ERR, MESH_TRAIN_T_PENDING, MESH_TRAIN_T_ROUND,
36 MESH_TRAIN_T_STEP_HI};
37 for (int i = 0; i < 5; i++) mesh_view_set_u32(v, zero[i], 0);
38 mesh_ctx_count(ctx, MESH_TRAIN_C_INITS, 1);
39}
40
41/* Round n with the cluster: publish own sums (unless resuming a round whose slot and counter were
42 * already published), wait bounded for every other system's, return slot0 + slot1 + ... + slot(N-1)
43 * folded in system order (in Z: every member updates bitwise identically). false on timeout. */
44static bool exchange(mesh_exchange ex, uint32_t k, uint32_t nsys, uint64_t n, const float own[5], bool resuming,
45 float out[5]) {
46 uint32_t par = (uint32_t)(n % 2);
47 if (!resuming) {
48 uint32_t w[MESH_EX_SLOT_WORDS] = {0};
49 for (int j = 0; j < 5; j++) w[j] = mesh_ftz_bits(own[j]);
50 mesh_exchange_set_slot(ex, k, par, w);
51 }
52 if (!mesh_exchange_round(ex, n)) return false;
53 uint32_t s[MESH_EX_SLOT_WORDS];
54 mesh_exchange_slot(ex, 0, par, s);
55 for (int j = 0; j < 5; j++) out[j] = mesh_ftz_from_bits(s[j]);
56 for (uint32_t m = 1; m < nsys; m++) {
57 mesh_exchange_slot(ex, m, par, s);
58 for (int j = 0; j < 5; j++) out[j] = mesh_ftz_add(out[j], mesh_ftz_from_bits(s[j]));
59 }
60 return true;
61}
62
63/* Arg s: the low half of the first step; the steps are base .. base + T_STEPS_PER_TASK - 1 with
64 * base = T_STEP_HI:s. Returns the steps completed. */
65static uint32_t run(const mesh_ctx *ctx, uint32_t s) {
66 mesh_view v = mesh_ctx_view(ctx);
67 uint64_t base = (uint64_t)mesh_view_u32(v, MESH_TRAIN_T_STEP_HI) << 32 | s;
68 uint32_t hdr[12];
69 mesh_view_read_u32s(v, 0, hdr, 12);
70 float p[4] = {mesh_ftz_from_bits(hdr[0]), mesh_ftz_from_bits(hdr[1]), mesh_ftz_from_bits(hdr[2]), mesh_ftz_from_bits(hdr[3])};
71 float lr = mesh_ftz_from_bits(hdr[4]), inv = mesh_ftz_from_bits(hdr[6]);
72 uint32_t steps = hdr[5];
73 uint32_t pending = hdr[9], round = hdr[10];
74 uint32_t sys = 0, nsys = 1;
75 mesh_exchange ex;
76 bool dual = peer_system(ctx, &sys, &nsys, &ex);
77 if (!dual) sys = 0, nsys = 1;
78 mesh_view_read_f32s(v, MESH_TRAIN_T_X, xs, MESH_TRAIN_SAMPLES);
79 mesh_view_read_f32s(v, MESH_TRAIN_T_Y, ys, MESH_TRAIN_SAMPLES);
80 /* A round left pending by a timed-out task is resumed, never recomputed. */
81 bool resume = dual && pending != 0 && round == s + 1u;
82 if (!resume && pending != 0) mesh_view_set_u32(v, MESH_TRAIN_T_PENDING, 0);
83 uint32_t done = 0;
84 bool timed_out = false;
85 while (done < steps) {
86 uint64_t t64 = base + done;
87 uint32_t t = (uint32_t)t64;
88 float sums[5] = {0, 0, 0, 0, 0};
89 bool resuming = resume;
90 resume = false;
91 /* Step t's batch: system k takes offset N t + k. */
92 if (!resuming) mesh_cubic_step_sums(p, xs, ys, t * nsys + sys, sums);
93 if (dual) {
94 float total[5];
95 if (!exchange(ex, sys, nsys, t64 + 1, sums, resuming, total)) {
96 timed_out = true;
97 break;
98 }
99 for (int j = 0; j < 5; j++) sums[j] = total[j];
100 }
101 mesh_cubic_apply(p, sums, lr, inv);
102 mesh_view_set_f32(v, MESH_TRAIN_T_LOSS + 4 * (t % 256), mesh_cubic_loss(sums, inv));
103 done++;
104 }
105 mesh_view_write_f32s(v, MESH_TRAIN_T_PARAMS, p, 4);
106 if (timed_out) {
107 mesh_view_set_u32(v, MESH_TRAIN_T_ROUND, s + done + 1u);
108 mesh_view_set_u32(v, MESH_TRAIN_T_PENDING, 1);
109 mesh_ctx_count(ctx, MESH_TRAIN_C_TIMEOUTS, 1);
110 } else if (dual) {
111 mesh_view_set_u32(v, MESH_TRAIN_T_PENDING, 0);
112 }
113 mesh_view_set_u32(v, MESH_TRAIN_T_ERR, timed_out);
114 uint64_t next = base + done;
115 mesh_view_set_u32(v, MESH_TRAIN_T_STEP_HI, (uint32_t)(next >> 32));
116 mesh_view_set_u32(v, MESH_TRAIN_T_NEXT_STEP, (uint32_t)next);
117 mesh_ctx_count(ctx, MESH_TRAIN_C_STEPS, done);
118 return done;
119}
120
121static uint32_t task(const mesh_ctx *ctx, uint32_t op, uint32_t arg) {
122 switch (op) {
123 case MESH_OP_TASK: return run(ctx, arg);
124 case MESH_TRAIN_OP_SET_LR: mesh_ctx_set_f32(ctx, MESH_TRAIN_T_LR, mesh_ftz_from_bits(arg)); return MESH_R_NONE;
125 default: return MESH_R_BAD_OP;
126 }
127}
128
129MESH_PROGRAM(init, task)