Jungfraujoch/common/ThreadAffinity.cpp

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

#include "ThreadAffinity.h"

#include <algorithm>
#include <set>

std::vector<int> CpusOfGpuPackages(const std::vector<int> &allowed, const std::vector<int> &package_of_cpu,
                                   const std::vector<int> &gpu_packages) {
    if (gpu_packages.empty()
        || std::find(gpu_packages.begin(), gpu_packages.end(), -1) != gpu_packages.end())
        return {};
    const auto package = [&](int c) {
        return c < static_cast<int>(package_of_cpu.size()) ? package_of_cpu[c] : -1;
    };
    std::set<int> packages;
    std::vector<int> out;
    for (const int c: allowed) {
        packages.insert(package(c));
        if (std::find(gpu_packages.begin(), gpu_packages.end(), package(c)) != gpu_packages.end())
            out.push_back(c);
    }
    if (packages.size() < 2 || out.empty() || out.size() == allowed.size())
        return {};
    return out;
}

#ifdef __linux__

#include <cmath>
#include <cstdio>
#include <filesystem>
#include <fstream>
#include <sstream>

#include <pthread.h>
#include <sched.h>

namespace {
    // The process's CPUs, taken before any thread has been pinned: static initialisation runs on the
    // main thread before main().
    cpu_set_t StartMask() {
        cpu_set_t mask;
        CPU_ZERO(&mask);
        if (sched_getaffinity(0, sizeof(mask), &mask) != 0)
            for (int c = 0; c < CPU_SETSIZE; c++)
                CPU_SET(c, &mask);
        return mask;
    }
    // Narrowed once by ConfineToGpuPackages, before the process starts any thread.
    cpu_set_t start_mask = StartMask();

    // The CPUs the cgroup's quota allows, rounded up; 0 where none is set. The process's own cgroup
    // first, then the root of the hierarchy as mounted, which is the container's own cgroup where the
    // container does not see the host's paths. cgroup v2 (cpu.max: "max 100000" or "200000 100000")
    // and v1 (cpu.cfs_quota_us = -1 or the quota, cpu.cfs_period_us).
    unsigned CgroupCpuQuota() {
        std::ifstream f("/proc/self/cgroup");
        std::string line;
        while (std::getline(f, line)) {
            const size_t c1 = line.find(':'), c2 = line.find(':', c1 + 1);
            if (c1 == std::string::npos || c2 == std::string::npos)
                continue;
            const std::string controllers = "," + line.substr(c1 + 1, c2 - c1 - 1) + ",";
            const std::string path = line.substr(c2 + 1);
            const bool v2 = line.compare(0, c2 + 1, "0::") == 0;
            if (!v2 && controllers.find(",cpu,") == std::string::npos)
                continue;
            for (const std::string &dir : {"/sys/fs/cgroup" + std::string(v2 ? "" : "/cpu") + path,
                                           "/sys/fs/cgroup" + std::string(v2 ? "" : "/cpu")}) {
                double quota = -1, period = 0;
                if (v2) {
                    std::ifstream q(dir + "/cpu.max");
                    std::string max;
                    if (!(q >> max >> period))
                        continue;
                    if (max != "max")
                        quota = std::atof(max.c_str());
                } else {
                    std::ifstream q(dir + "/cpu.cfs_quota_us"), p(dir + "/cpu.cfs_period_us");
                    if (!(q >> quota) || !(p >> period))
                        continue;
                }
                if (quota > 0 && period > 0)
                    return static_cast<unsigned>(std::ceil(quota / period));
                break;   // found, and no quota set
            }
        }
        return 0;
    }

    int NumaNodeCount() {
        int n = 0;
        std::error_code ec;
        for (const auto &e : std::filesystem::directory_iterator("/sys/devices/system/node", ec)) {
            const std::string name = e.path().filename().string();
            if (name.rfind("node", 0) == 0 && name.size() > 4
                && name.find_first_not_of("0123456789", 4) == std::string::npos)
                n++;
        }
        return n;
    }

    // "0-23,48-71" -> the CPUs it names.
    bool ParseCpuList(const std::string &list, cpu_set_t &out) {
        CPU_ZERO(&out);
        std::stringstream ss(list);
        std::string range;
        bool any = false;
        while (std::getline(ss, range, ',')) {
            int lo, hi;
            if (std::sscanf(range.c_str(), "%d-%d", &lo, &hi) == 2) {
            } else if (std::sscanf(range.c_str(), "%d", &lo) == 1) {
                hi = lo;
            } else
                continue;
            for (int c = lo; c <= hi && c < CPU_SETSIZE; c++) {
                CPU_SET(c, &out);
                any = true;
            }
        }
        return any;
    }

    int ReadInt(const std::string &path) {
        std::ifstream f(path);
        int v = -1;
        if (!(f >> v))
            return -1;
        return v;
    }

    int PackageOfCpu(int cpu) {
        return ReadInt("/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/topology/physical_package_id");
    }

    // The package of a NUMA node's first CPU; -1 for a node without CPUs or an unknown node.
    int PackageOfNode(int node) {
        if (node < 0)
            return -1;
        std::ifstream f("/sys/devices/system/node/node" + std::to_string(node) + "/cpulist");
        std::string list;
        cpu_set_t cpus;
        if (!std::getline(f, list) || !ParseCpuList(list, cpus))
            return -1;
        for (int c = 0; c < CPU_SETSIZE; c++)
            if (CPU_ISSET(c, &cpus))
                return PackageOfCpu(c);
        return -1;
    }
}

int NumaNodeOfPciDevice(const std::string &pci_bus_id) {
    if (NumaNodeCount() < 2)
        return -1;
    unsigned domain, bus, device, function;
    if (std::sscanf(pci_bus_id.c_str(), "%x:%x:%x.%x", &domain, &bus, &device, &function) != 4)
        return -1;
    char path[128];
    std::snprintf(path, sizeof(path), "/sys/bus/pci/devices/%04x:%02x:%02x.%x/numa_node",
                  domain, bus, device, function);
    std::ifstream f(path);
    int node = -1;
    if (!(f >> node))
        return -1;
    return node;
}

unsigned AvailableCpus() {
    static const unsigned quota = CgroupCpuQuota();
    const unsigned cpus = static_cast<unsigned>(std::max(1, CPU_COUNT(&start_mask)));
    return quota > 0 ? std::min(cpus, quota) : cpus;
}

std::vector<int> ConfineToGpuPackages(const std::vector<int> &gpu_nodes) {
    std::vector<int> allowed, package_of_cpu, gpu_packages;
    for (int c = 0; c < CPU_SETSIZE; c++)
        if (CPU_ISSET(c, &start_mask)) {
            allowed.push_back(c);
            package_of_cpu.resize(c + 1, -1);
            package_of_cpu[c] = PackageOfCpu(c);
        }
    for (const int node: gpu_nodes)
        gpu_packages.push_back(PackageOfNode(node));
    const std::vector<int> cpus = CpusOfGpuPackages(allowed, package_of_cpu, gpu_packages);
    if (cpus.empty())
        return {};
    CPU_ZERO(&start_mask);
    for (const int c: cpus)
        CPU_SET(c, &start_mask);
    RestoreThreadAffinity();
    return cpus;
}

void PinThreadToNumaNode(int node) {
    if (node < 0)
        return;
    std::ifstream f("/sys/devices/system/node/node" + std::to_string(node) + "/cpulist");
    std::string list;
    cpu_set_t node_cpus;
    if (!std::getline(f, list) || !ParseCpuList(list, node_cpus))
        return;
    cpu_set_t mask;
    CPU_AND(&mask, &node_cpus, &start_mask);
    if (CPU_COUNT(&mask) == 0)
        return;
    pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask);
}

void RestoreThreadAffinity() {
    pthread_setaffinity_np(pthread_self(), sizeof(start_mask), &start_mask);
}

std::vector<int> GetThreadCpus() {
    cpu_set_t mask;
    std::vector<int> cpus;
    if (pthread_getaffinity_np(pthread_self(), sizeof(mask), &mask) == 0)
        for (int c = 0; c < CPU_SETSIZE; c++)
            if (CPU_ISSET(c, &mask))
                cpus.push_back(c);
    return cpus;
}

void SetThreadCpus(const std::vector<int> &cpus) {
    if (cpus.empty())
        return;
    cpu_set_t mask;
    CPU_ZERO(&mask);
    for (const int c: cpus)
        CPU_SET(c, &mask);
    pthread_setaffinity_np(pthread_self(), sizeof(mask), &mask);
}

#else

#include <thread>

std::vector<int> ConfineToGpuPackages(const std::vector<int> &) { return {}; }
unsigned AvailableCpus() { return std::max(1u, std::thread::hardware_concurrency()); }
int NumaNodeOfPciDevice(const std::string &) { return -1; }
void PinThreadToNumaNode(int) {}
void RestoreThreadAffinity() {}
std::vector<int> GetThreadCpus() { return {}; }
void SetThreadCpus(const std::vector<int> &) {}

#endif