Jungfraujoch/common/CUDAWrapper.cu

// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <fstream>
#include <atomic>

#include "CUDAWrapper.h"
#include "JFJochException.h"

inline void cuda_err(cudaError_t val) {
    if (val != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}

int32_t get_gpu_count() {
    int device_count;
    cudaError_t val = cudaGetDeviceCount(&device_count);
    switch (val) {
        case cudaSuccess:
            return device_count;
        case cudaErrorNoDevice:
        case cudaErrorInsufficientDriver:
            return 0;
        default:
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
    }

}

void set_gpu(int32_t dev_id) {
    auto dev_count = get_gpu_count();

    // Ignore if no GPU present
    if (dev_count > 0) {
        if ((dev_id < 0) || (dev_id >= dev_count))
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Device ID cannot be negative");

        cuda_err(cudaSetDevice(dev_id));
    }
}

void pin_gpu() {
    static std::atomic<uint32_t> counter{0};
    auto dev_count = get_gpu_count();
    if (dev_count > 0)
        set_gpu(counter.fetch_add(1) % dev_count);
}