Files
Jungfraujoch/common/ThreadAffinity.h
T
leonarski_fandClaude Opus 5.5 3c5ff3f5df rugnux: on a machine of several sockets, keep to the sockets the GPUs are on
With -N left to choose, rugnux now confines the process (main thread, the CPUs
RestoreThreadAffinity hands back, and what AvailableCpus counts) to the CPUs of
the packages its cards hang off, and sizes its threads to them. Nothing changes
where the allowed CPUs sit in one package, there is no card or a card's NUMA
node is unknown, or the cards' packages already hold every allowed CPU. A mask
set from outside (taskset, numactl, a cpuset) is kept: the decision works
within it. An explicit -N keeps the old behaviour.

On a 48-CPU server with four cards on two of its four NUMA nodes, a rotation
run took 24.5-24.7 s confined to those nodes' 24 CPUs against 27.7-27.9 s on
all 48 (scaling and merging 8.1 -> 6.2 s), and two larger sets moved
95-99 -> 84-87 s and 131-135 -> 113-115 s. On a single-socket machine this is
a no-op.

The decision itself (CpusOfGpuPackages) is a pure function, tested on
synthetic 4-package, 2-package hyper-threaded, cards-on-both, no-card and
restricted-mask topologies.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
2026-10-11 09:27:32 +02:00

57 lines
3.4 KiB
C++

// 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);