1/* @ajax/mesh-inproc: a GPU program driven in-process (mesh-host inproc.rs):
2 * this process plays meshd's loader and publisher for one or more images,
3 * through meshd's backend -- compile, load, install a program-table slot in
4 * a mailbox of its own, launch, publish through the slot's ring, stop and
5 * relaunch.  mesh-kcheck checks every program this way; the gpt2 worker runs
6 * gpt2x's exact kernels this way. */
7#ifndef MESH_INPROC_H
8#define MESH_INPROC_H
9
10#include "mesh/backend.h"
11#include "mesh/manifest.h"
12#include "mesh/shm.h"
13
14#include "mesh/abi/words.h"
15
16#include <stdbool.h>
17#include <stddef.h>
18#include <stdint.h>
19
20double mesh_ms_since(uint64_t t0_ns);
21uint64_t mesh_inproc_now_ns(void);
22
23/* meshd's backend, as kcheck drives it. */
24typedef struct {
25    mesh_backend *b; /* owned */
26    const mesh_compiler *comp;
27    uint32_t kind; /* BACKEND_* */
28    void *ctx;     /* CUDA: the context (the module-load probe makes it current elsewhere); NULL elsewhere */
29    mesh_shm shm;  /* the mailbox the backend wraps */
30} mesh_dev;
31
32/* CUDA in blink mode (every program resident at once) over a shared-memory
33 * mailbox of its own, nvJitLink from MESH_NVJITLINK_DIR. */
34bool mesh_dev_open_cuda(mesh_dev *out, char **err);
35/* Any other backend (wasm: @ajax/mesh-exec-wasm's), already opened over shm; takes b. */
36void mesh_dev_from(mesh_backend *b, const mesh_shm *shm, mesh_dev *out);
37void mesh_dev_close(mesh_dev *d);
38static inline bool mesh_dev_cuda(const mesh_dev *d) { return d->kind == MESH_BACKEND_CUDA; }
39static inline bool mesh_dev_wasm(const mesh_dev *d) { return d->kind == MESH_BACKEND_WASM; }
40/* What the backend compiles an image to: "cubin", "MSL", "wasm". */
41const char *mesh_dev_code(const mesh_dev *d);
42uint32_t mesh_dev_sms_total(const mesh_dev *d);
43
44/* A mailbox (the backend's: pinned and device-mapped; zeroed) with word access. */
45typedef struct {
46    uint8_t *base;
47    uint64_t dev;
48} mesh_mb2;
49
50/* The backend's mailbox, with the header written (magic, version, peer none, SMs, backend). */
51bool mesh_mb2_alloc(const mesh_dev *d, mesh_mb2 *out, char **err);
52uint32_t mesh_mb2_r32(mesh_mb2 m, size_t off);     /* acquire */
53void mesh_mb2_w32(mesh_mb2 m, size_t off, uint32_t v); /* relaxed */
54void mesh_mb2_store32(mesh_mb2 m, size_t off, uint32_t v); /* release */
55uint32_t mesh_mb2_add32(mesh_mb2 m, size_t off, uint32_t v); /* AcqRel, the old value */
56uint64_t mesh_mb2_r64(mesh_mb2 m, size_t off);
57void mesh_mb2_w64(mesh_mb2 m, size_t off, uint64_t v);
58void mesh_mb2_store64(mesh_mb2 m, size_t off, uint64_t v);
59float mesh_mb2_rf(mesh_mb2 m, size_t off);
60void mesh_mb2_wf(mesh_mb2 m, size_t off, float v);
61void mesh_mb2_bytes(mesh_mb2 m, size_t off, uint8_t *out, size_t n); /* volatile byte reads */
62void mesh_mb2_write(mesh_mb2 m, size_t off, const uint8_t *b, size_t n);
63void mesh_mb2_zero(mesh_mb2 m, size_t off, size_t n);
64/* Pipes (mesh-abi pipe) from this side of the mailbox. */
65mesh_words mesh_mb2_words(mesh_mb2 m);
66
67/* A program image, compiled and loaded (not yet in a slot). */
68typedef struct {
69    mesh_manifest manifest;
70    mesh_kernel *kernel;
71    uint64_t hash;
72    uint32_t cubin_bytes, regs, rt_abi;
73    double jit_ms, load_ms;
74} mesh_image2;
75
76/* The backend's compiler on a catalog dir's bytes, then its load and the
77 * loader's checks (abi, MESH_RT_ABI).  CUDA: call it before launching anything. */
78bool mesh_inproc_load(const mesh_dev *d, const char *dir, mesh_image2 *out, char **err);
79void mesh_image2_free(const mesh_dev *d, mesh_image2 *im);
80
81/* A program in a slot of a mailbox, with the host side of its ring. */
82typedef struct {
83    mesh_mb2 mb;
84    uint32_t slot;
85    mesh_stream *stream;
86    mesh_image2 image;
87    uint32_t nodes;
88    uint32_t tail; /* next ticket to publish */
89} mesh_prog;
90
91/* Writes the slot's table entry and manifest extension under the seqlock
92 * (state LOADED), zeroes its control block, node records, trace and view.  Takes image. */
93bool mesh_prog_install(const mesh_dev *d, mesh_mb2 mb, uint32_t slot, mesh_image2 *image, uint32_t view_off, uint32_t reads_mask,
94                       mesh_prog *out, char **err);
95void mesh_prog_free(const mesh_dev *d, mesh_prog *p);
96size_t mesh_prog_ctl_off(const mesh_prog *p, size_t off);
97size_t mesh_prog_view(const mesh_prog *p);
98uint32_t mesh_prog_head(const mesh_prog *p);
99uint32_t mesh_prog_counter(const mesh_prog *p, uint32_t i);
100/* Rewinds the head to the tail, clears the shutdown, launches, waits (2 s) for every node online: *ms the time. */
101bool mesh_prog_launch(const mesh_dev *d, mesh_prog *p, double *ms, char **err);
102/* P_SHUTDOWN, then waits (3 s) until every node halted and the stream is done: *ms the time. */
103bool mesh_prog_stop(const mesh_dev *d, mesh_prog *p, double *ms, char **err);
104/* Publishes one task (waiting for the ticket that last used its ring entry): *ticket. */
105bool mesh_prog_publish(mesh_prog *p, uint32_t op, uint32_t arg, uint32_t *ticket, char **err);
106bool mesh_prog_done(const mesh_prog *p, uint32_t t);
107/* Ticket t's result (5 s). */
108bool mesh_prog_wait(const mesh_prog *p, uint32_t t, uint32_t *result, char **err);
109bool mesh_prog_run(mesh_prog *p, uint32_t op, uint32_t arg, uint32_t *result, char **err);
110/* One publish of op with args 0..n (consecutive tickets), then every result into results[n]. */
111bool mesh_prog_batch(mesh_prog *p, uint32_t op, uint32_t n, uint32_t *results, char **err);
112/* Appends bytes to the input queue (manifest `input`): data, then IQ_HEAD (release). */
113bool mesh_prog_input_push(const mesh_prog *p, const uint8_t *b, size_t n, char **err);
114/* One request through the window: the response bytes (*resp malloc'd, *rn). */
115bool mesh_prog_request(mesh_prog *p, const uint8_t *req, size_t n, uint8_t **resp, size_t *rn, char **err);
116
117/* First-fit-free sequential view placement: whole pages from VIEW_HEAP. */
118typedef struct {
119    uint32_t next;
120} mesh_heap;
121mesh_heap mesh_heap_new(void);
122bool mesh_heap_alloc(mesh_heap *h, uint32_t bytes, uint32_t *at, char **err);
123
124/* The manifest's `reads` as a slot mask over (name, slot) pairs; a name not loaded is an error. */
125typedef struct {
126    const char *name;
127    uint32_t slot;
128} mesh_slot_name;
129bool mesh_reads_mask(const mesh_manifest *m, const mesh_slot_name *slots, size_t n, uint32_t *mask, char **err);
130
131#endif