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

#include <algorithm>
#include <atomic>
#include <cctype>
#include <cstdlib>
#include <fstream>
#include <string>
#ifdef __linux__
#include <pthread.h>
#include <sched.h>
#endif

#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));
    }
}

// The NUMA node a GPU is attached to, from sysfs; -1 when unknown or not applicable. No libnuma:
// this is a file read. (This lookup existed before cc925b26 and is restored unchanged apart from
// lower-casing the bus id - CUDA reports hex digits in upper case on some drivers and sysfs paths
// are lower case.)
int get_gpu_numa_node(int32_t dev_id) {
#ifdef __linux__
    if (dev_id < 0 || dev_id >= get_gpu_count())
        return -1;
    char buf[64] = {};
    if (cudaDeviceGetPCIBusId(buf, static_cast<int>(sizeof(buf)), dev_id) != cudaSuccess)
        return -1;
    std::string bus_id(buf);
    std::transform(bus_id.begin(), bus_id.end(), bus_id.begin(),
                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
    std::ifstream f("/sys/bus/pci/devices/" + bus_id + "/numa_node");
    int node = -1;
    if (!(f >> node))
        return -1;
    return node;
#else
    (void) dev_id;
    return -1;
#endif
}

#ifdef __linux__
namespace {
    // Confine the calling thread to the cores of one NUMA node, read from sysfs. Its memory then
    // follows by first touch, which is what actually matters - the placement of what the thread
    // allocates, not the affinity itself.
    void run_on_numa_node(int node) {
        std::ifstream f("/sys/devices/system/node/node" + std::to_string(node) + "/cpulist");
        std::string list;
        if (!std::getline(f, list) || list.empty())
            return;
        cpu_set_t set;
        CPU_ZERO(&set);
        // cpulist is comma-separated singles and a-b ranges, e.g. "24-35" or "0,2,4-7".
        size_t pos = 0;
        while (pos < list.size()) {
            size_t comma = list.find(',', pos);
            const std::string item = list.substr(pos, comma == std::string::npos ? comma : comma - pos);
            const size_t dash = item.find('-');
            const int lo = std::atoi(item.c_str());
            const int hi = dash == std::string::npos ? lo : std::atoi(item.c_str() + dash + 1);
            for (int c = lo; c <= hi && c < CPU_SETSIZE; ++c)
                CPU_SET(c, &set);
            if (comma == std::string::npos) break;
            pos = comma + 1;
        }
        if (CPU_COUNT(&set) > 0)
            pthread_setaffinity_np(pthread_self(), sizeof(set), &set);
    }

    int numa_node_count() {
        int n = 0;
        while (std::ifstream("/sys/devices/system/node/node" + std::to_string(n) + "/cpulist").good())
            ++n;
        return n;
    }
}
#endif

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

#ifdef __linux__
    // Optionally also confine the thread to the cores local to that GPU. On a machine whose GPUs all
    // hang off some of the sockets, a worker can otherwise sit on a socket with no GPU of its own and
    // cross the interconnect for every transfer and every page it allocates. Off by default: it is a
    // real trade - it also denies the thread the cores of the other sockets - and only measurement on
    // a given box can say which way it goes. No-op unless there is more than one node and the GPU's
    // node is known.
    static const bool pin_cpu = [] {
        const char *e = std::getenv("JFJOCH_PIN_CPU_TO_GPU_NODE");
        return e && *e && *e != '0';
    }();
    if (!pin_cpu || numa_node_count() < 2)
        return;
    const int node = get_gpu_numa_node(dev);
    if (node >= 0)
        run_on_numa_node(node);
#endif
}