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

common/ThreadAffinity.cpp

@@ -3,9 +3,31 @@
 
 #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 <algorithm>
 #include <cmath>
 #include <cstdio>
 #include <filesystem>
@@ -26,7 +48,8 @@ namespace {
                 CPU_SET(c, &mask);
         return mask;
     }
-    const cpu_set_t start_mask = StartMask();
+    // 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
@@ -99,6 +122,33 @@ namespace {
         }
         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) {
@@ -118,13 +168,29 @@ int NumaNodeOfPciDevice(const std::string &pci_bus_id) {
 }
 
 unsigned AvailableCpus() {
-    static const unsigned n = [] {
-        unsigned cpus = static_cast<unsigned>(std::max(1, CPU_COUNT(&start_mask)));
-        if (const unsigned quota = CgroupCpuQuota(); quota > 0)
-            cpus = std::min(cpus, quota);
-        return cpus;
-    }();
-    return n;
+    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) {
@@ -168,9 +234,9 @@ void SetThreadCpus(const std::vector<int> &cpus) {
 
 #else
 
-#include <algorithm>
 #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) {}