From b43844ae064059fb26bee9547b3641f5765feed4 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Fri, 9 Oct 2026 12:51:17 +0200 Subject: [PATCH] rugnux: the merging engines sit on their card's NUMA node while they drive it The image-loop workers are kept on the CPUs of their card's NUMA node (pin_gpu), but the two threads that drive the merging tail were not: the quality guard took its card with set_gpu(1), and the main engine runs on the pipeline's own thread, which goes on to everything else. On a multi-socket server their pageable copies to the card went through a remote socket. RotationScaleMerge now holds a ScopedGpuPin (new, CUDAWrapper) for each IngestHost, IngestDevice and Run: pin_gpu(gpu_device) for the call, and afterwards the thread's previous device and CPU mask back (GetThreadCpus/SetThreadCpus, new in ThreadAffinity). So the binding covers whichever thread drives an engine, and does not stay on the pipeline thread. The threads the engine starts itself (ComputeAsuGroups, Combine, FitErrorModel, ApplyFrenchWilson use std::async) call RestoreThreadAffinity first, as the pool's workers do, so they are not squeezed onto one node. The guard thread, which does nothing but the guard's merge, uses pin_gpu(1) instead of set_gpu(1). Pool workers are untouched: they restore the start mask themselves; the new test checks both that and that a pinned caller keeps its own mask. No flag, no environment variable. On one NUMA node (this workstation) only the device setting happens, so nothing can move. Measured on the 5080 workstation (one card, one node), base = rc175-next 41fab6156, --model dropped, under the GPU lock: p.mtz md5 identical on myob_x10sa, cytc_x10sa, thau_x10sa_16keV, 8a1a, 11if, 7qis, 6h5t, kdp_x10sa_20keV at default -N (2 runs), -N 8 and CUDA_VISIBLE_DEVICES=0. Tail ("scale, merge and symmetry") unchanged within noise, e.g. 8a1a 30.06 -> 29.37 s (base runs 29.75/30.36), cytc 5.33 -> 5.35 s. The gain is for multi-socket boxes (T4 server: guard copies 1.7-2.1 vs 3.5-4.3 GB/s) and is not measured here. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- common/CUDAWrapper.cpp | 4 +++ common/CUDAWrapper.cu | 17 ++++++++++++ common/CUDAWrapper.h | 15 +++++++++++ common/ThreadAffinity.cpp | 22 ++++++++++++++++ common/ThreadAffinity.h | 8 ++++++ image_analysis/scale_merge/ErrorModel.cpp | 7 ++++- image_analysis/scale_merge/FrenchWilson.cpp | 6 ++++- .../scale_merge/RotationScaleMerge.cpp | 10 +++++++ rugnux/RugnuxScaleMerge.cpp | 2 +- tests/ParallelForTest.cpp | 26 +++++++++++++++++++ 10 files changed, 114 insertions(+), 3 deletions(-) diff --git a/common/CUDAWrapper.cpp b/common/CUDAWrapper.cpp index 3579a98e8..4c6f53f25 100644 --- a/common/CUDAWrapper.cpp +++ b/common/CUDAWrapper.cpp @@ -67,6 +67,10 @@ void pin_gpu() {} void pin_gpu(int32_t dev_id) {} +ScopedGpuPin::ScopedGpuPin(int32_t) {} + +ScopedGpuPin::~ScopedGpuPin() {} + void enable_gpu_numa_binding() {} void set_gpu_blocking_sync() {} diff --git a/common/CUDAWrapper.cu b/common/CUDAWrapper.cu index 2cfedf594..440f5eb60 100644 --- a/common/CUDAWrapper.cu +++ b/common/CUDAWrapper.cu @@ -91,6 +91,23 @@ void pin_gpu(int32_t dev_id) { PinThreadToNumaNode(GpuNumaNode(dev_id)); } +ScopedGpuPin::ScopedGpuPin(int32_t dev_id) { + if (get_gpu_count() == 0) + return; + cuda_err(cudaGetDevice(&prev_device)); + const int node = gpu_numa_binding ? GpuNumaNode(dev_id) : -1; + if (node >= 0) + prev_cpus = GetThreadCpus(); + set_gpu(dev_id); + PinThreadToNumaNode(node); +} + +ScopedGpuPin::~ScopedGpuPin() { + SetThreadCpus(prev_cpus); + if (prev_device >= 0) + cudaSetDevice(prev_device); // not set_gpu: a destructor must not throw +} + void pin_gpu() { static std::atomic counter{0}; auto dev_count = get_gpu_count(); diff --git a/common/CUDAWrapper.h b/common/CUDAWrapper.h index aa298ea7e..d713176f3 100644 --- a/common/CUDAWrapper.h +++ b/common/CUDAWrapper.h @@ -33,6 +33,21 @@ void enable_gpu_numa_binding(); // For a worker that takes its card by index rather than round-robin. void pin_gpu(int32_t dev_id); +// pin_gpu(dev_id) for as long as the object lives; then the calling thread gets back the device and +// the CPUs it had before. For a thread that drives a card only for a while - the merging engine, whose +// calls come from a thread that goes on to other work - where a pin for the rest of its life would +// carry the card's node onto everything that thread does, and starts, afterwards. +class ScopedGpuPin { +public: + explicit ScopedGpuPin(int32_t dev_id); + ~ScopedGpuPin(); + ScopedGpuPin(const ScopedGpuPin &) = delete; + ScopedGpuPin &operator=(const ScopedGpuPin &) = delete; +private: + int32_t prev_device = -1; // -1: nothing to give back + std::vector prev_cpus; // empty: the CPUs were not changed +}; + // Have every GPU's host threads BLOCK in a synchronisation (cudaDeviceScheduleBlockingSync) instead // of spinning on a core until the device finishes. Must be called before anything creates a CUDA // context; a device that already has one keeps the flags it was created with. No-op without CUDA. diff --git a/common/ThreadAffinity.cpp b/common/ThreadAffinity.cpp index cd4d9d52c..8e59d117d 100644 --- a/common/ThreadAffinity.cpp +++ b/common/ThreadAffinity.cpp @@ -146,6 +146,26 @@ void RestoreThreadAffinity() { pthread_setaffinity_np(pthread_self(), sizeof(start_mask), &start_mask); } +std::vector GetThreadCpus() { + cpu_set_t mask; + std::vector 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 &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 @@ -155,5 +175,7 @@ unsigned AvailableCpus() { return std::max(1u, std::thread::hardware_concurrency int NumaNodeOfPciDevice(const std::string &) { return -1; } void PinThreadToNumaNode(int) {} void RestoreThreadAffinity() {} +std::vector GetThreadCpus() { return {}; } +void SetThreadCpus(const std::vector &) {} #endif diff --git a/common/ThreadAffinity.h b/common/ThreadAffinity.h index 3af3232a4..aa7a02a8f 100644 --- a/common/ThreadAffinity.h +++ b/common/ThreadAffinity.h @@ -4,6 +4,7 @@ #pragma once #include +#include // 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 @@ -29,3 +30,10 @@ void PinThreadToNumaNode(int node); // 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 GetThreadCpus(); + +// Keep the calling thread on exactly these CPUs. Nothing happens for an empty list. +void SetThreadCpus(const std::vector &cpus); diff --git a/image_analysis/scale_merge/ErrorModel.cpp b/image_analysis/scale_merge/ErrorModel.cpp index 9ae291bbc..77999d826 100644 --- a/image_analysis/scale_merge/ErrorModel.cpp +++ b/image_analysis/scale_merge/ErrorModel.cpp @@ -9,6 +9,7 @@ #include #include "../../common/ParallelFor.h" +#include "../../common/ThreadAffinity.h" ErrorModelFit FitErrorModel(const std::vector &pool, std::vector &smp, size_t nthreads) { @@ -46,7 +47,11 @@ ErrorModelFit FitErrorModel(const std::vector &pool, std::vect std::nth_element(smp.begin() + lo, smp.begin() + k, smp.begin() + hi, by_snr); if (depth > 0 && hi - lo > SPLIT_MIN_PARALLEL) { auto left = std::async(std::launch::async, - [&, lo, k, b0, mid, depth] { split(lo, k, b0, mid, depth - 1); }); + [&, lo, k, b0, mid, depth] { + // Not the node of a merging engine that called this (ScopedGpuPin). + RestoreThreadAffinity(); + split(lo, k, b0, mid, depth - 1); + }); split(k, hi, mid + 1, b1, depth - 1); left.get(); } else { diff --git a/image_analysis/scale_merge/FrenchWilson.cpp b/image_analysis/scale_merge/FrenchWilson.cpp index c50c987d2..5971a4348 100644 --- a/image_analysis/scale_merge/FrenchWilson.cpp +++ b/image_analysis/scale_merge/FrenchWilson.cpp @@ -10,6 +10,7 @@ #include #include "../../common/ResolutionShells.h" +#include "../../common/ThreadAffinity.h" #include "gemmi/math.hpp" #include "gemmi/symmetry.hpp" @@ -230,7 +231,10 @@ void ApplyFrenchWilson(std::vector &merged, const gemmi::Space for (int t = 0; t < nt; ++t) { const int lo = t * chunk, hi = std::min(n, lo + chunk); if (lo >= hi) break; - futures.emplace_back(std::async(std::launch::async, [&do_chunk, lo, hi] { do_chunk(lo, hi); })); + futures.emplace_back(std::async(std::launch::async, [&do_chunk, lo, hi] { + RestoreThreadAffinity(); // not the node of a merging engine that called this (ScopedGpuPin) + do_chunk(lo, hi); + })); } for (auto &f : futures) f.get(); } diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index d346c4c74..57e35e070 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only +#include "../../common/CUDAWrapper.h" #include "../../common/ParallelFor.h" #include "RotationScaleMerge.h" #include "RotationScaleMergeGPU.h" // SurfaceTerm, the surface fit's term on both paths @@ -380,6 +381,10 @@ void RotationScaleMerge::Ingest() { } void RotationScaleMerge::IngestHost() { + // The engine's calls sit by its card, as the image loop's workers do: on the card's NUMA node, so + // the copies to it stay on that socket. Only for the call - the caller goes on to other work - and + // the threads the engine starts itself are given the process's CPUs back (RestoreThreadAffinity). + const ScopedGpuPin pin(gpu_device); TimingMark mark(*logger, "rsm ingest"); scale_partials_choice.reset(); n_frames = static_cast(partials_out.size()); @@ -702,6 +707,7 @@ void RotationScaleMerge::StageObservations() { #endif void RotationScaleMerge::IngestDevice() { + const ScopedGpuPin pin(gpu_device); // see IngestHost TimingMark mark(*logger, "rsm ingest (device)"); #ifdef JFJOCH_USE_CUDA if (resident_ingest) @@ -1398,6 +1404,7 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { for (int t = 0; t < n_hist; ++t) { const int lo = t * chunk, hi = std::min(n, lo + chunk); f1.emplace_back(std::async(std::launch::async, [&, t, lo, hi] { + RestoreThreadAffinity(); // not the engine's node (see IngestHost) int32_t *h = row(t); std::fill(h, h + n_groups, 0); for (int i = lo; i < hi; ++i) { const int g = group_ids[i]; if (g >= 0) ++h[g]; } @@ -1436,6 +1443,7 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { for (int t = 0; t < n_hist; ++t) { const int lo = t * chunk, hi = std::min(n, lo + chunk); f2.emplace_back(std::async(std::launch::async, [&, t, lo, hi] { + RestoreThreadAffinity(); // not the engine's node (see IngestHost) int32_t *fill = row(t); for (int i = lo; i < hi; ++i) { const int g = group_ids[i]; if (g >= 0) gperm[fill[g]++] = i; } })); @@ -4469,6 +4477,7 @@ void RotationScaleMerge::Combine() { const int r0 = t * chunk, r1 = std::min(n_run, r0 + chunk); if (r0 >= r1) break; futures.emplace_back(std::async(std::launch::async, [&, t, r0, r1] { + RestoreThreadAffinity(); // not the engine's node (see IngestHost) part[t].reserve(r1 - r0); for (int r = r0; r < r1; ++r) used[t] += process_rawrun(r, part[t], nullptr, overloaded[t]); })); @@ -6175,6 +6184,7 @@ namespace { RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_stats, bool measure_cc_before_corrections, const gemmi::SpaceGroup *group) { + const ScopedGpuPin pin(gpu_device); // see IngestHost run_group = group; struct ClearGroup { const gemmi::SpaceGroup *&g; ~ClearGroup() { g = nullptr; } } clear_group{run_group}; // Whether a frame's scale can be fitted on its partials depends on whether the partials tell a change diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index 3197f3765..f9207f744 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -759,7 +759,7 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { if (merge_first_part && !(postrefine_probe_ && !geometry_prepass && postrefine_probe_only_)) guard_ahead = std::async(std::launch::async, [&] { if (guard_on_second_gpu) - set_gpu(1); // for any device work of the merge that is not the engine's own + pin_gpu(1); // the card and its NUMA node, as an image-loop worker has them return merge_guard(); }); // The scaling engine's ingest reads the outcomes and nothing else either, so its host half runs diff --git a/tests/ParallelForTest.cpp b/tests/ParallelForTest.cpp index 74790440c..750cf5d99 100644 --- a/tests/ParallelForTest.cpp +++ b/tests/ParallelForTest.cpp @@ -8,9 +8,12 @@ #include #include #include +#include #include +#include "../common/CUDAWrapper.h" #include "../common/ParallelFor.h" +#include "../common/ThreadAffinity.h" namespace { const std::vector THREAD_COUNTS = {0, 1, 2, 3, 8, 64}; @@ -186,3 +189,26 @@ TEST_CASE("ParallelFor_PropagatesAnException", "[ParallelFor]") { if (i == 0) throw std::runtime_error("from an item"); }), std::runtime_error); } + +#ifdef __linux__ +// The merging engine's ScopedGpuPin keeps its thread by the card only while it lives: a thread kept to +// one CPU before is kept to that CPU after, and the pool's workers keep the process's CPUs while a +// narrowed thread hands them work. +TEST_CASE("ScopedGpuPin_GivesTheThreadItsCpusBack", "[ParallelFor]") { + const std::vector all = GetThreadCpus(); + REQUIRE(all.size() >= 2); + std::vector after_pin; + std::vector> seen_by_task(4); + std::thread([&] { + SetThreadCpus({all.front()}); + { + const ScopedGpuPin pin(0); + ParallelChunks(4, 4, [&](int lo, int) { seen_by_task[lo] = GetThreadCpus(); }); + } + after_pin = GetThreadCpus(); + }).join(); + CHECK(after_pin == std::vector{all.front()}); + for (int t = 1; t < 4; t++) + CHECK(seen_by_task[t] == all); +} +#endif