#include <nvJitLink.h>

#include <cstdlib>
#include <cstring>
#include <iostream>
#include <vector>

#define LINK_CHECK(call)                                                     \
    do {                                                                     \
        const nvJitLinkResult status = (call);                               \
        if (status != NVJITLINK_SUCCESS) {                                   \
            std::cerr << "nvJitLink status=" << static_cast<int>(status)     \
                      << '\n';                                               \
            std::exit(EXIT_FAILURE);                                         \
        }                                                                    \
    } while (0)

int main(int argc, char** argv) {
    if (argc != 2) {
        std::cerr << "usage: " << argv[0] << " sm_XX\n";
        return EXIT_FAILURE;
    }
    const std::string arch = std::string("-arch=") + argv[1];
    const char* options[] = {arch.c_str()};
    nvJitLinkHandle handle{};
    LINK_CHECK(nvJitLinkCreate(&handle, 1, options));

    static constexpr char ptx[] = R"(
.version 8.0
.target sm_80
.address_size 64
.visible .entry empty_kernel() { ret; }
)";
    LINK_CHECK(nvJitLinkAddData(handle, NVJITLINK_INPUT_PTX,
                                const_cast<char*>(ptx), std::strlen(ptx) + 1,
                                "empty.ptx"));
    LINK_CHECK(nvJitLinkComplete(handle));
    size_t cubin_bytes = 0;
    LINK_CHECK(nvJitLinkGetLinkedCubinSize(handle, &cubin_bytes));
    std::vector<char> cubin(cubin_bytes);
    LINK_CHECK(nvJitLinkGetLinkedCubin(handle, cubin.data()));
    LINK_CHECK(nvJitLinkDestroy(&handle));
    std::cout << "linked_cubin_bytes=" << cubin_bytes << '\n';
}