==> src/inproc.c <==1/* inproc.rs: Dev, Mb2, load, Prog, Heap, reads_mask. */
2#include "mesh/inproc.h"
3#include "mesh/loader.h"
4#include "mesh/manifest_encode.h"
5#include "mesh/mps.h"
6#include "mesh/nvjitlink.h"
7#include "mesh/words_host.h"
8
9#include "mesh/abi/backend.h"
10#include "mesh/abi/base.h"
11#include "mesh/abi/control_block.h"
12#include "mesh/abi/ext.h"
13#include "mesh/abi/header.h"
14#include "mesh/abi/input_queue.h"
15#include "mesh/abi/node.h"
16#include "mesh/abi/ops.h"
17#include "mesh/abi/program_table.h"
18#include "mesh/abi/ring.h"
19#include "mesh/abi/trace.h"
20#include "mesh/abi/view_heap.h"
21
22#include <stdatomic.h>
23#include <stdio.h>
24#include <stdlib.h>
25#include <string.h>
26#include <time.h>
27
28uint64_t mesh_inproc_now_ns(void) {
29 struct timespec t;
30 clock_gettime(CLOCK_MONOTONIC, &t);
31 return (uint64_t)t.tv_sec * 1000000000u + (uint64_t)t.tv_nsec;
32}
33
34double mesh_ms_since(uint64_t t0) { return (double)(mesh_inproc_now_ns() - t0) / 1e6; }
35
36bool mesh_dev_open_cuda(mesh_dev *out, char **err) {
37 *out = (mesh_dev){0};
38 if (!mesh_shm_new(&out->shm, err)) return false;
39 mesh_nvjitlink jit;
40 if (!mesh_nvjitlink_open(MESH_NVJITLINK_DIR, &jit, err)) return false;
41 if (!mesh_cuda_backend_open(MESH_ISOLATION_MODE_BLINK, 0, NULL, 0, &out->shm, &jit, &out->b, err)) return false;
42 out->comp = out->b->ops->compiler(out->b);
43 out->kind = MESH_BACKEND_CUDA;
44 out->ctx = mesh_cuda_backend_context(out->b);
45 return true;
46}
47
48void mesh_dev_from(mesh_backend *b, const mesh_shm *shm, mesh_dev *out) {
49 *out = (mesh_dev){b, b->ops->compiler(b), b->ops->caps(b)->kind, NULL, *shm};
50}
51
52void mesh_dev_close(mesh_dev *d) {
53 mesh_backend_destroy(d->b);
54 d->b = NULL;
55}
56
57const char *mesh_dev_code(const mesh_dev *d) {
58 return d->kind == MESH_BACKEND_CUDA ? "cubin" : d->kind == MESH_BACKEND_METAL ? "MSL" : "wasm";
59}
60
61uint32_t mesh_dev_sms_total(const mesh_dev *d) { return d->b->ops->caps(d->b)->sms_total; }
62
63static uint8_t *at(mesh_mb2 m, size_t off) {
64 if (off >= MESH_MAILBOX_BYTES) {
65 fprintf(stderr, "v2 offset 0x%zx\n", off);
66 abort();
67 }
68 return m.base + off;
69}
70static _Atomic uint32_t *a32(mesh_mb2 m, size_t off) {
71 if (off % 4) abort();
72 return (_Atomic uint32_t *)(void *)at(m, off);
73}
74static _Atomic uint64_t *a64(mesh_mb2 m, size_t off) {
75 if (off % 8) abort();
76 return (_Atomic uint64_t *)(void *)at(m, off);
77}
78
79uint32_t mesh_mb2_r32(mesh_mb2 m, size_t off) { return atomic_load_explicit(a32(m, off), memory_order_acquire); }
80void mesh_mb2_w32(mesh_mb2 m, size_t off, uint32_t v) { atomic_store_explicit(a32(m, off), v, memory_order_relaxed); }
81void mesh_mb2_store32(mesh_mb2 m, size_t off, uint32_t v) { atomic_store_explicit(a32(m, off), v, memory_order_release); }
82uint32_t mesh_mb2_add32(mesh_mb2 m, size_t off, uint32_t v) { return atomic_fetch_add_explicit(a32(m, off), v, memory_order_acq_rel); }
83uint64_t mesh_mb2_r64(mesh_mb2 m, size_t off) { return atomic_load_explicit(a64(m, off), memory_order_acquire); }
84void mesh_mb2_w64(mesh_mb2 m, size_t off, uint64_t v) { atomic_store_explicit(a64(m, off), v, memory_order_relaxed); }
85void mesh_mb2_store64(mesh_mb2 m, size_t off, uint64_t v) { atomic_store_explicit(a64(m, off), v, memory_order_release); }
86float mesh_mb2_rf(mesh_mb2 m, size_t off) {
87 uint32_t u = mesh_mb2_r32(m, off);
88 float f;
89 memcpy(&f, &u, 4);
90 return f;
91}
92void mesh_mb2_wf(mesh_mb2 m, size_t off, float v) {
93 uint32_t u;
94 memcpy(&u, &v, 4);
95 mesh_mb2_w32(m, off, u);
96}
97void mesh_mb2_bytes(mesh_mb2 m, size_t off, uint8_t *out, size_t n) {
98 for (size_t i = 0; i < n; i++) out[i] = *(volatile uint8_t *)at(m, off + i);
99}
100void mesh_mb2_write(mesh_mb2 m, size_t off, const uint8_t *b, size_t n) {
101 for (size_t i = 0; i < n; i++) *(volatile uint8_t *)at(m, off + i) = b[i];
102}
103void mesh_mb2_zero(mesh_mb2 m, size_t off, size_t n) {
104 for (size_t i = 0; i < n; i++) *(volatile uint8_t *)at(m, off + i) = 0;
105}
106mesh_words mesh_mb2_words(mesh_mb2 m) { return mesh_words_host(m.base); }
107
108bool mesh_mb2_alloc(const mesh_dev *d, mesh_mb2 *out, char **err) {
109 uint8_t *base;
110 uint64_t dev;
111 if (!d->b->ops->mailbox(d->b, &base, &dev, err)) return false;
112 mesh_mb2 m = {base, dev};
113 mesh_mb2_w32(m, MESH_H_MAGIC, MESH_MAGIC);
114 mesh_mb2_w32(m, MESH_H_VERSION, MESH_VERSION);
115 mesh_mb2_w32(m, MESH_H_PEER, MESH_PEER_NONE);
116 mesh_mb2_w32(m, MESH_H_SYSTEMS, 1);
117 mesh_mb2_w32(m, MESH_H_SMS, mesh_dev_sms_total(d));
118 mesh_mb2_w32(m, MESH_H_BACKEND, mesh_caps_word(d->b->ops->caps(d->b)));
119 *out = m;
120 return true;
121}
122
123bool mesh_inproc_load(const mesh_dev *d, const char *dir, mesh_image2 *out, char **err) {
124 *out = (mesh_image2){0};
125 const char *slash = strrchr(dir, '/');
126 char *root = slash ? mesh_strndup(dir, (size_t)(slash - dir)) : mesh_strdup(".");
127 const char *name = slash ? slash + 1 : dir;
128 if (!*root) {
129 free(root);
130 root = mesh_strdup("/");
131 }
132 mesh_payload payload;
133 mesh_prepared p;
134 if (!mesh_loader_read(root, name, d->comp->ops->needs_ptx(d->comp), &payload, err) ||
135 !d->comp->ops->prepare(d->comp, &payload, NULL, root, &p, err)) {
136 free(root);
137 return false;
138 }
139 free(root);
140 uint64_t t0 = mesh_inproc_now_ns();
141 mesh_kernel *k = NULL;
142 if (!d->b->ops->load(d->b, &p, &k, err)) {
143 mesh_err_prefix(err, "load: ");
144 mesh_prepared_free(&p);
145 return false;
146 }
147 double load_ms = mesh_ms_since(t0);
148 uint32_t rt_abi, regs;
149 bool ok = mesh_kernel_global_u32(k, MESH_RT_ABI_SYMBOL, &rt_abi, err);
150 if (ok && rt_abi != MESH_RT_ABI) ok = mesh_errf(err, "abi: %s = %u, meshd speaks %u", MESH_RT_ABI_SYMBOL, (unsigned)rt_abi, (unsigned)MESH_RT_ABI);
151 if (ok) ok = mesh_kernel_regs(k, ®s, err);
152 /* nodes = "fit" (as meshd's budget, alone on the device): 3/4 of what the occupancy allows. */
153 if (ok && p.manifest.nodes_fit && d->b->ops->resident(d->b)) {
154 uint32_t per_sm;
155 ok = mesh_kernel_max_blocks_per_sm(k, &per_sm, err);
156 if (ok) {
157 uint32_t n = per_sm * d->b->ops->caps(d->b)->sms * 3 / 4;
158 if (n > MESH_PROG_MAX_NODES) n = MESH_PROG_MAX_NODES;
159 p.manifest.nodes = n ? n : 1;
160 }
161 }
162 if (!ok) {
163 d->b->ops->unload(d->b, k, NULL);
164 mesh_prepared_free(&p);
165 return false;
166 }
167 out->manifest = p.manifest;
168 p.manifest = (mesh_manifest){0};
169 out->kernel = k;
170 out->hash = p.hash;
171 out->cubin_bytes = (uint32_t)p.code_len;
172 out->regs = regs;
173 out->rt_abi = rt_abi;
174 out->jit_ms = p.compile_ms;
175 out->load_ms = load_ms;
176 mesh_prepared_free(&p);
177 return true;
178}
179
180void mesh_image2_free(const mesh_dev *d, mesh_image2 *im) {
181 if (im->kernel) d->b->ops->unload(d->b, im->kernel, NULL);
182 mesh_manifest_free(&im->manifest);
183 *im = (mesh_image2){0};
184}
185
186bool mesh_prog_install(const mesh_dev *d, mesh_mb2 mb, uint32_t slot, mesh_image2 *image, uint32_t view_off, uint32_t reads_mask,
187 mesh_prog *out, char **err) {
188 const mesh_manifest *m = &image->manifest;
189 if ((m->view_bytes > 0 && view_off < MESH_VIEW_HEAP) || (size_t)view_off + m->view_bytes > MESH_VIEW_HEAP_END)
190 return mesh_errf(err, "view 0x%x+0x%x outside the view heap", (unsigned)view_off, (unsigned)m->view_bytes);
191 size_t e = mesh_prog_entry(slot);
192 uint32_t gen = mesh_mb2_r32(mb, e + MESH_PT_GEN);
193 mesh_mb2_store32(mb, e + MESH_PT_GEN, gen | 1);
194 atomic_thread_fence(memory_order_seq_cst);
195 mesh_image im = {view_off, reads_mask, image->hash, image->regs, image->cubin_bytes, image->rt_abi, 1};
196 uint8_t bytes[MESH_PT_BYTES], ext[MESH_PX_BYTES];
197 mesh_manifest_encode_entry(m, &im, bytes);
198 mesh_mb2_write(mb, e + 8, bytes + 8, MESH_PT_BYTES - 8); /* PT_STATE and PT_GEN stay ours */
199 mesh_manifest_encode_ext(m, ext);
200 uint64_t scratch = mesh_kernel_scratch(image->kernel);
201 for (int i = 0; i < 8; i++) ext[MESH_PX_SCRATCH + (size_t)i] = (uint8_t)(scratch >> (8 * i));
202 mesh_mb2_write(mb, mesh_prog_ext(slot), ext, MESH_PX_BYTES);
203 mesh_mb2_zero(mb, mesh_prog_ctl(slot), MESH_PROG_CTL_BYTES);
204 mesh_mb2_zero(mb, mesh_node_rec(slot, 0), MESH_PROG_NODES_BYTES);
205 mesh_mb2_zero(mb, mesh_trace_rec(slot, 0), (size_t)MESH_TRACE_BYTES * MESH_TRACE_PER_PROG);
206 mesh_mb2_zero(mb, view_off, m->view_bytes);
207 mesh_mb2_w32(mb, e + MESH_PT_STATE, MESH_ST_LOADED);
208 mesh_mb2_store32(mb, e + MESH_PT_GEN, (gen | 1) + 1);
209 mesh_mb2_add32(mb, MESH_H_PROG_SEQ, 1);
210 mesh_stream *st = NULL;
211 if (!d->b->ops->new_stream(d->b, slot, image->kernel, &st, err)) return false;
212 *out = (mesh_prog){mb, slot, st, *image, m->nodes, 0};
213 *image = (mesh_image2){0};
214 return true;
215}
216
217void mesh_prog_free(const mesh_dev *d, mesh_prog *p) {
218 mesh_stream_drop(p->stream);
219 mesh_image2_free(d, &p->image);
220}
221
222size_t mesh_prog_ctl_off(const mesh_prog *p, size_t off) { return mesh_prog_ctl(p->slot) + off; }
223size_t mesh_prog_view(const mesh_prog *p) { return mesh_mb2_r32(p->mb, mesh_prog_entry(p->slot) + MESH_PT_VIEW_OFF); }
224uint32_t mesh_prog_head(const mesh_prog *p) { return mesh_mb2_r32(p->mb, mesh_prog_ctl_off(p, MESH_P_RING_HEAD)); }
225uint32_t mesh_prog_counter(const mesh_prog *p, uint32_t i) { return mesh_mb2_r32(p->mb, mesh_prog_ctl_off(p, MESH_P_COUNTERS + 4 * i)); }
226
227bool mesh_prog_launch(const mesh_dev *d, mesh_prog *p, double *ms, char **err) {
228 mesh_mb2 mb = p->mb;
229 bool done = false;
230 if (!d->b->ops->query(d->b, p->stream, p->image.kernel, &done, err)) return false;
231 if (!done) return mesh_errf(err, "slot %u: still running", (unsigned)p->slot);
232 mesh_mb2_w32(mb, mesh_prog_ctl_off(p, MESH_P_RING_HEAD), p->tail);
233 mesh_mb2_w32(mb, mesh_prog_ctl_off(p, MESH_P_ONLINE), 0);
234 mesh_mb2_w32(mb, mesh_prog_ctl_off(p, MESH_P_HALTED), 0);
235 mesh_mb2_store32(mb, mesh_prog_ctl_off(p, MESH_P_SHUTDOWN), 0);
236 mesh_mb2_add32(mb, mesh_prog_ctl_off(p, MESH_P_LAUNCHES), 1);
237 atomic_thread_fence(memory_order_seq_cst);
238 uint64_t t0 = mesh_inproc_now_ns();
239 if (!d->b->ops->launch(d->b, p->stream, p->image.kernel, p->nodes, mb.dev, p->slot, err)) return false;
240 while (mesh_mb2_r32(mb, mesh_prog_ctl_off(p, MESH_P_ONLINE)) < p->nodes) {
241 if (mesh_ms_since(t0) > 2000) {
242 bool fin = false;
243 char *e2 = NULL;
244 const char *what = !d->b->ops->query(d->b, p->stream, p->image.kernel, &fin, &e2) ? e2 : fin ? "exited" : "running";
245 mesh_errf(err, "slot %u: %u of %u nodes online after 2 s (stream: %s)", (unsigned)p->slot,
246 (unsigned)mesh_mb2_r32(mb, mesh_prog_ctl_off(p, MESH_P_ONLINE)), (unsigned)p->nodes, what);
247 free(e2);
248 return false;
249 }
250 }
251 *ms = mesh_ms_since(t0);
252 mesh_mb2_w32(mb, mesh_prog_entry(p->slot) + MESH_PT_STATE, MESH_ST_RUNNING);
253 return true;
254}
255
256bool mesh_prog_stop(const mesh_dev *d, mesh_prog *p, double *ms, char **err) {
257 mesh_mb2 mb = p->mb;
258 uint64_t t0 = mesh_inproc_now_ns();
259 mesh_mb2_store32(mb, mesh_prog_ctl_off(p, MESH_P_SHUTDOWN), 1);
260 for (;;) {
261 bool done = false;
262 if (mesh_mb2_r32(mb, mesh_prog_ctl_off(p, MESH_P_HALTED)) >= p->nodes) {
263 if (!d->b->ops->query(d->b, p->stream, p->image.kernel, &done, err)) return false;
264 if (done) break;
265 }
266 if (mesh_ms_since(t0) > 3000)
267 return mesh_errf(err, "slot %u: %u of %u nodes halted in 3 s", (unsigned)p->slot,
268 (unsigned)mesh_mb2_r32(mb, mesh_prog_ctl_off(p, MESH_P_HALTED)), (unsigned)p->nodes);
269 }
270 *ms = mesh_ms_since(t0);
271 mesh_mb2_w32(mb, mesh_prog_entry(p->slot) + MESH_PT_STATE, MESH_ST_LOADED);
272 return true;
273}
274
275bool mesh_prog_done(const mesh_prog *p, uint32_t t) {
276 return mesh_mb2_r32(p->mb, mesh_prog_ctl_off(p, MESH_P_MARKS + 4 * (t % MESH_PROG_RESULTS))) == t + 1;
277}
278
279bool mesh_prog_wait(const mesh_prog *p, uint32_t t, uint32_t *result, char **err) {
280 uint64_t t0 = mesh_inproc_now_ns();
281 while (!mesh_prog_done(p, t)) {
282 if (mesh_ms_since(t0) > 5000)
283 return mesh_errf(err, "slot %u ticket %u not completed in 5 s (head %u, done %u, online %u, fault 0x%x)", (unsigned)p->slot,
284 (unsigned)t, (unsigned)mesh_prog_head(p), (unsigned)mesh_mb2_r32(p->mb, mesh_prog_ctl_off(p, MESH_P_DONE)),
285 (unsigned)mesh_mb2_r32(p->mb, mesh_prog_ctl_off(p, MESH_P_ONLINE)),
286 (unsigned)mesh_mb2_r32(p->mb, mesh_prog_ctl_off(p, MESH_P_FAULT)));
287 }
288 *result = mesh_mb2_r32(p->mb, mesh_prog_ctl_off(p, MESH_P_RESULTS + 4 * (t % MESH_PROG_RESULTS)));
289 return true;
290}
291
292bool mesh_prog_publish(mesh_prog *p, uint32_t op, uint32_t arg, uint32_t *ticket, char **err) {
293 uint32_t t = p->tail, r;
294 if (t >= MESH_PROG_RING && !mesh_prog_wait(p, t - MESH_PROG_RING, &r, err)) return false;
295 size_t e = mesh_prog_ctl_off(p, MESH_P_RING + MESH_ENTRY_BYTES * (t % MESH_PROG_RING));
296 mesh_mb2_w32(p->mb, e + MESH_E_OP, op);
297 mesh_mb2_w32(p->mb, e + MESH_E_ARG, arg);
298 mesh_mb2_w32(p->mb, e + MESH_E_SLOT, t % MESH_PROG_RESULTS);
299 mesh_mb2_store32(p->mb, e + MESH_E_SEQ, t + 1);
300 p->tail = t + 1;
301 mesh_mb2_w64(p->mb, mesh_prog_ctl_off(p, MESH_P_RING_TAIL), p->tail);
302 *ticket = t;
303 return true;
304}
305
306bool mesh_prog_run(mesh_prog *p, uint32_t op, uint32_t arg, uint32_t *result, char **err) {
307 uint32_t t;
308 return mesh_prog_publish(p, op, arg, &t, err) && mesh_prog_wait(p, t, result, err);
309}
310
311bool mesh_prog_batch(mesh_prog *p, uint32_t op, uint32_t n, uint32_t *results, char **err) {
312 uint32_t first = p->tail;
313 for (uint32_t a = 0; a < n; a++) {
314 uint32_t t;
315 if (!mesh_prog_publish(p, op, a, &t, err)) return false;
316 }
317 for (uint32_t a = 0; a < n; a++)
318 if (!mesh_prog_wait(p, first + a, &results[a], err)) return false;
319 return true;
320}
321
322bool mesh_prog_input_push(const mesh_prog *p, const uint8_t *b, size_t n, char **err) {
323 const mesh_manifest *m = &p->image.manifest;
324 if (m->input_len == 0) return mesh_errf(err, "slot %u: no input queue", (unsigned)p->slot);
325 size_t q = mesh_prog_view(p) + m->input_off;
326 uint32_t head = mesh_mb2_r32(p->mb, q + MESH_IQ_HEAD), tail = mesh_mb2_r32(p->mb, q + MESH_IQ_TAIL);
327 if ((size_t)(head - tail) + n > m->input_len)
328 return mesh_errf(err, "slot %u: input queue full (%u of %u)", (unsigned)p->slot, (unsigned)(head - tail), (unsigned)m->input_len);
329 for (size_t i = 0; i < n; i++) mesh_mb2_write(p->mb, q + MESH_IQ_DATA + (head + i) % m->input_len, &b[i], 1);
330 mesh_mb2_store32(p->mb, q + MESH_IQ_HEAD, head + (uint32_t)n);
331 return true;
332}
333
334bool mesh_prog_request(mesh_prog *p, const uint8_t *req, size_t n, uint8_t **resp, size_t *rn, char **err) {
335 if (!p->image.manifest.has_request) return mesh_errf(err, "slot %u: no request window", (unsigned)p->slot);
336 mesh_manifest_window w = p->image.manifest.request;
337 size_t v = mesh_prog_view(p);
338 if (n > w.req_max) n = w.req_max;
339 uint8_t *pad = mesh_alloc(w.req_max ? w.req_max : 1);
340 memcpy(pad, req, n);
341 mesh_mb2_write(p->mb, v + w.req_off, pad, w.req_max);
342 free(pad);
343 mesh_mb2_w32(p->mb, mesh_prog_ctl_off(p, MESH_P_REQ_LEN), (uint32_t)n);
344 uint64_t seq = mesh_mb2_r64(p->mb, mesh_prog_ctl_off(p, MESH_P_REQ_SEQ)) + 1;
345 mesh_mb2_store64(p->mb, mesh_prog_ctl_off(p, MESH_P_REQ_SEQ), seq);
346 uint32_t r;
347 if (!mesh_prog_run(p, MESH_OP_REQ, 0, &r, err)) return false;
348 uint64_t rs = mesh_mb2_r64(p->mb, mesh_prog_ctl_off(p, MESH_P_RESP_SEQ));
349 if (rs != seq) return mesh_errf(err, "slot %u: P_RESP_SEQ %llu != P_REQ_SEQ %llu", (unsigned)p->slot, (unsigned long long)rs, (unsigned long long)seq);
350 size_t len = mesh_mb2_r32(p->mb, mesh_prog_ctl_off(p, MESH_P_RESP_LEN));
351 if (len > w.resp_max) len = w.resp_max;
352 *resp = mesh_alloc(len ? len : 1);
353 mesh_mb2_bytes(p->mb, v + w.resp_off, *resp, len);
354 *rn = len;
355 return true;
356}
357
358mesh_heap mesh_heap_new(void) { return (mesh_heap){(uint32_t)MESH_VIEW_HEAP}; }
359
360bool mesh_heap_alloc(mesh_heap *h, uint32_t bytes, uint32_t *at_out, char **err) {
361 uint32_t a = h->next;
362 size_t end = a + ((size_t)bytes + MESH_VIEW_PAGE - 1) / MESH_VIEW_PAGE * MESH_VIEW_PAGE;
363 if (end > MESH_VIEW_HEAP_END) return mesh_errf(err, "view heap: 0x%x bytes do not fit", (unsigned)bytes);
364 h->next = (uint32_t)end;
365 *at_out = a;
366 return true;
367}
368
369bool mesh_reads_mask(const mesh_manifest *m, const mesh_slot_name *slots, size_t n, uint32_t *mask, char **err) {
370 uint32_t r = 0;
371 for (size_t i = 0; i < m->nreads; i++) {
372 if (strcmp(m->reads[i], "*") == 0) {
373 *mask = MESH_READS_ALL;
374 return true;
375 }
376 bool found = false;
377 for (size_t j = 0; j < n && !found; j++)
378 if (strcmp(slots[j].name, m->reads[i]) == 0) {
379 r |= 1u << slots[j].slot;
380 found = true;
381 }
382 if (!found) {
383 char *q = mesh_debug_str(m->reads[i]);
384 mesh_errf(err, "%s: reads %s, which is not loaded", m->name, q);
385 free(q);
386 return false;
387 }
388 }
389 *mask = r;
390 return true;
391}