From b6f34db258fd46163d8d05dbbaed4e50ecdc4be0 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Tue, 13 Oct 2026 16:07:33 +0200 Subject: [PATCH 1/3] rotation merge: helper buffers wait for the probe pass beside them, as the engine's own do The intermittent "out of GPU memory" on the 3600-frame cbf.gz set (about one run in three on a 16 GB card, base included) was caught with allocation backtraces: FitErrorModelGPU's 42.8 MB in MergeAndStats failed while the speculative geometry probe (GPUWorkBeside) held its FFT indexers on the same card. The engine's own arrays (RotationScaleMergeGPU Impl::Alloc) wait for such work to end and ask again; the merge's helpers - ErrorModelGPU, WilsonOutliersGPU, PairSortGPU and the cub scratch of the engine itself - allocated directly and failed at once. CudaAlloc::SynchronousAfterWorkBeside: Synchronous, and where the card is full after the pool is trimmed, it waits for the GPU work beside the main line (wait_for_gpu_work_beside, 10 min as for the engine) and asks once more. The four merge files use it. Only the probe passes register GPUWorkBeside and they index only, so none of them merges and nothing waits for itself. When a buffer is allocated is all that changes; what is computed is not. Catch2 CudaDevicePtr_AfterWorkBesideWaitsForItsMemory: with the card filled by work beside, a Synchronous buffer fails and a SynchronousAfterWorkBeside one gets its memory once the work ends. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- image_analysis/indexing/CUDAMemHelpers.h | 16 +++++++-- image_analysis/scale_merge/ErrorModelGPU.cu | 2 +- image_analysis/scale_merge/PairSortGPU.cu | 2 +- .../scale_merge/RotationScaleMergeGPU.cu | 2 +- .../scale_merge/WilsonOutliersGPU.cu | 2 +- tests/CUDAMemHelpersTest.cpp | 33 +++++++++++++++++++ 6 files changed, 51 insertions(+), 6 deletions(-) diff --git a/image_analysis/indexing/CUDAMemHelpers.h b/image_analysis/indexing/CUDAMemHelpers.h index 43ae2ca5d..0da8b4719 100644 --- a/image_analysis/indexing/CUDAMemHelpers.h +++ b/image_analysis/indexing/CUDAMemHelpers.h @@ -293,7 +293,11 @@ inline bool cuda_memory_available(size_t bytes, size_t &free_bytes) { // for a buffer that is freed by a thread other than the one whose work read it: cudaFreeAsync orders // the free on the freeing side's allocation stream, which says nothing about kernels queued // elsewhere, while cudaFree synchronises the whole device and therefore cannot be early. -enum class CudaAlloc { Pooled, Synchronous }; +// SynchronousAfterWorkBeside is Synchronous for a step of the main line - the merge's helpers - whose +// buffers may be asked for while a probe pass runs beside it on the same card (GPUWorkBeside): one that +// does not fit waits for that work to end and is asked for once more, as the merge's own arrays are +// (RotationScaleMergeGPU). Never for the work beside itself, which would wait for itself. +enum class CudaAlloc { Pooled, Synchronous, SynchronousAfterWorkBeside }; template class CudaDevicePtr { @@ -338,9 +342,17 @@ public: // stay behind as the last error (see above). cudaGetLastError(); cuda_trim_allocation_pool(); + cudaError_t err = cudaMalloc(&ptr_, count * sizeof(T)); + if (err == cudaErrorMemoryAllocation && alloc == CudaAlloc::SynchronousAfterWorkBeside) { + cudaGetLastError(); + // That work is a probe pass of seconds, so this long means it is stuck. + wait_for_gpu_work_beside(std::chrono::minutes(10)); + cuda_trim_allocation_pool(); + err = cudaMalloc(&ptr_, count * sizeof(T)); + } // With the reason: a context lost to another kernel's fault fails here too, and must not read // as a card that is merely full. - if (const cudaError_t err = cudaMalloc(&ptr_, count * sizeof(T)); err != cudaSuccess) + if (err != cudaSuccess) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, GPUMemoryFailure("Allocating GPU memory", count * sizeof(T), err == cudaErrorMemoryAllocation, cudaGetErrorString(err))); diff --git a/image_analysis/scale_merge/ErrorModelGPU.cu b/image_analysis/scale_merge/ErrorModelGPU.cu index 3443d8e09..cd8c5a4c4 100644 --- a/image_analysis/scale_merge/ErrorModelGPU.cu +++ b/image_analysis/scale_merge/ErrorModelGPU.cu @@ -16,7 +16,7 @@ namespace { constexpr int THREADS = 256; constexpr int CHUNKS = 64; // blocks per bin in the sum kernel; fixed, so the tree is too constexpr int NB = ERROR_MODEL_BINS; - constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + constexpr CudaAlloc ALLOC = CudaAlloc::SynchronousAfterWorkBeside; void cuda_err(cudaError_t val) { if (val != cudaSuccess) diff --git a/image_analysis/scale_merge/PairSortGPU.cu b/image_analysis/scale_merge/PairSortGPU.cu index 693821158..50ef96811 100644 --- a/image_analysis/scale_merge/PairSortGPU.cu +++ b/image_analysis/scale_merge/PairSortGPU.cu @@ -13,7 +13,7 @@ namespace { constexpr int THREADS = 256; - constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + constexpr CudaAlloc ALLOC = CudaAlloc::SynchronousAfterWorkBeside; void cuda_err(cudaError_t val) { if (val != cudaSuccess) diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 62805d660..4315e4b35 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -36,7 +36,7 @@ namespace { // free while unwinding from a failed call; nor can compute-sanitizer --track-stream-ordered-races // see it, and it reported every reassigned merge buffer as a use-after-free. cudaFree // synchronises the device first. - constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + constexpr CudaAlloc ALLOC = CudaAlloc::SynchronousAfterWorkBeside; __device__ __forceinline__ double SafeInvD(double x, double fallback) { return (isfinite(x) && x != 0.0) ? 1.0 / x : fallback; diff --git a/image_analysis/scale_merge/WilsonOutliersGPU.cu b/image_analysis/scale_merge/WilsonOutliersGPU.cu index aeb84cc10..b877d07a7 100644 --- a/image_analysis/scale_merge/WilsonOutliersGPU.cu +++ b/image_analysis/scale_merge/WilsonOutliersGPU.cu @@ -25,7 +25,7 @@ namespace { constexpr int THREADS = 256; constexpr size_t OBS_PER_SHELL = 2000; // WilsonOutliers.cpp - constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + constexpr CudaAlloc ALLOC = CudaAlloc::SynchronousAfterWorkBeside; void cuda_err(cudaError_t val) { if (val != cudaSuccess) diff --git a/tests/CUDAMemHelpersTest.cpp b/tests/CUDAMemHelpersTest.cpp index 2d0047319..c349d9d88 100644 --- a/tests/CUDAMemHelpersTest.cpp +++ b/tests/CUDAMemHelpersTest.cpp @@ -206,6 +206,39 @@ TEST_CASE("IndexerThreadPool_FFTIndexerWaitsForGPUWorkBeside", "[CUDAMemHelpers] } } +// A merge helper's buffer that finds the card full while a probe pass holds the rest of it waits for +// that work to end and its memory to come back (CudaAlloc::SynchronousAfterWorkBeside); a plain +// synchronous one fails at once. +TEST_CASE("CudaDevicePtr_AfterWorkBesideWaitsForItsMemory", "[CUDAMemHelpers]") { + if (get_gpu_count() == 0) + SKIP("no GPU"); + std::vector> held; + const auto fill = [&] { + for (size_t chunk = size_t(256) << 20; chunk >= (size_t(1) << 20); chunk /= 2) + while (true) { + try { + held.emplace_back(chunk, CudaAlloc::Synchronous); + } catch (const JFJochException &) { + break; + } + } + }; + constexpr size_t BYTES = size_t(64) << 20; + { + std::optional work(std::in_place); + fill(); + CHECK_THROWS_AS(CudaDevicePtr(BYTES, CudaAlloc::Synchronous), JFJochException); + std::thread ends([&] { + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + held.clear(); + work.reset(); + }); + CHECK_NOTHROW(CudaDevicePtr(BYTES, CudaAlloc::SynchronousAfterWorkBeside)); + ends.join(); + } + cuda_trim_allocation_pool(); +} + #endif // A merge that finds the card full waits for the GPU work running beside it: until the last one From df14a1e6312499572f0edbbccb8ed249b40c4214 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Tue, 13 Oct 2026 16:07:33 +0200 Subject: [PATCH 2/3] molecular replacement: the score engines keep within a quarter of the card together Molecular replacement's structure-factor score engines were made one per group, copy count and resolution, each admitted against a quarter of the card's total, and all kept to the end of the search. On a large cell they summed to well over the card: traced with nsys on the rotated 4nwv model (which failed on every run, base included), seven engines of 1.36 GB were held beside the two rigid-body pools (5.9 GB) when a 288 MB buffer did not fit. Each engine is still admitted, or the search scored on the CPU, exactly as before, on the card's total memory. Now the engines kept on a card also stay within that quarter together: an engine that does not fit beside them takes the place of the least recently used one no score is using, and waits for a score to finish where all are in use; one taken out is made again when next asked for. A score holds its engine for as long as it uses it (shared_ptr lease). Engines compute the same numbers whenever they are made, so only when the card holds them changes. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- .../MolecularReplacement.cpp | 147 +++++++++++++----- 1 file changed, 112 insertions(+), 35 deletions(-) diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 92a71b55c..8c8884a5c 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -277,6 +278,13 @@ PlacementScore r_factors(const gemmi::AsuData> &fmodel, cons // where the first placement of that kind is scored, each with the scale fit on the GPU beside it - or none, // where there is no GPU, the cell is too small for its gridding or the two would take over a quarter of the // card. Decided on the card's total memory, so the same input always takes the same path. +// +// Together they keep within that quarter too. A search asks for one per group, copy count and resolution, +// and kept to the end of the search they summed to well over the card on a large cell (seven of 1.4 GB +// beside the rigid-body engines). So an engine that does not fit beside the ones kept takes the place of +// the least recently used of those no score is using, and waits for a score to finish where every one is +// in use; one taken out is made again where it is next asked for. Every engine computes the same numbers +// whenever it is made, so this changes when the card holds them, not what a score gives. #ifdef JFJOCH_USE_CUDA class ScoreEngines { public: @@ -290,54 +298,123 @@ public: CudaDevicePtr> fcalc, fmask; }; - Engine *get(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, double d_min) { + // The engine for this kind of placement, kept on the card for as long as the returned pointer lives. + std::shared_ptr get(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, double d_min) { if (get_gpu_count() == 0) return nullptr; size_t atoms = 0; for (const gemmi::Chain &ch : st.models[0].chains) for (const gemmi::Residue &r : ch.residues) atoms += r.atoms.size(); - Entry *entry; - { - std::lock_guard lock(m_); - entry = &engines_[{sg.xhm(), atoms, d_min}]; + std::unique_lock lock(m_); + Entry &e = engines_[{sg.xhm(), atoms, d_min}]; + e.last_use = ++clock_; + // Made by the first to ask for it, without holding up those asking for other ones. + cv_.wait(lock, [&] { return !e.making; }); + if (e.engine || e.refused) + return Lease(e); + e.making = true; + lock.unlock(); + + std::unique_ptr sf; + size_t points = 0, bytes = 0, quarter = 0; + try { + sf = std::make_unique(RigidBodyGPUEngine::CurrentDevice(), st.models[0], + st.cell, sg, d_min); + points = sf->Reflections(); + bytes = sf->DeviceBytes() + ModelScaleGPU::DeviceBytes(points) + 2 * points * sizeof(std::complex); + quarter = sf->DeviceTotalMemory() / 4; + } catch (const JFJochException &) { + cuda_clear_error(); + sf.reset(); } - // Each made once, by the first to ask for it, without holding up the others. - Entry &e = *entry; - std::call_once(e.made, [&] { - try { - auto sf = std::make_unique(RigidBodyGPUEngine::CurrentDevice(), st.models[0], - st.cell, sg, d_min); - const size_t points = sf->Reflections(); - const size_t bytes = sf->DeviceBytes() + ModelScaleGPU::DeviceBytes(points) - + 2 * points * sizeof(std::complex); - if (sf->Supported() && bytes <= sf->DeviceTotalMemory() / 4) { - sf->Reserve(); - const int device = RigidBodyGPUEngine::CurrentDevice(); - set_gpu(sf->Device()); - e.engine = std::make_unique(); - e.engine->stream = std::make_unique(); - e.engine->scale = std::make_unique(*e.engine->stream, points); - e.engine->fcalc = CudaDevicePtr>(points); - e.engine->fmask = CudaDevicePtr>(points); - e.engine->sf = std::move(sf); - set_gpu(device); - } - } catch (const JFJochException &) { - e.engine.reset(); - cuda_clear_error(); - } - }); - return e.engine.get(); + lock.lock(); + if (!sf || !sf->Supported() || bytes > quarter) { + e.refused = true; + e.making = false; + cv_.notify_all(); + return nullptr; + } + const int card = sf->Device(); + cv_.wait(lock, [&] { return MakeRoom(card, bytes, quarter); }); + resident_[card] += bytes; + lock.unlock(); + + std::shared_ptr engine; + try { + sf->Reserve(); + const int device = RigidBodyGPUEngine::CurrentDevice(); + set_gpu(card); + engine = std::make_shared(); + engine->stream = std::make_unique(); + engine->scale = std::make_unique(*engine->stream, points); + engine->fcalc = CudaDevicePtr>(points); + engine->fmask = CudaDevicePtr>(points); + engine->sf = std::move(sf); + set_gpu(device); + } catch (const JFJochException &) { + engine.reset(); + cuda_clear_error(); + } + lock.lock(); + e.making = false; + if (engine) { + e.engine = std::move(engine); + e.card = card; + e.bytes = bytes; + } else { + resident_[card] -= bytes; + e.refused = true; + } + cv_.notify_all(); + return Lease(e); } private: struct Entry { - std::once_flag made; - std::unique_ptr engine; + std::shared_ptr engine; + bool making = false, refused = false; + int card = 0; + size_t bytes = 0; + uint64_t last_use = 0; }; + + // The entry's engine for a score: the copy it holds is what marks the engine in use, and handing it back + // is what may let an engine waiting for room in. + std::shared_ptr Lease(Entry &e) { + if (!e.engine) + return nullptr; + return std::shared_ptr(e.engine.get(), [this, held = e.engine](Engine *) mutable { + { + std::lock_guard lock(m_); + held.reset(); + } + cv_.notify_all(); + }); + } + + // Whether `bytes` more fit the card's quarter beside the engines kept there, once the least recently + // used of those no score is using have been taken out. Under m_. + bool MakeRoom(int card, size_t bytes, size_t quarter) { + while (resident_[card] + bytes > quarter) { + Entry *oldest = nullptr; + for (auto &[key, entry] : engines_) + if (entry.engine && entry.card == card && entry.engine.use_count() == 1 + && (oldest == nullptr || entry.last_use < oldest->last_use)) + oldest = &entry; + if (oldest == nullptr) + return false; + resident_[card] -= oldest->bytes; + oldest->engine.reset(); + } + return true; + } + std::map, Entry> engines_; + std::map resident_; + uint64_t clock_ = 0; std::mutex m_; + std::condition_variable cv_; }; // The scale of `scaling`'s points fitted on the engine's card: FitModelScale's grid over the solvent box, or @@ -464,7 +541,7 @@ PlacementScore score_placement(const gemmi::Structure &st, const gemmi::SpaceGro const gemmi::Op *law = nullptr, const gemmi::Fractional *tncs = nullptr) { gemmi::AsuData> fmask, fmodel; #ifdef JFJOCH_USE_CUDA - ScoreEngines::Engine *gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr; + const std::shared_ptr gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr; if (gpu) gpu->sf->Compute(st.models[0], fmodel, fmask); else From 7c622f9e7ffd917cc50075d01df10161c1ec6d33 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Tue, 13 Oct 2026 16:07:33 +0200 Subject: [PATCH 3/3] rugnux: the scaling engine's card memory goes back before molecular replacement The rotation merge engine is released before model validation decides whether its GPU path fits, but a predicted model (or --mr) skips that validation and went into molecular replacement with the engine still on the card - about 10 GB on the AF 9v16 run, which then failed with 0.0 GB free on some runs and on the others got fewer rigid-body engines. It is now released before the search too, on the same condition (no finalist-ledger merge still running on it). Card peak of that run 15.75 -> 14.05 GB; MR keys identical. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- rugnux/RugnuxScaleMerge.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index aab638267..f94b76f95 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -4124,6 +4124,10 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { phase("Molecular replacement"); Step("Molecular replacement"); TimingMark mark(logger, "model prior"); + // As before the validation above, which a predicted model or --mr does not run: the engine's card + // memory goes back before the search sizes its rigid-body engines from what the card has free. + if (!ledger_ahead.valid()) + rsm.reset(); std::vector merges{{"MERGE", &sm.merged, *result.consensus_cell, experiment_.GetSpaceGroupOrP1()}}; if (twin_file) merges.push_back({"TWIN_SUBGROUP_MERGE", &twin_file->merged, twin_file->cell, twin_file->space_group,