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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4

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 {

image_analysis/indexing/IndexerThreadPool.cpp

@@ -2,6 +2,10 @@
 // SPDX-License-Identifier: GPL-3.0-only
 
 #include "IndexerThreadPool.h"
+
+#include <map>
+#include <tuple>
+
 #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<int, float, float, float, float, int64_t, float, int64_t>;
+
+    struct KeptFFTIndexers {
+        std::mutex m;
+        int scopes = 0;
+        std::map<FFTIndexerKey, std::vector<std::unique_ptr<Indexer>>> 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<Indexer> 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> &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<FFTIndexerKey, std::vector<std::unique_ptr<Indexer>>> 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<Indexer> MakeFFTIndexerGPU(const IndexingSettings &settings, bool beside) {
+    if (auto kept = BorrowFFTIndexer(settings))
+        return kept;
     try {
         return std::make_unique<FFTIndexerGPU>(settings);
     } catch (const JFJochException &) {
@@ -58,6 +133,16 @@ static std::unique_ptr<Indexer> MakeFFTIndexerGPU(const IndexingSettings &settin
     cuda_trim_allocation_pool();
     return std::make_unique<FFTIndexerGPU>(settings);
 }
+
+void KeepFFTIndexers(const IndexingSettings &settings, int count) {
+    if (get_gpu_count() == 0)
+        return;
+    std::vector<std::unique_ptr<Indexer>> 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<std::mutex> 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();

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<CudaFFTPlanScope> 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;

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<double>(
+                        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;
         {

rugnux/RugnuxLive.cpp

@@ -3,7 +3,14 @@
 
 #include "RugnuxLive.h"
 
+#include <future>
+
 #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<const RugnuxModel> 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<std::future<void>> 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
+}

rugnux/RugnuxLive.h

@@ -19,3 +19,11 @@ struct RugnuxLiveSetup {
 
 RugnuxLiveSetup MakeRugnuxLiveSetup(const JFJochReaderDataset &dataset, int nthreads,
                                     std::shared_ptr<const RugnuxModel> 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);

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<Coord> 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