// Copyright (2019-2023) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include "NUMAHWPolicy.h"

#include "../common/CUDAWrapper.h"
#include "JFJochException.h"

#ifdef JFJOCH_USE_NUMA
#include <numa.h>
#endif

NUMAHWPolicy::NUMAHWPolicy(const std::string &policy) : name(policy) {
    if ((policy.empty()) || (policy == "none")) {
        name = "none";
    } else if (policy == "n2g2") {
        bindings.emplace_back(NUMABinding{.cpu_node = 0, .mem_node = 0, .gpu = 0});
        bindings.emplace_back(NUMABinding{.cpu_node = 1, .mem_node = 1, .gpu = 1});
    } else if (policy == "n2g4") {
        bindings.emplace_back(NUMABinding{.cpu_node = 0, .mem_node = 0, .gpu = 0});
        bindings.emplace_back(NUMABinding{.cpu_node = 1, .mem_node = 1, .gpu = 2});
        bindings.emplace_back(NUMABinding{.cpu_node = 0, .mem_node = 0, .gpu = 1});
        bindings.emplace_back(NUMABinding{.cpu_node = 1, .mem_node = 1, .gpu = 3});
    } else if (policy == "n2g4_hbm") {
        bindings.emplace_back(NUMABinding{.cpu_node = 0, .mem_node = 2, .gpu = 0});
        bindings.emplace_back(NUMABinding{.cpu_node = 1, .mem_node = 3, .gpu = 2});
        bindings.emplace_back(NUMABinding{.cpu_node = 0, .mem_node = 2, .gpu = 1});
        bindings.emplace_back(NUMABinding{.cpu_node = 1, .mem_node = 3, .gpu = 3});
    } else if (policy == "n8g4") {
        for (int32_t i = 0; i < 8; i++)
            bindings.emplace_back(NUMABinding{.cpu_node = i, .mem_node = i, .gpu = i/2});
    } else if (policy == "n8g4_hbm") {
        for (int32_t i = 0; i < 8; i++)
            bindings.emplace_back(NUMABinding{.cpu_node = i, .mem_node = i + 8, .gpu = i / 2});
    } else if (policy == "g2") {
        bindings.emplace_back(NUMABinding{.cpu_node = -1, .mem_node = -1, .gpu = 0});
        bindings.emplace_back(NUMABinding{.cpu_node = -1, .mem_node = -1, .gpu = 1});
    } else if (policy == "g4") {
        bindings.emplace_back(NUMABinding{.cpu_node = -1, .mem_node = -1, .gpu = 0});
        bindings.emplace_back(NUMABinding{.cpu_node = -1, .mem_node = -1, .gpu = 1});
        bindings.emplace_back(NUMABinding{.cpu_node = -1, .mem_node = -1, .gpu = 2});
        bindings.emplace_back(NUMABinding{.cpu_node = -1, .mem_node = -1, .gpu = 3});
    } else
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Unknown NUMA policy");
}

NUMAHWPolicy::NUMAHWPolicy(const NUMAHWPolicy &other) : bindings(other.bindings), name(other.name), curr_thread(0) {}

NUMAHWPolicy &NUMAHWPolicy::operator=(const NUMAHWPolicy &other) {
    bindings = other.bindings;
    name = other.name;
    curr_thread = 0;
    return *this;
}

NUMABinding NUMAHWPolicy::GetBinding(uint32_t thread) {
    if (bindings.empty())
        return NUMABinding{.cpu_node = -1, .mem_node = -1, .gpu = -1};
    else
        return bindings.at(thread % bindings.size());
}

NUMABinding NUMAHWPolicy::GetBinding() {
    return GetBinding(curr_thread++);
}

void NUMAHWPolicy::Bind() {
    Bind(GetBinding());
}

void NUMAHWPolicy::Bind(uint32_t thread) {
    Bind(GetBinding(thread));
}

void NUMAHWPolicy::Bind(const NUMABinding &binding) {
    RunOnNode(binding.cpu_node);
    MemOnNode(binding.mem_node);
    SelectGPU(binding.gpu);
}

void NUMAHWPolicy::RunOnNode(int32_t cpu_node) {
#ifdef JFJOCH_USE_NUMA
    if (numa_available() != -1) {
        auto max_nodes = numa_num_configured_nodes();

        if (cpu_node >= 0) {
            if (cpu_node < max_nodes)
                numa_run_on_node(cpu_node);
            else
                throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "CPU NUMA node out of bounds");
        }
    }
#endif
}

void NUMAHWPolicy::MemOnNode(int32_t mem_node) {
#ifdef JFJOCH_USE_NUMA
    if (numa_available() != -1) {
        auto max_nodes = numa_num_configured_nodes();

        if (mem_node >= 0) {
            if (mem_node < max_nodes) {
                struct bitmask *mask = numa_allocate_nodemask();
                numa_bitmask_setbit(mask, mem_node);
                numa_set_membind(mask);
                numa_bitmask_free(mask);
            } else
                throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Memory NUMA node out of bounds");
        }
    }
#endif
}

void NUMAHWPolicy::SelectGPU(int32_t gpu) {
#ifdef JFJOCH_USE_CUDA
    if (gpu > 0) {
        if (gpu < get_gpu_count())
            set_gpu(gpu);
        else
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "GPU device out of bounds");
    }
#endif
}

const std::string &NUMAHWPolicy::GetName() const {
    return name;
}