==> src/mps.c <==1/* mps.rs's workers (Mps, Proc, Worker, the pool), the dual-mode GPU and the CUDA backend. */
2#include "mps_internal.h"
3
4#include "mesh/loader.h"
5
6#include "mesh/abi/base.h"
7#include "mesh/abi/backend.h"
8#include "mesh/abi/host_status.h"
9
10#include <errno.h>
11#include <signal.h>
12#include <stdio.h>
13#include <stdlib.h>
14#include <string.h>
15#include <sys/wait.h>
16#include <unistd.h>
17
18/* ---- globals: the daemon, live workers, the server, native GPU clients, the pool ---------- */
19
20typedef struct mps mps;
21
22static pthread_mutex_t g_mu = PTHREAD_MUTEX_INITIALIZER;
23static char *daemon_dir;
24static bool daemon_temporary;
25static pid_t *workers;
26static size_t nworkers;
27static mps *server;
28static pid_t *native_clients;
29static size_t nnative;
30
31void mps_set_daemon(const char *dir, bool temporary) {
32 pthread_mutex_lock(&g_mu);
33 free(daemon_dir);
34 daemon_dir = mesh_strdup(dir);
35 daemon_temporary = temporary;
36 pthread_mutex_unlock(&g_mu);
37}
38
39static void pids_add(pid_t **v, size_t *n, pid_t p) {
40 *v = mesh_realloc(*v, (*n + 1) * sizeof **v);
41 (*v)[(*n)++] = p;
42}
43
44static void pids_remove(pid_t *v, size_t *n, pid_t p) {
45 for (size_t i = 0; i < *n; i++)
46 if (v[i] == p) {
47 v[i] = v[--*n];
48 return;
49 }
50}
51
52static bool pids_has(const pid_t *v, size_t n, pid_t p) {
53 for (size_t i = 0; i < n; i++)
54 if (v[i] == p) return true;
55 return false;
56}
57
58static char *pids_debug(const pid_t *v, size_t n) {
59 mesh_text t = {0};
60 mesh_text_putc(&t, '[');
61 for (size_t i = 0; i < n; i++) mesh_text_printf(&t, "%s%d", i ? ", " : "", (int)v[i]);
62 mesh_text_putc(&t, ']');
63 return mesh_text_take(&t);
64}
65
66/* ---- worker processes -------------------------------------------------------------------- */
67
68typedef struct reply {
69 uint8_t op;
70 uint8_t *body;
71 size_t n;
72 struct reply *next;
73} reply;
74
75/* A worker process. Refcounted: its owner and its reader thread. */
76typedef struct proc {
77 atomic_int refs;
78 pid_t pid;
79 int in_fd, out_fd;
80 pthread_mutex_t in_mu; /* one request at a time */
81 pthread_mutex_t ask_mu; /* the replies */
82 pthread_mutex_t rq_mu;
83 pthread_cond_t rq_cv;
84 reply *rq_head, *rq_tail;
85 bool rq_closed;
86 /* The last launch's generation, and the last one whose stream finished. */
87 _Atomic uint32_t gen, done;
88 atomic_bool dead;
89 pthread_mutex_t err_mu;
90 char *err;
91 pthread_mutex_t reap_mu;
92 bool reaped;
93} proc;
94
95static void proc_unref(proc *p) {
96 if (!p || atomic_fetch_sub(&p->refs, 1) != 1) return;
97 for (reply *r = p->rq_head; r;) {
98 reply *n = r->next;
99 free(r->body);
100 free(r);
101 r = n;
102 }
103 if (p->in_fd >= 0) close(p->in_fd);
104 free(p->err);
105 free(p);
106}
107
108static bool proc_reaped(proc *p) {
109 pthread_mutex_lock(&p->reap_mu);
110 if (!p->reaped) {
111 int st;
112 if (waitpid(p->pid, &st, WNOHANG) == p->pid) p->reaped = true;
113 }
114 bool r = p->reaped;
115 pthread_mutex_unlock(&p->reap_mu);
116 return r;
117}
118
119static void *proc_reader(void *a) {
120 proc *p = a;
121 uint8_t op, *b;
122 size_t n;
123 while (mesh_mps_read_frame(p->out_fd, &op, &b, &n)) {
124 if (op == MESH_W_DONE) {
125 if (mesh_mps_u32_at(b, n, 4) == 0) {
126 pthread_mutex_lock(&p->err_mu);
127 free(p->err);
128 p->err = mesh_fmt("worker %d: %.*s", (int)p->pid, (int)(n > 8 ? n - 8 : 0), n > 8 ? (const char *)b + 8 : "");
129 pthread_mutex_unlock(&p->err_mu);
130 }
131 atomic_store(&p->done, mesh_mps_u32_at(b, n, 0));
132 free(b);
133 continue;
134 }
135 reply *r = mesh_alloc(sizeof *r);
136 *r = (reply){op, b, n, NULL};
137 pthread_mutex_lock(&p->rq_mu);
138 if (p->rq_tail) p->rq_tail->next = r;
139 else p->rq_head = r;
140 p->rq_tail = r;
141 pthread_cond_signal(&p->rq_cv);
142 pthread_mutex_unlock(&p->rq_mu);
143 }
144 close(p->out_fd);
145 atomic_store(&p->dead, true);
146 pthread_mutex_lock(&p->rq_mu);
147 p->rq_closed = true;
148 pthread_cond_broadcast(&p->rq_cv);
149 pthread_mutex_unlock(&p->rq_mu);
150 proc_unref(p);
151 return NULL;
152}
153
154static bool proc_send(proc *p, uint8_t op, const uint8_t *body, size_t n, char **err) {
155 pthread_mutex_lock(&p->in_mu);
156 bool ok = p->in_fd >= 0 && mesh_mps_write_frame(p->in_fd, op, body, n);
157 int e = errno;
158 pthread_mutex_unlock(&p->in_mu);
159 if (!ok) {
160 char *s = mps_os_error(e ? e : EPIPE);
161 mesh_errf(err, "worker %d: %s", (int)p->pid, s);
162 free(s);
163 }
164 return ok;
165}
166
167/* Rust's Duration Debug for whole seconds: "30s", "2s". */
168static bool proc_ask(proc *p, uint8_t op, const uint8_t *body, size_t n, uint8_t want, unsigned within_ms, uint8_t **out, size_t *on,
169 char **err) {
170 pthread_mutex_lock(&p->ask_mu);
171 if (!proc_send(p, op, body, n, err)) {
172 pthread_mutex_unlock(&p->ask_mu);
173 return false;
174 }
175 struct timespec dl;
176 clock_gettime(CLOCK_REALTIME, &dl);
177 dl.tv_sec += within_ms / 1000;
178 dl.tv_nsec += (long)(within_ms % 1000) * 1000000L;
179 if (dl.tv_nsec >= 1000000000L) {
180 dl.tv_sec++;
181 dl.tv_nsec -= 1000000000L;
182 }
183 pthread_mutex_lock(&p->rq_mu);
184 int rc = 0;
185 while (!p->rq_head && !p->rq_closed && rc == 0) rc = pthread_cond_timedwait(&p->rq_cv, &p->rq_mu, &dl);
186 reply *r = p->rq_head;
187 if (r) {
188 p->rq_head = r->next;
189 if (!p->rq_head) p->rq_tail = NULL;
190 }
191 bool closed = p->rq_closed;
192 pthread_mutex_unlock(&p->rq_mu);
193 pthread_mutex_unlock(&p->ask_mu);
194 bool ok = false;
195 if (r) {
196 if (r->op == want && mesh_mps_u32_at(r->body, r->n, 0) == 1) {
197 size_t k = r->n > 4 ? r->n - 4 : 0;
198 *out = mesh_alloc(k ? k : 1);
199 if (k) memcpy(*out, r->body + 4, k);
200 *on = k;
201 ok = true;
202 } else if (r->op == want) {
203 mesh_errf(err, "%.*s", (int)(r->n > 4 ? r->n - 4 : 0), r->n > 4 ? (const char *)r->body + 4 : "");
204 } else {
205 mesh_errf(err, "worker %d: reply %u, expected %u", (int)p->pid, r->op, want);
206 }
207 free(r->body);
208 free(r);
209 } else if (closed) {
210 mesh_errf(err, "worker %d exited", (int)p->pid);
211 } else {
212 mesh_errf(err, "worker %d: no reply in %us", (int)p->pid, within_ms / 1000);
213 }
214 return ok;
215}
216
217/* *done = the last launch's stream finished; false with *err on a stream error or a death. */
218static bool proc_query(proc *p, bool *done, char **err) {
219 pthread_mutex_lock(&p->err_mu);
220 char *e = p->err ? mesh_strdup(p->err) : NULL;
221 pthread_mutex_unlock(&p->err_mu);
222 if (e) {
223 mesh_errf(err, "%s", e);
224 free(e);
225 return false;
226 }
227 bool d = atomic_load(&p->done) == atomic_load(&p->gen);
228 if (!d && atomic_load(&p->dead)) return mesh_errf(err, "worker %d exited while its kernel ran", (int)p->pid);
229 *done = d;
230 return true;
231}
232
233/* Alive, no stream error, nothing running: fit for another image. */
234static bool proc_idle(proc *p) {
235 bool done = false;
236 return !atomic_load(&p->dead) && proc_query(p, &done, NULL) && done;
237}
238
239static void *retire_thread(void *a) {
240 proc *p = a;
241 uint64_t t0 = mps_now_ns();
242 while (mps_now_ns() - t0 < 2000000000u) {
243 if (proc_reaped(p)) {
244 mesh_info("mps: worker %d exited", (int)p->pid);
245 proc_unref(p);
246 return NULL;
247 }
248 mps_sleep_us(1000);
249 }
250 mesh_error("mps: worker %d did not exit in 2 s; killing it", (int)p->pid);
251 kill(p->pid, SIGKILL);
252 uint64_t t1 = mps_now_ns();
253 while (!proc_reaped(p)) {
254 if (mps_now_ns() - t1 > 2000000000u) {
255 mesh_error("mps: worker %d not reaped 2 s after SIGKILL; left behind", (int)p->pid);
256 break;
257 }
258 mps_sleep_us(1000);
259 }
260 proc_unref(p);
261 return NULL;
262}
263
264/* W_QUIT, then reaped off the driver thread (killed if it lingers). Takes the reference. */
265static void proc_retire(proc *p) {
266 proc_send(p, MESH_W_QUIT, NULL, 0, NULL);
267 pthread_mutex_lock(&g_mu);
268 pids_remove(workers, &nworkers, p->pid);
269 pthread_mutex_unlock(&g_mu);
270 pthread_t t;
271 if (pthread_create(&t, NULL, retire_thread, p) == 0) pthread_detach(t);
272 else retire_thread(p);
273}
274
275/* Idle workers, kept with their context (module unloaded) for the next load. */
276static pthread_mutex_t pool_mu = PTHREAD_MUTEX_INITIALIZER;
277static proc **pool;
278static size_t npool;
279static atomic_size_t pool_cap;
280
281/* ---- Mps ------------------------------------------------------------------------------------ */
282
283struct mps {
284 char *dir;
285 char *mailbox; /* the path workers open the mailbox's memfd by */
286 uint32_t sms; /* 0: no green context */
287 char *exe;
288 pthread_mutex_t server_mu;
289 char *server; /* the server's pid once a worker was seen in its client list */
290};
291
292typedef struct {
293 uint32_t rt_abi, regs, per_sm, local, sms;
294 uint64_t scratch;
295} winfo;
296
297/* One image in its own MPS client process. */
298typedef struct {
299 proc *p; /* until dropped */
300 char *name;
301 winfo info;
302} worker;
303
304static bool mps_check_clients(mps *m, const pid_t *want, size_t nwant, char **err) {
305 char *servers = mps_control(m->dir, "get_server_list", err);
306 if (!servers) return false;
307 char *nl = strchr(servers, '\n');
308 if (nl) *nl = 0;
309 const char *sv = servers;
310 size_t sn = strlen(sv);
311 mesh_trim(&sv, &sn);
312 if (sn == 0) {
313 char *w = pids_debug(want, nwant);
314 mesh_errf(err, "mps: workers %s are up but our daemon runs no server: they are not MPS clients; refusing", w);
315 free(w);
316 free(servers);
317 return false;
318 }
319 char *server_pid = mesh_strndup(sv, sn);
320 free(servers);
321 char *req = mesh_fmt("get_client_list %s", server_pid);
322 char *clients = mps_control(m->dir, req, err);
323 free(req);
324 if (!clients) {
325 free(server_pid);
326 return false;
327 }
328 pid_t *pids = NULL;
329 size_t np = 0;
330 for (char *line = clients; line && *line;) {
331 char *next = strchr(line, '\n');
332 if (next) *next++ = 0;
333 const char *s = line;
334 size_t n = strlen(s);
335 mesh_trim(&s, &n);
336 uint32_t v;
337 if (mesh_parse_u32(s, n, 10, &v)) pids_add(&pids, &np, (pid_t)v);
338 line = next;
339 }
340 free(clients);
341 bool ok = true;
342 for (size_t i = 0; i < nwant && ok; i++)
343 if (!pids_has(pids, np, want[i])) {
344 char *d = pids_debug(pids, np);
345 ok = mesh_errf(err, "mps: worker %d is not among server %s's clients %s; refusing", (int)want[i], server_pid, d);
346 free(d);
347 }
348 if (ok) {
349 pthread_mutex_lock(&g_mu);
350 pid_t *foreign = NULL;
351 size_t nf = 0;
352 for (size_t i = 0; i < np; i++)
353 if (!pids_has(workers, nworkers, pids[i]) && !pids_has(native_clients, nnative, pids[i])) pids_add(&foreign, &nf, pids[i]);
354 pthread_mutex_unlock(&g_mu);
355 if (nf) {
356 char *d = pids_debug(foreign, nf);
357 mesh_error("mps: server %s has clients that are not meshd's workers: %s", server_pid, d);
358 free(d);
359 }
360 free(foreign);
361 pthread_mutex_lock(&m->server_mu);
362 if (!m->server || strcmp(m->server, server_pid) != 0) {
363 char *d = pids_debug(pids, np);
364 mesh_info("mps: server pid %s, clients %s", server_pid, d);
365 free(d);
366 free(m->server);
367 m->server = mesh_strdup(server_pid);
368 }
369 pthread_mutex_unlock(&m->server_mu);
370 }
371 free(pids);
372 free(server_pid);
373 return ok;
374}
375
376static bool mps_spawn(mps *m, proc **out, char **err) {
377 int in[2], out_p[2];
378 if (!mesh_pipe_cloexec(in, err) || !mesh_pipe_cloexec(out_p, err)) return false;
379 char *sms = m->sms ? mesh_fmt("%u", (unsigned)m->sms) : NULL;
380 char *argv[] = {m->exe, "worker", "--mailbox", m->mailbox, sms ? "--sms" : NULL, sms, NULL};
381 char *env[] = {mesh_fmt("%s=%s/pipe", MESH_MPS_PIPE_ENV, m->dir), mesh_fmt("%s=%s/log", MESH_MPS_LOG_ENV, m->dir), NULL};
382 mesh_spawn_fd fds[] = {{in[0], 0}, {out_p[1], 1}, {2, 2}};
383 pid_t pid;
384 char *why = NULL;
385 /* No PR_SET_PDEATHSIG: it fires when the spawning (JIT) thread exits. A
386 * worker exits at EOF on stdin instead, which meshd's death also gives. */
387 bool ok = mesh_spawn(argv, env, fds, 3, &pid, &why);
388 free(sms);
389 free(env[0]);
390 free(env[1]);
391 close(in[0]);
392 close(out_p[1]);
393 if (!ok) {
394 const char *e = strstr(why, ": ");
395 mesh_errf(err, "load: mps worker: %s", e ? e + 2 : why);
396 free(why);
397 close(in[1]);
398 close(out_p[0]);
399 return false;
400 }
401 proc *p = mesh_alloc(sizeof *p);
402 atomic_init(&p->refs, 2);
403 p->pid = pid;
404 p->in_fd = in[1];
405 p->out_fd = out_p[0];
406 pthread_mutex_init(&p->in_mu, NULL);
407 pthread_mutex_init(&p->ask_mu, NULL);
408 pthread_mutex_init(&p->rq_mu, NULL);
409 pthread_cond_init(&p->rq_cv, NULL);
410 pthread_mutex_init(&p->err_mu, NULL);
411 pthread_mutex_init(&p->reap_mu, NULL);
412 pthread_mutex_lock(&g_mu);
413 pids_add(&workers, &nworkers, pid);
414 pthread_mutex_unlock(&g_mu);
415 pthread_t t;
416 if (pthread_create(&t, NULL, proc_reader, p) != 0) abort();
417 pthread_detach(t);
418 *out = p;
419 return true;
420}
421
422static void worker_drop(worker *w);
423
424static bool mps_preload(mps *m, const char *name, const uint8_t *cubin, size_t n, uint32_t threads, uint32_t smem, uint32_t scratch,
425 worker **out, char **err) {
426 uint64_t t0 = mps_now_ns();
427 proc *p = NULL;
428 for (;;) {
429 pthread_mutex_lock(&pool_mu);
430 proc *q = npool ? pool[--npool] : NULL;
431 pthread_mutex_unlock(&pool_mu);
432 if (!q) break;
433 if (proc_idle(q)) {
434 p = q;
435 break;
436 }
437 proc_retire(q);
438 }
439 bool fresh = p == NULL;
440 if (!p && !mps_spawn(m, &p, err)) return false;
441 pid_t pid = p->pid;
442 worker *w = mesh_alloc(sizeof *w);
443 w->p = p;
444 w->name = mesh_strdup(name);
445 uint8_t *body = NULL;
446 size_t bn = 0;
447 mesh_mps_put_u32(&body, &bn, threads);
448 mesh_mps_put_u32(&body, &bn, smem);
449 mesh_mps_put_u32(&body, &bn, scratch);
450 mesh_mps_put_u32(&body, &bn, (uint32_t)strlen(name));
451 mesh_mps_put_bytes(&body, &bn, name, strlen(name));
452 mesh_mps_put_bytes(&body, &bn, cubin, n);
453 uint8_t *b = NULL;
454 size_t k = 0;
455 char *why = NULL;
456 bool ok = proc_ask(p, MESH_W_LOAD, body, bn, MESH_W_LOADED, MESH_MPS_WORKER_LOAD_MS, &b, &k, &why);
457 free(body);
458 if (!ok) {
459 mesh_errf(err, "load: %s", why);
460 free(why);
461 worker_drop(w);
462 return false;
463 }
464 w->info = (winfo){mesh_mps_u32_at(b, k, 0), mesh_mps_u32_at(b, k, 4), mesh_mps_u32_at(b, k, 8), mesh_mps_u32_at(b, k, 12),
465 mesh_mps_u32_at(b, k, 16), mesh_mps_u64_at(b, k, 28)};
466 uint64_t dev = mesh_mps_u64_at(b, k, 20);
467 free(b);
468 if (m->sms && m->sms != w->info.sms) {
469 mesh_errf(err, "load: the worker's green context has %u SMs, not %u", (unsigned)w->info.sms, (unsigned)m->sms);
470 worker_drop(w);
471 return false;
472 }
473 if (fresh && !mps_check_clients(m, &pid, 1, err)) {
474 worker_drop(w);
475 return false;
476 }
477 mesh_info("mps: worker %d (%s) for %s: module loaded, mailbox at device 0x%llx, %u SMs, in %.1f ms", (int)pid, fresh ? "new" : "pooled",
478 name, (unsigned long long)dev, (unsigned)w->info.sms, (double)(mps_now_ns() - t0) / 1e6);
479 *out = w;
480 return true;
481}
482
483static bool mps_prewarm(mps *m, uint32_t n, char **err) {
484 atomic_store(&pool_cap, n);
485 if (n == 0) return true;
486 uint64_t t0 = mps_now_ns();
487 proc **procs = mesh_alloc(n * sizeof *procs);
488 pid_t *pids = mesh_alloc(n * sizeof *pids);
489 uint32_t made = 0;
490 bool ok = true;
491 for (; made < n && ok; made++) {
492 ok = mps_spawn(m, &procs[made], err);
493 if (!ok) break;
494 pids[made] = procs[made]->pid;
495 }
496 for (uint32_t i = 0; i < made && ok; i++) {
497 uint8_t *b = NULL;
498 size_t k;
499 char *why = NULL;
500 ok = proc_ask(procs[i], MESH_W_READY, NULL, 0, MESH_W_READY, MESH_MPS_WORKER_LOAD_MS, &b, &k, &why);
501 free(b);
502 if (!ok) {
503 mesh_errf(err, "mps: pool: %s", why);
504 free(why);
505 }
506 }
507 if (ok) ok = mps_check_clients(m, pids, made, err);
508 if (ok) {
509 char *d = pids_debug(pids, made);
510 mesh_info("mps: pool of %u idle workers %s ready in %.0f ms", (unsigned)n, d, (double)(mps_now_ns() - t0) / 1e6);
511 free(d);
512 pthread_mutex_lock(&pool_mu);
513 pool = mesh_realloc(pool, (npool + made) * sizeof *pool);
514 for (uint32_t i = 0; i < made; i++) pool[npool++] = procs[i];
515 pthread_mutex_unlock(&pool_mu);
516 } else {
517 for (uint32_t i = 0; i < made; i++) proc_retire(procs[i]);
518 }
519 free(procs);
520 free(pids);
521 return ok;
522}
523
524static bool worker_launch(worker *w, uint32_t slot, uint32_t nodes, char **err) {
525 proc *p = w->p;
526 uint32_t gen = atomic_load(&p->gen) + 1;
527 atomic_store(&p->gen, gen);
528 uint8_t *body = NULL;
529 size_t bn = 0;
530 mesh_mps_put_u32(&body, &bn, gen);
531 mesh_mps_put_u32(&body, &bn, slot);
532 mesh_mps_put_u32(&body, &bn, nodes);
533 uint8_t *b = NULL;
534 size_t k;
535 char *why = NULL;
536 bool ok = proc_ask(p, MESH_W_LAUNCH, body, bn, MESH_W_LAUNCHED, MESH_MPS_WORKER_LAUNCH_MS, &b, &k, &why);
537 free(body);
538 free(b);
539 if (!ok) {
540 mesh_errf(err, "%s (%s)", w->name, why);
541 free(why);
542 }
543 return ok;
544}
545
546static bool worker_query(worker *w, bool *done, char **err) {
547 char *why = NULL;
548 if (proc_query(w->p, done, &why)) return true;
549 mesh_errf(err, "%s: %s", w->name, why);
550 free(why);
551 return false;
552}
553
554/* Back to the pool with its module unloaded (W_UNLOAD), unless it is unfit or
555 * the pool is full: then the process exits. */
556static void worker_drop(worker *w) {
557 proc *p = w->p;
558 if (p) {
559 pthread_mutex_lock(&pool_mu);
560 bool keep = proc_idle(p) && npool < atomic_load(&pool_cap) && proc_send(p, MESH_W_UNLOAD, NULL, 0, NULL);
561 if (keep) {
562 pool = mesh_realloc(pool, (npool + 1) * sizeof *pool);
563 pool[npool++] = p;
564 }
565 pthread_mutex_unlock(&pool_mu);
566 if (!keep) proc_retire(p);
567 }
568 free(w->name);
569 free(w);
570}
571
572/* ---- globals' accessors ------------------------------------------------------------------- */
573
574char *mesh_mps_native_pipe_dir(void) {
575 pthread_mutex_lock(&g_mu);
576 char *d = server ? mesh_fmt("%s/pipe", server->dir) : NULL;
577 pthread_mutex_unlock(&g_mu);
578 return d;
579}
580
581bool mesh_mps_check_native(pid_t pid, char **err) {
582 pthread_mutex_lock(&g_mu);
583 mps *m = server;
584 if (m) pids_add(&native_clients, &nnative, pid);
585 pthread_mutex_unlock(&g_mu);
586 if (!m) return mesh_errf(err, "mps: no server (blink mode)");
587 bool ok = mps_check_clients(m, &pid, 1, err);
588 if (!ok) {
589 pthread_mutex_lock(&g_mu);
590 pids_remove(native_clients, &nnative, pid);
591 pthread_mutex_unlock(&g_mu);
592 }
593 return ok;
594}
595
596void mesh_mps_cleanup(void) {
597 /* Idle ones see EOF and exit; then every worker is killed and reaped. */
598 pthread_mutex_lock(&pool_mu);
599 proc **idle = pool;
600 size_t nidle = npool;
601 pool = NULL;
602 npool = 0;
603 pthread_mutex_unlock(&pool_mu);
604 for (size_t i = 0; i < nidle; i++) {
605 pthread_mutex_lock(&idle[i]->in_mu);
606 if (idle[i]->in_fd >= 0) close(idle[i]->in_fd);
607 idle[i]->in_fd = -1;
608 pthread_mutex_unlock(&idle[i]->in_mu);
609 }
610 pthread_mutex_lock(&g_mu);
611 pid_t *kids = workers;
612 size_t nk = nworkers;
613 workers = NULL;
614 nworkers = 0;
615 char *dir = daemon_dir;
616 bool temporary = daemon_temporary;
617 daemon_dir = NULL;
618 pthread_mutex_unlock(&g_mu);
619 for (size_t i = 0; i < nk; i++) kill(kids[i], SIGKILL);
620 /* Reaped within a shared deadline: one stuck in the kernel is left behind. */
621 uint64_t t0 = mps_now_ns();
622 size_t left = nk;
623 while (left && mps_now_ns() - t0 < 500000000u) {
624 left = 0;
625 for (size_t i = 0; i < nk; i++) {
626 if (kids[i] == 0) continue;
627 int st;
628 pid_t r = waitpid(kids[i], &st, WNOHANG);
629 if (r == kids[i] || (r < 0 && errno == ECHILD)) kids[i] = 0;
630 else left++;
631 }
632 if (left) mps_sleep_us(1000);
633 }
634 for (size_t i = 0; i < nk; i++)
635 if (kids[i]) mesh_error("mps: worker %d not reaped 500 ms after SIGKILL; left behind", (int)kids[i]);
636 free(kids);
637 for (size_t i = 0; i < nidle; i++) proc_unref(idle[i]);
638 free(idle);
639 if (!dir) return;
640 mps_quit_daemon(dir, temporary);
641 free(dir);
642}
643
644/* ---- the GPU in either mode, and the backend -------------------------------------------- */
645
646typedef enum { K_ONE, K_WORKER } kind;
647
648typedef struct {
649 kind kind;
650 mesh_cuda_kernel one;
651 worker *w;
652 struct cuda_backend *b;
653} ckernel;
654
655typedef struct cuda_backend {
656 bool mps_mode;
657 mesh_cuda_gpu gpu; /* blink: the context */
658 uint64_t dev; /* blink: the mailbox's device address */
659 mps *m; /* mps: the server */
660 bool host_atomics;
661 mesh_cuda_limits limits;
662 mesh_shm shm;
663 mesh_caps caps;
664 mesh_nvjitlink jit;
665 char *arch;
666 mesh_compiler comp;
667} cuda_backend;
668
669static cuda_backend *cb(const mesh_backend *b) { return b->self; }
670
671static bool k_regs(const mesh_kernel *k, uint32_t *out, char **err) {
672 ckernel *c = k->self;
673 if (c->kind == K_WORKER) {
674 *out = c->w->info.regs;
675 return true;
676 }
677 return mesh_cuda_kernel_regs(&c->one, out, err);
678}
679static bool k_local(const mesh_kernel *k, uint32_t *out, char **err) {
680 ckernel *c = k->self;
681 if (c->kind == K_WORKER) {
682 *out = c->w->info.local;
683 return true;
684 }
685 return mesh_cuda_kernel_local_bytes(&c->one, out, err);
686}
687static bool k_per_sm(const mesh_kernel *k, uint32_t *out, char **err) {
688 ckernel *c = k->self;
689 if (c->kind == K_WORKER) {
690 *out = c->w->info.per_sm;
691 return true;
692 }
693 return mesh_cuda_kernel_max_blocks_per_sm(&c->one, out, err);
694}
695static uint64_t k_scratch(const mesh_kernel *k) {
696 ckernel *c = k->self;
697 return c->kind == K_WORKER ? c->w->info.scratch : c->one.scratch;
698}
699static bool k_global(const mesh_kernel *k, const char *name, uint32_t *out, char **err) {
700 ckernel *c = k->self;
701 if (c->kind == K_ONE) return mesh_cuda_kernel_global_u32(&c->one, name, out, err);
702 if (strcmp(name, MESH_RT_ABI_SYMBOL) == 0) {
703 *out = c->w->info.rt_abi;
704 return true;
705 }
706 return mesh_errf(err, "global %s: a worker reports only %s", name, MESH_RT_ABI_SYMBOL);
707}
708static void k_drop(mesh_kernel *k) {
709 ckernel *c = k->self;
710 /* A context module dropped without its unload stays loaded (as the Rust's). */
711 if (c->kind == K_WORKER) worker_drop(c->w);
712 free(c);
713}
714
715static const mesh_kernel_ops kernel_ops = {MESH_BACKEND_CUDA, k_regs, k_local, k_per_sm, k_scratch, k_global, k_drop};
716
717static mesh_kernel *new_kernel(ckernel *c) {
718 mesh_kernel *k = mesh_alloc(sizeof *k);
719 *k = (mesh_kernel){&kernel_ops, c};
720 return k;
721}
722
723/* PTX -> cubin with nvJitLink for this GPU's arch (and, mps, the worker preload). */
724static bool c_needs_ptx(const mesh_compiler *c) {
725 (void)c;
726 return true;
727}
728
729static bool c_prepare(const mesh_compiler *c, mesh_payload *src, const uint64_t *expect, const char *root, mesh_prepared *out, char **err) {
730 (void)root;
731 cuda_backend *b = c->self;
732 *out = (mesh_prepared){0};
733 mesh_manifest m;
734 if (!mesh_check_manifest(src, MESH_BACKEND_CUDA, &m, err)) {
735 mesh_payload_free(src);
736 return false;
737 }
738 uint64_t hash = mesh_image_hash(src->ptx, src->ptx_len, m.param);
739 uint32_t reg_cap = mesh_cuda_reg_cap(&b->limits, m.warps);
740 uint8_t *cubin = NULL;
741 size_t cn = 0;
742 if (!mesh_check_hash(m.name, "PTX", hash, expect, err) ||
743 !mesh_cuda_limits_check(&b->limits, m.name, m.warps, mesh_manifest_smem(&m), err))
744 goto fail;
745 uint64_t t0 = mps_now_ns();
746 char *why = NULL;
747 if (!mesh_nvjitlink_ptx_to_cubin(&b->jit, src->ptx, src->ptx_len, b->arch, reg_cap, &cubin, &cn, &why)) {
748 for (char *s = why; *s; s++)
749 if (*s == '\n') {
750 /* replace '\n' with " | " */
751 size_t at = (size_t)(s - why);
752 char *r = mesh_fmt("%.*s | %s", (int)at, why, s + 1);
753 free(why);
754 why = r;
755 s = why + at + 2;
756 }
757 mesh_errf(err, "jit: %s", why);
758 free(why);
759 goto fail;
760 }
761 double compile_ms = (double)(mps_now_ns() - t0) / 1e6;
762 mesh_kernel *k = NULL;
763 if (b->mps_mode) {
764 worker *w;
765 if (!mps_preload(b->m, m.name, cubin, cn, mesh_manifest_threads(&m), mesh_manifest_smem(&m), m.scratch_bytes, &w, err)) goto fail;
766 ckernel *ck = mesh_alloc(sizeof *ck);
767 *ck = (ckernel){K_WORKER, {0}, w, b};
768 k = new_kernel(ck);
769 }
770 out->manifest = m;
771 out->hash = hash;
772 out->code = cubin;
773 out->code_len = cn;
774 out->compile_ms = compile_ms;
775 out->src = *src;
776 *src = (mesh_payload){0};
777 out->kernel = k;
778 out->exec = MESH_BACKEND_CUDA;
779 out->reg_cap = reg_cap;
780 return true;
781fail:
782 free(cubin);
783 mesh_manifest_free(&m);
784 mesh_payload_free(src);
785 return false;
786}
787
788static const mesh_compiler_ops compiler_ops = {c_needs_ptx, c_prepare};
789
790static const mesh_caps *b_caps(const mesh_backend *b) { return &cb(b)->caps; }
791static bool b_isolated(const mesh_backend *b) { return cb(b)->mps_mode; }
792static bool b_resident(const mesh_backend *b) {
793 (void)b;
794 return true;
795}
796static uint32_t b_isolation(const mesh_backend *b) { return cb(b)->mps_mode ? MESH_ISOLATION_MPS : MESH_ISOLATION_BLINK; }
797static bool b_mailbox(const mesh_backend *b, uint8_t **host, uint64_t *dev, char **err) {
798 (void)err;
799 *host = cb(b)->shm.base;
800 *dev = cb(b)->mps_mode ? 0 : cb(b)->dev;
801 return true;
802}
803static bool b_make_current(const mesh_backend *b, char **err) { return cb(b)->mps_mode ? true : mesh_cuda_make_current(&cb(b)->gpu, err); }
804static const mesh_compiler *b_compiler(const mesh_backend *b) { return &cb(b)->comp; }
805
806static bool b_load(const mesh_backend *b, mesh_prepared *p, mesh_kernel **out, char **err) {
807 cuda_backend *c = cb(b);
808 mesh_kernel *pre = p->kernel;
809 p->kernel = NULL;
810 if (pre && pre->ops != &kernel_ops) {
811 mesh_kernel_drop(pre);
812 return mesh_errf(err, "internal: a non-CUDA image on the CUDA backend");
813 }
814 if (!c->mps_mode && !pre) {
815 ckernel *k = mesh_alloc(sizeof *k);
816 k->kind = K_ONE;
817 k->b = c;
818 if (!mesh_cuda_load_cubin(&c->gpu, p->code, MESH_PROGRAM_ENTRY, &k->one, err)) {
819 free(k);
820 return false;
821 }
822 if (!mesh_cuda_kernel_set_shape(&k->one, mesh_manifest_threads(&p->manifest), mesh_manifest_smem(&p->manifest), err) ||
823 !mesh_cuda_kernel_alloc_scratch(&k->one, p->manifest.scratch_bytes, err)) {
824 mesh_cuda_unload(&c->gpu, &k->one, NULL);
825 free(k);
826 return false;
827 }
828 *out = new_kernel(k);
829 return true;
830 }
831 if (c->mps_mode && pre && ((ckernel *)pre->self)->kind == K_WORKER) {
832 *out = pre;
833 return true;
834 }
835 mesh_kernel_drop(pre);
836 return mesh_errf(err, "internal: an image prepared for the other isolation mode");
837}
838
839static void stream_drop(mesh_stream *s) {
840 if (s->self) {
841 mesh_cuda_stream_destroy(s->self);
842 free(s->self);
843 }
844}
845
846static bool b_new_stream(const mesh_backend *b, uint32_t slot, const mesh_kernel *k, mesh_stream **out, char **err) {
847 (void)k;
848 mesh_stream *s = mesh_alloc(sizeof *s);
849 *s = (mesh_stream){MESH_BACKEND_CUDA, slot, NULL, stream_drop};
850 if (!cb(b)->mps_mode) {
851 mesh_cuda_stream *cs = mesh_alloc(sizeof *cs);
852 if (!mesh_cuda_new_stream(&cb(b)->gpu, cs, err)) {
853 free(cs);
854 free(s);
855 return false;
856 }
857 s->self = cs;
858 }
859 *out = s;
860 return true;
861}
862
863static bool b_launch(const mesh_backend *b, const mesh_stream *st, const mesh_kernel *k, uint32_t nodes, uint64_t dev, uint32_t slot, char **err) {
864 cuda_backend *c = cb(b);
865 ckernel *ck = k->self;
866 if (!c->mps_mode && st->self && ck->kind == K_ONE) {
867 uint64_t d = dev;
868 uint32_t s = slot;
869 void *params[] = {&d, &s};
870 return mesh_cuda_launch_on(&c->gpu, st->self, &ck->one, nodes, params, err);
871 }
872 if (c->mps_mode && ck->kind == K_WORKER) return worker_launch(ck->w, slot, nodes, err);
873 return mesh_errf(err, "internal: stream/image of the other isolation mode");
874}
875
876static bool b_query(const mesh_backend *b, const mesh_stream *st, const mesh_kernel *k, bool *done, char **err) {
877 cuda_backend *c = cb(b);
878 ckernel *ck = k->self;
879 if (!c->mps_mode && st->self) return mesh_cuda_query(&c->gpu, st->self, done, err);
880 if (c->mps_mode && ck->kind == K_WORKER) return worker_query(ck->w, done, err);
881 return mesh_errf(err, "internal: stream/image of the other isolation mode");
882}
883
884static bool b_unload(const mesh_backend *b, mesh_kernel *k, char **err) {
885 cuda_backend *c = cb(b);
886 ckernel *ck = k->self;
887 bool ok = true;
888 if (ck->kind == K_ONE) {
889 if (c->mps_mode) ok = mesh_errf(err, "internal: a context module in mps mode");
890 else ok = mesh_cuda_unload(&c->gpu, &ck->one, err);
891 }
892 mesh_kernel_drop(k); /* a worker goes back to the pool */
893 return ok;
894}
895
896static bool b_stack_limit(const mesh_backend *b, size_t *out, char **err) {
897 if (cb(b)->mps_mode) {
898 *out = SIZE_MAX;
899 return true;
900 }
901 return mesh_cuda_stack_limit(&cb(b)->gpu, out, err);
902}
903static bool b_set_stack_limit(const mesh_backend *b, size_t bytes, char **err) {
904 return cb(b)->mps_mode ? true : mesh_cuda_set_stack_limit(&cb(b)->gpu, bytes, err);
905}
906static bool b_module_scope(const mesh_backend *b, bool has_slot, uint32_t slot, uint32_t *scope) {
907 if (!cb(b)->mps_mode) return false;
908 *scope = has_slot ? slot : MESH_NO_SLOT;
909 return true;
910}
911static void b_destroy(mesh_backend *b) {
912 cuda_backend *c = cb(b);
913 mesh_caps_free(&c->caps);
914 mesh_cuda_gpu_free(&c->gpu);
915 mesh_nvjitlink_free(&c->jit);
916 free(c->arch);
917 free(c);
918 free(b);
919}
920
921static const mesh_backend_ops backend_ops = {
922 b_caps, b_isolated, b_resident, b_isolation, b_mailbox, b_make_current, b_compiler, b_load, b_new_stream, b_launch, b_query,
923 b_unload, b_stack_limit, b_set_stack_limit, b_module_scope, NULL, NULL, NULL, NULL, NULL, b_destroy,
924};
925
926#define CKD(call, what) \
927 do { \
928 if (!mesh_cu_ck_debug((call), (what), err)) \
929 return false; \
930 } while (0)
931
932/* (name, compute capability, SMs, blocks per SM, host-native atomics, block limits), no context. */
933static bool device_info(cuda_backend *c, char **err) {
934 if (!mesh_cuda() && !mesh_cuda_load(err)) return false;
935 const mesh_cuda_fns *cu = mesh_cuda();
936 CKD(cu->cuInit(0), "cuInit");
937 if (!mesh_cuda_injected()) {
938 char *lib = mesh_cuda_check_lib(err);
939 if (!lib) return false;
940 free(lib);
941 }
942 CUdevice dev = 0;
943 CKD(cu->cuDeviceGet(&dev, 0), "cuDeviceGet");
944 char name[128] = {0};
945 CKD(cu->cuDeviceGetName(name, 128, dev), "cuDeviceGetName");
946 int a[5];
947 const int attrs[5] = {CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
948 CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, CU_DEVICE_ATTRIBUTE_MAX_BLOCKS_PER_MULTIPROCESSOR,
949 CU_DEVICE_ATTRIBUTE_HOST_NATIVE_ATOMIC_SUPPORTED};
950 for (int i = 0; i < 5; i++) CKD(cu->cuDeviceGetAttribute(&a[i], attrs[i], dev), "cuDeviceGetAttribute");
951 if (!mesh_cuda_limits_of(dev, &c->limits, err)) return false;
952 c->gpu.name = mesh_strdup(name);
953 c->gpu.cc_major = a[0];
954 c->gpu.cc_minor = a[1];
955 c->gpu.sms_total = (uint32_t)a[2];
956 c->gpu.blocks_per_sm = (uint32_t)a[3];
957 c->host_atomics = a[4] != 0;
958 return true;
959}
960
961bool mesh_cuda_backend_open(mesh_isolation_mode iso, uint32_t sms, const char *mps_dir, uint32_t pool_n, const mesh_shm *shm,
962 mesh_nvjitlink *jit, mesh_backend **out, char **err) {
963 cuda_backend *c = mesh_alloc(sizeof *c);
964 c->shm = *shm;
965 c->jit = *jit;
966 *jit = (mesh_nvjitlink){0};
967 if (iso == MESH_ISOLATION_MODE_BLINK) {
968 if (!mesh_cuda_gpu_open(sms, &c->gpu, err) || !mesh_cuda_make_current(&c->gpu, err) ||
969 !mesh_cuda_host_register(&c->gpu, shm->base, shm->len, &c->dev, err))
970 goto fail;
971 c->limits = c->gpu.limits;
972 c->host_atomics = c->gpu.host_atomics;
973 } else {
974 c->mps_mode = true;
975 char *dir = mps_dir ? mesh_strdup(mps_dir) : mesh_fmt("/tmp/meshd-mps.%d", (int)getpid());
976 if (!mps_start_daemon(dir, mps_dir == NULL, err)) {
977 free(dir);
978 goto fail;
979 }
980 /* meshd's own cuInit (device attributes, no context) must join no daemon. */
981 char *none = mesh_fmt("%s/none", dir);
982 setenv(MESH_MPS_PIPE_ENV, none, 1);
983 free(none);
984 if (!device_info(c, err)) {
985 free(dir);
986 goto fail;
987 }
988 if (sms > c->gpu.sms_total) {
989 mesh_errf(err, "--sms %u: the device has %u", (unsigned)sms, (unsigned)c->gpu.sms_total);
990 free(dir);
991 goto fail;
992 }
993 c->gpu.sms = sms ? sms : c->gpu.sms_total;
994 c->gpu.green = sms != 0;
995 mps *m = mesh_alloc(sizeof *m);
996 char exe[4096];
997 ssize_t n = readlink("/proc/self/exe", exe, sizeof exe - 1);
998 if (n < 0) {
999 char *e = mps_os_error(errno);
1000 mesh_errf(err, "mps: /proc/self/exe: %s", e);
1001 free(e);
1002 free(dir);
1003 free(m);
1004 goto fail;
1005 }
1006 exe[n] = 0;
1007 m->dir = dir;
1008 m->mailbox = mesh_shm_path(shm);
1009 m->sms = sms;
1010 m->exe = mesh_strdup(exe);
1011 pthread_mutex_init(&m->server_mu, NULL);
1012 mesh_info("mps: mailbox %d bytes in a memfd (%s); workers run %s", MESH_MAILBOX_BYTES, m->mailbox, m->exe);
1013 c->m = m;
1014 mesh_backend_on_cleanup(MESH_CLEANUP_MPS, mesh_mps_cleanup);
1015 if (!mps_prewarm(m, pool_n, err)) goto fail;
1016 pthread_mutex_lock(&g_mu);
1017 server = m;
1018 pthread_mutex_unlock(&g_mu);
1019 }
1020 c->arch = mesh_fmt("sm_%d%d", c->gpu.cc_major, c->gpu.cc_minor);
1021 c->caps = (mesh_caps){MESH_BACKEND_CUDA, mesh_strdup(c->gpu.name), c->gpu.cc_major, c->gpu.cc_minor, c->gpu.sms_total, c->gpu.sms,
1022 c->gpu.green, c->gpu.blocks_per_sm, 0, c->limits.max_threads, c->limits.max_shared, mesh_strdup(c->arch)};
1023 c->comp = (mesh_compiler){&compiler_ops, c};
1024 mesh_backend *b = mesh_alloc(sizeof *b);
1025 *b = (mesh_backend){&backend_ops, c};
1026 *out = b;
1027 return true;
1028fail:
1029 mesh_cuda_gpu_free(&c->gpu);
1030 mesh_nvjitlink_free(&c->jit);
1031 free(c);
1032 return false;
1033}
1034
1035CUcontext mesh_cuda_backend_context(const mesh_backend *b) { return cb(b)->mps_mode ? NULL : cb(b)->gpu.ctx; }
1036bool mesh_cuda_backend_mps(const mesh_backend *b) { return cb(b)->mps_mode; }
1037bool mesh_cuda_backend_host_atomics(const mesh_backend *b) { return cb(b)->host_atomics; }