1/* nvjitlink.rs: open, version check, one handle per link (destroyed on every path). */
2#include "mesh/nvjitlink.h"
3#include "mesh/text.h"
4
5#include <dlfcn.h>
6#include <stdlib.h>
7#include <string.h>
8
9#define INPUT_PTX 2 /* NVJITLINK_INPUT_PTX */
10
11static void *sym(void *lib, const char *path, const char *name, char **err) {
12    void *p = dlsym(lib, name);
13    if (!p) mesh_errf(err, "nvJitLink: %s has no %s", path, name);
14    return p;
15}
16
17bool mesh_nvjitlink_open(const char *dir, mesh_nvjitlink *out, char **err) {
18    *out = (mesh_nvjitlink){0};
19    size_t dn = strlen(dir);
20    while (dn > 0 && dir[dn - 1] == '/') dn--;
21    char *path = mesh_fmt("%.*s/%s", (int)dn, dir, MESH_NVJITLINK_LIB);
22    void *lib = dlopen(path, RTLD_NOW | RTLD_LOCAL);
23    if (!lib) {
24        const char *e = dlerror();
25        mesh_errf(err, "nvJitLink: cannot open %s: %s (set --nvjitlink DIR; there is no fallback)", path, e ? e : "");
26        free(path);
27        return false;
28    }
29    void *ver = sym(lib, path, "nvJitLinkVersion", err);
30    if (!ver) {
31        free(path);
32        return false;
33    }
34    int (*version)(uint32_t *, uint32_t *);
35    memcpy(&version, &ver, sizeof ver);
36    uint32_t major = 0, minor = 0;
37    if (version(&major, &minor) != 0) {
38        mesh_errf(err, "nvJitLink: nvJitLinkVersion failed (%s)", path);
39        free(path);
40        return false;
41    }
42    if (major != 13) {
43        mesh_errf(err, "nvJitLink: %s is version %u.%u; CUDA 13 needed", path, (unsigned)major, (unsigned)minor);
44        free(path);
45        return false;
46    }
47    static const char *const names[] = {"__nvJitLinkCreate_13_0", "__nvJitLinkDestroy_13_0", "__nvJitLinkAddData_13_0",
48                                         "__nvJitLinkComplete_13_0", "__nvJitLinkGetLinkedCubinSize_13_0",
49                                         "__nvJitLinkGetLinkedCubin_13_0", "__nvJitLinkGetErrorLogSize_13_0",
50                                         "__nvJitLinkGetErrorLog_13_0"};
51    void *ptrs[8];
52    for (int i = 0; i < 8; i++)
53        if (!(ptrs[i] = sym(lib, path, names[i], err))) {
54            free(path);
55            return false;
56        }
57    memcpy(&out->fns.create, &ptrs[0], sizeof(void *));
58    memcpy(&out->fns.destroy, &ptrs[1], sizeof(void *));
59    memcpy(&out->fns.add_data, &ptrs[2], sizeof(void *));
60    memcpy(&out->fns.complete, &ptrs[3], sizeof(void *));
61    memcpy(&out->fns.cubin_size, &ptrs[4], sizeof(void *));
62    memcpy(&out->fns.cubin, &ptrs[5], sizeof(void *));
63    memcpy(&out->fns.err_size, &ptrs[6], sizeof(void *));
64    memcpy(&out->fns.err_log, &ptrs[7], sizeof(void *));
65    out->path = path;
66    out->major = major;
67    out->minor = minor;
68    return true;
69}
70
71void mesh_nvjitlink_from(const mesh_nvjitlink_fns *fns, mesh_nvjitlink *out) {
72    *out = (mesh_nvjitlink){mesh_strdup("(test)"), 13, 0, *fns};
73}
74
75void mesh_nvjitlink_free(mesh_nvjitlink *j) {
76    free(j->path);
77    *j = (mesh_nvjitlink){0};
78}
79
80/* The error log, trimmed of NULs and whitespace ("" when there is none). */
81static char *error_log(const mesh_nvjitlink *j, void *h) {
82    size_t n = 0;
83    if (j->fns.err_size(h, &n) != 0 || n == 0) return mesh_strdup("");
84    char *buf = mesh_alloc(n + 1);
85    if (j->fns.err_log(h, buf) != 0) {
86        free(buf);
87        return mesh_strdup("");
88    }
89    size_t len = n;
90    while (len > 0 && buf[len - 1] == 0) len--;
91    const char *s = buf;
92    mesh_trim(&s, &len);
93    char *r = mesh_strndup(s, len);
94    free(buf);
95    return r;
96}
97
98bool mesh_nvjitlink_ptx_to_cubin(const mesh_nvjitlink *j, const uint8_t *ptx, size_t n, const char *arch, uint32_t reg_cap,
99                                 uint8_t **cubin, size_t *cubin_len, char **err) {
100    char *opt_arch = mesh_fmt("-arch=%s", arch);
101    char *opt_reg = reg_cap ? mesh_fmt("-maxrregcount=%u", (unsigned)reg_cap) : NULL;
102    const char *opts[2] = {opt_arch, opt_reg};
103    void *h = NULL;
104    int r = j->fns.create(&h, opt_reg ? 2 : 1, opts);
105    free(opt_arch);
106    free(opt_reg);
107    if (r != 0) return mesh_errf(err, "nvJitLinkCreate(%s) failed: code %d", arch, r);
108    bool ok = false;
109    /* PTX is text: hand it NUL-terminated. */
110    bool nul = n > 0 && ptx[n - 1] == 0;
111    uint8_t *src = mesh_alloc(n + 1);
112    memcpy(src, ptx, n);
113    size_t sn = nul ? n : n + 1;
114    r = j->fns.add_data(h, INPUT_PTX, src, sn, "program.ptx");
115    if (r != 0) {
116        char *log = error_log(j, h);
117        mesh_errf(err, "nvJitLinkAddData failed: code %d: %s", r, log);
118        free(log);
119        goto done;
120    }
121    r = j->fns.complete(h);
122    if (r != 0) {
123        char *log = error_log(j, h);
124        mesh_errf(err, "nvJitLinkComplete failed: code %d: %s", r, log);
125        free(log);
126        goto done;
127    }
128    size_t len = 0;
129    r = j->fns.cubin_size(h, &len);
130    if (r != 0 || len == 0) {
131        mesh_errf(err, "nvJitLinkGetLinkedCubinSize failed: code %d", r);
132        goto done;
133    }
134    uint8_t *out = mesh_alloc(len);
135    r = j->fns.cubin(h, out);
136    if (r != 0) {
137        free(out);
138        mesh_errf(err, "nvJitLinkGetLinkedCubin failed: code %d", r);
139        goto done;
140    }
141    *cubin = out;
142    *cubin_len = len;
143    ok = true;
144done:
145    free(src);
146    j->fns.destroy(&h);
147    return ok;
148}