==> tests/train.c <==1/* train on the simulator against @ajax/mesh-ref-cubic, bitwise: the samples
2 * (sin in f64), every step's params and loss alone; then a lockstep group of
3 * two (two mailboxes, the exchange pages mirrored by a thread as meshd's peer
4 * relay does) whose members end bitwise identical and equal to the reference
5 * of a 64-sample batch. */
6#include "mesh/check.h"
7#include "mesh/ref/cubic.h"
8#include "mesh/rt/sim_mailbox.h"
9
10#include "../src/train.c"
11
12#include <pthread.h>
13#include <stdatomic.h>
14#include <string.h>
15#include <time.h>
16
17#define TASKS 6
18
19static float xs[MESH_TRAIN_SAMPLES], ys[MESH_TRAIN_SAMPLES];
20
21static void check_samples(uint8_t *view) {
22 for (uint32_t i = 0; i < MESH_TRAIN_SAMPLES; i++) {
23 CHECK_U(mesh_sim_u32(view + MESH_TRAIN_T_X + 4 * i), mesh_f32_bits(xs[i]));
24 CHECK_U(mesh_sim_u32(view + MESH_TRAIN_T_Y + 4 * i), mesh_f32_bits(ys[i]));
25 }
26}
27
28static void alone(void) {
29 uint8_t *mb = mesh_sim_mailbox();
30 mesh_sim_prog p = {.name = "train", .view_bytes = MESH_TRAIN_VIEW_BYTES};
31 uint8_t *view = mesh_sim_install(mb, 0, &p);
32 mesh_sim_launch *l = mesh_sim_launch_prog(mb, 0, &p, mesh_main);
33 CHECK_U(mesh_sim_call(mb, 0, MESH_OP_INIT, 0), MESH_R_NONE);
34 check_samples(view);
35 CHECK_U(mesh_sim_u32(view + MESH_TRAIN_T_INV_BATCH), mesh_f32_bits(1.0f / 32.0f));
36 float p_ref[4] = {0, 0, 0, 0}, loss[256] = {0};
37 for (uint32_t task = 0; task < TASKS; task++) {
38 uint32_t s = mesh_sim_u32(view + MESH_TRAIN_T_NEXT_STEP);
39 CHECK_U(mesh_sim_call(mb, 0, MESH_OP_TASK, s), MESH_TRAIN_STEPS_PER_TASK);
40 for (uint32_t k = 0; k < MESH_TRAIN_STEPS_PER_TASK; k++) {
41 uint32_t t = s + k;
42 float sums[5];
43 mesh_cubic_step_sums(p_ref, xs, ys, t, sums);
44 mesh_cubic_apply(p_ref, sums, MESH_TRAIN_LR, 1.0f / 32.0f);
45 loss[t % 256] = mesh_cubic_loss(sums, 1.0f / 32.0f);
46 }
47 for (int k = 0; k < 4; k++) CHECK_U(mesh_sim_u32(view + MESH_TRAIN_T_PARAMS + 4 * k), mesh_f32_bits(p_ref[k]));
48 }
49 for (uint32_t j = 0; j < 256; j++) CHECK_U(mesh_sim_u32(view + MESH_TRAIN_T_LOSS + 4 * j), mesh_f32_bits(loss[j]));
50 CHECK_U(mesh_sim_u32(view + MESH_TRAIN_T_NEXT_STEP), TASKS * MESH_TRAIN_STEPS_PER_TASK);
51 CHECK_U(mesh_sim_u32(mesh_sim_ctl(mb, 0) + MESH_P_COUNTERS + 4 * MESH_TRAIN_C_STEPS), TASKS * MESH_TRAIN_STEPS_PER_TASK);
52 /* OP_SET_LR. */
53 CHECK_U(mesh_sim_call(mb, 0, MESH_TRAIN_OP_SET_LR, mesh_f32_bits(0.25f)), MESH_R_NONE);
54 CHECK_U(mesh_sim_u32(view + MESH_TRAIN_T_LR), mesh_f32_bits(0.25f));
55 mesh_sim_halt(mb, 0, l);
56 mesh_sim_mailbox_free(mb);
57}
58
59/* The peer relay: member k's round counter and slot (of the round's parity)
60 * copied into every other member's exchange page, the counter released last. */
61typedef struct {
62 uint8_t *mb[2];
63 _Atomic int stop;
64} relay;
65
66static void *mirror(void *arg) {
67 relay *r = arg;
68 while (!atomic_load(&r->stop)) {
69 for (uint32_t k = 0; k < 2; k++) {
70 uint8_t *src = r->mb[k], *dst = r->mb[1 - k];
71 uint64_t n = mesh_sim_u64(src + mesh_ex_counter(k));
72 if (n > mesh_sim_u64(dst + mesh_ex_counter(k))) {
73 size_t par = MESH_EX_SLOT_BYTES * (size_t)(n % 2);
74 memcpy(dst + mesh_ex_slots(k) + par, src + mesh_ex_slots(k) + par, MESH_EX_SLOT_BYTES);
75 mesh_sim_set_u64(dst + mesh_ex_counter(k), n);
76 }
77 }
78 struct timespec ts = {0, 10000};
79 nanosleep(&ts, NULL);
80 }
81 return NULL;
82}
83
84/* One member's driver, as meshd's periodic publisher: tasks with arg = its
85 * own T_NEXT_STEP until it reaches `until`. A round the partner was late for
86 * times out (bounded) and is resumed by the next task, never recomputed. */
87typedef struct {
88 uint8_t *mb, *view;
89 uint32_t until, tasks, timeouts;
90} driver;
91
92static void *drive(void *arg) {
93 driver *d = arg;
94 while (mesh_sim_u32(d->view + MESH_TRAIN_T_NEXT_STEP) < d->until && d->tasks < 1000) {
95 uint32_t s = mesh_sim_u32(d->view + MESH_TRAIN_T_NEXT_STEP), r = 0;
96 uint32_t t = mesh_sim_publish(d->mb, 0, MESH_OP_TASK, s);
97 if (!mesh_sim_wait(d->mb, 0, t, &r, 120000)) break;
98 d->tasks++;
99 d->timeouts += mesh_sim_u32(d->view + MESH_TRAIN_T_ERR);
100 }
101 return NULL;
102}
103
104static void lockstep(void) {
105 relay r = {{mesh_sim_mailbox(), mesh_sim_mailbox()}, 0};
106 mesh_sim_prog p = {.name = "train", .view_bytes = MESH_TRAIN_VIEW_BYTES, .peer_mode = MESH_PEER_MODE_LOCKSTEP};
107 mesh_sim_launch *l[2];
108 uint8_t *view[2];
109 for (uint32_t k = 0; k < 2; k++) {
110 mesh_sim_set_u32(r.mb[k] + MESH_H_GROUP_RANK, k);
111 mesh_sim_set_u32(r.mb[k] + MESH_H_GROUP_N, 2);
112 view[k] = mesh_sim_install(r.mb[k], 0, &p);
113 l[k] = mesh_sim_launch_prog(r.mb[k], 0, &p, mesh_main);
114 CHECK_U(mesh_sim_call(r.mb[k], 0, MESH_OP_INIT, 0), MESH_R_NONE);
115 check_samples(view[k]);
116 CHECK_U(mesh_sim_u32(view[k] + MESH_TRAIN_T_INV_BATCH), mesh_f32_bits(1.0f / 64.0f));
117 }
118 pthread_t th, dt[2];
119 pthread_create(&th, NULL, mirror, &r);
120 driver d[2] = {{r.mb[0], view[0], TASKS * MESH_TRAIN_STEPS_PER_TASK, 0, 0}, {r.mb[1], view[1], TASKS * MESH_TRAIN_STEPS_PER_TASK, 0, 0}};
121 for (int k = 0; k < 2; k++) pthread_create(&dt[k], NULL, drive, &d[k]);
122 for (int k = 0; k < 2; k++) pthread_join(dt[k], NULL);
123 /* A member cannot complete a round without the other's sums, so both stop at the same step. */
124 uint32_t steps = mesh_sim_u32(view[0] + MESH_TRAIN_T_NEXT_STEP);
125 CHECK_U(mesh_sim_u32(view[1] + MESH_TRAIN_T_NEXT_STEP), steps);
126 CHECK(steps >= TASKS * MESH_TRAIN_STEPS_PER_TASK);
127 float p_ref[4] = {0, 0, 0, 0};
128 for (uint32_t t = 0; t < steps; t++) {
129 float a[5], b[5], sum[5];
130 mesh_cubic_step_sums(p_ref, xs, ys, 2 * t, a);
131 mesh_cubic_step_sums(p_ref, xs, ys, 2 * t + 1, b);
132 for (int j = 0; j < 5; j++) sum[j] = a[j] + b[j];
133 mesh_cubic_apply(p_ref, sum, MESH_TRAIN_LR, 1.0f / 64.0f);
134 }
135 for (int k = 0; k < 4; k++) {
136 CHECK_U(mesh_sim_u32(view[0] + MESH_TRAIN_T_PARAMS + 4 * k), mesh_f32_bits(p_ref[k]));
137 CHECK_U(mesh_sim_u32(view[1] + MESH_TRAIN_T_PARAMS + 4 * k), mesh_f32_bits(p_ref[k]));
138 }
139 for (int j = 0; j < 256; j++) CHECK_U(mesh_sim_u32(view[0] + MESH_TRAIN_T_LOSS + 4 * j), mesh_sim_u32(view[1] + MESH_TRAIN_T_LOSS + 4 * j));
140 if (d[0].timeouts + d[1].timeouts) fprintf(stderr, "lockstep: %u and %u rounds timed out and were resumed\n", d[0].timeouts, d[1].timeouts);
141 atomic_store(&r.stop, 1);
142 pthread_join(th, NULL);
143 for (uint32_t k = 0; k < 2; k++) {
144 mesh_sim_halt(r.mb[k], 0, l[k]);
145 mesh_sim_mailbox_free(r.mb[k]);
146 }
147}
148
149int main(void) {
150 mesh_cubic_samples(xs, ys);
151 for (uint32_t i = 0; i < MESH_TRAIN_SAMPLES; i++) {
152 CHECK_U(mesh_f32_bits(train_sample_x(i)), mesh_f32_bits(xs[i]));
153 CHECK_U(mesh_f32_bits(train_sin_pi(xs[i])), mesh_f32_bits(mesh_cubic_sin_pi(xs[i])));
154 }
155 alone();
156 lockstep();
157 return mesh_check_done();
158}