#include "gpr/device.hpp"

#include <cstdio>

namespace gpr {

void cu_check(CUresult r, const char* what) {
    if (r != CUDA_SUCCESS) {
        const char* s = nullptr;
        cuGetErrorString(r, &s);
        throw std::runtime_error(std::string("CUDA error in ") + what + ": " + (s ? s : "?"));
    }
}

CudaContext::CudaContext() {
    cu_check(cuInit(0), "cuInit");
    cu_check(cuDeviceGet(&dev, 0), "cuDeviceGet");
    cu_check(cuDeviceGetAttribute(&cc_major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev),
             "cc_major");
    cu_check(cuDeviceGetAttribute(&cc_minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev),
             "cc_minor");
    cu_check(cuDevicePrimaryCtxRetain(&ctx, dev), "primaryCtxRetain");
    cu_check(cuCtxSetCurrent(ctx), "ctxSetCurrent");
}

CudaContext::~CudaContext() {
    if (ctx) cuDevicePrimaryCtxRelease(dev);
}

std::string CudaContext::arch_flag() const {
    return "--gpu-architecture=sm_" + std::to_string(cc_major) + std::to_string(cc_minor);
}

}  // namespace gpr