GPU memory: beam-centre point transforms in two spectra and one work area; a probe on a small card indexes with one FFT indexer

- BeamCenterFFTGPU's point surfaces: each spectrum product goes in the place
  of a spectrum that is finished with (a.a where m's spectrum goes before it
  is there, a.m and m.m and a^2.m in a's), so the third spectrum buffer goes,
  and the forward and inverse plans share one work area - they never run at
  once. The same transforms and products in the same order, so the surfaces
  are unchanged. 2.6 -> 2.05 GB on a 16 Mpx detector, the largest single
  holder at the first pass's peak; DeviceMemoryNeeded counts it so.
- Rugnux::FirstPassIndexerThreads: a probe pass beside the main line on a card
  below the take-turns size (gpu_takes_turns(), 12 GB on the first card's
  total) builds one FFT indexer for its pools instead of two, and indexes its
  two schemes one after the other. Which indexer indexes a frame does not
  change what it finds.

Measured with the RTX 5080's total read as 8.55 GB: first-pass peak 6.50 ->
5.92 GB (8qaw), 5.76 -> 4.84 GB (4nwv) together with the batched FFT indexer;
the speculative probe after the output 6.30 -> 4.35 GB (8qaw). Output
unchanged: p.hkl byte-identical on cytc_x10sa, myob_x10sa, 4nwv and 8qaw on
this card and on the 8.55 GB simulation; run times unchanged within noise.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

common/CUDAWrapper.cpp

@@ -56,6 +56,10 @@ namespace {
     }
 }
 
+bool gpu_takes_turns() {
+    return take_turns();
+}
+
 GPUWorkBeside::GPUWorkBeside(bool after_merges) {
     std::unique_lock lock(gpu_work_m);
     if (take_turns())

common/CUDAWrapper.h

@@ -154,6 +154,10 @@ public:
     GPUMerge &operator=(const GPUMerge &) = delete;
 };
 
+// Whether the card is too small for a probe pass beside the main line (12 GB, decided on the first card's
+// total memory), so that they take turns (GPUWorkBeside). Build-independent.
+bool gpu_takes_turns();
+
 // 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();

image_analysis/geom_refinement/BeamCenterFFTGPU.cu

@@ -274,10 +274,13 @@ PointSurfacesDevice TransformPoint(const std::vector<float> &a, const std::vecto
     PointSurfacesDevice out;
     out.nx = nx;
     out.real = CudaDevicePtr<float>(nreal, ALLOC);
-    CudaDevicePtr<cufftComplex> d_a(ncomplex, ALLOC), d_m(ncomplex, ALLOC), d_prod(ncomplex, ALLOC);
-
-    CudaFFTPlan fwd(PointForward(ny, nx));
-    CudaFFTPlan inv(PointInverse(ny, nx));
+    // Two spectra, a product taking the place of one that is finished with (see below); and one work
+    // area for both plans, which never run at the same time.
+    CudaDevicePtr<cufftComplex> d_a(ncomplex, ALLOC), d_m(ncomplex, ALLOC);
+    CudaDevicePtr<char> work(std::max<size_t>({CudaFFTPlan::WorkBytes(PointForward(ny, nx)),
+                                               CudaFFTPlan::WorkBytes(PointInverse(ny, nx)), 1}), ALLOC);
+    CudaFFTPlan fwd(PointForward(ny, nx), nullptr, work.get());
+    CudaFFTPlan inv(PointInverse(ny, nx), nullptr, work.get());
 
     auto upload = [&](const std::vector<float> &src) {
         Check(cudaMemset(out.real.get(), 0, nreal * sizeof(float)),
@@ -287,17 +290,14 @@ PointSurfacesDevice TransformPoint(const std::vector<float> &a, const std::vecto
               "BeamCenterFFT: image upload failed");
     };
 
-    upload(a);
-    Check(cufftExecR2C(fwd, out.real, d_a), "BeamCenterFFT: forward transform failed");
-    upload(m);
-    Check(cufftExecR2C(fwd, out.real, d_m), "BeamCenterFFT: forward transform failed");
-
     const float norm = 1.0f / (static_cast<float>(ny) * static_cast<float>(nx));
-    // The inverse leaves an unnormalised ny x nx result whose 2h x 2w corner is the surface.
-    auto inverse_product = [&](const cufftComplex *u, const cufftComplex *v) {
-        MultiplySpectra<<<Blocks(ncomplex), BLOCK>>>(u, v, d_prod.get(), ncomplex);
+    // The product goes to `dst`, which may be u itself, and the inverse - which leaves an unnormalised
+    // ny x nx result whose 2h x 2w corner is the surface - may overwrite it, as cuFFT's C2R does with
+    // its input.
+    auto inverse_product = [&](const cufftComplex *u, const cufftComplex *v, cufftComplex *dst) {
+        MultiplySpectra<<<Blocks(ncomplex), BLOCK>>>(u, v, dst, ncomplex);
         CheckLastKernel("BeamCenterFFT: spectrum product failed");
-        Check(cufftExecC2R(inv, d_prod, out.real), "BeamCenterFFT: inverse transform failed");
+        Check(cufftExecC2R(inv, dst, out.real), "BeamCenterFFT: inverse transform failed");
     };
     auto take_corner = [&](CudaDevicePtr<float> &dst) {
         dst = CudaDevicePtr<float>(ncrop, ALLOC);
@@ -306,18 +306,23 @@ PointSurfacesDevice TransformPoint(const std::vector<float> &a, const std::vecto
         CheckLastKernel("BeamCenterFFT: surface scaling failed");
     };
 
-    inverse_product(d_a, d_a);
+    // a.a goes where m's spectrum will be, before it is there; a.m in the place of a's spectrum, the last
+    // to read it; m.m there too; and the spectrum of a^2 there after that.
+    upload(a);
+    Check(cufftExecR2C(fwd, out.real, d_a), "BeamCenterFFT: forward transform failed");
+    inverse_product(d_a, d_a, d_m);
     take_corner(out.C);
-    inverse_product(d_a, d_m);
+    upload(m);
+    Check(cufftExecR2C(fwd, out.real, d_m), "BeamCenterFFT: forward transform failed");
+    inverse_product(d_a, d_m, d_a);
     take_corner(out.S);
-    inverse_product(d_m, d_m);
+    inverse_product(d_m, d_m, d_a);
     take_corner(out.D);
-    // The spectrum of a is finished with, so a^2 goes in its place rather than in a fourth buffer.
     upload(a);
     SquareInPlace<<<Blocks(nreal), BLOCK>>>(out.real.get(), nreal);
     CheckLastKernel("BeamCenterFFT: squaring failed");
     Check(cufftExecR2C(fwd, out.real, d_a), "BeamCenterFFT: forward transform failed");
-    inverse_product(d_a, d_m);
+    inverse_product(d_a, d_m, d_a);
     ScaleCorner<<<Blocks(ncrop), BLOCK>>>(out.real.get(), nx, 2 * h, 2 * w, norm);
     CheckLastKernel("BeamCenterFFT: surface scaling failed");
     return out;
@@ -511,11 +516,11 @@ size_t BeamCenterFFTGPU::DeviceMemoryNeeded(int64_t width, int64_t height) {
     const int64_t nxc = nx / 2 + 1;
     const size_t work_fwd = CudaFFTPlan::WorkBytes(PointForward(ny, nx));
     const size_t work_inv = CudaFFTPlan::WorkBytes(PointInverse(ny, nx));
-    // One transform buffer, three spectra, and three of the four convolution surfaces - the
-    // fourth stays in the transform buffer, where the last inverse leaves it.
-    size_t needed = ny * nx * sizeof(float) + 3 * ny * nxc * sizeof(cufftComplex)
-                    + 3 * 4 * static_cast<size_t>(width) * height * sizeof(float) + work_fwd
-                    + work_inv;
+    // One transform buffer, two spectra, three of the four convolution surfaces - the fourth stays in
+    // the transform buffer, where the last inverse leaves it - and one work area for both plans.
+    size_t needed = ny * nx * sizeof(float) + 2 * ny * nxc * sizeof(cufftComplex)
+                    + 3 * 4 * static_cast<size_t>(width) * height * sizeof(float)
+                    + std::max(work_fwd, work_inv);
 
     // The two 1D phases, the larger of them: they run after the point surfaces, with its buffers
     // already freed.

rugnux/Rugnux.cpp

@@ -536,6 +536,10 @@ void Rugnux::AddFirstPassMemo(const FirstPassMemo &memo) {
     first_pass_memo_.push_back(memo);
 }
 
+int Rugnux::FirstPassIndexerThreads() const {
+    return gpu_work_beside_ && gpu_takes_turns() ? 1 : 2;
+}
+
 void Rugnux::StartSpeculativeGeometryProbe(std::shared_ptr<Rugnux> run, const std::array<float, 5> &geometry) {
     // What the walk's index_at_geometry sets up on its copy of the run before it indexes at this geometry.
     run->prepass_detector_geometry_ = prepass_detector_geometry_;
@@ -924,7 +928,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
         // those two so a rotation run doesn't pay to initialise (cuFFT plan + GPU alloc) indexers it
         // never uses. Stills fire the FFT once per image across all workers, so keep the default there.
         if (config_.rotation_indexing && config_.two_pass_rotation && experiment_.IsRotationIndexing()) {
-            indexing_settings.IndexingThreads(2);
+            indexing_settings.IndexingThreads(FirstPassIndexerThreads());
             // Two indexer threads leave the rest of the machine idle for the whole first pass, which is
             // a third of the run. Split each scheme's candidate-cell refinement over -N/2 threads so the
             // two together use -N. The split is by candidate cell, so the result does not depend on it.

rugnux/Rugnux.h

@@ -1274,6 +1274,11 @@ class Rugnux {
     // How many workers with GPU engines of their own - at most `most` - one card takes: decided from
     // the card's total memory and the detector, never from what is free (RugnuxImagePass.cpp).
     [[nodiscard]] int GPUWorkersPerCard(int most) const;
+    // The FFT indexers of a first pass's pool: two, for its two schemes at once; one for a probe pass on a
+    // card too small for it beside the main line, which then indexes the schemes one after the other.
+    // An FFT indexer holds 0.3-0.6 GB of the card, and which one indexes a frame does not change what it
+    // finds.
+    [[nodiscard]] int FirstPassIndexerThreads() const;
     // The pre-scan's workers (RugnuxPreScan.cpp).
     [[nodiscard]] size_t PrescanWorkers() const;
     // After the image pass of a rotation sweep: the decisions taken on the sweep's first part

rugnux/RugnuxFirstPass.cpp

@@ -1475,7 +1475,7 @@ void Rugnux::FirstPassRun::PrefetchCertainFrames() {
 IndexingSettings Rugnux::FirstPassRun::ShortAxisSettings(const DiffractionExperiment &x) const {
     auto settings = x.GetIndexingSettings();
     settings.FFT_MinUnitCell_A(IndexingSettings::fft_min_unit_cell_limit_A);
-    settings.IndexingThreads(2);
+    settings.IndexingThreads(rx.FirstPassIndexerThreads());
     settings.RefineThreads(std::clamp(rx.config_.nthreads / 2, 1, MAX_REFINE_THREADS));
     return settings;
 }

rugnux/RugnuxFirstPassRescues.cpp

@@ -182,7 +182,7 @@ auto Rugnux::FirstPassRun::LongAxisRescue(const FirstPass &standing) -> std::opt
     constexpr float RESCUE_MAX_UNIT_CELL_A = 1000.0f;
     if (coarse_settings.GetFFT_MaxUnitCell_A() < RESCUE_MAX_UNIT_CELL_A)
         coarse_settings.FFT_MaxUnitCell_A(RESCUE_MAX_UNIT_CELL_A);
-    coarse_settings.IndexingThreads(2);
+    coarse_settings.IndexingThreads(rx.FirstPassIndexerThreads());
     coarse_settings.RefineThreads(std::clamp(rx.config_.nthreads / 2, 1, MAX_REFINE_THREADS));
     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-
@@ -208,7 +208,7 @@ auto Rugnux::FirstPassRun::LongAxisRescue(const FirstPass &standing) -> std::opt
         // CountIndexed forces the lattice and never touches the pool, so the existing
         // IndexAndRefine is reused; only the FFT needs the wider one.
         auto wide_settings = rx.experiment_.GetIndexingSettings();
-        wide_settings.IndexingThreads(2);
+        wide_settings.IndexingThreads(rx.FirstPassIndexerThreads());
         wide_settings.RefineThreads(std::clamp(rx.config_.nthreads / 2, 1, MAX_REFINE_THREADS));
         RaiseFFTBoundForKnownCell(wide_settings, rx.experiment_);
         FirstPass constrained;