#include <cuda.h>
#include <nvrtc.h>

#include <cstdlib>
#include <iostream>
#include <string>
#include <vector>

#define NVRTC_CHECK(call)                                                     \
    do {                                                                      \
        const nvrtcResult status = (call);                                    \
        if (status != NVRTC_SUCCESS) {                                        \
            std::cerr << nvrtcGetErrorString(status) << '\n';                 \
            std::exit(EXIT_FAILURE);                                          \
        }                                                                     \
    } while (0)

int main() {
    static constexpr char source[] = R"(
extern "C" __global__ void scale(float* x, float value) {
    x[threadIdx.x] *= value;
})";

    nvrtcProgram program{};
    NVRTC_CHECK(nvrtcCreateProgram(&program, source, "scale.cu", 0, nullptr, nullptr));
    const char* options[] = {"--std=c++17"};
    const nvrtcResult compile_status = nvrtcCompileProgram(program, 1, options);

    size_t log_bytes = 0;
    NVRTC_CHECK(nvrtcGetProgramLogSize(program, &log_bytes));
    std::string log(log_bytes, '\0');
    if (log_bytes > 1) {
        NVRTC_CHECK(nvrtcGetProgramLog(program, log.data()));
        std::cerr << log;
    }
    if (compile_status != NVRTC_SUCCESS) {
        return EXIT_FAILURE;
    }

    size_t ptx_bytes = 0;
    NVRTC_CHECK(nvrtcGetPTXSize(program, &ptx_bytes));
    std::vector<char> ptx(ptx_bytes);
    NVRTC_CHECK(nvrtcGetPTX(program, ptx.data()));
    NVRTC_CHECK(nvrtcDestroyProgram(&program));
    std::cout << "generated_ptx_bytes=" << ptx_bytes << '\n';
}