From 3befddd6493340a8072114a8f8cea5e94dd0b220 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 11 Oct 2026 11:11:42 +0200 Subject: [PATCH] rugnux: an FFT indexer that does not fit waits for the GPU work beside it, then is built again Several pipeline copies of one run build FFT indexer pools on the card at once - the rotation-scale walk's indexing probes start as the output pass begins its image loop, and the speculative geometry probe runs beside the merge - and each indexer is a cuFFT plan plus 0.6-0.9 GB of buffers. Construction had no admission control: a cufftPlanMany or device allocation that found the card full failed straight through, and since IndexerResult carries the exception that ends the run. With 7.2 GB of the card held by another process, a small-molecule rotation set failed 4 of 4 runs this way although each pass fits on its own. Construction now does what the merge's allocations do: hand back the pool's unused memory, wait for the work beside it to end, and build again. Who waits for whom is fixed so that nothing waits in a circle: - the main line waits for the probes (GPUWorkBeside, as the merge already does); - a probe waits for the main line's image pass (new GPUImagePass), never for a probe, itself included. The image pass waits for no probe, so it ends. The walk's indexing probes now hold a GPUWorkBeside too (they are exactly that work), and a run carries gpu_work_beside_ so its pools know which side they are on. Only when an indexer is built changes, never what it computes; with nothing to wait for it still fails at once. Same run with the 7.2 GB hog: 5/5 complete (also at 7.6 and 8.0 GB), p.hkl identical to the unhogged run; a card that cannot hold the run still fails in 0.5 s. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- common/CUDAWrapper.cpp | 47 +++++++++++---- common/CUDAWrapper.h | 15 +++++ image_analysis/indexing/IndexerThreadPool.cpp | 50 ++++++++++++---- image_analysis/indexing/IndexerThreadPool.h | 9 ++- rugnux/Rugnux.cpp | 9 ++- rugnux/Rugnux.h | 4 ++ rugnux/RugnuxFirstPass.cpp | 2 +- rugnux/RugnuxFirstPassRescues.cpp | 6 +- rugnux/RugnuxPasses.cpp | 7 +++ rugnux/RugnuxStills.cpp | 2 +- tests/CUDAMemHelpersTest.cpp | 60 +++++++++++++++++++ 11 files changed, 181 insertions(+), 30 deletions(-) diff --git a/common/CUDAWrapper.cpp b/common/CUDAWrapper.cpp index dc4d31102..3eb2ba1b8 100644 --- a/common/CUDAWrapper.cpp +++ b/common/CUDAWrapper.cpp @@ -28,27 +28,50 @@ std::string get_gpu_description() { } namespace { - std::mutex gpu_work_beside_mutex; - std::condition_variable gpu_work_beside_done; - int gpu_work_beside = 0; + // How many of one kind of GPU work are running, for a thread to wait until none is. + struct GPUWorkCount { + std::mutex m; + std::condition_variable none; + int n = 0; + void Add() { + std::lock_guard lock(m); + ++n; + } + void Remove() { + { + std::lock_guard lock(m); + --n; + } + none.notify_all(); + } + }; + GPUWorkCount gpu_work_beside, gpu_image_passes; } GPUWorkBeside::GPUWorkBeside() { - std::lock_guard lock(gpu_work_beside_mutex); - ++gpu_work_beside; + gpu_work_beside.Add(); } GPUWorkBeside::~GPUWorkBeside() { - { - std::lock_guard lock(gpu_work_beside_mutex); - --gpu_work_beside; - } - gpu_work_beside_done.notify_all(); + gpu_work_beside.Remove(); } bool wait_for_gpu_work_beside(std::chrono::seconds timeout) { - std::unique_lock lock(gpu_work_beside_mutex); - return gpu_work_beside_done.wait_for(lock, timeout, [] { return gpu_work_beside == 0; }); + std::unique_lock lock(gpu_work_beside.m); + return gpu_work_beside.none.wait_for(lock, timeout, [] { return gpu_work_beside.n == 0; }); +} + +GPUImagePass::GPUImagePass() { + gpu_image_passes.Add(); +} + +GPUImagePass::~GPUImagePass() { + gpu_image_passes.Remove(); +} + +void wait_for_gpu_image_passes() { + std::unique_lock lock(gpu_image_passes.m); + gpu_image_passes.none.wait(lock, [] { return gpu_image_passes.n == 0; }); } #ifndef JFJOCH_USE_CUDA diff --git a/common/CUDAWrapper.h b/common/CUDAWrapper.h index 83e02bca8..3d29f349b 100644 --- a/common/CUDAWrapper.h +++ b/common/CUDAWrapper.h @@ -96,3 +96,18 @@ public: // Wait until no GPUWorkBeside is held, for at most `timeout`. True once none is - at once if none was. bool wait_for_gpu_work_beside(std::chrono::seconds timeout); + +// The image pass of a run's main line - not of a probe holding a GPUWorkBeside: the stretch that holds +// the most of the card, and that never waits for a probe, so a probe that does not fit beside it can +// wait for it to end. Hold one for as long as the image pass runs. Build-independent - it only counts. +class GPUImagePass { +public: + GPUImagePass(); + ~GPUImagePass(); + GPUImagePass(const GPUImagePass &) = delete; + GPUImagePass &operator=(const GPUImagePass &) = delete; +}; + +// Wait until no GPUImagePass is held - at once if none was. With no timeout: an image pass waits for no +// probe, so it ends. +void wait_for_gpu_image_passes(); diff --git a/image_analysis/indexing/IndexerThreadPool.cpp b/image_analysis/indexing/IndexerThreadPool.cpp index 4afe38a97..ec9e750e3 100644 --- a/image_analysis/indexing/IndexerThreadPool.cpp +++ b/image_analysis/indexing/IndexerThreadPool.cpp @@ -24,12 +24,40 @@ void WarmUpCuFFT() { #endif } +#ifdef JFJOCH_USE_CUDA +// An FFT indexer is a cuFFT plan and ~0.6-0.9 GB of device buffers, and several pipeline copies of one +// process run on the card at the same time: probe passes (GPUWorkBeside) beside the main line, for a +// few seconds each. So one that does not fit hands back what the allocation pool holds unused and waits +// for the work beside it to end, then is built again - as the merge's allocations do. The main line +// waits for the probes; a probe waits for the main line's image pass (GPUImagePass), never for a probe, +// itself included. Neither of those waits for the other, so nothing waits in a circle. Only when the +// indexer is built depends on this, never what it computes; one that does not fit the card fails as +// before. +static std::unique_ptr MakeFFTIndexerGPU(const IndexingSettings &settings, bool beside) { + try { + return std::make_unique(settings); + } catch (const JFJochException &) { + cuda_throw_if_context_lost(); // not memory: waiting would only hide the real error + cuda_clear_error(); + } + if (beside) + wait_for_gpu_image_passes(); + else + // That work is a probe pass of seconds, so this long means it is stuck: the indexer fails as + // it would have without it. + wait_for_gpu_work_beside(std::chrono::minutes(10)); + cuda_trim_allocation_pool(); + return std::make_unique(settings); +} +#endif + // The indexer for one RESOLVED algorithm, or nullptr if this build/host cannot serve it. -static std::unique_ptr MakeIndexer(IndexingAlgorithmEnum algorithm, const IndexingSettings &settings) { +static std::unique_ptr MakeIndexer(IndexingAlgorithmEnum algorithm, const IndexingSettings &settings, + bool beside) { #ifdef JFJOCH_USE_CUDA if (get_gpu_count() > 0) { if (algorithm == IndexingAlgorithmEnum::FFT) - return std::make_unique(settings); + return MakeFFTIndexerGPU(settings, beside); if (algorithm == IndexingAlgorithmEnum::FFBIDX) return std::make_unique(); } @@ -41,8 +69,9 @@ static std::unique_ptr MakeIndexer(IndexingAlgorithmEnum algorithm, con return nullptr; } -IndexerThread::IndexerThread(const IndexingSettings &settings, int threadid, IndexerConstruction construction) - : settings_(settings), construction_(construction) { +IndexerThread::IndexerThread(const IndexingSettings &settings, int threadid, IndexerConstruction construction, + bool beside) + : settings_(settings), construction_(construction), beside_(beside) { std::unique_lock lock(m); state = TaskState::STARTING; worker_thread = std::thread(&IndexerThread::Worker, this, threadid); @@ -74,14 +103,14 @@ void IndexerThread::Worker(int threadid) { const auto requested = settings_.GetAlgorithm(); if (requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFT) - fft_indexer = MakeIndexer(IndexingAlgorithmEnum::FFT, settings_); + fft_indexer = MakeIndexer(IndexingAlgorithmEnum::FFT, settings_, beside_); if (requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFBIDX) - ffbidx_indexer = MakeIndexer(IndexingAlgorithmEnum::FFBIDX, settings_); + ffbidx_indexer = MakeIndexer(IndexingAlgorithmEnum::FFBIDX, settings_, beside_); if ((requested == IndexingAlgorithmEnum::Auto && get_gpu_count() == 0) || requested == IndexingAlgorithmEnum::FFTW) - fftw_indexer = MakeIndexer(IndexingAlgorithmEnum::FFTW, settings_); + fftw_indexer = MakeIndexer(IndexingAlgorithmEnum::FFTW, settings_, beside_); } catch (const std::exception &e) { spdlog::error("Failed to initialize indexer: {}", e.what()); { @@ -134,7 +163,7 @@ void IndexerThread::Worker(int threadid) { // A preconstructing worker already holds it; an OnFirstUse worker builds it here, // on the first frame that resolves to this algorithm. if (slot && !*slot) - *slot = MakeIndexer(algorithm, settings_); + *slot = MakeIndexer(algorithm, settings_, beside_); if (!slot || !*slot) { // Algorithm is already resolved here (never Auto/None - see @@ -208,14 +237,15 @@ IndexerThread::~IndexerThread() { Finalize(); } -IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings, IndexerConstruction construction) +IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings, IndexerConstruction construction, + bool beside) : worker_busy(settings.GetIndexingThreads(), 0), worker_free_count(settings.GetIndexingThreads()), viable_cell_min_spots(settings.GetViableCellMinSpots()), blocking(settings.GetBlockingBehavior()), settings_(settings) { for (size_t i = 0; i < settings.GetIndexingThreads(); ++i) - tasks.emplace_back(std::make_unique(std::cref(settings), i, construction)); + tasks.emplace_back(std::make_unique(std::cref(settings), i, construction, beside)); } int IndexerThreadPool::GetFreeWorker() { diff --git a/image_analysis/indexing/IndexerThreadPool.h b/image_analysis/indexing/IndexerThreadPool.h index 35f4beccb..709f984a3 100644 --- a/image_analysis/indexing/IndexerThreadPool.h +++ b/image_analysis/indexing/IndexerThreadPool.h @@ -52,6 +52,7 @@ class IndexerThread { // IndexerThreadPool(experiment.GetIndexingSettings()), which returns by value. const IndexingSettings settings_; const IndexerConstruction construction_; + const bool beside_; bool stop = false; enum class TaskState {STARTING, IDLE, READY, RUNNING, COMPLETED, ERROR} state = TaskState::STARTING; @@ -65,7 +66,7 @@ class IndexerThread { void Worker(int threadid); public: - IndexerThread(const IndexingSettings& settings, int threadid, IndexerConstruction construction); + IndexerThread(const IndexingSettings& settings, int threadid, IndexerConstruction construction, bool beside); ~IndexerThread(); std::unique_ptr Run(const DiffractionExperiment &experiment, const std::vector &recip, bool severity_only = false); @@ -83,8 +84,12 @@ class IndexerThreadPool { const IndexingSettings settings_; int GetFreeWorker(); public: + // beside: the pool serves a probe pass that holds GPUWorkBeside, so an FFT indexer that does not fit + // the card waits for the main line's image pass to end, not for the probes - it is one of them (see + // MakeFFTIndexerGPU). IndexerThreadPool(const IndexingSettings& settings, - IndexerConstruction construction = IndexerConstruction::Preconstruct); + IndexerConstruction construction = IndexerConstruction::Preconstruct, + bool beside = false); // The settings the pool's indexers were built with. [[nodiscard]] const IndexingSettings &Settings() const { return settings_; } // severity_only skips indexing and produces just the spindle severity - see Indexer::Run. diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 4d5e408e2..2338f2efd 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -435,6 +435,7 @@ void Rugnux::StartNearTieRecheck(std::shared_ptr run, bool geometry_prep run->speculative_probe_.reset(); run->speculate_geometry_probe_ = false; run->rotation_scale_check_.reset(); + run->gpu_work_beside_ = false; // it runs after the run, on its own std::lock_guard lock(near_tie_rechecks_->m); near_tie_rechecks_->pending.push_back({std::move(run), geometry_prepass, first_pass_outcome_}); } @@ -503,6 +504,7 @@ void Rugnux::StartSpeculativeGeometryProbe(std::shared_ptr run, const st run->postrefine_probe_ = false; run->postrefine_probe_only_ = false; run->rerun_if_starved_ = false; + run->gpu_work_beside_ = true; // The walk's other probe, at the geometry in hand, on a copy of its own: where this pass's own // first pass has not left its memo the walk would index there too, after the pass. auto hand = std::make_shared(*run); @@ -873,7 +875,8 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b RaiseFFTBoundForKnownCell(indexing_settings, experiment_); } indexing_bound_A = indexing_settings.GetFFT_MaxUnitCell_A(); - indexer_pool = std::make_unique(indexing_settings, IndexerConstruction::OnFirstUse); + indexer_pool = std::make_unique(indexing_settings, IndexerConstruction::OnFirstUse, + gpu_work_beside_); indexer = std::make_unique(experiment_, indexer_pool.get()); // With no per-image file to write, nothing reads the message's reflection list - except the // geometry pre-pass, which measures the spots against it (SpotFootprint.h). @@ -986,7 +989,11 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b } TimingMark image_pass_mark(logger, "image pass"); + std::optional gpu_image_pass; // what a probe that does not fit beside it waits for + if (!gpu_work_beside_) + gpu_image_pass.emplace(); ImagePass(p); + gpu_image_pass.reset(); image_pass_mark.Stop(); // Scaling and merging (full analysis only). diff --git a/rugnux/Rugnux.h b/rugnux/Rugnux.h index a476554ac..682ef13af 100644 --- a/rugnux/Rugnux.h +++ b/rugnux/Rugnux.h @@ -1045,6 +1045,10 @@ class Rugnux { }; std::shared_ptr speculative_probe_; bool speculate_geometry_probe_ = false; + // Whether this run is a probe pass running beside the main line under a GPUWorkBeside - the + // speculative probe, the rotation-scale walk's indexing probes. Its indexer pools then wait for the + // main line's image pass, never for that work, being that work themselves (IndexerThreadPool). + bool gpu_work_beside_ = false; // run: the copy of this run the probe runs on, taken by the caller on the thread that owns the // run - the post-refinement that decides whether to start it runs beside the merge, which does // not leave the run's members alone. diff --git a/rugnux/RugnuxFirstPass.cpp b/rugnux/RugnuxFirstPass.cpp index e1bd20d49..9afa3cacd 100644 --- a/rugnux/RugnuxFirstPass.cpp +++ b/rugnux/RugnuxFirstPass.cpp @@ -1552,7 +1552,7 @@ void Rugnux::FirstPassRun::StartShortAxisSpeculation() { sp.spots_per_image = first_pass_spots_per_image; sp.experiment = std::make_unique(rx.experiment_); sp.pool = std::make_unique(ShortAxisSettings(*sp.experiment), - IndexerConstruction::OnFirstUse); + IndexerConstruction::OnFirstUse, rx.gpu_work_beside_); for (const auto &[name, ordinals] : schemes) sp.sris.push_back(FeedScheme(*sp.pool, ordinals, sp.experiment.get())); for (auto &ri : sp.sris) { diff --git a/rugnux/RugnuxFirstPassRescues.cpp b/rugnux/RugnuxFirstPassRescues.cpp index 765007839..c48a3b68e 100644 --- a/rugnux/RugnuxFirstPassRescues.cpp +++ b/rugnux/RugnuxFirstPassRescues.cpp @@ -184,7 +184,7 @@ auto Rugnux::FirstPassRun::LongAxisRescue(const FirstPass &standing) -> std::opt coarse_settings.FFT_MaxUnitCell_A(RESCUE_MAX_UNIT_CELL_A); coarse_settings.IndexingThreads(2); coarse_settings.RefineThreads(std::clamp(rx.config_.nthreads / 2, 1, MAX_REFINE_THREADS)); - IndexerThreadPool coarse_pool(coarse_settings, IndexerConstruction::OnFirstUse); + IndexerThreadPool coarse_pool(coarse_settings, IndexerConstruction::OnFirstUse, rx.gpu_work_beside_); // Coarse first pass: keep the recovered cell with the LONGEST axis directly. Its full- // resolution per-frame validation would be low (the coarse cell is metrically right but // imprecise), so do NOT score it here - it is only the reference for the constrained re-index. @@ -213,7 +213,7 @@ auto Rugnux::FirstPassRun::LongAxisRescue(const FirstPass &standing) -> std::opt RaiseFFTBoundForKnownCell(wide_settings, rx.experiment_); FirstPass constrained; if (wide_settings.GetFFT_MaxUnitCell_A() > indexing_bound_A) { - IndexerThreadPool wide_pool(wide_settings, IndexerConstruction::OnFirstUse); + IndexerThreadPool wide_pool(wide_settings, IndexerConstruction::OnFirstUse, rx.gpu_work_beside_); constrained = PickBest(wide_pool, *indexer); } else { constrained = PickBest(*indexer_pool, *indexer); @@ -808,7 +808,7 @@ auto Rugnux::FirstPassRun::ShortAxisHypothesis(FirstPass best) -> FirstPass { } else { short_spec.reset(); short_pool = std::make_unique(ShortAxisSettings(rx.experiment_), - IndexerConstruction::OnFirstUse); + IndexerConstruction::OnFirstUse, rx.gpu_work_beside_); sris = FeedSchemes(*short_pool); RunIndexing(sris); } diff --git a/rugnux/RugnuxPasses.cpp b/rugnux/RugnuxPasses.cpp index 93bcbb6bb..9ed951fab 100644 --- a/rugnux/RugnuxPasses.cpp +++ b/rugnux/RugnuxPasses.cpp @@ -666,6 +666,13 @@ RotationScaleWalk Rugnux::WalkRotationScaleBeside(const ProcessResult &pass1, Ro // to take, which the validation frames of the whole sweep then judge - so it is fitted on // the frames pass 1 fitted the first one on. r->prepass_frames_only_ = !index_only && !decide_late_; + // An indexing probe runs beside the refined passes, on the same card, and gives its memory + // back in seconds: what of theirs does not fit beside it waits for it to end. A refit + // runs only once they have ended, so it is nobody's work beside. + std::optional beside; + if (index_only) + beside.emplace(); + r->gpu_work_beside_ = index_only; ProcessResult res; try { res = r->RunPipeline(nullptr, /*write_output=*/false, /*geometry_prepass=*/false); diff --git a/rugnux/RugnuxStills.cpp b/rugnux/RugnuxStills.cpp index 2c938876d..b710ffcbb 100644 --- a/rugnux/RugnuxStills.cpp +++ b/rugnux/RugnuxStills.cpp @@ -117,7 +117,7 @@ void Rugnux::RefineStillsGeometry(int start_image, int end_image, int images_to_ // First-pass engines, built from the current (nominal) geometry: an index-only pass (no // integration) to obtain each frame's spots + assigned HKL + orientation. AzimuthalIntegrationMapping mapping(experiment_, pixel_mask_); - IndexerThreadPool pool(experiment_.GetIndexingSettings(), IndexerConstruction::OnFirstUse); + IndexerThreadPool pool(experiment_.GetIndexingSettings(), IndexerConstruction::OnFirstUse, gpu_work_beside_); IndexAndRefine indexer(experiment_, &pool, /*retain_outcomes=*/false); auto pass_settings = config_.spot_finding; diff --git a/tests/CUDAMemHelpersTest.cpp b/tests/CUDAMemHelpersTest.cpp index dad8a2f3a..87415c1ec 100644 --- a/tests/CUDAMemHelpersTest.cpp +++ b/tests/CUDAMemHelpersTest.cpp @@ -15,6 +15,7 @@ #include #include "../image_analysis/indexing/CUDAMemHelpers.h" +#include "../image_analysis/indexing/IndexerThreadPool.h" // The broker starts fresh worker threads for every data collection. A stream per thread ever started // is a device-memory leak that ends a long-running broker, so threads that follow one another have to @@ -103,6 +104,65 @@ TEST_CASE("CudaDevicePtr_HandledPoolFailureLeavesNoError", "[CUDAMemHelpers]") { cudaMemPoolDestroy(capped); } +// An FFT indexer that finds the card full waits for the GPU work beside it to give its memory back, and +// is built then: the main line's for the probes, a probe's for the main line's image pass - never for a +// probe, itself included. With nothing to wait for it fails at once, as one that does not fit always has. +TEST_CASE("IndexerThreadPool_FFTIndexerWaitsForGPUWorkBeside", "[CUDAMemHelpers]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + + IndexingSettings settings; + settings.Algorithm(IndexingAlgorithmEnum::FFT).IndexingThreads(1); + DiffractionExperiment x; + x.ImportIndexingSettings(settings); + const std::vector recip{{0.1, 0, 0}, {0, 0.1, 0}, {0, 0, 0.1}}; + + // The whole card, held by the work beside. + std::optional>> held; + const auto fill = [&held] { + held.emplace(); + 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; + } + } + }; + IndexerThreadPool probe(settings, IndexerConstruction::OnFirstUse, /*beside=*/true); + IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse); + + fill(); + { + GPUWorkBeside work; + CHECK(probe.Run(x, recip).error != nullptr); + } + { + std::optional image_pass(std::in_place); + std::thread ends([&] { + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + held.reset(); + image_pass.reset(); + }); + GPUWorkBeside work; + CHECK(probe.Run(x, recip).error == nullptr); + ends.join(); + } + + fill(); + { + std::optional work(std::in_place); + std::thread ends([&] { + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + held.reset(); + work.reset(); + }); + CHECK(pool.Run(x, recip).error == nullptr); + ends.join(); + } +} + #endif // A merge that finds the card full waits for the GPU work running beside it: until the last one