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

#pragma once

#include <string>
#include <vector>

// Keeping a GPU's worker threads on the socket the card hangs off, on a machine with more than one
// NUMA node, so the pinned host buffers they allocate and the copies to the card stay on that
// socket's memory controller and PCIe root. Read from /sys - no libnuma - and Linux only: elsewhere,
// and on a machine with a single node, every call is a no-op.

// The NUMA node of a PCI device, from its bus id as CUDA gives it ("0000:41:00.0"); -1 where it is
// not known or the machine has a single node.
int NumaNodeOfPciDevice(const std::string &pci_bus_id);

// How many CPUs this process may use: the CPUs it was started on (its affinity mask - what taskset,
// a cpuset or a Slurm job gives it), and no more than its cgroup's CPU quota where one is set (a
// container started with --cpus). std::thread::hardware_concurrency() counts every CPU of the
// machine, so a job given 16 of a node's 192 would otherwise start 192 threads. Read once, at the
// first call. Outside Linux it is hardware_concurrency(). At least 1.
unsigned AvailableCpus();

// The CPUs a process driving GPUs keeps to on a machine of several packages (sockets): those of the
// packages its cards hang off. `allowed` is what the process may use, `package_of_cpu[c]` the package
// of CPU c (-1 unknown) and `gpu_packages` the package of each card (-1 unknown). Empty where nothing
// changes: the allowed CPUs sit in one package, there is no card or one card's package is unknown, the
// cards' packages hold every allowed CPU, or none of them.
std::vector<int> CpusOfGpuPackages(const std::vector<int> &allowed, const std::vector<int> &package_of_cpu,
                                   const std::vector<int> &gpu_packages);

// Keep the process to CpusOfGpuPackages for cards on the NUMA nodes `gpu_nodes` (-1 unknown): the
// calling thread, the CPUs RestoreThreadAffinity gives back and what AvailableCpus counts all become
// those, within the CPUs it was started on. Threads, and the memory they first touch, spread over
// sockets the cards are not on make the latency-bound scaling slower: on a 48-CPU server with its cards
// on two of four NUMA nodes, a rotation run kept to those two nodes' 24 CPUs took 12 % less wall time
// than on all 48, and the 12 CPUs of one node without a card still beat all 48. Call it on the main thread before any thread
// is started. Returns the CPUs it keeps to; empty where nothing changed. Linux only.
std::vector<int> ConfineToGpuPackages(const std::vector<int> &gpu_nodes);

// Restrict the calling thread to the CPUs of `node` that the process may run on. Nothing happens for
// node < 0, or where that would leave the thread no CPU.
void PinThreadToNumaNode(int node);

// Give the calling thread back the CPUs the process started with. A thread inherits the affinity of
// the thread that creates it, so a thread pool first reached from a pinned worker would otherwise
// run every later parallel pass of the process on one socket.
void RestoreThreadAffinity();

// The calling thread's CPUs now, for SetThreadCpus to give back after a while on one node. Empty
// outside Linux.
std::vector<int> GetThreadCpus();

// Keep the calling thread on exactly these CPUs. Nothing happens for an empty list.
void SetThreadCpus(const std::vector<int> &cpus);