1/* @ajax/mesh-rt: the device runtime every mesh program is built on (contract v2). A program writes
2 *
3 *     #include "mesh/rt.h"
4 *     static void init(const mesh_ctx *ctx) { .. }                          // op 1 (OP_INIT)
5 *     static uint32_t task(const mesh_ctx *ctx, uint32_t op, uint32_t arg) // op 0 and 2..
6 *     { .. }
7 *     MESH_PROGRAM(init, task);
8 *
9 * which expands to the kernel `mesh_main(mb, slot)` -- the node loop over the
10 * slot's control block and ring (@ajax/mesh-rt-loop) -- and the global
11 * `MESH_RT_ABI` meshd checks at load. Every task runs on one node, a block of
12 * 32 x `warps` threads (the manifest's `warps`, default 1): every thread calls
13 * init/task; the task's result is thread 0's return value. The node has the
14 * runtime's shared memory and the program's `shared_bytes`
15 * (mesh_ctx_shared) and the block collectives (@ajax/mesh-rt-block). Every
16 * runtime function is inlined into the kernel, so one PTX per program.
17 *
18 * Compiled for nvptx64 (clang --target=nvptx64-nvidia-cuda) it is the GPU
19 * program; compiled for the host, the same source runs on @ajax/mesh-rt-sim's
20 * threads (mesh_sim_start(mesh_main, ...)) for tests.
21 * Ported from alpha-rust-mesh crates/mesh-rt/src/lib.rs. */
22#ifndef MESH_RT_H
23#define MESH_RT_H
24
25#include <stdbool.h>
26#include <stdint.h>
27
28#include "mesh/abi.h"
29#include "mesh/rt/block.h"
30#include "mesh/rt/ctx.h"
31#include "mesh/rt/loop.h"
32#include "mesh/rt/prim.h"
33#include "mesh/rt/view.h"
34
35/* The program's kernel: `nodes` blocks of 32 x `warps` threads (PT_NODES, PT_WARPS). */
36#if defined(__NVPTX__)
37#define MESH_PROGRAM(INIT, TASK)                                                                                    \
38    __attribute__((nvptx_kernel)) void mesh_main(uint8_t *mb, uint32_t slot) { mesh_rt_run(mb, slot, INIT, TASK); } \
39    /* The runtime contract this image was built against (meshd refuses a mismatch). */                           \
40    __attribute__((used)) const uint32_t mesh_rt_abi_symbol_ __asm__("MESH_RT_ABI") = MESH_RT_ABI
41#else
42/* On the host the kernel is an ordinary function for @ajax/mesh-rt-sim's
43 * launcher, named mesh_main unless MESH_SIM_KERNEL says otherwise (a test that
44 * runs two programs includes each with its own name). */
45#ifndef MESH_SIM_KERNEL
46#define MESH_SIM_KERNEL mesh_main
47#endif
48#define MESH_PROGRAM(INIT, TASK)                                                                   \
49    void MESH_SIM_KERNEL(uint8_t *mb, uint32_t slot);                                              \
50    void MESH_SIM_KERNEL(uint8_t *mb, uint32_t slot) { mesh_rt_run(mb, slot, INIT, TASK); }        \
51    _Static_assert(MESH_RT_ABI > 0, "the runtime contract")
52#endif
53
54#endif