==> tests/mps.c <==1/* The CUDA backend without a GPU, in both isolation modes, over
2 * @ajax/cuda-driver-fake (whose kernels play the program's node loop on the
3 * real mailbox) and a fake nvJitLink.
4 *
5 * This one binary is three programs: the test; `nvidia-cuda-mps-control`
6 * (argv[0], via a symlink on PATH: a fake control daemon that knows its
7 * clients from a file); and `meshd worker` (argv[1]: an MPS worker on the
8 * fake device, which registers itself as a client). So --isolation mps runs
9 * for real: a private daemon, a pool of worker processes, preloads, launches
10 * through the wire format, halts, the pool's reuse, and the exit cleanup. */
11#include "mesh/check.h"
12#include "mesh/cuda_fake.h"
13#include "mesh/manifest_encode.h"
14#include "mesh/mps.h"
15#include "mesh/mps_worker.h"
16#include "mesh/worker_node.h"
17
18#include "mesh/abi/control_block.h"
19#include "mesh/abi/ext.h"
20#include "mesh/abi/header.h"
21#include "mesh/abi/host_status.h"
22#include "mesh/abi/ops.h"
23#include "mesh/abi/program_table.h"
24#include "mesh/abi/ring.h"
25
26#include <dirent.h>
27#include <fcntl.h>
28#include <signal.h>
29#include <stdio.h>
30#include <stdlib.h>
31#include <string.h>
32#include <sys/stat.h>
33#include <sys/wait.h>
34#include <time.h>
35#include <unistd.h>
36
37/* ---- the program the fake GPU runs: hello (arg * 3 + 7), with a view magic ---------------- */
38
39static void hello_init(const mesh_worker_ctx *c, void *u) {
40 (void)u;
41 mesh_ctx_set_u32(c, 0, 0x4f4c4548);
42}
43static uint32_t hello_task(const mesh_worker_ctx *c, uint32_t op, uint32_t arg, void *u) {
44 (void)u;
45 if (op != MESH_OP_TASK) return MESH_R_BAD_OP;
46 mesh_ctx_count(c, 0, 1);
47 return arg * 3 + 7;
48}
49static const mesh_worker_program hello = {hello_init, hello_task, NULL};
50
51static void kernel(const char *entry, uint64_t mailbox, uint32_t slot, uint32_t nodes, uint32_t threads, void *user) {
52 (void)entry, (void)nodes, (void)threads, (void)user;
53 mesh_worker_run_nodes(mesh_mailbox_new((uint8_t *)(uintptr_t)mailbox, 0), slot, &hello, NULL);
54}
55
56/* ---- a fake nvJitLink: any PTX becomes a fake cubin ------------------------------------------ */
57
58static int j_create(void **h, uint32_t n, const char *const *o) {
59 (void)n, (void)o;
60 *h = (void *)1;
61 return 0;
62}
63static int j_destroy(void **h) {
64 *h = NULL;
65 return 0;
66}
67static int j_add(void *h, int k, const void *d, size_t n, const char *name) {
68 (void)h, (void)k, (void)name;
69 return n > 0 && ((const char *)d)[n - 1] == 0 && strstr(d, ".entry mesh_main") ? 0 : 5;
70}
71static int j_complete(void *h) {
72 (void)h;
73 return 0;
74}
75static const char cubin[] = "FAKECUBIN entry=mesh_main regs=40 local=64 rt_abi=6";
76static int j_size(void *h, size_t *n) {
77 (void)h;
78 *n = sizeof cubin;
79 return 0;
80}
81static int j_cubin(void *h, void *out) {
82 (void)h;
83 memcpy(out, cubin, sizeof cubin);
84 return 0;
85}
86static int j_err_size(void *h, size_t *n) {
87 (void)h;
88 *n = 0;
89 return 0;
90}
91static int j_err_log(void *h, void *o) {
92 (void)h, (void)o;
93 return 0;
94}
95static const mesh_nvjitlink_fns jit_fns = {j_create, j_destroy, j_add, j_complete, j_size, j_cubin, j_err_size, j_err_log};
96
97/* ---- the fake control daemon -------------------------------------------------------------- */
98
99static char *pipe_file(const char *name) {
100 const char *d = getenv("CUDA_MPS_PIPE_DIRECTORY");
101 return mesh_fmt("%s/%s", d ? d : ".", name);
102}
103
104static int fake_control(int argc, char **argv) {
105 if (argc > 1 && strcmp(argv[1], "-d") == 0) {
106 pid_t p = fork();
107 if (p != 0) return p > 0 ? 0 : 1;
108 setsid();
109 char *pidf = pipe_file("nvidia-cuda-mps-control.pid"), *ctl = pipe_file("control"), *quit = pipe_file("quit");
110 FILE *f = fopen(pidf, "w");
111 fprintf(f, "%d\n", (int)getpid());
112 fclose(f);
113 close(open(ctl, O_CREAT | O_WRONLY, 0600));
114 for (int i = 0; i < 600000 && access(quit, F_OK) != 0; i++) usleep(1000);
115 unlink(ctl);
116 _exit(0);
117 }
118 char line[128] = {0};
119 if (!fgets(line, sizeof line, stdin)) return 1;
120 line[strcspn(line, "\n")] = 0;
121 char *clients = pipe_file("clients"), *quit = pipe_file("quit");
122 int rc = 0;
123 if (strcmp(line, "get_server_list") == 0) {
124 if (access(clients, F_OK) == 0) puts("4242");
125 } else if (strncmp(line, "get_client_list ", 16) == 0) {
126 FILE *f = fopen(clients, "r");
127 char buf[64];
128 while (f && fgets(buf, sizeof buf, f)) fputs(buf, stdout);
129 if (f) fclose(f);
130 } else if (strcmp(line, "quit") == 0) {
131 close(open(quit, O_CREAT | O_WRONLY, 0600));
132 } else {
133 fprintf(stderr, "fake control: unknown request %s\n", line);
134 rc = 1;
135 }
136 free(clients);
137 free(quit);
138 return rc;
139}
140
141static void fake_worker(int argc, char **argv) {
142 mesh_fake_cuda_install(&(mesh_fake_cuda_cfg){.run = kernel});
143 char *clients = pipe_file("clients");
144 FILE *f = fopen(clients, "a");
145 fprintf(f, "%d\n", (int)getpid());
146 fclose(f);
147 free(clients);
148 mesh_mps_worker_main(argc, argv);
149}
150
151/* ---- the test -------------------------------------------------------------------------- */
152
153static uint64_t now_ms(void) {
154 struct timespec t;
155 clock_gettime(CLOCK_MONOTONIC, &t);
156 return (uint64_t)t.tv_sec * 1000 + (uint64_t)t.tv_nsec / 1000000;
157}
158
159static const char *hello_ptx = ".version 8.8\n.target sm_121\n.address_size 64\n.visible .entry mesh_main() { ret; }\n";
160
161static mesh_payload payload(const char *name) {
162 char *toml = mesh_fmt("name = \"%s\"\nabi = 6\nview_bytes = 0x1000\nnodes = 2\ncounters = [\"tasks\"]\n", name);
163 mesh_payload p = {mesh_strdup(name), toml, (uint8_t *)mesh_strdup(hello_ptx), strlen(hello_ptx)};
164 return p;
165}
166
167static void install(mesh_mailbox mb, uint32_t slot, const mesh_prepared *pr, uint32_t view_off) {
168 mesh_image im = {view_off, 0, pr->hash, 40, (uint32_t)pr->code_len, 6, 1};
169 uint8_t e[MESH_PT_BYTES], x[MESH_PX_BYTES];
170 mesh_manifest_encode_entry(&pr->manifest, &im, e);
171 mesh_manifest_encode_ext(&pr->manifest, x);
172 for (size_t i = 2; i < MESH_PT_BYTES / 4; i++) {
173 uint32_t w;
174 memcpy(&w, e + 4 * i, 4);
175 mesh_mailbox_set(mb, mesh_prog_entry(slot) + 4 * i, w);
176 }
177 for (size_t i = 0; i < MESH_PX_BYTES / 4; i++) {
178 uint32_t w;
179 memcpy(&w, x + 4 * i, 4);
180 mesh_mailbox_set(mb, mesh_prog_ext(slot) + 4 * i, w);
181 }
182 mesh_mailbox_set(mb, mesh_prog_entry(slot) + MESH_PT_STATE, MESH_ST_LOADED);
183}
184
185static void prep_ring(mesh_mailbox mb, uint32_t slot, uint32_t tail) {
186 size_t c = mesh_prog_ctl(slot);
187 mesh_mailbox_set(mb, c + MESH_P_RING_HEAD, tail);
188 mesh_mailbox_set(mb, c + MESH_P_ONLINE, 0);
189 mesh_mailbox_set(mb, c + MESH_P_HALTED, 0);
190 mesh_mailbox_publish(mb, c + MESH_P_SHUTDOWN, 0);
191}
192
193static bool ticket(mesh_mailbox mb, uint32_t slot, uint32_t t, uint32_t op, uint32_t arg, uint32_t *res) {
194 size_t c = mesh_prog_ctl(slot), e = c + MESH_P_RING + MESH_ENTRY_BYTES * (t % MESH_PROG_RING), k = t % MESH_PROG_RESULTS;
195 mesh_mailbox_set(mb, e + MESH_E_OP, op);
196 mesh_mailbox_set(mb, e + MESH_E_ARG, arg);
197 mesh_mailbox_set(mb, e + MESH_E_SLOT, k);
198 mesh_mailbox_publish(mb, e + MESH_E_SEQ, t + 1);
199 uint64_t t0 = now_ms();
200 while (mesh_mailbox_get(mb, c + MESH_P_MARKS + 4 * k) != t + 1)
201 if (now_ms() - t0 > 10000) return false;
202 *res = mesh_mailbox_get(mb, c + MESH_P_RESULTS + 4 * k);
203 return true;
204}
205
206/* Runs one program through a backend: prepare, load, launch, tickets, halt, unload. */
207static void run_one(mesh_backend *b, mesh_mailbox mb, uint32_t slot, const char *name, uint32_t view_off) {
208 char *err = NULL;
209 mesh_payload p = payload(name);
210 mesh_prepared pr;
211 const mesh_compiler *comp = b->ops->compiler(b);
212 CHECK(comp->ops->needs_ptx(comp));
213 if (!CHECK(comp->ops->prepare(comp, &p, NULL, "/", &pr, &err))) {
214 fprintf(stderr, "prepare: %s\n", err);
215 return;
216 }
217 CHECK(pr.exec == MESH_BACKEND_CUDA);
218 install(mb, slot, &pr, view_off);
219 mesh_kernel *k = NULL;
220 if (!CHECK(b->ops->load(b, &pr, &k, &err))) {
221 fprintf(stderr, "load: %s\n", err);
222 mesh_prepared_free(&pr);
223 return;
224 }
225 uint32_t v;
226 CHECK(mesh_kernel_global_u32(k, "MESH_RT_ABI", &v, &err) && v == 6);
227 CHECK(mesh_kernel_regs(k, &v, &err) && v == 40);
228 CHECK(mesh_kernel_local_bytes(k, &v, &err) && v == 64);
229 CHECK(mesh_kernel_max_blocks_per_sm(k, &v, &err) && v == 16);
230 mesh_stream *st = NULL;
231 CHECK(b->ops->new_stream(b, slot, k, &st, &err));
232 uint8_t *host;
233 uint64_t dev;
234 CHECK(b->ops->mailbox(b, &host, &dev, &err));
235 prep_ring(mb, slot, 0);
236 if (!CHECK(b->ops->launch(b, st, k, 2, dev, slot, &err))) fprintf(stderr, "launch: %s\n", err);
237 uint64_t t0 = now_ms();
238 while (mesh_mailbox_get(mb, mesh_prog_ctl(slot) + MESH_P_ONLINE) < 2 && now_ms() - t0 < 10000) usleep(100);
239 CHECK_U(mesh_mailbox_get(mb, mesh_prog_ctl(slot) + MESH_P_ONLINE), 2);
240 uint32_t r;
241 CHECK(ticket(mb, slot, 0, MESH_OP_INIT, 0, &r) && r == MESH_R_NONE);
242 CHECK_U(mesh_mailbox_get(mb, view_off), 0x4f4c4548);
243 for (uint32_t t = 1; t < 100; t++) CHECK(ticket(mb, slot, t, MESH_OP_TASK, t, &r) && r == t * 3 + 7);
244 CHECK_U(mesh_mailbox_get(mb, mesh_prog_ctl(slot) + MESH_P_COUNTERS), 99);
245 bool done = true;
246 CHECK(b->ops->query(b, st, k, &done, &err) && !done);
247 /* A halt: the nodes leave, the stream (or the worker's) goes idle. */
248 mesh_mailbox_publish(mb, mesh_prog_ctl(slot) + MESH_P_SHUTDOWN, MESH_SHUTDOWN_HALT);
249 t0 = now_ms();
250 do {
251 CHECK(b->ops->query(b, st, k, &done, &err));
252 } while (!done && now_ms() - t0 < 10000);
253 CHECK(done);
254 CHECK_U(mesh_mailbox_get(mb, mesh_prog_ctl(slot) + MESH_P_HALTED), 2);
255 mesh_stream_drop(st);
256 CHECK(b->ops->unload(b, k, &err));
257 mesh_prepared_free(&pr);
258 free(err);
259}
260
261static void blink(mesh_shm *shm) {
262 mesh_fake_cuda_install(&(mesh_fake_cuda_cfg){.run = kernel});
263 mesh_nvjitlink jit;
264 mesh_nvjitlink_from(&jit_fns, &jit);
265 mesh_backend *b;
266 char *err = NULL;
267 if (!CHECK(mesh_cuda_backend_open(MESH_ISOLATION_MODE_BLINK, 0, NULL, 0, shm, &jit, &b, &err))) {
268 fprintf(stderr, "open: %s\n", err);
269 return;
270 }
271 const mesh_caps *c = b->ops->caps(b);
272 CHECK_S(c->arch, "sm_121");
273 CHECK(c->kind == MESH_BACKEND_CUDA && c->sms_total == 48 && c->blocks_per_sm == 24 && c->max_tg_threads == 1024);
274 CHECK(!b->ops->isolated(b) && b->ops->resident(b));
275 CHECK_U(b->ops->isolation(b), MESH_ISOLATION_BLINK);
276 uint32_t scope;
277 CHECK(!b->ops->module_scope(b, true, 3, &scope));
278 CHECK(mesh_cuda_backend_context(b) != NULL);
279 uint8_t *host;
280 uint64_t dev;
281 CHECK(b->ops->mailbox(b, &host, &dev, &err) && host == shm->base && dev == (uintptr_t)shm->base);
282 run_one(b, mesh_mailbox_new(shm->base, 0), 3, "hello", 0x100000);
283 mesh_fake_cuda_stats fs = mesh_fake_cuda_stats_now();
284 CHECK(fs.modules_live == 0 && fs.launches == 1 && fs.last_grid == 2 && fs.last_block == 32 && fs.last_smem == 0x200);
285 /* A refusal from the compiler: the node does not fit a block. */
286 mesh_payload p = payload("wide");
287 free(p.manifest);
288 p.manifest = mesh_strdup("name = \"wide\"\nabi = 6\nshared_bytes = 0x30000\n");
289 mesh_prepared pr;
290 const mesh_compiler *comp = b->ops->compiler(b);
291 CHECK(!comp->ops->prepare(comp, &p, NULL, "/", &pr, &err));
292 CHECK_S(err, "occupancy: wide: 197120 bytes of shared memory per node (the runtime's 512 + shared_bytes); the device allows 101376 per block");
293 /* And the hash a peer announced. */
294 p = payload("hello");
295 uint64_t wrong = 1;
296 CHECK(!comp->ops->prepare(comp, &p, &wrong, "/", &pr, &err));
297 CHECK(strncmp(err, "peer: hello: the PTX received hashes to ", 40) == 0);
298 free(err);
299 mesh_backend_destroy(b);
300}
301
302static void mps_mode(mesh_shm *shm, const char *self) {
303 /* nvidia-cuda-mps-control on PATH is this binary. */
304 char bin[] = "/tmp/mesh-mps-bin-XXXXXX";
305 CHECK(mkdtemp(bin) != NULL);
306 char *link = mesh_fmt("%s/nvidia-cuda-mps-control", bin);
307 CHECK(symlink(self, link) == 0);
308 char *path = mesh_fmt("%s:%s", bin, getenv("PATH") ? getenv("PATH") : "/usr/bin:/bin");
309 setenv("PATH", path, 1);
310 char *dir = mesh_fmt("/tmp/mesh-mps-test-%d", (int)getpid());
311 mesh_fake_cuda_install(&(mesh_fake_cuda_cfg){.run = kernel});
312 mesh_nvjitlink jit;
313 mesh_nvjitlink_from(&jit_fns, &jit);
314 mesh_backend *b;
315 char *err = NULL;
316 if (!CHECK(mesh_cuda_backend_open(MESH_ISOLATION_MODE_MPS, 0, dir, 1, shm, &jit, &b, &err))) {
317 fprintf(stderr, "open mps: %s\n", err);
318 return;
319 }
320 CHECK(b->ops->isolated(b) && mesh_cuda_backend_mps(b));
321 CHECK_U(b->ops->isolation(b), MESH_ISOLATION_MPS);
322 uint32_t scope;
323 CHECK(b->ops->module_scope(b, true, 4, &scope) && scope == 4);
324 CHECK(b->ops->module_scope(b, false, 0, &scope) && scope == MESH_NO_SLOT);
325 char *pd = mesh_mps_native_pipe_dir(), *want = mesh_fmt("%s/pipe", dir);
326 CHECK_S(pd, want);
327 free(pd);
328 free(want);
329 /* Two programs, one after the other: the second reuses the pooled worker. */
330 run_one(b, mesh_mailbox_new(shm->base, 0), 4, "hello", 0x100000);
331 run_one(b, mesh_mailbox_new(shm->base, 0), 5, "again", 0x101000);
332 /* The cleanup kills the workers, quits the daemon and keeps the logs of a --mps-dir. */
333 mesh_backend_cleanup();
334 char *ctl = mesh_fmt("%s/pipe", dir), *logs = mesh_fmt("%s/log", dir);
335 CHECK(access(ctl, F_OK) != 0);
336 CHECK(access(logs, F_OK) == 0);
337 rmdir(logs);
338 rmdir(dir);
339 free(ctl);
340 free(logs);
341 mesh_backend_destroy(b);
342 unlink(link);
343 rmdir(bin);
344 free(link);
345 free(path);
346 free(dir);
347 free(err);
348}
349
350int main(int argc, char **argv) {
351 const char *base = strrchr(argv[0], '/');
352 base = base ? base + 1 : argv[0];
353 if (strcmp(base, "nvidia-cuda-mps-control") == 0) return fake_control(argc, argv);
354 if (argc > 1 && strcmp(argv[1], "worker") == 0) fake_worker(argc, argv);
355 char self[4096];
356 ssize_t n = readlink("/proc/self/exe", self, sizeof self - 1);
357 if (n <= 0) return 1;
358 self[n] = 0;
359 setenv("CUDA_MODULE_LOADING", "EAGER", 1);
360 mesh_shm shm;
361 char *err = NULL;
362 if (!mesh_shm_new(&shm, &err)) return 1;
363 mesh_mailbox mb = mesh_mailbox_new(shm.base, 0);
364 mesh_mailbox_set(mb, MESH_H_MAGIC, MESH_MAGIC);
365 mesh_mailbox_set(mb, MESH_H_VERSION, MESH_VERSION);
366 blink(&shm);
367 mps_mode(&shm, self);
368 return mesh_check_done();
369}