1/* nvjitlink.rs's link_and_refuse, with a fake nvJitLink (the real one needs the
2 * CUDA 13 toolkit, which this machine lacks; on a DGX the open succeeds). */
3#include "mesh/check.h"
4#include "mesh/nvjitlink.h"
5
6#include <stdlib.h>
7#include <string.h>
8
9static int created, destroyed;
10static char opts_seen[128];
11static size_t added;
12static int fail_complete;
13
14static int f_create(void **h, uint32_t n, const char *const *opts) {
15    created++;
16    opts_seen[0] = 0;
17    for (uint32_t i = 0; i < n; i++) {
18        strcat(opts_seen, opts[i]);
19        strcat(opts_seen, " ");
20    }
21    *h = &created;
22    return 0;
23}
24static int f_destroy(void **h) {
25    destroyed++;
26    *h = NULL;
27    return 0;
28}
29static int f_add(void *h, int kind, const void *data, size_t n, const char *name) {
30    (void)h;
31    CHECK_I(kind, 2);
32    CHECK_S(name, "program.ptx");
33    CHECK_U(((const char *)data)[n - 1], 0);
34    added = n;
35    return 0;
36}
37static int f_complete(void *h) {
38    (void)h;
39    return fail_complete;
40}
41static int f_size(void *h, size_t *n) {
42    (void)h;
43    *n = 8;
44    return 0;
45}
46static int f_cubin(void *h, void *out) {
47    (void)h;
48    memcpy(out, "\x7f" "ELFcubn", 8);
49    return 0;
50}
51static int f_err_size(void *h, size_t *n) {
52    (void)h;
53    *n = 24;
54    return 0;
55}
56static int f_err_log(void *h, void *out) {
57    (void)h;
58    memcpy(out, "ptxas: bogus opcode \n\0\0\0", 24);
59    return 0;
60}
61
62int main(void) {
63    mesh_nvjitlink j;
64    char *err = NULL;
65    CHECK(!mesh_nvjitlink_open("/nonexistent", &j, &err));
66    CHECK(strstr(err, "cannot open /nonexistent/libnvJitLink.so.13") != NULL);
67    CHECK(strstr(err, "(set --nvjitlink DIR; there is no fallback)") != NULL);
68    free(err);
69    err = NULL;
70    mesh_nvjitlink_fns f = {f_create, f_destroy, f_add, f_complete, f_size, f_cubin, f_err_size, f_err_log};
71    mesh_nvjitlink_from(&f, &j);
72    const char *ptx = ".version 8.8\n.target sm_121\n.address_size 64\n.visible .entry k() { ret; }\n";
73    uint8_t *cubin;
74    size_t n;
75    CHECK(mesh_nvjitlink_ptx_to_cubin(&j, (const uint8_t *)ptx, strlen(ptx), "sm_121", 0, &cubin, &n, &err));
76    CHECK(n == 8 && memcmp(cubin, "\x7f" "ELF", 4) == 0);
77    CHECK_S(opts_seen, "-arch=sm_121 ");
78    CHECK_U(added, strlen(ptx) + 1);
79    free(cubin);
80    CHECK(mesh_nvjitlink_ptx_to_cubin(&j, (const uint8_t *)ptx, strlen(ptx), "sm_86", 64, &cubin, &n, &err));
81    CHECK_S(opts_seen, "-arch=sm_86 -maxrregcount=64 ");
82    free(cubin);
83    fail_complete = 7;
84    CHECK(!mesh_nvjitlink_ptx_to_cubin(&j, (const uint8_t *)ptx, strlen(ptx), "sm_121", 0, &cubin, &n, &err));
85    CHECK_S(err, "nvJitLinkComplete failed: code 7: ptxas: bogus opcode");
86    free(err);
87    CHECK_I(created, destroyed);
88    mesh_nvjitlink_free(&j);
89    return mesh_check_done();
90}