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
This commit is contained in:
2026-10-11 09:27:32 +02:00
co-authored by Claude Opus 5.5
parent 54f8e07591
commit 3c5ff3f5df
7 changed files with 190 additions and 12 deletions
+65
View File
@@ -212,3 +212,68 @@ TEST_CASE("ScopedGpuPin_GivesTheThreadItsCpusBack", "[ParallelFor]") {
CHECK(seen_by_task[t] == all);
}
#endif
namespace {
std::vector<int> CpuRange(int lo, int hi) { // [lo, hi]
std::vector<int> v(hi - lo + 1);
std::iota(v.begin(), v.end(), lo);
return v;
}
std::vector<int> Concat(std::vector<int> a, const std::vector<int> &b) {
a.insert(a.end(), b.begin(), b.end());
return a;
}
}
// Synthetic machines: which CPUs a GPU run keeps to, given the CPUs it may use, the package of every
// CPU and the package of every card.
TEST_CASE("CpusOfGpuPackages_KeepsToTheCardsPackages", "[ParallelFor]") {
// Four packages of 12 CPUs, four cards on packages 2 and 3.
std::vector<int> four(48);
for (int c = 0; c < 48; c++)
four[c] = c / 12;
CHECK(CpusOfGpuPackages(CpuRange(0, 47), four, {2, 2, 3, 3}) == CpuRange(24, 47));
// Two packages of 12 cores with hyper-threading, numbered as Linux does: CPUs 0-11 and 24-35 on
// package 0, 12-23 and 36-47 on package 1.
std::vector<int> two(48);
for (int c = 0; c < 48; c++)
two[c] = (c / 12) % 2;
const std::vector<int> package1 = Concat(CpuRange(12, 23), CpuRange(36, 47));
CHECK(CpusOfGpuPackages(CpuRange(0, 47), two, {1}) == package1);
CHECK(CpusOfGpuPackages(CpuRange(0, 47), two, {1, 1}) == package1);
// Cards on both packages: nothing to keep to.
CHECK(CpusOfGpuPackages(CpuRange(0, 47), two, {0, 1}).empty());
// No card (a CPU build, or none visible), or a card whose package is not known.
CHECK(CpusOfGpuPackages(CpuRange(0, 47), two, {}).empty());
CHECK(CpusOfGpuPackages(CpuRange(0, 47), two, {1, -1}).empty());
// One package, as on a desktop.
CHECK(CpusOfGpuPackages(CpuRange(0, 31), std::vector<int>(32, 0), {0}).empty());
}
// A process started on fewer CPUs (taskset, numactl, a cpuset) keeps within them.
TEST_CASE("CpusOfGpuPackages_RespectsARestrictedMask", "[ParallelFor]") {
std::vector<int> four(48);
for (int c = 0; c < 48; c++)
four[c] = c / 12;
// Already on one package - with the cards or without them - stays where it is.
CHECK(CpusOfGpuPackages(CpuRange(0, 11), four, {2, 3}).empty());
CHECK(CpusOfGpuPackages(CpuRange(24, 35), four, {2, 3}).empty());
// Already on the cards' packages only.
CHECK(CpusOfGpuPackages(CpuRange(24, 47), four, {2, 3}).empty());
// Over packages 1 and 2, cards on 2 and 3: the part on package 2.
CHECK(CpusOfGpuPackages(CpuRange(18, 29), four, {2, 3}) == CpuRange(24, 29));
}
// Without a card nothing is confined, so the process keeps the CPUs it started with.
TEST_CASE("ConfineToGpuPackages_DoesNothingWithoutACard", "[ParallelFor]") {
const unsigned before = AvailableCpus();
const std::vector<int> cpus = GetThreadCpus();
CHECK(ConfineToGpuPackages({}).empty());
CHECK(AvailableCpus() == before);
CHECK(GetThreadCpus() == cpus);
}