1/* @ajax/mesh-backend: the executor interface meshd's driver sees (mesh-host
2 * backend.rs). Everything above it -- rings, publisher, lifecycle,
3 * placement, peers -- is backend-neutral mailbox words.
4 *
5 * - CUDA (@ajax/mesh-mps over @ajax/mesh-cuda, Linux): one resident kernel per
6 * program, PTX compiled by nvJitLink, the mailbox in pinned host memory.
7 * - wasm (@ajax/mesh-exec-wasm): CPU programs on lease lanes.
8 * - native (@ajax/mesh-native): a program that is an executable runs as its own
9 * worker process on the shared-memory mailbox.
10 * - Metal is not ported (Metal Shading Language is C++).
11 *
12 * The Rust's traits are tables of functions here: a backend, its compiler, a
13 * loaded image (mesh_kernel) and where a slot's work runs (mesh_stream) each
14 * carry `ops` and their implementation's `self`. Errors are `char **err`:
15 * on failure the function returns false and sets *err to a malloc'd message
16 * (the Rust's text; load refusals start with a LOAD_REASONS word). */
17#ifndef MESH_BACKEND_H
18#define MESH_BACKEND_H
1920#include "mesh/loader.h"
21#include "mesh/manifest.h"
2223#include "mesh/abi/backend.h"
24#include "mesh/abi/control_block.h"
2526#include <stdbool.h>
27#include <stddef.h>
28#include <stdint.h>
2930/* CUDA's isolation modes (`--isolation`): one context that blinks, or MPS workers. */
31typedef enum { MESH_ISOLATION_MODE_BLINK, MESH_ISOLATION_MODE_MPS } mesh_isolation_mode;
3233/* nvJitLink's default dir (`--nvjitlink`). */
34#define MESH_NVJITLINK_DIR "/usr/local/cuda/lib64"
35/* A blink scope that halts nothing (no program sits in it): PROG_SLOTS. */
36#define MESH_NO_SLOT 32u
3738/* This machine's GPU (or executor) as the mailbox header and F_JOIN describe it. */
39typedef struct {
40 uint32_t kind; /* BACKEND_* */
41 char *name;
42 int32_t cc_major, cc_minor; /* CUDA compute capability; 0, 0 elsewhere */
43 uint32_t sms_total; /* SMs (CUDA) / cores of the device */
44 uint32_t sms; /* those programs run on */
45 bool green;
46 uint32_t blocks_per_sm; /* CUDA resident blocks per SM; 0 elsewhere */
47 uint32_t family; /* Metal's MTLGPUFamilyAppleN; 0 on CUDA */
48 uint32_t max_tg_threads; /* the most threads a node may have (1024); 0: no GPU */
49 uint32_t max_shared; /* the most shared memory a node may have, the runtime's included; 0: no GPU */
50 char *arch; /* what code is compiled for: "sm_121" */
51} mesh_caps;
5253uint32_t mesh_caps_word(const mesh_caps *c); /* backend_word(kind, family) */
54const char *mesh_caps_backend_name(const mesh_caps *c); /* BACKENDS[kind] */
55void mesh_caps_copy(mesh_caps *dst, const mesh_caps *src);
56void mesh_caps_free(mesh_caps *c);
57/* A backend label for logs and the control port: "cuda", "wasm", "metal-apple9", "backend7" (malloc'd). */
58char *mesh_backend_label(uint32_t word);
5960/* ---- a loaded program image (Rust: backend::Kernel) ------------------------------------ */
6162typedef struct mesh_kernel mesh_kernel;
63typedef struct {
64 uint32_t exec; /* the executor that runs it (BACKEND_*) */
65 /* Registers per thread (CUDA); 0 elsewhere. */
66 bool (*regs)(const mesh_kernel *k, uint32_t *out, char **err);
67 /* Local (stack) bytes per thread; 0 on the CPU executors. */
68 bool (*local_bytes)(const mesh_kernel *k, uint32_t *out, char **err);
69 /* CUDA: resident blocks of the node's shape per SM; the CPU executors: 1. */
70 bool (*max_blocks_per_sm)(const mesh_kernel *k, uint32_t *out, char **err);
71 /* The program's scratch: its device address (PX_SCRATCH); 0: none. */
72 uint64_t (*scratch)(const mesh_kernel *k);
73 /* A u32 global of the image, e.g. MESH_RT_ABI. */
74 bool (*global_u32)(const mesh_kernel *k, const char *name, uint32_t *out, char **err);
75 /* Drops the image without the backend's unload (Rust's Drop: an image a
76 * failed load or a superseded prepare leaves behind). */
77 void (*drop)(mesh_kernel *k);
78} mesh_kernel_ops;
7980struct mesh_kernel {
81 const mesh_kernel_ops *ops;
82 void *self;
83};
8485static inline uint32_t mesh_kernel_exec(const mesh_kernel *k) { return k->ops->exec; }
86/* The native worker image behind k, or NULL (Rust: Kernel::native). */
87static inline void *mesh_kernel_native(const mesh_kernel *k) {
88 return k && k->ops->exec == MESH_BACKEND_NATIVE ? k->self : NULL;
89}
90bool mesh_kernel_regs(const mesh_kernel *k, uint32_t *out, char **err);
91bool mesh_kernel_local_bytes(const mesh_kernel *k, uint32_t *out, char **err);
92bool mesh_kernel_max_blocks_per_sm(const mesh_kernel *k, uint32_t *out, char **err);
93uint64_t mesh_kernel_scratch(const mesh_kernel *k);
94bool mesh_kernel_global_u32(const mesh_kernel *k, const char *name, uint32_t *out, char **err);
95/* Drops and frees a heap-allocated kernel (NULL: nothing). */
96void mesh_kernel_drop(mesh_kernel *k);
9798/* ---- where a slot's work runs (Rust: backend::Stream) --------------------------------- */
99100typedef struct mesh_stream {
101 uint32_t exec; /* BACKEND_* */
102 uint32_t slot;
103 void *self; /* the backend's (a CUDA stream, ...); NULL for the CPU executors */
104 void (*drop)(struct mesh_stream *s); /* NULL: nothing to release */
105} mesh_stream;
106107/* Releases the stream's resources and frees it (NULL: nothing). */
108void mesh_stream_drop(mesh_stream *s);
109110/* ---- a compiled image, ready for the blink (Rust: backend::Prepared) ------------------ */
111112typedef struct {
113 mesh_manifest manifest;
114 /* The image's identity in the program table and on the peer link. */
115 uint64_t hash;
116 /* CUDA: the cubin; native: the binary. */
117 uint8_t *code;
118 size_t code_len;
119 double compile_ms;
120 /* The bytes it came from (replicated to the peers). */
121 mesh_payload src;
122 /* mps: its worker, started with the module loaded; native: the worker
123 * image; wasm: the compiled module. NULL: the blink loads the code. */
124 mesh_kernel *kernel;
125 /* The executor that compiled it (BACKEND_*): the one that runs it. */
126 uint32_t exec;
127 /* CUDA: the register cap it was compiled with (PX_REG_CAP); 0: none. */
128 uint32_t reg_cap;
129} mesh_prepared;
130131static inline const char *mesh_prepared_dir(const mesh_prepared *p) { return p->src.dir; }
132/* Frees everything, dropping a kernel still held. */
133void mesh_prepared_free(mesh_prepared *p);
134135/* ---- compilers (Rust: trait Compiler) -------------------------------------------------- */
136137typedef struct mesh_compiler mesh_compiler;
138typedef struct {
139 /* Whether the catalog must hold the PTX (CUDA compiles it). */
140 bool (*needs_ptx)(const mesh_compiler *c);
141 /* Manifest checks, the hash (checked against *expect, a peer's F_LOAD;
142 * expect NULL: none), compile. root is --programs. Takes ownership of
143 * *src (it moves into out->src on success, is freed on failure). Errors
144 * start with a LOAD_REASONS word. Runs on the JIT thread: thread-safe. */
145 bool (*prepare)(const mesh_compiler *c, mesh_payload *src, const uint64_t *expect, const char *root, mesh_prepared *out,
146 char **err);
147} mesh_compiler_ops;
148149struct mesh_compiler {
150 const mesh_compiler_ops *ops;
151 void *self;
152};
153154/* ---- the backend (Rust: trait Backend) ------------------------------------------------- */
155156typedef struct mesh_backend mesh_backend;
157typedef struct {
158 const mesh_caps *(*caps)(const mesh_backend *b);
159 /* A module change halts only the slot it changes (mps, the CPU executors);
160 * false: every program of the context blinks (CUDA's default). */
161 bool (*isolated)(const mesh_backend *b);
162 /* Programs are resident kernels sharing the device's block slots: the occupancy budget applies (CUDA). */
163 bool (*resident)(const mesh_backend *b);
164 /* S_ISOLATION (ISOLATION_*). */
165 uint32_t (*isolation)(const mesh_backend *b);
166 /* The mailbox: its host address and the device address the kernels get. */
167 bool (*mailbox)(const mesh_backend *b, uint8_t **host, uint64_t *dev, char **err);
168 bool (*make_current)(const mesh_backend *b, char **err);
169 /* The compiler (owned by the backend, valid while it lives). */
170 const mesh_compiler *(*compiler)(const mesh_backend *b);
171 /* The image's module (blink: loaded here; mps, native, wasm: prepared
172 * already, taken out of p->kernel). *out is a new heap kernel. */
173 bool (*load)(const mesh_backend *b, mesh_prepared *p, mesh_kernel **out, char **err);
174 bool (*new_stream)(const mesh_backend *b, uint32_t slot, const mesh_kernel *k, mesh_stream **out, char **err);
175 /* Starts slot's nodes (returns at once; P_ONLINE tells when they are up). */
176 bool (*launch)(const mesh_backend *b, const mesh_stream *st, const mesh_kernel *k, uint32_t nodes, uint64_t dev,
177 uint32_t slot, char **err);
178 /* *done = true once the slot's last launch finished. */
179 bool (*query)(const mesh_backend *b, const mesh_stream *st, const mesh_kernel *k, bool *done, char **err);
180 /* Unloads and frees k (whatever happens). */
181 bool (*unload)(const mesh_backend *b, mesh_kernel *k, char **err);
182 bool (*stack_limit)(const mesh_backend *b, size_t *out, char **err);
183 bool (*set_stack_limit)(const mesh_backend *b, size_t bytes, char **err);
184 /* The halt scope of a module change of slot (has_slot false: a load or a
185 * free). Returns true with *scope set, or false for None (the whole context). */
186 bool (*module_scope)(const mesh_backend *b, bool has_slot, uint32_t slot, uint32_t *scope);
187 /* The slot's work ended abnormally (a native worker died): why (malloc'd), once; NULL: nothing. NULL op: never. */
188 char *(*crashed)(const mesh_backend *b, const mesh_stream *st, const mesh_kernel *k);
189 /* The slot's last launch is over by the executor's own account (a native worker). NULL op: false. */
190 bool (*owns_halt)(const mesh_backend *b);
191 /* Running nodes re-read their reads mask at every peer view. NULL op: false. */
192 bool (*reads_live)(const mesh_backend *b);
193 /* About once a second: refresh what the executor measures. NULL op: nothing. */
194 void (*sample)(const mesh_backend *b);
195 /* A launch's nodes must be online within this many ms. NULL op: 2000. */
196 uint32_t (*online_timeout_ms)(const mesh_backend *b);
197 /* Frees the backend (its exit paths have run cleanup already). */
198 void (*destroy)(mesh_backend *b);
199} mesh_backend_ops;
200201struct mesh_backend {
202 const mesh_backend_ops *ops;
203 void *self;
204};
205206/* The trait's default methods, applied when an op is NULL. */
207char *mesh_backend_crashed(const mesh_backend *b, const mesh_stream *st, const mesh_kernel *k);
208bool mesh_backend_owns_halt(const mesh_backend *b);
209bool mesh_backend_reads_live(const mesh_backend *b);
210void mesh_backend_sample(const mesh_backend *b);
211uint32_t mesh_backend_online_timeout_ms(const mesh_backend *b);
212void mesh_backend_destroy(mesh_backend *b);
213214/* ---- free functions of backend.rs ------------------------------------------------------- */
215216/* The executor for a manifest: the first of its `backends` this meshd runs
217 * (have: BACKEND_* kinds, the GPU's first). False with *err ("unsupported: ..."). */
218bool mesh_pick_exec(const mesh_manifest *m, const uint32_t *have, size_t nhave, uint32_t *exec, char **err);
219/* Program m in catalog dir has an image for one of have: the executor
220 * pick_exec names, and for a native program its binary for this target. */
221bool mesh_runs_here(const char *dir, const mesh_manifest *m, const uint32_t *have, size_t nhave);
222223/* The executors' compilers, the GPU's first (what the JIT thread gets). */
224typedef struct {
225 uint32_t kind;
226 const mesh_compiler *comp;
227} mesh_compiler_entry;
228229/* Picks the executor for the payload's manifest and prepares the image with
230 * its compiler (on the JIT thread). Takes ownership of *src. */
231bool mesh_prepare_any(const mesh_compiler_entry *comps, size_t n, mesh_payload *src, const uint64_t *expect, const char *root,
232 mesh_prepared *out, char **err);
233/* The manifest and abi checks every backend makes first, including
234 * `backends`: a program that does not list this backend is refused as `unsupported`. */
235bool mesh_check_manifest(const mesh_payload *src, uint32_t kind, mesh_manifest *out, char **err);
236237/* Every exit path calls it: native workers first, then the CUDA MPS daemon
238 * and workers (as backend.rs). Executors register their cleanup when they
239 * start; cleanups run in ascending `order` (MESH_CLEANUP_NATIVE before
240 * MESH_CLEANUP_MPS), each once per registration. */
241#define MESH_CLEANUP_NATIVE 10
242#define MESH_CLEANUP_WASM 15
243#define MESH_CLEANUP_MPS 20
244void mesh_backend_cleanup(void);
245void mesh_backend_on_cleanup(int order, void (*fn)(void));
246247/* The executable a native program has for this target: <dir>/<target triple>/<name> (malloc'd). */
248char *mesh_native_binary(const char *dir, const char *name);
249/* The target triple this was built for ("x86_64-unknown-linux-gnu"). */
250const char *mesh_native_target(void);
251252#endif