From 1c3a2e202cc9e3ddda38f4e4928177e7a6a09325 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 11:39:19 +0200 Subject: [PATCH] GPU memory over the run: FFT indexers given back after the last indexing; speculative probe takes turns with the merge on a small card - IndexerThreadPool::ReleaseIndexers(): each worker drops the indexers it has built (once its current frame is done) and the allocation pool is trimmed, so their memory goes back to the card, not only to the pool. A later frame builds them again, as on first use. A rotation pass calls it once its first pass is done and again when the leftover-lattice census - the last of its indexing - ends. Measured on the RTX 5080: the 1.24 GB (pre-pass) / 0.62 GB (later passes) the pass's pool held through the image pass and the scaling is gone; 4nwv's pre-pass image pass peaks 5.97 instead of 7.18 GB of device memory. - GPUMerge: the main line holds one through ScaleMergeAndSymmetry. The speculative geometry probe (GPUWorkBeside(after_merges=true)) waits for it on a card below the take-turns size (12 GB), where the probe's FFT indexers and image workers beside the merge's arrays would overfill the card - on a Windows (WDDM) card that does not fail an allocation but pages it and crawls. The merge waits for nothing, so nothing the merge waits on (the rotation-scale walk's probes) can wait in a circle with it. No change on a 16 GB card. Output unchanged: p.hkl byte-identical on cytc_x10sa, myob_x10sa, 4nwv and 8qaw on this card; also identical with the card's total read as 8.55 GB (10 / 9 image workers instead of 16) and with 4 image workers, so the worker count does not change the result. Co-Authored-By: Claude Opus 5.5 --- common/CUDAWrapper.cpp | 17 +++- common/CUDAWrapper.h | 19 ++++- image_analysis/indexing/IndexerThreadPool.cpp | 42 +++++++++- image_analysis/indexing/IndexerThreadPool.h | 8 ++ rugnux/Rugnux.cpp | 13 ++- rugnux/RugnuxFirstPassAccept.cpp | 2 + tests/CUDAMemHelpersTest.cpp | 83 +++++++++++++++++++ 7 files changed, 174 insertions(+), 10 deletions(-) diff --git a/common/CUDAWrapper.cpp b/common/CUDAWrapper.cpp index 6b5eb1272..3a71878f4 100644 --- a/common/CUDAWrapper.cpp +++ b/common/CUDAWrapper.cpp @@ -33,7 +33,7 @@ namespace { // thread to wait until none of one kind is. std::mutex gpu_work_m; std::condition_variable gpu_work_changed; - int gpu_work_beside_n = 0, gpu_image_passes_n = 0; + int gpu_work_beside_n = 0, gpu_image_passes_n = 0, gpu_merges_n = 0; // A card too small for a probe pass beside the main line's image pass, or beside another probe pass: // measured on a 16 Mpx detector, two of them together held 8.6 GB. On such a card a probe pass waits, @@ -56,10 +56,12 @@ namespace { } } -GPUWorkBeside::GPUWorkBeside() { +GPUWorkBeside::GPUWorkBeside(bool after_merges) { std::unique_lock lock(gpu_work_m); if (take_turns()) - gpu_work_changed.wait(lock, [] { return gpu_image_passes_n == 0 && gpu_work_beside_n == 0; }); + gpu_work_changed.wait(lock, [after_merges] { + return gpu_image_passes_n == 0 && gpu_work_beside_n == 0 && (!after_merges || gpu_merges_n == 0); + }); ++gpu_work_beside_n; } @@ -83,6 +85,15 @@ GPUImagePass::~GPUImagePass() { remove(gpu_image_passes_n); } +GPUMerge::GPUMerge() { + std::lock_guard lock(gpu_work_m); + ++gpu_merges_n; +} + +GPUMerge::~GPUMerge() { + remove(gpu_merges_n); +} + void wait_for_gpu_image_passes() { std::unique_lock lock(gpu_work_m); gpu_work_changed.wait(lock, [] { return gpu_image_passes_n == 0; }); diff --git a/common/CUDAWrapper.h b/common/CUDAWrapper.h index b3a4cbf1f..f772cb81f 100644 --- a/common/CUDAWrapper.h +++ b/common/CUDAWrapper.h @@ -115,10 +115,11 @@ void cuda_throw_if_context_lost(); // GPU work that runs beside the main line of a process and gives all of its device memory back when it // ends: a probe pass run on a copy of the run while the run itself scales and merges. Hold one for as // long as that work runs. On a small card it waits first for the main line's image passes (GPUImagePass) -// and for other work beside to end. Build-independent - it only counts. +// and for other work beside to end - and, with after_merges, for the main line's merges (GPUMerge). +// Build-independent - it only counts. class GPUWorkBeside { public: - GPUWorkBeside(); + explicit GPUWorkBeside(bool after_merges = false); ~GPUWorkBeside(); GPUWorkBeside(const GPUWorkBeside &) = delete; GPUWorkBeside &operator=(const GPUWorkBeside &) = delete; @@ -139,6 +140,20 @@ public: GPUImagePass &operator=(const GPUImagePass &) = delete; }; +// The scaling, merging and model validation of a run's main line - not of a probe. A probe that the main +// line never waits for (the speculative geometry probe) waits for it on a small card: the merge holds as +// much of the card as the image pass, and on a card that pages what does not fit rather than failing the +// allocation, the two together would not fail but crawl. It waits for nothing itself, so nothing that the +// main line waits for inside its merge - the rotation-scale walk's probes - can wait in a circle with it. +// Build-independent - it only counts. +class GPUMerge { +public: + GPUMerge(); + ~GPUMerge(); + GPUMerge(const GPUMerge &) = delete; + GPUMerge &operator=(const GPUMerge &) = 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 40845cd77..820020e15 100644 --- a/image_analysis/indexing/IndexerThreadPool.cpp +++ b/image_analysis/indexing/IndexerThreadPool.cpp @@ -147,14 +147,28 @@ void IndexerThread::Worker(int threadid) { while (true) { std::unique_ptr input; + bool drop = false; // Look for task + handle stop { std::unique_lock lock(m); - c_start.wait(lock, [this] { return stop || state == TaskState::READY; }); + c_start.wait(lock, [this] { return stop || drop_indexers || state == TaskState::READY; }); if (stop && (state != TaskState::READY)) return; - state = TaskState::RUNNING; - input = std::move(task_input); + if (state == TaskState::READY) { + state = TaskState::RUNNING; + input = std::move(task_input); + } else + drop = true; + } + if (drop) { + fft_indexer.reset(); + ffbidx_indexer.reset(); + fftw_indexer.reset(); + { + std::unique_lock lock(m); + drop_indexers = false; + } + c_done.notify_all(); } if (input) { std::unique_ptr tmp_result; @@ -210,6 +224,15 @@ void IndexerThread::Worker(int threadid) { } } +void IndexerThread::Release() { + std::unique_lock lock(m); + if (stop) + return; + drop_indexers = true; + c_start.notify_all(); + c_done.wait(lock, [this] { return !drop_indexers; }); +} + void IndexerThread::Finalize() { { std::unique_lock lock(m); @@ -257,6 +280,19 @@ IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings, IndexerCo tasks.emplace_back(std::make_unique(std::cref(settings), i, construction, beside)); } +void IndexerThreadPool::ReleaseIndexers() { + for (auto &t : tasks) + t->Release(); +#ifdef JFJOCH_USE_CUDA + // Their buffers went back to the allocation pool, which would keep them reserved for the next + // buffer of whoever asks; hand them to the card. + for (int32_t d = 0; d < get_gpu_count(); d++) { + ScopedGpuPin pin(d); + cuda_trim_allocation_pool(); + } +#endif +} + int IndexerThreadPool::GetFreeWorker() { std::unique_lock lock(m); diff --git a/image_analysis/indexing/IndexerThreadPool.h b/image_analysis/indexing/IndexerThreadPool.h index 7415ea7e6..ea011f9cb 100644 --- a/image_analysis/indexing/IndexerThreadPool.h +++ b/image_analysis/indexing/IndexerThreadPool.h @@ -68,6 +68,7 @@ class IndexerThread { const bool beside_; bool stop = false; + bool drop_indexers = false; // see Release() enum class TaskState {STARTING, IDLE, READY, RUNNING, COMPLETED, ERROR} state = TaskState::STARTING; std::mutex m; std::condition_variable c_running; @@ -83,6 +84,9 @@ public: ~IndexerThread(); std::unique_ptr Run(const DiffractionExperiment &experiment, const std::vector &recip, bool severity_only = false); + // The thread drops the indexers it has built - once the frame it is on, if any, is done - and builds + // them again on the next frame that needs one. Returns when they are gone. + void Release(); void Finalize(); }; @@ -108,6 +112,10 @@ public: // severity_only skips indexing and produces just the spindle severity - see Indexer::Run. IndexerResult Run(const DiffractionExperiment& experiment, const std::vector& recip, bool severity_only = false); + // Give back to the card the device memory of the indexers built so far - an FFT indexer holds + // 0.6-1.2 GB of it - for a stretch with nothing to index. A frame that comes after all builds its + // indexer again, as on first use, and the indexer it gets computes what the dropped one would. + void ReleaseIndexers(); }; diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 53f284e53..33483f7a0 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -556,9 +556,10 @@ void Rugnux::StartSpeculativeGeometryProbe(std::shared_ptr run, const st sp->hand = hand; // A probe that fails leaves no memo, and the walk's own probe then runs it again and meets the // failure itself. It runs beside the merge of the pass that started it, on the same card: a merge - // that does not fit beside it waits for it to end (GPUWorkBeside). + // that does not fit beside it waits for it to end (GPUWorkBeside) - and on a small card it waits for + // that merge to end instead (GPUMerge). sp->done = std::async(std::launch::async, [run, hand] { - GPUWorkBeside beside; + GPUWorkBeside beside(/*after_merges=*/true); for (const auto &r : {run, hand}) { try { r->RunPipeline(nullptr, /*write_output=*/false, /*geometry_prepass=*/false); @@ -1025,6 +1026,10 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b if (probe_done) return result; } + // A rotation pass indexes nothing after its first pass but the leftover-lattice census, which gives + // them back again when it ends; the image pass and the merge get the card without them. + if (indexer_pool && experiment_.IsRotationIndexing()) + indexer_pool->ReleaseIndexers(); // The rotation-scale walk beside this pass starts once its frames are indexed (see RunAllPasses). if (rotation_scale_check_) rotation_scale_check_->Set(rotation_scale_check_->output_indexed); @@ -1061,7 +1066,11 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b if (full && !cancelled_ && (result.indexing_rate.value_or(0.0f) > 0 || result.rotation_lattice_found) && (config_.run_scaling || !config_.reference_data.empty())) { TimingMark mark(logger, "scale, merge and symmetry"); + std::optional gpu_merge; // what the speculative probe waits for on a small card + if (!gpu_work_beside_) + gpu_merge.emplace(); const bool finished = ScaleMergeAndSymmetry(p); + gpu_merge.reset(); if (p.postrefine_ahead.valid()) { p.postrefine_ahead.get(); p.postrefine_log.ReplayInto(logger); diff --git a/rugnux/RugnuxFirstPassAccept.cpp b/rugnux/RugnuxFirstPassAccept.cpp index 1df384d0b..f0eb897c6 100644 --- a/rugnux/RugnuxFirstPassAccept.cpp +++ b/rugnux/RugnuxFirstPassAccept.cpp @@ -395,6 +395,8 @@ void Rugnux::FirstPassRun::StartLeftoverCensus(const FirstPass &best) { census.primary_spots_pct, n_all, census.primary_intensity_pct, census.frames, census.lattices.size(), std::chrono::duration(std::chrono::steady_clock::now() - census_start).count()); + // The last of the pass's indexing (see RunPipeline). + indexer_pool->ReleaseIndexers(); return census; }); } diff --git a/tests/CUDAMemHelpersTest.cpp b/tests/CUDAMemHelpersTest.cpp index ae944a2cd..33b59cb0f 100644 --- a/tests/CUDAMemHelpersTest.cpp +++ b/tests/CUDAMemHelpersTest.cpp @@ -4,6 +4,7 @@ #include #include "../common/CUDAWrapper.h" +#include #include #include @@ -329,6 +330,65 @@ TEST_CASE("GetGpuFreeMemory_IntegratedCountsThePageCache", "[CUDAMemHelpers]") { CHECK(double(free_bytes) == Catch::Approx(double(cuda_free)).margin(SLACK)); } + +// The pool's indexers go back to the card - not only to the allocation pool - when it is asked, and the +// next frame builds them again. +TEST_CASE("IndexerThreadPool_ReleaseIndexersGivesTheCardBack", "[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}}; + + // What the allocation pools of all cards hand out: an indexer worker takes its card round-robin. + const auto pooled = [] { + uint64_t sum = 0; + for (int d = 0; d < get_gpu_count(); d++) { + cudaMemPool_t pool = nullptr; + uint64_t used = 0; + if (cudaDeviceGetDefaultMemPool(&pool, d) == cudaSuccess + && cudaMemPoolGetAttribute(pool, cudaMemPoolAttrUsedMemCurrent, &used) == cudaSuccess) + sum += used; + } + return sum; + }; + const auto reserved = [] { + uint64_t sum = 0; + for (int d = 0; d < get_gpu_count(); d++) { + cudaMemPool_t pool = nullptr; + uint64_t r = 0; + if (cudaDeviceGetDefaultMemPool(&pool, d) == cudaSuccess + && cudaMemPoolGetAttribute(pool, cudaMemPoolAttrReservedMemCurrent, &r) == cudaSuccess) + sum += r; + } + return sum; + }; + const uint64_t before = pooled(); + const uint64_t reserved_before = reserved(); + { + IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse); + const IndexerResult first = pool.Run(x, recip); + REQUIRE(first.error == nullptr); + REQUIRE(pooled() > before + (uint64_t(100) << 20)); + + pool.ReleaseIndexers(); + CHECK(pooled() <= before + (uint64_t(16) << 20)); + CHECK(reserved() <= reserved_before + (uint64_t(64) << 20)); + + const IndexerResult again = pool.Run(x, recip); + CHECK(again.error == nullptr); + CHECK(again.executed == first.executed); + CHECK(again.lattice.size() == first.lattice.size()); + } + for (int d = 0; d < get_gpu_count(); d++) { + ScopedGpuPin pin(d); + cuda_trim_allocation_pool(); + } +} + #endif // A merge that finds the card full waits for the GPU work running beside it: until the last one @@ -347,3 +407,26 @@ TEST_CASE("GPUWorkBeside_WaitEndsWithTheWork", "[CUDAMemHelpers]") { CHECK(wait_for_gpu_work_beside(std::chrono::seconds(60))); ends.join(); } + +// The speculative geometry probe waits for the main line's merge on a card too small for both, and on a +// larger one runs beside it; the rest of the work beside never waits for a merge. +TEST_CASE("GPUWorkBeside_AfterMergesTakesTurnsOnASmallCard", "[CUDAMemHelpers]") { + // As take_turns() in CUDAWrapper.cpp. + const bool small = get_gpu_count() > 0 && get_gpu_total_memory(0) < 12'000'000'000ull; + + std::optional merge(std::in_place); + { GPUWorkBeside other; } // does not wait + + std::atomic started = false; + std::thread probe([&started] { + GPUWorkBeside beside(/*after_merges=*/true); + started = true; + }); + if (small) { + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + CHECK_FALSE(started); + } + merge.reset(); + probe.join(); + CHECK(started); +}