sass_fma/test_kernel.cu at master · jonpry/sass_fma · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
#include <iostream>
#include <vector>
#include <fstream>
#include <cuda.h>
#include <cuda_runtime.h>

#define CHECK_CUDA(call) \
    do { \
        CUresult res = call; \
        if (res != CUDA_SUCCESS) { \
            const char* err_str; \
            cuGetErrorString(res, &err_str); \
            std::cerr << "CUDA Error: " << err_str << " at line " << __LINE__ << std::endl; \
            exit(1); \
        } \
    } while (0)

int main() {
    int n = 1024;

    // 1. Initialize Driver API
    CHECK_CUDA(cuInit(0));
    CUdevice device;
    CHECK_CUDA(cuDeviceGet(&device, 0));
    CUcontext context;
    CHECK_CUDA(cuCtxCreate(&context, 0, device));

    // 2. Load the CUBIN (The interceptor will hook this)
    // We expect discovery.cubin to exist after ./build.sh
    std::ifstream ifs("discovery.cubin", std::ios::binary | std::ios::ate);
    if (!ifs.is_open()) {
        std::cerr << "Error: discovery.cubin not found. Run ./build.sh first." << std::endl;
        return 1;
    }
    std::streamsize cubin_size = ifs.tellg();
    ifs.seekg(0, std::ios::beg);
    std::vector<char> cubin_data(cubin_size);
    ifs.read(cubin_data.data(), cubin_size);

    CUmodule module;
    // THIS IS THE HOOK POINT
    CHECK_CUDA(cuModuleLoadDataEx(&module, cubin_data.data(), 0, NULL, NULL));

    CUfunction function;
    CHECK_CUDA(cuModuleGetFunction(&function, module, "_Z16discovery_kernelPfS_S_S_"));

    // 3. Setup Data
    std::vector<float> h_a(1, 2.0f), h_b(1, 3.0f), h_c(1, 4.0f), h_d(9, 0.0f);
    CUdeviceptr d_a, d_b, d_c, d_d;

    CHECK_CUDA(cuMemAlloc(&d_a, sizeof(float)));
    CHECK_CUDA(cuMemAlloc(&d_b, sizeof(float)));
    CHECK_CUDA(cuMemAlloc(&d_c, sizeof(float)));
    CHECK_CUDA(cuMemAlloc(&d_d, 9 * sizeof(float)));

    CHECK_CUDA(cuMemcpyHtoD(d_a, h_a.data(), sizeof(float)));
    CHECK_CUDA(cuMemcpyHtoD(d_b, h_b.data(), sizeof(float)));
    CHECK_CUDA(cuMemcpyHtoD(d_c, h_c.data(), sizeof(float)));

    // 4. Launch
    void* args[] = { &d_a, &d_b, &d_c, &d_d };
    CHECK_CUDA(cuLaunchKernel(function, 1, 1, 1, 1, 1, 1, 0, NULL, args, NULL));
    CHECK_CUDA(cuCtxSynchronize());

    // 5. Verify
    CHECK_CUDA(cuMemcpyDtoH(h_d.data(), d_d, 9 * sizeof(float)));
    std::cout << "Result: " << (h_d[6] == 10.0f ? "SUCCESS" : "FAILURE") << " (Got " << h_d[6] << ")" << std::endl;

    return 0;
}