1/* @ajax/mesh-cuda: the CUDA driver calls meshd (and kcheck, and the MPS
2 * workers) need (mesh-host cuda.rs): a primary or green context, streams (one
3 * per program), the pinned mailbox, and loading (cubins only: nvJitLink
4 * compiles, never the driver) / launching / unloading resident kernels.
5 *
6 * Module loading is EAGER, always: under LAZY (the driver default) a kernel's
7 * first launch is itself a module load, which blocks behind every resident
8 * kernel of the context. mesh_cuda_force_eager runs before cuInit. */
9#ifndef MESH_CUDA_H
10#define MESH_CUDA_H
1112#include "mesh/cuda_driver.h"
1314#include <stdbool.h>
15#include <stddef.h>
16#include <stdint.h>
1718#define MESH_CUDA_MODULE_LOADING "CUDA_MODULE_LOADING"
19/* The per-machine CUDA driver library dir (a machine whose kernel module and
20 * userland libcuda differ keeps a matching private libcuda.so.1 there). */
21#define MESH_CUDA_LIB_ENV "MESH_CUDA_LIB"
2223/* Sets CUDA_MODULE_LOADING=EAGER, or refuses a different setting. Call it
24 * first in main, before any thread exists. */
25bool mesh_cuda_force_eager(char **err);
26/* `--cuda-lib DIR` from argv, else $MESH_CUDA_LIB; NULL: the system libcuda (malloc'd). */
27char *mesh_cuda_lib_dir(int argc, char **argv);
28/* Makes this process (and every child) load libcuda.so.1 from dir: when dir is
29 * not LD_LIBRARY_PATH's first entry yet, re-executes this program with dir
30 * prepended (and does not return). Call first in main. dir NULL: nothing. */
31bool mesh_cuda_use_lib(const char *dir, int argc, char **argv, char **err);
32/* After cuInit: the libcuda this process mapped (malloc'd); with
33 * MESH_CUDA_LIB set it must be the one in that dir. */
34char *mesh_cuda_check_lib(char **err);
3536/* The device's per-block limits a multi-warp node must fit. */
37typedef struct {
38 uint32_t max_threads; /* MAX_THREADS_PER_BLOCK (1024) */
39 uint32_t regs_per_block; /* MAX_REGISTERS_PER_BLOCK (65536) */
40 uint32_t max_shared; /* MAX_SHARED_MEMORY_PER_BLOCK_OPTIN (sm_86 and GB10: 99 KiB) */
41} mesh_cuda_limits;
4243bool mesh_cuda_limits_of(CUdevice dev, mesh_cuda_limits *out, char **err);
44/* The register cap nvJitLink must compile a node of `warps` warps with so
45 * that one block fits the register file: 0 when 255 per thread fit already. */
46uint32_t mesh_cuda_reg_cap(const mesh_cuda_limits *l, uint32_t warps);
47/* A node of warps warps and smem dynamic shared bytes fits a block here. */
48bool mesh_cuda_limits_check(const mesh_cuda_limits *l, const char *name, uint32_t warps, uint32_t smem, char **err);
4950typedef struct {
51 CUcontext ctx;
52 CUgreenCtx green_ctx; /* the green context behind ctx, if any */
53 char *name;
54 int32_t cc_major, cc_minor;
55 uint32_t sms_total;
56 uint32_t sms; /* SMs the mesh may use: the green context's, else all */
57 bool green;
58 uint32_t blocks_per_sm; /* the resident-block budget per SM every program's nodes share */
59 bool host_atomics; /* HOST_NATIVE_ATOMIC_SUPPORTED (0 on an RTX 3060 over PCIe) */
60 mesh_cuda_limits limits;
61} mesh_cuda_gpu;
6263/* A loaded kernel module and its entry, with the launch shape of its nodes. */
64typedef struct {
65 CUmodule module;
66 CUfunction func;
67 uint32_t threads; /* per block (32 x warps) */
68 uint32_t smem; /* dynamic shared bytes per block (RT_SHARED_BYTES + shared_bytes) */
69 uint64_t scratch; /* the program's scratch, zeroed at load; 0: none */
70} mesh_cuda_kernel;
7172/* A non-blocking stream of the context (one per program). */
73typedef struct {
74 CUstream s;
75} mesh_cuda_stream;
7677/* Device 0; sms > 0: a green context of that many SMs. Needs
78 * CUDA_MODULE_LOADING=EAGER (mesh_cuda_force_eager first). */
79bool mesh_cuda_gpu_open(uint32_t sms, mesh_cuda_gpu *out, char **err);
80void mesh_cuda_gpu_free(mesh_cuda_gpu *g);
81bool mesh_cuda_make_current(const mesh_cuda_gpu *g, char **err);
82/* The per-thread stack the context reserves (raise it only while nothing runs: in a blink). */
83bool mesh_cuda_stack_limit(const mesh_cuda_gpu *g, size_t *out, char **err);
84bool mesh_cuda_set_stack_limit(const mesh_cuda_gpu *g, size_t bytes, char **err);
85/* Host memory (the mailbox's shared file) pinned and mapped into the device: its device address. */
86bool mesh_cuda_host_register(const mesh_cuda_gpu *g, void *p, size_t bytes, uint64_t *dev, char **err);
87/* A cubin and its entry (quick; but like every module operation it blocks while any kernel runs). */
88bool mesh_cuda_load_cubin(const mesh_cuda_gpu *g, const uint8_t *cubin, const char *entry, mesh_cuda_kernel *out, char **err);
89bool mesh_cuda_new_stream(const mesh_cuda_gpu *g, mesh_cuda_stream *out, char **err);
90void mesh_cuda_stream_destroy(mesh_cuda_stream *s);
91/* nodes blocks of k's shape on st; params point at the kernel's parameters. Returns at once. */
92bool mesh_cuda_launch_on(const mesh_cuda_gpu *g, const mesh_cuda_stream *st, const mesh_cuda_kernel *k, uint32_t nodes,
93 void **params, char **err);
94/* *done = true once everything on st finished. */
95bool mesh_cuda_query(const mesh_cuda_gpu *g, const mesh_cuda_stream *st, bool *done, char **err);
96/* The module and the scratch (both block while a kernel of the context runs: call it halted). */
97bool mesh_cuda_unload(const mesh_cuda_gpu *g, mesh_cuda_kernel *k, char **err);
9899bool mesh_cuda_kernel_regs(const mesh_cuda_kernel *k, uint32_t *out, char **err);
100bool mesh_cuda_kernel_local_bytes(const mesh_cuda_kernel *k, uint32_t *out, char **err);
101/* Sets the launch shape, opting the entry in above 48 KiB of shared memory;
102 * refuses a shape the compiled entry cannot launch with. */
103bool mesh_cuda_kernel_set_shape(mesh_cuda_kernel *k, uint32_t threads, uint32_t smem, char **err);
104/* bytes of device memory, zeroed before this returns (nothing when 0). */
105bool mesh_cuda_kernel_alloc_scratch(mesh_cuda_kernel *k, uint32_t bytes, char **err);
106bool mesh_cuda_kernel_max_blocks_per_sm(const mesh_cuda_kernel *k, uint32_t *out, char **err);
107/* A u32 module global (MESH_RT_ABI): a synchronous copy, call it while no kernel runs. */
108bool mesh_cuda_kernel_global_u32(const mesh_cuda_kernel *k, const char *name, uint32_t *out, char **err);
109110#endif