From bf8f0a482d22f008965fe2f018292ca23c06edf1 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 16:44:01 +0200 Subject: [PATCH] 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 --- common/CUDAWrapper.cpp | 4 ++ common/CUDAWrapper.h | 4 ++ .../geom_refinement/BeamCenterFFTGPU.cu | 51 ++++++++++--------- rugnux/Rugnux.cpp | 6 ++- rugnux/Rugnux.h | 5 ++ rugnux/RugnuxFirstPass.cpp | 2 +- rugnux/RugnuxFirstPassRescues.cpp | 4 +- 7 files changed, 49 insertions(+), 27 deletions(-) diff --git a/common/CUDAWrapper.cpp b/common/CUDAWrapper.cpp index 3a71878f4..722bdc9ae 100644 --- a/common/CUDAWrapper.cpp +++ b/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()) diff --git a/common/CUDAWrapper.h b/common/CUDAWrapper.h index f772cb81f..1dd596a82 100644 --- a/common/CUDAWrapper.h +++ b/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(); diff --git a/image_analysis/geom_refinement/BeamCenterFFTGPU.cu b/image_analysis/geom_refinement/BeamCenterFFTGPU.cu index b71831cb4..f4d4be1fe 100644 --- a/image_analysis/geom_refinement/BeamCenterFFTGPU.cu +++ b/image_analysis/geom_refinement/BeamCenterFFTGPU.cu @@ -274,10 +274,13 @@ PointSurfacesDevice TransformPoint(const std::vector &a, const std::vecto PointSurfacesDevice out; out.nx = nx; out.real = CudaDevicePtr(nreal, ALLOC); - CudaDevicePtr 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 d_a(ncomplex, ALLOC), d_m(ncomplex, ALLOC); + CudaDevicePtr work(std::max({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 &src) { Check(cudaMemset(out.real.get(), 0, nreal * sizeof(float)), @@ -287,17 +290,14 @@ PointSurfacesDevice TransformPoint(const std::vector &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(ny) * static_cast(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<<>>(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<<>>(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 &dst) { dst = CudaDevicePtr(ncrop, ALLOC); @@ -306,18 +306,23 @@ PointSurfacesDevice TransformPoint(const std::vector &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<<>>(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<<>>(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(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(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. diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index d52b589ba..e6e6eb076 100644 --- a/rugnux/Rugnux.cpp +++ b/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 run, const std::array &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. diff --git a/rugnux/Rugnux.h b/rugnux/Rugnux.h index 2b679e8c3..e5cca61d9 100644 --- a/rugnux/Rugnux.h +++ b/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 diff --git a/rugnux/RugnuxFirstPass.cpp b/rugnux/RugnuxFirstPass.cpp index 588b37e49..2f4840c13 100644 --- a/rugnux/RugnuxFirstPass.cpp +++ b/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; } diff --git a/rugnux/RugnuxFirstPassRescues.cpp b/rugnux/RugnuxFirstPassRescues.cpp index c48a3b68e..ccc616dfe 100644 --- a/rugnux/RugnuxFirstPassRescues.cpp +++ b/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;