==> src/cuda.c <==1/* cuda.rs over the @ajax/cuda-driver table. */
2#include "mesh/cuda.h"
3#include "mesh/text.h"
4
5#include "mesh/abi/base.h"
6
7#include <errno.h>
8#include <limits.h>
9#include <stdio.h>
10#include <stdlib.h>
11#include <string.h>
12#include <unistd.h>
13
14#define CU (mesh_cuda())
15#define CK(call, what) \
16 do { \
17 if (!mesh_cu_ck((call), (what), err)) \
18 return false; \
19 } while (0)
20
21static char *os_error(int e) { return mesh_fmt("%s (os error %d)", strerror(e), e); }
22
23bool mesh_cuda_force_eager(char **err) {
24 const char *v = getenv(MESH_CUDA_MODULE_LOADING);
25 if (!v) return setenv(MESH_CUDA_MODULE_LOADING, "EAGER", 1) == 0;
26 if (strcmp(v, "EAGER") == 0) return true;
27 char *q = mesh_debug_str(v);
28 mesh_errf(err, "%s=%s: refusing to run; module loading must be EAGER (a lazy load blocks behind resident kernels)",
29 MESH_CUDA_MODULE_LOADING, q);
30 free(q);
31 return false;
32}
33
34char *mesh_cuda_lib_dir(int argc, char **argv) {
35 for (int i = 0; i < argc; i++)
36 if (strcmp(argv[i], "--cuda-lib") == 0) return i + 1 < argc ? mesh_strdup(argv[i + 1]) : NULL;
37 const char *e = getenv(MESH_CUDA_LIB_ENV);
38 return e && *e ? mesh_strdup(e) : NULL;
39}
40
41bool mesh_cuda_use_lib(const char *dir, int argc, char **argv, char **err) {
42 if (!dir) return true;
43 char canon[PATH_MAX];
44 if (!realpath(dir, canon)) {
45 char *e = os_error(errno);
46 mesh_errf(err, "%s / --cuda-lib %s: %s", MESH_CUDA_LIB_ENV, dir, e);
47 free(e);
48 return false;
49 }
50 char *lib = mesh_fmt("%s/libcuda.so.1", canon);
51 bool exists = access(lib, F_OK) == 0;
52 free(lib);
53 if (!exists) return mesh_errf(err, "%s / --cuda-lib %s: no libcuda.so.1 there", MESH_CUDA_LIB_ENV, dir);
54 const char *ld = getenv("LD_LIBRARY_PATH");
55 if (!ld) ld = "";
56 size_t first = strcspn(ld, ":");
57 if (first == strlen(canon) && strncmp(ld, canon, first) == 0) return setenv(MESH_CUDA_LIB_ENV, canon, 1) == 0;
58 if (getenv("MESH_CUDA_LIB_EXEC"))
59 return mesh_errf(err, "%s: re-executed, but LD_LIBRARY_PATH still does not start with %s", MESH_CUDA_LIB_ENV, canon);
60 char *path = *ld ? mesh_fmt("%s:%s", canon, ld) : mesh_strdup(canon);
61 setenv("LD_LIBRARY_PATH", path, 1);
62 free(path);
63 setenv(MESH_CUDA_LIB_ENV, canon, 1);
64 setenv("MESH_CUDA_LIB_EXEC", "1", 1);
65 /* The resolved path, not /proc/self/exe: the process keeps its name. */
66 char exe[PATH_MAX];
67 ssize_t n = readlink("/proc/self/exe", exe, sizeof exe - 1);
68 if (n < 0) {
69 char *e = os_error(errno);
70 mesh_errf(err, "%s: current_exe: %s", MESH_CUDA_LIB_ENV, e);
71 free(e);
72 return false;
73 }
74 exe[n] = 0;
75 char **args = mesh_alloc(((size_t)argc + 1) * sizeof *args);
76 for (int i = 0; i < argc; i++) args[i] = argv[i];
77 execv(exe, args);
78 char *e = os_error(errno);
79 mesh_errf(err, "%s: re-exec with LD_LIBRARY_PATH=%s: %s", MESH_CUDA_LIB_ENV, canon, e);
80 free(e);
81 free(args);
82 return false;
83}
84
85char *mesh_cuda_check_lib(char **err) {
86 FILE *f = fopen("/proc/self/maps", "r");
87 if (!f) {
88 char *e = os_error(errno);
89 mesh_errf(err, "/proc/self/maps: %s", e);
90 free(e);
91 return NULL;
92 }
93 char line[4096];
94 char *lib = NULL;
95 while (!lib && fgets(line, sizeof line, f)) {
96 /* the 6th whitespace-separated field is the path */
97 char *save = NULL, *tok = strtok_r(line, " \t\n", &save);
98 for (int i = 0; tok && i < 5; i++) tok = strtok_r(NULL, " \t\n", &save);
99 if (!tok) continue;
100 const char *slash = strrchr(tok, '/');
101 if (strncmp(slash ? slash + 1 : tok, "libcuda.so", 10) == 0) lib = mesh_strdup(tok);
102 }
103 fclose(f);
104 if (!lib) {
105 mesh_errf(err, "no libcuda mapped after cuInit");
106 return NULL;
107 }
108 const char *dir = getenv(MESH_CUDA_LIB_ENV);
109 if (dir && *dir) {
110 char canon[PATH_MAX];
111 const char *c = realpath(dir, canon) ? canon : dir;
112 const char *slash = strrchr(lib, '/');
113 size_t pn = slash ? (size_t)(slash - lib) : 0;
114 if (!(slash && strlen(c) == pn && strncmp(lib, c, pn) == 0)) {
115 mesh_errf(err, "%s=%s, but the driver library mapped is %s", MESH_CUDA_LIB_ENV, dir, lib);
116 free(lib);
117 return NULL;
118 }
119 }
120 return lib;
121}
122
123static bool check_eager(char **err) {
124 int mode = CU_MODULE_LAZY_LOADING;
125 CK(CU->cuModuleGetLoadingMode(&mode), "cuModuleGetLoadingMode");
126 if (mode != CU_MODULE_EAGER_LOADING) {
127 const char *name = mode == CU_MODULE_LAZY_LOADING ? "CU_MODULE_LAZY_LOADING" : "CUmoduleLoadingMode_enum(?)";
128 return mesh_errf(err, "the driver loads modules %s; EAGER is required (%s must be EAGER before cuInit)", name,
129 MESH_CUDA_MODULE_LOADING);
130 }
131 return true;
132}
133
134static bool attr(CUdevice dev, int a, int *v, char **err) {
135 CK(CU->cuDeviceGetAttribute(v, a, dev), "cuDeviceGetAttribute");
136 return true;
137}
138
139bool mesh_cuda_limits_of(CUdevice dev, mesh_cuda_limits *out, char **err) {
140 int t, r, s;
141 if (!attr(dev, CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK, &t, err) ||
142 !attr(dev, CU_DEVICE_ATTRIBUTE_MAX_REGISTERS_PER_BLOCK, &r, err) ||
143 !attr(dev, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, &s, err))
144 return false;
145 *out = (mesh_cuda_limits){(uint32_t)t, (uint32_t)r, (uint32_t)s};
146 return true;
147}
148
149uint32_t mesh_cuda_reg_cap(const mesh_cuda_limits *l, uint32_t warps) {
150 uint32_t threads = MESH_NODE_THREADS * (warps ? warps : 1);
151 uint32_t cap = (l->regs_per_block / threads) / 8 * 8;
152 return cap < 255 ? (cap > 16 ? cap : 16) : 0;
153}
154
155bool mesh_cuda_limits_check(const mesh_cuda_limits *l, const char *name, uint32_t warps, uint32_t smem, char **err) {
156 uint32_t threads = MESH_NODE_THREADS * warps;
157 if (threads > l->max_threads)
158 return mesh_errf(err, "occupancy: %s: warps %u = %u threads per node; the device allows %u per block", name, (unsigned)warps,
159 (unsigned)threads, (unsigned)l->max_threads);
160 if (smem > l->max_shared)
161 return mesh_errf(err,
162 "occupancy: %s: %u bytes of shared memory per node (the runtime's %u + shared_bytes); the device allows %u per block",
163 name, (unsigned)smem, (unsigned)MESH_RT_SHARED_BYTES, (unsigned)l->max_shared);
164 return true;
165}
166
167bool mesh_cuda_gpu_open(uint32_t sms, mesh_cuda_gpu *out, char **err) {
168 *out = (mesh_cuda_gpu){0};
169 const char *ml = getenv(MESH_CUDA_MODULE_LOADING);
170 if (!ml || strcmp(ml, "EAGER") != 0)
171 return mesh_errf(err, "%s is not EAGER: call cuda::force_eager() first", MESH_CUDA_MODULE_LOADING);
172 if (!CU && !mesh_cuda_load(err)) return false;
173 CK(CU->cuInit(0), "cuInit");
174 if (!check_eager(err)) return false;
175 if (!mesh_cuda_injected()) {
176 char *lib = mesh_cuda_check_lib(err);
177 if (!lib) return false;
178 free(lib);
179 }
180 CUdevice dev = 0;
181 CK(CU->cuDeviceGet(&dev, 0), "cuDeviceGet");
182 char name[128] = {0};
183 CK(CU->cuDeviceGetName(name, 128, dev), "cuDeviceGetName");
184 int major, minor, smt, bps, ha;
185 if (!attr(dev, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, &major, err) ||
186 !attr(dev, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, &minor, err) ||
187 !attr(dev, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, &smt, err) ||
188 !attr(dev, CU_DEVICE_ATTRIBUTE_MAX_BLOCKS_PER_MULTIPROCESSOR, &bps, err) ||
189 !attr(dev, CU_DEVICE_ATTRIBUTE_HOST_NATIVE_ATOMIC_SUPPORTED, &ha, err))
190 return false;
191 mesh_cuda_limits limits;
192 if (!mesh_cuda_limits_of(dev, &limits, err)) return false;
193 CUcontext ctx = NULL;
194 CUgreenCtx green_ctx = NULL;
195 uint32_t used = (uint32_t)smt;
196 if (sms == 0) {
197 CK(CU->cuDevicePrimaryCtxRetain(&ctx, dev), "cuDevicePrimaryCtxRetain");
198 CK(CU->cuCtxSetCurrent(ctx), "cuCtxSetCurrent");
199 } else {
200 CUdevResource all, group, rest;
201 memset(&all, 0, sizeof all);
202 memset(&group, 0, sizeof group);
203 memset(&rest, 0, sizeof rest);
204 CK(CU->cuDeviceGetDevResource(dev, &all, CU_DEV_RESOURCE_TYPE_SM), "cuDeviceGetDevResource");
205 unsigned int n = 1;
206 /* IGNORE_SM_COSCHEDULING: 2-SM granularity. */
207 CK(CU->cuDevSmResourceSplitByCount(&group, &n, &all, &rest, CU_DEV_SM_RESOURCE_SPLIT_IGNORE_SM_COSCHEDULING, sms),
208 "cuDevSmResourceSplitByCount");
209 if (n != 1) return mesh_errf(err, "green context: cannot split %u SMs off %d", (unsigned)sms, smt);
210 used = group.u.sm.smCount;
211 CUdevResourceDesc desc = NULL;
212 CK(CU->cuDevResourceGenerateDesc(&desc, &group, 1), "cuDevResourceGenerateDesc");
213 CK(CU->cuGreenCtxCreate(&green_ctx, desc, dev, CU_GREEN_CTX_DEFAULT_STREAM), "cuGreenCtxCreate");
214 CK(CU->cuCtxFromGreenCtx(&ctx, green_ctx), "cuCtxFromGreenCtx");
215 CK(CU->cuCtxSetCurrent(ctx), "cuCtxSetCurrent");
216 }
217 *out = (mesh_cuda_gpu){ctx, green_ctx, mesh_strdup(name), major, minor, (uint32_t)smt, used, sms != 0, (uint32_t)bps, ha != 0, limits};
218 return true;
219}
220
221void mesh_cuda_gpu_free(mesh_cuda_gpu *g) {
222 free(g->name);
223 g->name = NULL;
224}
225
226bool mesh_cuda_make_current(const mesh_cuda_gpu *g, char **err) {
227 CK(CU->cuCtxSetCurrent(g->ctx), "cuCtxSetCurrent");
228 return true;
229}
230
231bool mesh_cuda_stack_limit(const mesh_cuda_gpu *g, size_t *out, char **err) {
232 (void)g;
233 *out = 0;
234 CK(CU->cuCtxGetLimit(out, CU_LIMIT_STACK_SIZE), "cuCtxGetLimit(STACK_SIZE)");
235 return true;
236}
237
238bool mesh_cuda_set_stack_limit(const mesh_cuda_gpu *g, size_t bytes, char **err) {
239 (void)g;
240 CK(CU->cuCtxSetLimit(CU_LIMIT_STACK_SIZE, bytes), "cuCtxSetLimit(STACK_SIZE)");
241 return true;
242}
243
244bool mesh_cuda_host_register(const mesh_cuda_gpu *g, void *p, size_t bytes, uint64_t *dev, char **err) {
245 (void)g;
246 CK(CU->cuMemHostRegister_v2(p, bytes, CU_MEMHOSTREGISTER_PORTABLE | CU_MEMHOSTREGISTER_DEVICEMAP), "cuMemHostRegister");
247 CUdeviceptr d = 0;
248 CK(CU->cuMemHostGetDevicePointer_v2(&d, p, 0), "cuMemHostGetDevicePointer");
249 *dev = d;
250 return true;
251}
252
253bool mesh_cuda_load_cubin(const mesh_cuda_gpu *g, const uint8_t *cubin, const char *entry, mesh_cuda_kernel *out, char **err) {
254 (void)g;
255 CUmodule module = NULL;
256 CK(CU->cuModuleLoadData(&module, cubin), "cuModuleLoadData(cubin)");
257 CUfunction func = NULL;
258 char *what = mesh_fmt("cuModuleGetFunction(%s)", entry);
259 bool ok = mesh_cu_ck(CU->cuModuleGetFunction(&func, module, entry), what, err);
260 free(what);
261 if (!ok) {
262 CU->cuModuleUnload(module);
263 return false;
264 }
265 *out = (mesh_cuda_kernel){module, func, MESH_NODE_THREADS, MESH_RT_SHARED_BYTES, 0};
266 return true;
267}
268
269static void stream_drop(mesh_cuda_stream *s) {
270 if (s->s) CU->cuStreamDestroy_v2(s->s);
271 s->s = NULL;
272}
273
274bool mesh_cuda_new_stream(const mesh_cuda_gpu *g, mesh_cuda_stream *out, char **err) {
275 CUstream s = NULL;
276 if (g->green_ctx) CK(CU->cuGreenCtxStreamCreate(&s, g->green_ctx, CU_STREAM_NON_BLOCKING, 0), "cuGreenCtxStreamCreate");
277 else CK(CU->cuStreamCreate(&s, CU_STREAM_NON_BLOCKING), "cuStreamCreate");
278 out->s = s;
279 return true;
280}
281
282void mesh_cuda_stream_destroy(mesh_cuda_stream *s) { stream_drop(s); }
283
284bool mesh_cuda_launch_on(const mesh_cuda_gpu *g, const mesh_cuda_stream *st, const mesh_cuda_kernel *k, uint32_t nodes,
285 void **params, char **err) {
286 (void)g;
287 CK(CU->cuLaunchKernel(k->func, nodes, 1, 1, k->threads, 1, 1, k->smem, st->s, params, NULL), "cuLaunchKernel");
288 return true;
289}
290
291bool mesh_cuda_query(const mesh_cuda_gpu *g, const mesh_cuda_stream *st, bool *done, char **err) {
292 (void)g;
293 CUresult r = CU->cuStreamQuery(st->s);
294 if (r == CUDA_SUCCESS) {
295 *done = true;
296 return true;
297 }
298 if (r == CUDA_ERROR_NOT_READY) {
299 *done = false;
300 return true;
301 }
302 return mesh_cu_ck(r, "program stream", err);
303}
304
305bool mesh_cuda_unload(const mesh_cuda_gpu *g, mesh_cuda_kernel *k, char **err) {
306 (void)g;
307 if (k->scratch) {
308 CK(CU->cuMemFree_v2(k->scratch), "cuMemFree(scratch)");
309 k->scratch = 0;
310 }
311 CK(CU->cuModuleUnload(k->module), "cuModuleUnload");
312 return true;
313}
314
315static bool func_attr(const mesh_cuda_kernel *k, int a, int *v, char **err) {
316 CK(CU->cuFuncGetAttribute(v, a, k->func), "cuFuncGetAttribute");
317 return true;
318}
319
320bool mesh_cuda_kernel_regs(const mesh_cuda_kernel *k, uint32_t *out, char **err) {
321 int v = 0;
322 if (!func_attr(k, CU_FUNC_ATTRIBUTE_NUM_REGS, &v, err)) return false;
323 *out = (uint32_t)v;
324 return true;
325}
326
327bool mesh_cuda_kernel_local_bytes(const mesh_cuda_kernel *k, uint32_t *out, char **err) {
328 int v = 0;
329 if (!func_attr(k, CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, &v, err)) return false;
330 *out = (uint32_t)v;
331 return true;
332}
333
334bool mesh_cuda_kernel_set_shape(mesh_cuda_kernel *k, uint32_t threads, uint32_t smem, char **err) {
335 int max = 0;
336 if (!func_attr(k, CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, &max, err)) return false;
337 if (threads > (uint32_t)max) {
338 uint32_t regs = 0;
339 if (!mesh_cuda_kernel_regs(k, ®s, err)) return false;
340 return mesh_errf(err, "regs: %u registers per thread allow %d threads per block; the node has %u", (unsigned)regs, max,
341 (unsigned)threads);
342 }
343 int stat = 0;
344 if (!func_attr(k, CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, &stat, err)) return false;
345 if (smem + (uint32_t)stat > 48u << 10)
346 CK(CU->cuFuncSetAttribute(k->func, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, (int)smem),
347 "cuFuncSetAttribute(MAX_DYNAMIC_SHARED_SIZE_BYTES)");
348 k->threads = threads;
349 k->smem = smem;
350 return true;
351}
352
353bool mesh_cuda_kernel_alloc_scratch(mesh_cuda_kernel *k, uint32_t bytes, char **err) {
354 if (bytes == 0 || k->scratch != 0) return true;
355 CUdeviceptr d = 0;
356 char *what = mesh_fmt("cuMemAlloc(scratch %u bytes)", (unsigned)bytes);
357 bool ok = mesh_cu_ck(CU->cuMemAlloc_v2(&d, bytes), what, err);
358 free(what);
359 if (!ok) return false;
360 /* On a stream of its own, waited for: the program streams are non-blocking. */
361 CUstream st = NULL;
362 ok = mesh_cu_ck(CU->cuStreamCreate(&st, CU_STREAM_NON_BLOCKING), "cuStreamCreate", err);
363 if (ok) {
364 ok = mesh_cu_ck(CU->cuMemsetD8Async(d, 0, bytes, st), "cuMemsetD8Async(scratch)", err) &&
365 mesh_cu_ck(CU->cuStreamSynchronize(st), "cuStreamSynchronize", err);
366 CU->cuStreamDestroy_v2(st);
367 }
368 if (!ok) {
369 CU->cuMemFree_v2(d);
370 return false;
371 }
372 k->scratch = d;
373 return true;
374}
375
376bool mesh_cuda_kernel_max_blocks_per_sm(const mesh_cuda_kernel *k, uint32_t *out, char **err) {
377 int v = 0;
378 CK(CU->cuOccupancyMaxActiveBlocksPerMultiprocessor(&v, k->func, (int)k->threads, k->smem),
379 "cuOccupancyMaxActiveBlocksPerMultiprocessor");
380 *out = (uint32_t)v;
381 return true;
382}
383
384bool mesh_cuda_kernel_global_u32(const mesh_cuda_kernel *k, const char *name, uint32_t *out, char **err) {
385 CUdeviceptr d = 0;
386 size_t bytes = 0;
387 char *what = mesh_fmt("cuModuleGetGlobal(%s)", name);
388 bool ok = mesh_cu_ck(CU->cuModuleGetGlobal_v2(&d, &bytes, k->module, name), what, err);
389 free(what);
390 if (!ok) return false;
391 if (bytes != 4) return mesh_errf(err, "global %s is %zu bytes, not 4", name, bytes);
392 uint32_t v = 0;
393 CK(CU->cuMemcpyDtoH_v2(&v, d, 4), "cuMemcpyDtoH");
394 *out = v;
395 return true;
396}