/* SPDX-License-Identifier: BSD-3-Clause */
/* Copyright (c) 2026 Open-IOV contributors */
/* cuda_drv_test.c - minimal CUDA driver-API smoke test (dlopen, no toolkit).
 *
 * Run inside the guest. Exercises the paths CUDA needs beyond NVML:
 *   cuInit -> cuDeviceGetCount/Get -> cuCtxCreate -> cuMemAlloc ->
 *   cuMemcpyHtoD/DtoH -> cuMemFree -> cuCtxDestroy
 *
 * Expected output on success: "ALL OK".
 *
 * Status (2026-09-23): cuInit returns 3 (CUDA_ERROR_NOT_INITIALIZED) on the
 * emulated GPU. strace shows every nvidia.ko ioctl succeeding (RM status 0),
 * the driver completing RmInitAdapter and then closing the device, and
 * libcuda opening /dev/nvidia-uvm before /dev/nvidia0 - i.e. the next blocker
 * is the UVM driver (nvidia-uvm.ko), which needs the GPU MMU/VAS model that
 * the GSP-RM emulation does not provide yet. Build in the guest with:
 *   gcc -O1 -o /tmp/cdt cuda_drv_test.c -ldl
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <dlfcn.h>
#include <stdint.h>

typedef int CUresult;
typedef int CUdevice;
typedef void *CUcontext;
typedef uint64_t CUdeviceptr;

/* On failure: print cuGetErrorName/String and, when CDT_HOLD=<seconds> is
 * set, keep the process alive so the live state (trackers, RM error record)
 * can be inspected with gdb after cuCtxCreate has given up. */
static void fail_info(CUresult r)
{
    const char *(*errname)(CUresult, const char **) = dlsym(RTLD_DEFAULT, "cuGetErrorName");
    const char *(*errstr)(CUresult, const char **) = dlsym(RTLD_DEFAULT, "cuGetErrorString");
    const char *s = NULL;
    const char *hold = getenv("CDT_HOLD");

    if (errname && errname(r, &s) == 0 && s)
        printf("cuGetErrorName: %s\n", s);
    if (errstr && errstr(r, &s) == 0 && s)
        printf("cuGetErrorString: %s\n", s);
    if (hold && atoi(hold) > 0) {
        printf("CDT_HOLD: sleeping %d s for inspection\n", atoi(hold));
        fflush(stdout);
        sleep((unsigned)atoi(hold));
    }
}

#define CHECK(expr) do { \
    CUresult r = (expr); \
    if (r) { printf("%-24s FAILED rc=%d\n", #expr, r); fail_info(r); return 1; } \
    printf("%-24s OK\n", #expr); \
} while (0)

/* Print before entering libcuda so a crash identifies the call. */
#define MARK(name) do { printf("-> %s\n", name); } while (0)

int main(void)
{
    setvbuf(stdout, NULL, _IONBF, 0);

    void *h = dlopen("libcuda.so.1", RTLD_NOW);
    if (!h) { printf("dlopen libcuda.so.1 failed: %s\n", dlerror()); return 1; }

    CUresult (*cuInit)(unsigned int) = dlsym(h, "cuInit");
    CUresult (*cuDeviceGetCount)(int *) = dlsym(h, "cuDeviceGetCount");
    CUresult (*cuDeviceGet)(CUdevice *, int) = dlsym(h, "cuDeviceGet");
    CUresult (*cuCtxCreate)(CUcontext *, unsigned int, CUdevice) = dlsym(h, "cuCtxCreate_v2");
    CUresult (*cuMemAlloc)(CUdeviceptr *, size_t) = dlsym(h, "cuMemAlloc_v2");
    CUresult (*cuMemcpyHtoD)(CUdeviceptr, const void *, size_t) = dlsym(h, "cuMemcpyHtoD_v2");
    CUresult (*cuMemcpyDtoH)(void *, CUdeviceptr, size_t) = dlsym(h, "cuMemcpyDtoH_v2");
    CUresult (*cuMemFree)(CUdeviceptr) = dlsym(h, "cuMemFree_v2");
    CUresult (*cuCtxDestroy)(CUcontext) = dlsym(h, "cuCtxDestroy_v2");

    if (!cuInit || !cuDeviceGetCount || !cuDeviceGet || !cuCtxCreate ||
        !cuMemAlloc || !cuMemcpyHtoD || !cuMemcpyDtoH || !cuMemFree || !cuCtxDestroy) {
        printf("dlsym failed\n");
        return 1;
    }

    int count = 0;
    CUdevice dev;
    CUcontext ctx = NULL;
    CUdeviceptr dptr = 0;
    unsigned int buf[64];
    unsigned int back[64];
    size_t i;

    MARK("cuInit");
    CHECK(cuInit(0));
    MARK("cuDeviceGetCount");
    CHECK(cuDeviceGetCount(&count));
    printf("device count = %d\n", count);
    if (count < 1) return 1;

    MARK("cuDeviceGet");
    CHECK(cuDeviceGet(&dev, 0));
    printf("device 0 = %d\n", dev);

    MARK("cuCtxCreate");
    {
        const char *pause = getenv("CDT_PAUSE");
        if (pause && atoi(pause) > 0) {
            printf("CDT_PAUSE: sleeping %d s before cuCtxCreate\n", atoi(pause));
            fflush(stdout);
            sleep((unsigned)atoi(pause));
        }
    }
    CHECK(cuCtxCreate(&ctx, 0, dev));
    printf("context = %p\n", ctx);

    MARK("cuMemAlloc");
    CHECK(cuMemAlloc(&dptr, sizeof(buf)));
    printf("device ptr = 0x%llx (%zu bytes)\n", (unsigned long long)dptr, sizeof(buf));

    for (i = 0; i < 64; i++) buf[i] = (unsigned int)(i * 3 + 1);
    MARK("cuMemcpyHtoD");
    CHECK(cuMemcpyHtoD(dptr, buf, sizeof(buf)));
    memset(back, 0, sizeof(back));
    MARK("cuMemcpyDtoH");
    CHECK(cuMemcpyDtoH(back, dptr, sizeof(buf)));
    {
        size_t matched = 0;
        for (i = 0; i < 64; i++)
            if (buf[i] == back[i]) matched++;
        printf("roundtrip %s matched=%zu/64 back[0..7]=%u,%u,%u,%u,%u,%u,%u,%u expected 1,4,7,10,13,16,19,22\n",
               matched == 64 ? "match" : "MISMATCH", matched,
               back[0], back[1], back[2], back[3],
               back[4], back[5], back[6], back[7]);
    }

    MARK("cuMemFree");
    CHECK(cuMemFree(dptr));
    MARK("cuCtxDestroy");
    CHECK(cuCtxDestroy(ctx));
    printf("ALL OK\n");
    {
        const char *wait = getenv("CDT_WAIT");
        if (wait && atoi(wait) > 0) {
            printf("CDT_WAIT: parked %d s (buf=%p back=%p)\n",
                   atoi(wait), (void *)buf, (void *)back);
            fflush(stdout);
            sleep((unsigned)atoi(wait));
        }
    }
    return 0;
}