From 19ee20fbe83a6e06bfb20f25090f70ac8d6396c2 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 19 Oct 2026 11:56:28 +0200 Subject: [PATCH] rugnux live: keep FFT indexers and cuFFT plans for the broker's lifetime, built at START jfjoch_broker opened no FFTPlanReuse, so every cuFFT plan of a live run was made on first use and destroyed after it, and every pass (first pass, measured-centre arm, short-axis pool, pass 2) built its GPU FFT indexers afresh. Now: - FFTIndexerReuse: while one lives, a pool worker borrows a kept GPU FFT indexer of its card and construction settings instead of building one, and gives it back when it ends or is released. - The rugnux receiver holds an FFTIndexerReuse and an FFTPlanReuse for the broker's lifetime. - At START, beside the run (which is waiting for its first frames), WarmUpRugnuxLiveGPU builds on every card the first-pass and short-axis indexers and the beam-centre capture's plan pair. The CLI opens no FFTIndexerReuse and is unchanged. Results identical: CLI p.hkl md5 equal on two sets (CLI and --simulate-live --fast-live); broker replay, 3 collections each, MTZ data identical. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- image_analysis/indexing/CUDAMemHelpers.h | 4 +- image_analysis/indexing/IndexerThreadPool.cpp | 94 ++++++++++++++++++- image_analysis/indexing/IndexerThreadPool.h | 23 ++++- receiver/JFJochReceiverRugnux.cpp | 21 +++++ rugnux/RugnuxLive.cpp | 34 +++++++ rugnux/RugnuxLive.h | 8 ++ tests/CUDAMemHelpersTest.cpp | 58 ++++++++++++ 7 files changed, 237 insertions(+), 5 deletions(-) diff --git a/image_analysis/indexing/CUDAMemHelpers.h b/image_analysis/indexing/CUDAMemHelpers.h index f0b2128bf..5c9d680d7 100644 --- a/image_analysis/indexing/CUDAMemHelpers.h +++ b/image_analysis/indexing/CUDAMemHelpers.h @@ -357,8 +357,8 @@ struct CudaFFTShape { // same shape. A plan is never shared by two users at once - cuFFT is thread safe only across distinct // plans. // -// The scope is one processing run (rugnux), so nothing is kept from one run for the next, nor past the -// device reset at the end of a run. A kept plan holds no work area: the user supplies one (CudaFFTPlan), +// In `rugnux` the scope is one processing run, so nothing is kept from one run for the next, nor past the +// device reset at the end of a run; jfjoch_broker, which resets no device, keeps one for its lifetime. A kept plan holds no work area: the user supplies one (CudaFFTPlan), // so an idle plan holds only its tables and modules. namespace jfjoch_cufft_plans { struct Plan { diff --git a/image_analysis/indexing/IndexerThreadPool.cpp b/image_analysis/indexing/IndexerThreadPool.cpp index 2707482ad..53f0b2d82 100644 --- a/image_analysis/indexing/IndexerThreadPool.cpp +++ b/image_analysis/indexing/IndexerThreadPool.cpp @@ -2,6 +2,10 @@ // SPDX-License-Identifier: GPL-3.0-only #include "IndexerThreadPool.h" + +#include +#include + #include "../common/CUDAWrapper.h" #include "../common/Logger.h" @@ -33,7 +37,76 @@ FFTPlanReuse::FFTPlanReuse() = default; FFTPlanReuse::~FFTPlanReuse() = default; +#ifndef JFJOCH_USE_CUDA +FFTIndexerReuse::FFTIndexerReuse() = default; +FFTIndexerReuse::~FFTIndexerReuse() = default; +void KeepFFTIndexers(const IndexingSettings &, int) {} +#endif + #ifdef JFJOCH_USE_CUDA +namespace { + // The card and what FFTIndexer's constructor reads from the settings. + using FFTIndexerKey = std::tuple; + + struct KeptFFTIndexers { + std::mutex m; + int scopes = 0; + std::map>> idle; + }; + + KeptFFTIndexers &kept_fft_indexers() { + static auto *k = new KeptFFTIndexers; + return *k; + } + + FFTIndexerKey KeyOnCurrentCard(const IndexingSettings &s) { + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Failed to read the CUDA device"); + return {device, s.GetFFT_MaxUnitCell_A(), s.GetFFT_MinUnitCell_A(), s.GetFFT_MinAngle_deg(), + s.GetFFT_MaxAngle_deg(), s.GetFFT_NumVectors(), s.GetFFT_HighResolution_A(), s.GetRefineThreads()}; + } + + std::unique_ptr BorrowFFTIndexer(const IndexingSettings &settings) { + auto &k = kept_fft_indexers(); + std::lock_guard lock(k.m); + auto it = k.idle.find(KeyOnCurrentCard(settings)); + if (it == k.idle.end() || it->second.empty()) + return nullptr; + auto ret = std::move(it->second.back()); + it->second.pop_back(); + return ret; + } + + // On the card the indexer was built on, which is the calling worker's. + void GiveBackFFTIndexer(const IndexingSettings &settings, std::unique_ptr &indexer) { + if (!indexer) + return; + auto &k = kept_fft_indexers(); + std::lock_guard lock(k.m); + if (k.scopes > 0) + k.idle[KeyOnCurrentCard(settings)].push_back(std::move(indexer)); + } +} + +FFTIndexerReuse::FFTIndexerReuse() { + std::lock_guard lock(kept_fft_indexers().m); + kept_fft_indexers().scopes++; +} + +FFTIndexerReuse::~FFTIndexerReuse() { + std::map>> idle; + { + std::lock_guard lock(kept_fft_indexers().m); + if (--kept_fft_indexers().scopes == 0) + idle.swap(kept_fft_indexers().idle); + } + for (auto &[key, list] : idle) { + ScopedGpuPin pin(std::get<0>(key)); + list.clear(); + } +} + // 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 @@ -43,6 +116,8 @@ FFTPlanReuse::~FFTPlanReuse() = default; // 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) { + if (auto kept = BorrowFFTIndexer(settings)) + return kept; try { return std::make_unique(settings); } catch (const JFJochException &) { @@ -58,6 +133,16 @@ static std::unique_ptr MakeFFTIndexerGPU(const IndexingSettings &settin cuda_trim_allocation_pool(); return std::make_unique(settings); } + +void KeepFFTIndexers(const IndexingSettings &settings, int count) { + if (get_gpu_count() == 0) + return; + std::vector> indexers; + for (int i = 0; i < count; i++) + indexers.push_back(MakeFFTIndexerGPU(settings, false)); + for (auto &i : indexers) + GiveBackFFTIndexer(settings, i); +} #endif // The indexer for one RESOLVED algorithm, or nullptr if this build/host cannot serve it. @@ -152,8 +237,12 @@ void IndexerThread::Worker(int threadid) { { std::unique_lock lock(m); c_start.wait(lock, [this] { return stop || drop_indexers || state == TaskState::READY; }); - if (stop && (state != TaskState::READY)) + if (stop && (state != TaskState::READY)) { +#ifdef JFJOCH_USE_CUDA + GiveBackFFTIndexer(settings_, fft_indexer); +#endif return; + } if (state == TaskState::READY) { state = TaskState::RUNNING; input = std::move(task_input); @@ -161,6 +250,9 @@ void IndexerThread::Worker(int threadid) { drop = true; } if (drop) { +#ifdef JFJOCH_USE_CUDA + GiveBackFFTIndexer(settings_, fft_indexer); +#endif fft_indexer.reset(); ffbidx_indexer.reset(); fftw_indexer.reset(); diff --git a/image_analysis/indexing/IndexerThreadPool.h b/image_analysis/indexing/IndexerThreadPool.h index ea011f9cb..61b6a145a 100644 --- a/image_analysis/indexing/IndexerThreadPool.h +++ b/image_analysis/indexing/IndexerThreadPool.h @@ -41,8 +41,8 @@ enum class IndexerConstruction { Preconstruct, OnFirstUse }; void WarmUpCuFFT(); // cuFFT plans are kept and reused for as long as one of these lives, and destroyed when the last one -// ends (CudaFFTPlanScope): one per processing run, so no plan outlives the run that made it. Nothing -// without CUDA. +// ends (CudaFFTPlanScope): one per processing run in `rugnux`, so no plan outlives the run that made it; +// one for its lifetime in jfjoch_broker. Nothing without CUDA. class CudaFFTPlanScope; class FFTPlanReuse { std::unique_ptr scope; @@ -53,6 +53,25 @@ public: FFTPlanReuse &operator=(const FFTPlanReuse &) = delete; }; +// GPU FFT indexers are kept for reuse while one of these is alive, by card and the settings they are +// built from, and destroyed when the last one ends. A pool's worker then borrows a kept indexer instead +// of building one, and gives it back - not to the card - when it ends or is released (ReleaseIndexers): +// the device memory stays taken in between. For jfjoch_broker's rugnux live runs, which build the same +// pools for every pass of every collection; an indexer computes the same with any earlier frames behind +// it, as it does within a pool. Nothing without CUDA. +class FFTIndexerReuse { +public: + FFTIndexerReuse(); + ~FFTIndexerReuse(); + FFTIndexerReuse(const FFTIndexerReuse &) = delete; + FFTIndexerReuse &operator=(const FFTIndexerReuse &) = delete; +}; + +// Builds `count` GPU FFT indexers of these settings on the current card at the same time, and keeps them +// for the pools that ask for them later (inside an FFTIndexerReuse; outside one they are only built and +// destroyed). Nothing without a GPU. +void KeepFFTIndexers(const IndexingSettings &settings, int count); + class IndexerThread { struct TaskInput { const DiffractionExperiment &experiment; diff --git a/receiver/JFJochReceiverRugnux.cpp b/receiver/JFJochReceiverRugnux.cpp index 17b707d42..4e1da466d 100644 --- a/receiver/JFJochReceiverRugnux.cpp +++ b/receiver/JFJochReceiverRugnux.cpp @@ -16,6 +16,7 @@ #include "../common/ParallelFor.h" #include "../frame_serialize/CBORStream2Deserializer.h" #include "../frame_serialize/CBORStream2Serializer.h" +#include "../image_analysis/indexing/IndexerThreadPool.h" #include "../rugnux/RugnuxLive.h" #include "../rugnux/ResultReport.h" #include "../writer/RugnuxResultBlocks.h" @@ -139,6 +140,12 @@ JFJochReceiverRugnux::JFJochReceiverRugnux(const DiffractionExperiment &in_exper live_status(in_live_status), measurement_started(1), forward_started(forward_nthreads) { + // GPU FFT indexers and cuFFT plans are kept for reuse for the life of the broker, not of one run as + // in `rugnux` (FFTIndexerReuse, FFTPlanReuse): every pass of every collection builds the same pools, + // and making a plan takes a process-wide lock and took over a second on a card busy with the online + // analysis. Never destroyed - at exit the CUDA runtime may already be gone. + [[maybe_unused]] static const FFTIndexerReuse *fft_indexers = new FFTIndexerReuse; + [[maybe_unused]] static const FFTPlanReuse *fft_plans = new FFTPlanReuse; // The ParallelFor pool is made by the first thread to use it and its workers take that thread's // priority. Made here, at the broker's own priority, before the rugnux thread lowers its own - so // the pool is lowered for the analysis only where it can be raised back (see RugnuxThread). @@ -304,6 +311,20 @@ void JFJochReceiverRugnux::RugnuxThread() { const auto dataset = reader.GetDataset(); RugnuxLiveSetup setup = MakeRugnuxLiveSetup(*dataset, rugnux_nthreads, model); setup.config.host_memory_budget_bytes = host_budget; + // The cards are idle until the first images arrive: the indexers and plans the passes will + // need are built then, beside the run, which meanwhile waits for its first frames. Waited for + // when the run ends. + auto gpu_warm_up = std::async(std::launch::async, [this, setup] { + try { + const auto t0 = std::chrono::steady_clock::now(); + WarmUpRugnuxLiveGPU(setup); + logger.Info("rugnux live: FFT indexers and plans ready in {:.2f} s", std::chrono::duration( + std::chrono::steady_clock::now() - t0).count()); + } catch (const std::exception &e) { + logger.Warning("rugnux live: building the FFT indexers ahead failed, they are built when used - {}", + e.what()); + } + }); LiveObserver observer(plots, scan_result, *az_int_mapping, live_status, logger); ProcessResult result; { diff --git a/rugnux/RugnuxLive.cpp b/rugnux/RugnuxLive.cpp index c5902a9e0..886db699c 100644 --- a/rugnux/RugnuxLive.cpp +++ b/rugnux/RugnuxLive.cpp @@ -3,7 +3,14 @@ #include "RugnuxLive.h" +#include + #include "RugnuxDefaults.h" +#include "../common/CUDAWrapper.h" +#include "../image_analysis/indexing/IndexerThreadPool.h" +#ifdef JFJOCH_USE_CUDA +#include "../image_analysis/geom_refinement/BeamCenterFFTGPU.h" +#endif RugnuxLiveSetup MakeRugnuxLiveSetup(const JFJochReaderDataset &dataset, int nthreads, std::shared_ptr model) { @@ -55,3 +62,30 @@ RugnuxLiveSetup MakeRugnuxLiveSetup(const JFJochReaderDataset &dataset, int nthr c.adaptive_integration_radius = false; return s; } + +void WarmUpRugnuxLiveGPU(const RugnuxLiveSetup &setup) { +#ifdef JFJOCH_USE_CUDA + // As Rugnux sets them for its first-pass pool (the measured centre's indexing and pass 2 use it as + // well) and its short-axis pool. Should they part, the run builds its own, as without this. + IndexingSettings first_pass = setup.experiment.GetIndexingSettings(); + first_pass.RefineThreads(std::clamp(setup.config.nthreads / 2, 1, MAX_REFINE_THREADS)); + IndexingSettings short_axis = first_pass; + short_axis.FFT_MinUnitCell_A(IndexingSettings::fft_min_unit_cell_limit_A); + + // Each pool has two workers, on consecutive cards. + const int gpus = get_gpu_count(); + const int per_card = gpus == 1 ? 2 : 1; + std::vector> cards; + for (int32_t i = 0; i < gpus; i++) + cards.push_back(std::async(std::launch::async, [&, i] { + ScopedGpuPin pin(i); + KeepFFTIndexers(first_pass, per_card); + KeepFFTIndexers(short_axis, per_card); + // Sizing the capture's work area makes its two plans (CudaFFTPlan::WorkBytes). + BeamCenterFFTGPU::DeviceMemoryNeeded(setup.experiment.GetXPixelsNumConv(), + setup.experiment.GetYPixelsNumConv()); + })); + for (auto &c: cards) + c.get(); +#endif +} diff --git a/rugnux/RugnuxLive.h b/rugnux/RugnuxLive.h index c24df0e26..cda7d1fb7 100644 --- a/rugnux/RugnuxLive.h +++ b/rugnux/RugnuxLive.h @@ -19,3 +19,11 @@ struct RugnuxLiveSetup { RugnuxLiveSetup MakeRugnuxLiveSetup(const JFJochReaderDataset &dataset, int nthreads, std::shared_ptr model); + +// Builds on every card, at the same time, the GPU FFT indexers a live run of this setup is known to +// need - its first pass's and its short-axis pass's (the settings Rugnux gives them) - and the cuFFT +// plans of its beam-centre capture, and with them the card's CUDA context and cuFFT's one-time +// initialisation. Meant for the time between the start message and the first image, while the cards +// are idle, under an FFTIndexerReuse and an FFTPlanReuse that outlive it: they are kept for the +// passes to borrow, and nothing the run computes depends on it. No-op without a GPU. +void WarmUpRugnuxLiveGPU(const RugnuxLiveSetup &setup); diff --git a/tests/CUDAMemHelpersTest.cpp b/tests/CUDAMemHelpersTest.cpp index b1aa58319..9c70f0c0f 100644 --- a/tests/CUDAMemHelpersTest.cpp +++ b/tests/CUDAMemHelpersTest.cpp @@ -389,6 +389,64 @@ TEST_CASE("IndexerThreadPool_ReleaseIndexersGivesTheCardBack", "[CUDAMemHelpers] } } +// jfjoch_broker keeps the FFT indexers from one rugnux pass and collection to the next: a pool's worker +// borrows the one kept on its card instead of building its own, gives it back when the pool ends, and +// the card gets the memory back only with the last FFTIndexerReuse. +TEST_CASE("IndexerThreadPool_FFTIndexerKeptAcrossPools", "[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}}; + + 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; + }; + + IndexerResult fresh; + { + IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse); + fresh = pool.Run(x, recip); + REQUIRE(fresh.error == nullptr); + } + + const uint64_t before = pooled(); + { + FFTIndexerReuse reuse; + for (int d = 0; d < get_gpu_count(); d++) { + ScopedGpuPin pin(d); + KeepFFTIndexers(settings, 1); + } + const uint64_t kept = pooled(); + REQUIRE(kept > before + (uint64_t(100) << 20)); + for (int i = 0; i < 2; i++) { + IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse); + const IndexerResult r = pool.Run(x, recip); + CHECK(pooled() == kept); + CHECK(r.error == nullptr); + CHECK(r.executed == fresh.executed); + CHECK(r.lattice.size() == fresh.lattice.size()); + } + CHECK(pooled() == kept); + } + CHECK(pooled() <= before + (uint64_t(16) << 20)); + 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