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