From 3ebf085e4c82cca8de19bc509054c7ce9b1db1a8 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 17 Oct 2026 09:31:04 +0200 Subject: [PATCH 1/7] Host memory: the machine's total, a budget decided from it, and a pool for large per-worker buffers get_host_total_memory (sysconf / hw.memsize / GlobalMemoryStatusEx), DefaultHostMemoryBudget at 75% of the total - decided from the total, never from what is free, as the GPU plans are - and HostWorkerPool, a process-wide pool (a third of the budget) that workers reserve their large buffers against before making them, so how many run at once is decided by bytes rather than threads. Only scheduling reads any of it. [HostMemory] tests. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- common/CMakeLists.txt | 1 + common/HostMemory.cpp | 77 ++++++++++++++++++++++++++++++++++++++++ common/HostMemory.h | 60 +++++++++++++++++++++++++++++++ tests/CMakeLists.txt | 1 + tests/HostMemoryTest.cpp | 48 +++++++++++++++++++++++++ 5 files changed, 187 insertions(+) create mode 100644 common/HostMemory.cpp create mode 100644 common/HostMemory.h create mode 100644 tests/HostMemoryTest.cpp diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index f8a8ff5c9..b4711ecf3 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -77,6 +77,7 @@ ADD_LIBRARY(JFJochCommon STATIC DetectorModuleGeometry.cpp DetectorModuleGeometry.h DetectorSetup.h DetectorSetup.cpp ZeroCopyReturnValue.h Histogram.h DiffractionGeometry.h CUDAWrapper.cpp CUDAWrapper.h + HostMemory.cpp HostMemory.h ThreadAffinity.cpp ThreadAffinity.h ADUHistogram.cpp ADUHistogram.h RawToConvertedGeometryCore.h diff --git a/common/HostMemory.cpp b/common/HostMemory.cpp new file mode 100644 index 000000000..54c4a8b66 --- /dev/null +++ b/common/HostMemory.cpp @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "HostMemory.h" + +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#elif defined(__APPLE__) +#include +#include +#include +#else +#include +#endif + +size_t get_host_total_memory() { +#if defined(_WIN32) + MEMORYSTATUSEX status{}; + status.dwLength = sizeof(status); + if (!GlobalMemoryStatusEx(&status)) + return 0; + return static_cast(status.ullTotalPhys); +#elif defined(__APPLE__) + uint64_t bytes = 0; + size_t len = sizeof(bytes); + if (sysctlbyname("hw.memsize", &bytes, &len, nullptr, 0) != 0) + return 0; + return static_cast(bytes); +#else + const long pages = sysconf(_SC_PHYS_PAGES), page = sysconf(_SC_PAGE_SIZE); + if (pages <= 0 || page <= 0) + return 0; + return static_cast(pages) * static_cast(page); +#endif +} + +size_t DefaultHostMemoryBudget() { + return static_cast(HOST_MEMORY_BUDGET_FRACTION * static_cast(get_host_total_memory())); +} + +HostWorkerPool &HostWorkerPool::Instance() { + static HostWorkerPool pool; + return pool; +} + +HostWorkerPool::HostWorkerPool() { + SetSize(0); +} + +void HostWorkerPool::SetSize(size_t bytes) { + std::lock_guard lock(m); + size = bytes > 0 ? bytes + : static_cast(HOST_WORKER_POOL_FRACTION * static_cast(DefaultHostMemoryBudget())); + cv.notify_all(); +} + +size_t HostWorkerPool::Size() { + std::lock_guard lock(m); + return size; +} + +HostWorkerPool::Reservation::Reservation(HostWorkerPool &pool, size_t bytes) : pool(pool), bytes(bytes) { + std::unique_lock lock(pool.m); + pool.cv.wait(lock, [&] { return pool.reserved == 0 || pool.reserved + bytes <= pool.size; }); + pool.reserved += bytes; +} + +HostWorkerPool::Reservation::~Reservation() { + { + std::lock_guard lock(pool.m); + pool.reserved -= bytes; + } + pool.cv.notify_all(); +} diff --git a/common/HostMemory.h b/common/HostMemory.h new file mode 100644 index 000000000..d016ddb58 --- /dev/null +++ b/common/HostMemory.h @@ -0,0 +1,60 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include + +// The machine's physical memory in bytes - its total, not what is free at the moment. 0 where it +// cannot be read. +size_t get_host_total_memory(); + +// The host memory a run plans against: this fraction of the machine's physical memory, decided from +// the total, never from what happens to be free (as what runs on a GPU is decided on the card's total). +// The plan decides only what runs at the same time - two merges side by side or one after the other, +// how many large per-worker buffers are held at once - and never what is computed, so no result +// depends on it. +constexpr double HOST_MEMORY_BUDGET_FRACTION = 0.75; +size_t DefaultHostMemoryBudget(); + +// The share of the budget that the large per-worker buffers of parallel steps - a translation function's +// grids, a null replicate's model maps - may hold at once. The rest is what the run holds beside them: +// its reflections, its merge, the detector-sized arrays. +constexpr double HOST_WORKER_POOL_FRACTION = 1.0 / 3.0; + +// One pool for those buffers, process-wide. A worker reserves its bytes before it makes them and hands +// them back when it is done; one that would take the pool past its size waits until others have handed +// theirs back - except that a reservation made while nothing is reserved is always let through, however +// large. So how many such workers run at once is decided by their bytes, not by the thread count, and +// however the steps holding them nest or run side by side; what each worker computes, and where its +// result goes, does not change. A worker must not reserve again while it holds a reservation, and none +// may reserve on ParallelFor's pool: a holder's own parallel passes need those threads. Work spread over +// that pool is bounded by its count instead (Size() / bytes each). +class HostWorkerPool { +public: + static HostWorkerPool &Instance(); + + // The pool's size in bytes; 0 sets it back to HOST_WORKER_POOL_FRACTION of DefaultHostMemoryBudget(). + void SetSize(size_t bytes); + [[nodiscard]] size_t Size(); + + class Reservation { + public: + Reservation(HostWorkerPool &pool, size_t bytes); + ~Reservation(); + Reservation(const Reservation &) = delete; + Reservation &operator=(const Reservation &) = delete; + private: + HostWorkerPool &pool; + size_t bytes; + }; + +private: + HostWorkerPool(); + std::mutex m; + std::condition_variable cv; + size_t size = 0; + size_t reserved = 0; +}; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c6dd46391..e5fb8f6bc 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -102,6 +102,7 @@ ADD_EXECUTABLE(jfjoch_test HotPixelFinderTest.cpp BandwidthEstimateTest.cpp ParallelForTest.cpp + HostMemoryTest.cpp RugnuxTest.cpp ResultReportTest.cpp DiagnosticOutputTest.cpp diff --git a/tests/HostMemoryTest.cpp b/tests/HostMemoryTest.cpp new file mode 100644 index 000000000..e58d17ebd --- /dev/null +++ b/tests/HostMemoryTest.cpp @@ -0,0 +1,48 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#include +#include +#include +#include +#include +#include + +#include "../common/HostMemory.h" + +TEST_CASE("HostMemory_BudgetFromTheTotal", "[HostMemory]") { + const size_t total = get_host_total_memory(); + REQUIRE(total > 0); + CHECK(DefaultHostMemoryBudget() == static_cast(HOST_MEMORY_BUDGET_FRACTION * static_cast(total))); +} + +// As many reservations at once as the pool holds, and one at a time where a single one is larger than it. +TEST_CASE("HostWorkerPool_BoundsWorkersByBytes", "[HostMemory]") { + HostWorkerPool &pool = HostWorkerPool::Instance(); + for (const auto &[size, bytes, at_most] : std::vector>{ + {100, 40, 2}, {100, 100, 1}, {100, 1000, 1}, {1000, 10, 8}}) { + CAPTURE(size, bytes); + pool.SetSize(size); + std::atomic running{0}, most{0}, done{0}; + std::vector workers; + for (int i = 0; i < 8; i++) + workers.emplace_back([&] { + const HostWorkerPool::Reservation reserved(pool, bytes); + const int now = ++running; + int seen = most.load(); + while (now > seen && !most.compare_exchange_weak(seen, now)) {} + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + --running; + ++done; + }); + for (auto &w : workers) + w.join(); + CHECK(done == 8); + CHECK(most >= 1); + CHECK(most <= at_most); + } + pool.SetSize(0); + CHECK(pool.Size() == static_cast(HOST_WORKER_POOL_FRACTION * static_cast(DefaultHostMemoryBudget()))); +} From 64a87cdf52d7c30a66746724a393d378da65bdd3 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 17 Oct 2026 09:31:04 +0200 Subject: [PATCH 2/7] CPU beam-centre FFT: two transforms at a time PointSurfaces held four padded (2h x 2w) r2c workspaces and three spectra at once - ~4.2 GiB on an 18 Mpx detector at any thread count, the Mac's peak on small proteins. Two workspaces now, the image's spectrum handed back before the squared image's is made: 24 instead of 44 bytes per padded point. Each transform is built and run as before (own buffers, own plan), so the surfaces are bit-identical. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../geom_refinement/BeamCenterFFTCPU.cpp | 53 ++++++++++--------- 1 file changed, 29 insertions(+), 24 deletions(-) diff --git a/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp b/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp index 659c5c0ca..1496752f5 100644 --- a/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp +++ b/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp @@ -83,34 +83,39 @@ public: BeamCenterConvSurfaces2D BeamCenterFFTCPU::PointSurfaces(const std::vector &a, const std::vector &m, int64_t h, int64_t w) { - // The three forward transforms are independent of each other, and so are the four inverse ones, - // so each group runs at the same time, one workspace per transform. A workspace is built exactly - // as a single one was - its own buffers and its own plan for them - so every transform is the one - // it was, bit for bit; the four inverse ones reuse the three forward workspaces and one more. + // Two transforms at a time, each on a workspace of its own. A workspace is padded to twice the + // detector each way, so on a large detector the four workspaces and three spectra that the seven + // transforms used to hold at once were gigabytes. In order: the forward transforms of the image + // and the mask; the two inverse ones that read the image's spectrum, which is then handed back; + // the squared image's forward transform beside the mask's own inverse; and the last inverse. + // A workspace is built exactly as a single one was - its own buffers and its own plan for them - + // so every transform is the one it was, bit for bit. const int64_t ny = BeamCenterFFTPadSize(2 * h), nx = BeamCenterFFTPadSize(2 * w); - std::vector> fft; - for (int i = 0; i < 4; i++) - fft.push_back(std::make_unique(ny, nx)); + Rfft2 ws0(ny, nx), ws1(ny, nx); + + std::vector> A, M; + { + auto fA = std::async(std::launch::async, [&] { return ws0.Forward(a, h, w); }); + M = ws1.Forward(m, h, w); + A = fA.get(); + } + + BeamCenterConvSurfaces2D out; + { + auto fC = std::async(std::launch::async, [&] { return ws0.InverseProduct(A, A, 2 * h, 2 * w); }); + out.S = ws1.InverseProduct(A, M, 2 * h, 2 * w); + out.C = fC.get(); + } + std::vector>().swap(A); std::vector a2(a.size()); for (size_t i = 0; i < a2.size(); i++) a2[i] = a[i] * a[i]; - auto fA = std::async(std::launch::async, [&] { return fft[0]->Forward(a, h, w); }); - auto fM = std::async(std::launch::async, [&] { return fft[1]->Forward(m, h, w); }); - auto fA2 = std::async(std::launch::async, [&] { return fft[2]->Forward(a2, h, w); }); - const auto A = fA.get(); - const auto M = fM.get(); - const auto A2 = fA2.get(); - - BeamCenterConvSurfaces2D out; - auto fC = std::async(std::launch::async, [&] { return fft[0]->InverseProduct(A, A, 2 * h, 2 * w); }); - auto fS = std::async(std::launch::async, [&] { return fft[1]->InverseProduct(A, M, 2 * h, 2 * w); }); - auto fQ = std::async(std::launch::async, [&] { return fft[2]->InverseProduct(A2, M, 2 * h, 2 * w); }); - auto fD = std::async(std::launch::async, [&] { return fft[3]->InverseProduct(M, M, 2 * h, 2 * w); }); - out.C = fC.get(); - out.S = fS.get(); - out.Q = fQ.get(); - out.D = fD.get(); + { + auto fD = std::async(std::launch::async, [&] { return ws1.InverseProduct(M, M, 2 * h, 2 * w); }); + const auto A2 = ws0.Forward(a2, h, w); + out.Q = ws0.InverseProduct(A2, M, 2 * h, 2 * w); + out.D = fD.get(); + } return out; } - From 43e836573d4a73601eac00dcfae4df42127c1e2d Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 17 Oct 2026 09:31:05 +0200 Subject: [PATCH 3/7] MR and model validation on the CPU: workers bounded by their bytes, not the thread count - The CPU translation function runs as many rotations at once as the host worker pool holds (8 B per grid point each, 550 MiB on a 400 A cell; 33 at once at 32 threads was 18 GiB), at most nthreads. Counted, not reserved: the rotations run on ParallelFor's pool, whose workers must not wait on memory. - Null replicates (model validation's and MR's) whose structure factors are made on the CPU reserve 32 B per point of their grid (ModelGridPoints) against the pool before they start, so the nine replicates, and nulls run side by side, hold at most the pool between them. Every replicate and rotation computes what it did and is stored by index: no result depends on how many run at once. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../structure_refinement/ModelGrid.cpp | 7 +++++++ .../structure_refinement/ModelGrid.h | 9 ++++++++ .../structure_refinement/ModelValidation.cpp | 11 ++++++++++ .../MolecularReplacement.cpp | 21 +++++++++++++++++-- .../MolecularReplacement.h | 3 ++- 5 files changed, 48 insertions(+), 3 deletions(-) diff --git a/image_analysis/structure_refinement/ModelGrid.cpp b/image_analysis/structure_refinement/ModelGrid.cpp index 2dde37fe2..c10ea95fd 100644 --- a/image_analysis/structure_refinement/ModelGrid.cpp +++ b/image_analysis/structure_refinement/ModelGrid.cpp @@ -317,3 +317,10 @@ void PutMaskOnGrid(gemmi::Grid &grid, const gemmi::Model &model, const st if (shrink_has_offsets) masker.shrink(grid); } + +size_t ModelGridPoints(const gemmi::UnitCell &cell, const gemmi::SpaceGroup *sg, double d_min, double rate) { + const double spacing = d_min / (2 * rate); + const std::array m = gemmi::good_grid_size({{cell.a / spacing, cell.b / spacing, cell.c / spacing}}, + gemmi::GridSizeRounding::Up, sg); + return static_cast(m[0]) * static_cast(m[1]) * static_cast(m[2]); +} diff --git a/image_analysis/structure_refinement/ModelGrid.h b/image_analysis/structure_refinement/ModelGrid.h index 4998f2226..96f647292 100644 --- a/image_analysis/structure_refinement/ModelGrid.h +++ b/image_analysis/structure_refinement/ModelGrid.h @@ -27,3 +27,12 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator, float> & // `grid` sized, with its cell and group set. void PutMaskOnGrid(gemmi::Grid &grid, const gemmi::Model &model, const std::vector &orbit_leaders, size_t nthreads); + +// The number of points of the grid a model's structure factors to d_min are computed on at DensityCalculator's +// `rate` (a spacing of d_min / (2 rate)) in this cell and space group: the grid's size, without making it. +size_t ModelGridPoints(const gemmi::UnitCell &cell, const gemmi::SpaceGroup *sg, double d_min, double rate = 1.5); + +// What computing a model's structure factors on the CPU holds at its peak, per point of that grid: the density +// and the bulk-solvent mask grids, and for each of them its transform's input copy, spectrum and coefficients, +// with the two made at once (ModelValidation's compute_model_factors). +constexpr size_t MODEL_FACTORS_BYTES_PER_GRID_POINT = 32; diff --git a/image_analysis/structure_refinement/ModelValidation.cpp b/image_analysis/structure_refinement/ModelValidation.cpp index 7be413fe1..b82008986 100644 --- a/image_analysis/structure_refinement/ModelValidation.cpp +++ b/image_analysis/structure_refinement/ModelValidation.cpp @@ -35,6 +35,7 @@ #include "../../common/CorrelationCoefficient.h" #include "../../common/JFJochMath.h" // PI (M_PI is not standard, and MSVC does not define it) +#include "../../common/HostMemory.h" #include "../../common/Logger.h" #include "../../common/ParallelFor.h" #include "../scale_merge/ReindexAmbiguity.h" // ReindexReflections @@ -1247,7 +1248,17 @@ ModelValidationResult Validate(const std::vector &merged, // the order the concurrent replicates finish in. std::vector null_r_work(NULL_REPLICATES), null_margin(NULL_REPLICATES); std::atomic placed_redraws{0}; + // A replicate whose structure factors are made on the CPU holds a gigabyte or more of grids on a large + // cell, so it reserves them against the host's worker pool (HostWorkerPool) first: as many replicates + // run at once as it holds, whatever the thread count, and each computes what it computes either way. + size_t replicate_bytes = MODEL_FACTORS_BYTES_PER_GRID_POINT + * ModelGridPoints(null_model.cell, null_model.find_spacegroup(), null_d_min); +#ifdef JFJOCH_USE_CUDA + if ((sf_gpu != nullptr && null_d_min == sf_gpu->DMin()) || (sf_null != nullptr && null_d_min == sf_null->DMin())) + replicate_bytes = 0; +#endif auto run_replicate = [&](int i) { + const HostWorkerPool::Reservation reserved(HostWorkerPool::Instance(), replicate_bytes); std::mt19937 rng(NULL_SEED + 1 + i); gemmi::Mat33 rot = null_rotation[i]; for (int redraw = 0;; redraw++) { diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 7e5fdc3b1..64d0ca0c1 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -33,6 +33,8 @@ #include "../../common/JFJochMath.h" // PI #include "../../common/Logger.h" +#include "../../common/CUDAWrapper.h" // get_gpu_count +#include "../../common/HostMemory.h" #include "../../common/ParallelFor.h" #include "ModelFFT.h" #include "ModelGrid.h" @@ -45,7 +47,6 @@ #include "RigidBodyGPUEngine.h" #include "ModelStructureFactorsGPU.h" #include "ModelScaleGPU.h" -#include "../../common/CUDAWrapper.h" #include "../../common/JFJochException.h" #include "../indexing/CUDAMemHelpers.h" // cuda_trim_allocation_pool #endif @@ -766,11 +767,18 @@ double null_sigma(const gemmi::Structure &placed, int copies, const gemmi::Mat33 where = translate(turned); } std::vector r_work(NULL_REPLICATES); + // Without a GPU each replicate's structure factors are made on the CPU, a gigabyte or more of grids on a + // large cell, so it reserves them against the host's worker pool (HostWorkerPool) first: as many run at + // once as it holds - with the nulls of other candidates beside it - and each computes the same either way. + const size_t replicate_bytes = get_gpu_count() == 0 + ? MODEL_FACTORS_BYTES_PER_GRID_POINT * ModelGridPoints(placed.cell, placed.find_spacegroup(), null_d_min) + : 0; std::atomic next{0}; std::vector> workers; for (size_t t = 0; t < std::min(std::max(nthreads, 1), NULL_REPLICATES); t++) workers.push_back(std::async(std::launch::async, [&] { for (int i = next++; i < NULL_REPLICATES; i = next++) { + const HostWorkerPool::Reservation reserved(HostWorkerPool::Instance(), replicate_bytes); gemmi::Structure rep = placed; std::vector &chains = rep.models[0].chains; for (int k = 0; k < copies; k++) { @@ -1561,7 +1569,16 @@ std::vector TranslationFunction(const MolecularTransform &ft, c const std::vector> *known, size_t nthreads) { const TranslationGather gather = make_gather(setup, known != nullptr); std::vector out(rotations.size()); - ParallelFor(static_cast(rotations.size()), nthreads, [&](int i) { + // A rotation's buffers - the grid's coefficients and the function, its terms' factors - are hundreds of + // megabytes on a large cell, so as many rotations run at once as the host's worker pool holds + // (HostWorkerPool), at most nthreads. Counted here rather than reserved in each: the rotations run on + // ParallelFor's pool, whose workers must not wait on one another's memory. + const size_t nu = setup.nu, nv = setup.nv, nw = setup.nw; + const size_t bytes = nu * nv * (nw / 2 + 1) * sizeof(std::complex) + nu * nv * nw * sizeof(float) + + setup.Terms() * (setup.ops * sizeof(std::complex) + sizeof(double)); + const size_t workers = std::clamp(HostWorkerPool::Instance().Size() / std::max(bytes, 1), 1, + std::max(nthreads, 1)); + ParallelFor(static_cast(rotations.size()), workers, [&](int i) { out[i] = translation_function(ft, setup, gather, rotations[i], known); }); return out; diff --git a/image_analysis/structure_refinement/MolecularReplacement.h b/image_analysis/structure_refinement/MolecularReplacement.h index f60ab15dc..750eac3af 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.h +++ b/image_analysis/structure_refinement/MolecularReplacement.h @@ -181,7 +181,8 @@ struct TranslationPeak { TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, const gemmi::Mat33 &rotation, const std::vector> *known = nullptr); -// The same at each of `rotations`, `nthreads` of them at once, the grid's gather worked out once for all. +// The same at each of `rotations`, `nthreads` of them at once - fewer where their buffers do not all fit the host's +// worker pool (HostWorkerPool) - the grid's gather worked out once for all. std::vector TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, const std::vector &rotations, const std::vector> *known, size_t nthreads); From 71afe748b3f7c6cf89c64d93aa5a83abb4d532e3 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 17 Oct 2026 09:31:15 +0200 Subject: [PATCH 4/7] Rotation scaling on the host reads the reflections in place: 24 B per partial instead of 84 Without a GPU (and for the observation dump) RotationScaleMerge kept a full Obs copy of every partial - 84 B, 9.8-10.9 GiB on a 123 M-partial set, 2.6-2.9 GiB on 4nwv - beside the reflections it was copied from, which the run keeps anyway. A partial is now a PartialObs: a pointer to its reflection plus frame, prescaling_corr (with the flux divided out), the recomputed partiality and the working corr. Every other field is read from the reflection, which nothing changes after integration but h,k,l (read only at the ingest). Fields the CPU never read on a partial (h,k,l, delta_phi, var_per_I, capture, group) are gone. Same values in the same arithmetic order: output unchanged. The geometry pre-pass handed each image's reflections back right after the ingest; an engine that reads them (ReadsReflections) now has the post-refinement view gathered once it is reset, after the pre-pass merge - the gather reads only reflection fields, so it is the same gather. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../scale_merge/RotationScaleMerge.cpp | 144 +++++++++--------- .../scale_merge/RotationScaleMerge.h | 42 +++-- rugnux/RugnuxScaleMerge.cpp | 15 +- 3 files changed, 116 insertions(+), 85 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 86dc97067..5a73a377a 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -596,10 +596,8 @@ void RotationScaleMerge::IngestHost() { SmoothMosaicityAndPartiality(); ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - partials[i].delta_phi = ingest_delta_phi[i]; + for (int i = lo; i < hi; ++i) partials[i].partiality = ingest_partiality[i]; - } }); // The corr every pass starts from (see the header) was filled by the build, straight from the @@ -831,12 +829,13 @@ void RotationScaleMerge::IngestDevice(const RotationScaleMerge *observations_fro const int cnt = std::min(n - at, slice); ParallelChunks(cnt, nthreads, [&](int lo, int hi) { for (int i = lo; i < hi; ++i) { - const auto &o = partials[at + i]; - fbuf[0][i] = o.I; fbuf[1][i] = o.sigma; fbuf[2][i] = o.prescaling_corr; - fbuf[3][i] = o.partiality; fbuf[4][i] = o.zeta; fbuf[5][i] = o.corr; - fbuf[6][i] = o.bkg; fbuf[7][i] = o.var_bkg; fbuf[8][i] = o.image_number; - fbuf[9][i] = o.d; - flags[i] = static_cast((o.on_ice ? 1 : 0) | (o.clipped << 1)); // SetObsFlags + const PartialObs &o = partials[at + i]; + const Reflection &r = *o.r; + fbuf[0][i] = r.I; fbuf[1][i] = r.sigma; fbuf[2][i] = o.prescaling_corr; + fbuf[3][i] = o.partiality; fbuf[4][i] = r.zeta; fbuf[5][i] = o.corr; + fbuf[6][i] = r.bkg; fbuf[7][i] = r.var_bkg; fbuf[8][i] = r.image_number; + fbuf[9][i] = r.d; + flags[i] = static_cast((r.on_ice_ring ? 1 : 0) | (ClippedFlags(r) << 1)); // SetObsFlags } }); constexpr F fields[NFLOAT] = {F::I, F::Sigma, F::PrescalingCorr, F::Partiality, @@ -849,7 +848,7 @@ void RotationScaleMerge::IngestDevice(const RotationScaleMerge *observations_fro auto pos = std::make_shared(); pos->px.resize(n); pos->py.resize(n); ParallelChunks(n, nthreads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { pos->px[i] = partials[i].px; pos->py[i] = partials[i].py; } + for (int i = lo; i < hi; ++i) { pos->px[i] = partials[i].r->predicted_x; pos->py[i] = partials[i].r->predicted_y; } }); obs_position = std::move(pos); gpu_->SetRawRuns(static_cast(rawrun_start.size()), static_cast(perm.size()), perm.data(), @@ -1000,25 +999,18 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr &ke for (const auto &r : partials_out[o].reflections) { if (!keep[src]) { ++src; continue; } ++src; - Obs &obs = partials[at]; - obs.h = r.h; obs.k = r.k; obs.l = r.l; - obs.I = r.I; obs.sigma = r.sigma; obs.d = r.d; + PartialObs &obs = partials[at]; + obs.r = &r; // The whole deterministic correction - Lorentz x polarization x sensor efficiency x air path - is the // numerator of corr, so the two halves are multiplied back together here. The incident flux is // divided out at the same time so the fits below see only what is left; it is the one term in // the product that is measured from the data rather than computed from the geometry. obs.prescaling_corr = static_cast(r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]); - obs.partiality = r.partiality; obs.zeta = r.zeta; obs.delta_phi = r.delta_phi_deg; obs.bkg = r.bkg; obs.var_bkg = r.var_bkg; - obs.var_per_I = 0.0f; obs.capture = 0.0f; - obs.px = r.predicted_x; obs.py = r.predicted_y; - obs.image_number = r.image_number; + obs.partiality = r.partiality; obs.frame = o; - obs.on_ice = r.on_ice_ring ? 1 : 0; - obs.clipped = (r.clipped ? OBS_CLIPPED : 0) | (r.overloaded ? OBS_OVERLOADED : 0); obs.corr = r.image_scale_corr; - obs.group = -1; - finite_ok[at] = (std::isfinite(obs.I) && std::isfinite(obs.prescaling_corr) && obs.prescaling_corr != 0.0f - && std::isfinite(obs.sigma) && obs.sigma > 0.0f) ? 1 : 0; + finite_ok[at] = (std::isfinite(r.I) && std::isfinite(obs.prescaling_corr) && obs.prescaling_corr != 0.0f + && std::isfinite(r.sigma) && r.sigma > 0.0f) ? 1 : 0; ingest_delta_phi[at] = SmoothedDeltaPhi(sg, o, r.h, r.k, r.l, r.delta_phi_deg, resolved); ingest_partiality[at] = r.partiality; ingest_image_number[at] = r.image_number; @@ -1562,8 +1554,8 @@ void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vectorI, o.r->sigma, min_partiality, I_corr); if (w <= 0.0) continue; sw += w; swI += w * I_corr; @@ -2813,11 +2805,11 @@ int RotationScaleMerge::RockingEventFrames(int64_t *n_events) const { } return RockingEventFramesOver( [&](int i) { - const Obs &o = partials[i]; - return std::isfinite(o.corr) && o.corr > 0.0f && std::isfinite(o.I) - && std::isfinite(o.sigma) && o.sigma > 0.0f; + const PartialObs &o = partials[i]; + return std::isfinite(o.corr) && o.corr > 0.0f && std::isfinite(o.r->I) + && std::isfinite(o.r->sigma) && o.r->sigma > 0.0f; }, - [&](int i) { return partials[i].image_number; }, n_events); + [&](int i) { return partials[i].r->image_number; }, n_events); } template @@ -4389,15 +4381,23 @@ void RotationScaleMerge::Combine() { // A raw hkl's partials lie a frame apart in memory, and the work on one rocking event is long // enough that the CPU does not reach the next runs' loads on its own - the combine waited on a // cache miss for nearly every partial. So each run asks for the partials of the run - // COMBINE_PREFETCH_RUNS ahead. + // COMBINE_PREFETCH_RUNS ahead, and for the reflections of the run half as far ahead, whose partials + // - which say where the reflections are - were asked for that many runs before. constexpr int COMBINE_PREFETCH_RUNS = 8; auto prefetch_run = [&](int r) { - if (r >= n_run) return; - for (int p = rawrun_start[r]; p < rawrun_start[r] + rawrun_count[r]; ++p) { - const char *o = reinterpret_cast(&partials[perm[p]]); - Prefetch(o); - Prefetch(o + sizeof(Obs) - 1); - } + if (r < n_run) + for (int p = rawrun_start[r]; p < rawrun_start[r] + rawrun_count[r]; ++p) { + const char *o = reinterpret_cast(&partials[perm[p]]); + Prefetch(o); + Prefetch(o + sizeof(PartialObs) - 1); + } + const int rr = r - COMBINE_PREFETCH_RUNS / 2; + if (rr < n_run) + for (int p = rawrun_start[rr]; p < rawrun_start[rr] + rawrun_count[rr]; ++p) { + const char *o = reinterpret_cast(partials[perm[p]].r); + Prefetch(o); + Prefetch(o + sizeof(Reflection) - 1); + } }; // Combine one raw-hkl run into fulls, appended to `out`. Independent per run (the events of one hkl @@ -4409,9 +4409,9 @@ void RotationScaleMerge::Combine() { std::vector ev; // usable perm-indices of this raw hkl, in image-number order ev.reserve(hi - lo); for (int p = lo; p < hi; ++p) { - const auto &o = partials[perm[p]]; + const PartialObs &o = partials[perm[p]]; if (!std::isfinite(o.corr) || o.corr <= 0.0f) continue; - if (!std::isfinite(o.I) || !std::isfinite(o.sigma) || o.sigma <= 0.0f) continue; + if (!std::isfinite(o.r->I) || !std::isfinite(o.r->sigma) || o.r->sigma <= 0.0f) continue; ev.push_back(perm[p]); } const int group = rawrun_group[r]; @@ -4419,9 +4419,9 @@ void RotationScaleMerge::Combine() { size_t i = 0; while (i < ev.size()) { size_t kk = i + 1; - float last_frame = partials[ev[i]].image_number; + float last_frame = partials[ev[i]].r->image_number; while (kk < ev.size()) { - const float frame = partials[ev[kk]].image_number; + const float frame = partials[ev[kk]].r->image_number; if (frame - last_frame > max_frame_gap) break; last_frame = frame; ++kk; @@ -4430,11 +4430,11 @@ void RotationScaleMerge::Combine() { double pooled_bkg = 0.0; int n_pool = 0; for (size_t m = i; m < kk; ++m) { - const float b = partials[ev[m]].bkg; + const float b = partials[ev[m]].r->bkg; if (std::isfinite(b)) { pooled_bkg += b; ++n_pool; } } pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0; - auto pooled_I = [&](const Obs &r2) { + auto pooled_I = [&](const Reflection &r2) { const double n_bkg = std::max(0.0, static_cast(r2.sigma) * r2.sigma - r2.I) / std::max(r2.bkg, 1.0f); return static_cast(r2.I) + n_bkg * (static_cast(r2.bkg) - pooled_bkg); @@ -4443,12 +4443,12 @@ void RotationScaleMerge::Combine() { double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0; float d = NAN; int peak_outcome = partials[ev[i]].frame; - float peak_frame = partials[ev[i]].image_number; - float peak_px = partials[ev[i]].px, peak_py = partials[ev[i]].py; + float peak_frame = partials[ev[i]].r->image_number; + float peak_px = partials[ev[i]].r->predicted_x, peak_py = partials[ev[i]].r->predicted_y; float peak_partiality = -1.0f; - const bool on_ice = partials[ev[i]].on_ice; + const bool on_ice = partials[ev[i]].r->on_ice_ring; uint8_t flags = 0; - for (size_t m = i; m < kk; ++m) flags |= partials[ev[m]].clipped; + for (size_t m = i; m < kk; ++m) flags |= ClippedFlags(*partials[ev[m]].r); // A saturated pixel anywhere in the rocking curve means its brightest part was not measured: // the event is not a measurement of the reflection, and summing what is left - or scaling // the flanks up by the partiality model - reads it 2-3x low. Dropped whole, as XDS drops an @@ -4460,17 +4460,18 @@ void RotationScaleMerge::Combine() { } const bool clipped = flags & OBS_CLIPPED; for (size_t m = i; m < kk; ++m) { - const auto &r2 = partials[ev[m]]; - const double sigma_corr = static_cast(r2.sigma) * r2.corr; + const PartialObs &p2 = partials[ev[m]]; + const Reflection &r2 = *p2.r; + const double sigma_corr = static_cast(r2.sigma) * p2.corr; const double w = 1.0 / (sigma_corr * sigma_corr); sum_w += w; - sum_wI += w * pooled_I(r2) * r2.corr; - sum_partiality += r2.partiality; - if (r2.partiality > peak_partiality) { - peak_partiality = r2.partiality; - peak_outcome = r2.frame; + sum_wI += w * pooled_I(r2) * p2.corr; + sum_partiality += p2.partiality; + if (p2.partiality > peak_partiality) { + peak_partiality = p2.partiality; + peak_outcome = p2.frame; peak_frame = r2.image_number; - peak_px = r2.px; peak_py = r2.py; + peak_px = r2.predicted_x; peak_py = r2.predicted_y; } if (!std::isfinite(d) && std::isfinite(r2.d) && r2.d > 0.0f) d = r2.d; } @@ -4482,8 +4483,8 @@ void RotationScaleMerge::Combine() { for (int iter = 0; iter < 3; ++iter) { sum_w = 0.0; sum_wI = 0.0; sum_wb = 0.0; sum_cwb = 0.0; for (size_t m = i; m < kk; ++m) { - const auto &r2 = partials[ev[m]]; - const double corr = r2.corr; + const Reflection &r2 = *partials[ev[m]].r; + const double corr = partials[ev[m]].corr; const double I_corr = pooled_I(r2) * corr; const double sigma_corr = static_cast(r2.sigma) * corr; // The non-signal variance as the integrator measured it, which is where it belongs: @@ -4661,16 +4662,17 @@ void RotationScaleMerge::ComputePerFrameCC(const std::vector &partial_gr size_t n = 0; const int lo = frame_start[f], hi = frame_start[f] + frame_count[f]; for (int i = lo; i < hi; ++i) { - const auto &o = partials[i]; + const PartialObs &o = partials[i]; + const Reflection &r = *o.r; const int g = partial_group[i]; - if (o.on_ice) continue; + if (r.on_ice_ring) continue; if (g < 0) continue; if (o.partiality < min_partiality) continue; - if (!std::isfinite(o.I) || !std::isfinite(o.corr) || o.corr <= 0.0f) continue; - if (!std::isfinite(o.sigma) || o.sigma <= 0.0f) continue; + if (!std::isfinite(r.I) || !std::isfinite(o.corr) || o.corr <= 0.0f) continue; + if (!std::isfinite(r.sigma) || r.sigma <= 0.0f) continue; const double mean = partial_group_mean[g]; if (!std::isfinite(mean)) continue; - const double img = static_cast(o.I) * o.corr; + const double img = static_cast(r.I) * o.corr; sx += img; sy += mean; sx2 += img * img; sy2 += mean * mean; sxy += img * mean; ++n; } @@ -6435,7 +6437,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool #endif // The CPU partial scaling below writes every partial's corr when it is done, from a copy it starts // off corr_ingested itself, and nothing between here and there reads corr - so there the restore - // is not made, which saves a pass over the eighty-byte records to set four bytes of each. + // is not made, which saves a pass over the partials to set four bytes of each. if (!restored_on_gpu && !scale_partials) ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), [&](int lo, int hi) { @@ -6568,8 +6570,8 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool FillVector scaling_corr(n_partials); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { for (int q = lo; q < hi; ++q) { - const Obs &o = partials[scaling_index[q]]; - scaling_refs[q] = {o.frame, o.I, o.sigma, o.partiality, o.prescaling_corr}; + const PartialObs &o = partials[scaling_index[q]]; + scaling_refs[q] = {o.frame, o.r->I, o.r->sigma, o.partiality, o.prescaling_corr}; scaling_corr[q] = corr_ingested[scaling_index[q]]; // the restore (see the top) } }); @@ -6580,9 +6582,9 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // fitted on them is fitted on noise that the reference does not answer to. The Rotation // model needs zeta > 0. auto fitted_obs = [&](int i) { - const Obs &o = partials[i]; - return partial_group[i] >= 0 && !o.on_ice && o.partiality >= min_partiality - && std::isfinite(o.zeta) && o.zeta > 0.0f; + const PartialObs &o = partials[i]; + return partial_group[i] >= 0 && !o.r->on_ice_ring && o.partiality >= min_partiality + && std::isfinite(o.r->zeta) && o.r->zeta > 0.0f; }; const int n_fit_frames = static_cast(frame_start.size()); std::vector fit_start(n_fit_frames + 1, 0); @@ -6597,8 +6599,8 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool int32_t k = fit_start[f]; for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) { if (!fitted_obs(i)) continue; - const Obs &o = partials[i]; - scaling_fit[k++] = {partial_group[i], o.I, o.sigma, o.partiality, o.prescaling_corr}; + const PartialObs &o = partials[i]; + scaling_fit[k++] = {partial_group[i], o.r->I, o.r->sigma, o.partiality, o.prescaling_corr}; } }); @@ -6679,7 +6681,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), [&](int lo, int hi) { for (int i = lo; i < hi; ++i) { - Obs &o = partials[i]; + PartialObs &o = partials[i]; if (apply[o.frame] && std::isfinite(o.corr)) o.corr = static_cast(o.corr * ratio[o.frame]); } @@ -6738,8 +6740,8 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool [&](int lo, int hi) { int64_t n = 0; for (int i = lo; i < hi; ++i) { - Obs &o = partials[i]; - if (!(std::isfinite(o.zeta) && o.zeta >= search_min_zeta)) { + PartialObs &o = partials[i]; + if (!(std::isfinite(o.r->zeta) && o.r->zeta >= search_min_zeta)) { if (std::isfinite(o.corr) && o.corr > 0.0f) ++n; o.corr = 0.0f; } diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index b1aa6d658..fafd1c892 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -128,7 +128,10 @@ public: // experiment: read live (its space group is changed by the caller between Run() calls). // partial_outcomes: the per-frame partials; the final per-frame scale (G, CC, mosaicity) is written - // back onto them so the offline per-image scaling table is still exported. + // back onto them so the offline per-image scaling table is still exported. Where the engine scales on + // the host (ReadsReflections), it reads the reflections themselves from the ingest on, so they + // must stay where they are, as they are, for as long as it merges: not freed, moved or changed, + // but for their h, k and l, which it does not read after the ingest. // reference_cell: the consensus cell (for the completeness count and the cell-consistency mask). RotationScaleMerge(const DiffractionExperiment &experiment, std::vector &partial_outcomes, @@ -259,6 +262,9 @@ public: [[nodiscard]] int GetGpuDevice() const { return gpu_device; } // Whether this engine scales on the GPU (decided by the ingest). [[nodiscard]] bool OnGPU() const; + // Whether the engine still reads the source reflections in later Runs (see the constructor): where it + // scales on the host, or keeps its partials for the observation dump. + [[nodiscard]] bool ReadsReflections() const { return !partials.empty(); } // Whether Run() also hands back the merge's fulls unmerged (Result::scaled_fulls). Off by default: // it is a copy as long as the fulls, wanted only by a caller that writes them. @@ -296,10 +302,10 @@ private: static constexpr uint8_t OBS_CLIPPED = 1; static constexpr uint8_t OBS_OVERLOADED = 2; - // One integrated observation - a per-frame partial during scaling/combine, or a combined full during - // scale-fulls/merge. Flat (not nested per image); a POD so the arrays translate straight to CUDA. - // No member initialisers, so that `partials` and `fulls` (FillVector) are first written by the - // parallel passes that fill them: whoever makes one sets every field. + // One combined full during scale-fulls/merge (a per-frame partial is a PartialObs, below). Flat (not + // nested per image); a POD so the arrays translate straight to CUDA. No member initialisers - here or in + // PartialObs - so that `fulls` and `partials` (FillVector) are first written by the parallel passes + // that fill them: whoever makes one sets every field. struct Obs { int32_t h, k, l; float I, sigma, d, prescaling_corr, partiality, zeta, delta_phi, bkg, var_bkg; @@ -321,6 +327,24 @@ private: int32_t group; // dense ASU-group id for the current space group; <0 = never mergeable }; + // One per-frame partial as the host holds it - without a GPU, or for the observation dump: where its + // reflection is, and the four values the ingest works out for it or the scaling changes. Every other + // field is the reflection's own, read from it: the run keeps the reflections for as long as the engine + // merges, and nothing changes in them after the images are integrated but the Miller indices of a + // reindexing, which only the ingest reads (from the reflections, into the raw-hkl runs). Twenty-four + // bytes where a copy of the fields was eighty-four - on a long axis, a hundred million partials. + struct PartialObs { + const Reflection *r; + int32_t frame; // index of the outcome whose per-frame scale G applies to this obs + float prescaling_corr; // the reflection's whole deterministic correction over the frame's flux + float partiality; // recomputed from the frame-order-smoothed mosaicity (Ingest) + float corr; // image_scale_corr (working; updated by scaling) + }; + // A reflection's Obs::clipped. + static uint8_t ClippedFlags(const Reflection &r) { + return (r.clipped ? OBS_CLIPPED : 0) | (r.overloaded ? OBS_OVERLOADED : 0); + } + // One leverage-corrected error-model sample per usable full: its counting variance at the group's // mean, that mean squared, its squared deviation from the mean - and the resolution it sits at, // because the fit is re-run over the samples that survive the automatic resolution cutoff. See @@ -414,7 +438,7 @@ private: // built only when the device pipeline is NOT resident (no GPU, or an observation dump was asked // for); on the resident path (resident_ingest) the device holds the observations and the host // keeps only the ingest arrays below, until the end of Ingest. - FillVector partials; // all per-frame partials, grouped by frame + FillVector partials; // all per-frame partials, grouped by frame bool resident_ingest = false; // What the ingest-time geometry smoothing (SmoothMosaicityAndPartiality) needs, one array per field // in `partials` order: the two fields it rewrites, the two the rocking-event walks read in raw-hkl @@ -696,13 +720,13 @@ private: std::vector partial_group_start, partial_group_count; NotCopied partial_group_perm; // Without a GPU: each partial's ASU group (-1 = in none), in observation order - where the CPU - // passes read a partial's group from. A partial's own Obs::group is not stamped: written into the - // eighty-byte records at random, once per grouping, it cost more than the grouping itself. + // passes read a partial's group from. It is not kept on the partial: written into the records at + // random, once per grouping, it cost more than the grouping itself. NotCopied partial_group; // The CPU scaling loops' working copy of the partials, or of the fulls. A loop passes over them // several times per iteration and reads a handful of fields each time, so it works on these rather - // than on the eighty-byte Obs. ScalingRef holds what the group means and the corr update read, in the CSR's + // than on the partials and their reflections. ScalingRef holds what the group means and the corr update read, in the CSR's // group order with the partials in no group after them, and corr is kept beside it in the same // order - so neither pass gathers. ScalingFit holds, frame by frame, the observations // FitScalingFrames fits: the fit's filters, which do not change over the loop, applied once, diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index 658fadb27..8706399d0 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -839,14 +839,17 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // the scaling engine is finished - nothing after this point reads it on that pass // (the search and the reference re-merge run only on the canonical one) - so it is // freed BEFORE the fit: the engine holds about twice the partial payload again, - // and with it gone the fit no longer sits on the pass's memory peak. - if (prepass_postrefine_obs) + // and with it gone the fit no longer sits on the pass's memory peak. An engine that + // read the reflections in its merges left the gather to here, and from here nothing + // on the pass reads them either: each image's is handed back as it is gathered. + const bool release = geometry_prepass && rsm && config_.reference_data.empty(); + if (prepass_postrefine_obs || release) rsm.reset(); auto pr_obs = prepass_postrefine_obs ? std::move(*prepass_postrefine_obs) : gather ? GatherPostRefineObservations(indexer->GetIntegrationOutcome(), config_.nthreads, - /*release_reflections=*/false, logger, + /*release_reflections=*/release, logger, max_images) : PostRefineObservations{}; prepass_postrefine_obs.reset(); @@ -1201,8 +1204,10 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // per-frame fields - so take the post-refine's view NOW and hand each image's vector back as // it is read: on a fine-sliced long axis the vectors are gigabytes, held to the end of the // pass for those 32 bytes a partial. A run with a reference MTZ keeps them: resolving the - // indexing ambiguity re-reads them (and re-ingests). - if (geometry_prepass && rsm && config_.reference_data.empty()) + // indexing ambiguity re-reads them (and re-ingests). So does an engine that scales on the host, + // which reads them in its merges (RotationScaleMerge::ReadsReflections): there the view is taken + // once the engine is done with them (prepare_post_refine). + if (geometry_prepass && rsm && config_.reference_data.empty() && !rsm->ReadsReflections()) prepass_postrefine_obs = GatherPostRefineObservations( indexer->GetIntegrationOutcome(), config_.nthreads, /*release_reflections=*/true, logger); From 46c44e04d9478c5dd256a1a71e730824d1a84b19 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 17 Oct 2026 09:31:24 +0200 Subject: [PATCH 5/7] Rotation scaling: a host-memory budget decides what runs side by side The run's budget (ProcessConfig::host_memory_budget_bytes, else 75% of the machine's RAM) is decided once and logged, and sizes the host worker pool. Three concurrencies of the scaling chain are decided on it from counts known before each starts (RotationScaleMerge::HostIngestBytes, HostMergeBytes, HostBytesBound), each logged where it serialises: - the quality guard's merge beside the main ingest: where they would not fit together, the guard merges first; - the canonical pass's post-refinement beside the merge chain: done before the first merge; - the all-observation P1 search merge beside the main engine's (GPU path): started where the search first reads it, after the main engine's merges and the release of their scratch. None of these reads what the other writes while it runs, so only the order changes: a [large] test runs the rotation set at a 1 TB and a 1 MB budget and compares the .hkl byte for byte. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../scale_merge/RotationScaleMerge.cpp | 37 +++++++ .../scale_merge/RotationScaleMerge.h | 9 ++ rugnux/Rugnux.cpp | 15 +++ rugnux/Rugnux.h | 7 ++ rugnux/RugnuxScaleMerge.cpp | 96 +++++++++++++++---- tests/RugnuxLargeTest.cpp | 64 +++++++++++++ 6 files changed, 210 insertions(+), 18 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 5a73a377a..120dc0753 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -779,6 +779,43 @@ size_t RotationScaleMerge::DeviceBytes() const { return std::max(bytes + bytes / 5, ingest); } +namespace { + // The host bytes the ingest holds per partial from then on: the corr every Run starts from and, with + // the GPU resident, the predicted position the fulls are given; on the host, the partial itself, the + // raw-hkl order, the finiteness mask and the ASU grouping's three per-partial arrays. Per raw-hkl run: + // its start, count, h, k, l, d and group. + constexpr size_t HOST_INGEST_BYTES_PER_PARTIAL_GPU = 3 * sizeof(float); + constexpr size_t HOST_INGEST_BYTES_PER_RUN = 7 * sizeof(int32_t); + // What one merge holds on the host at its peak, per rocking event (a full each at most: the full, its + // download staging and the merge's per-full fields), per distinct hkl (an ASU group and a merged + // reflection each at most: the merge's per-group accumulators, R_meas sums and their staging, the ISa + // scatter, the merged reflection), and per partial where the engine scales on the host (the scaling + // loop's copies of the partials and the memo of their corr). + constexpr size_t HOST_MERGE_BYTES_PER_EVENT = 190, HOST_MERGE_BYTES_PER_HKL = 420; + constexpr size_t HOST_MERGE_BYTES_PER_PARTIAL_CPU = 65; +} + +size_t RotationScaleMerge::HostIngestBytes() const { + constexpr size_t PER_PARTIAL_CPU = sizeof(PartialObs) + sizeof(float) + sizeof(uint8_t) + 4 * sizeof(int32_t); + const size_t per_partial = OnGPU() ? HOST_INGEST_BYTES_PER_PARTIAL_GPU : PER_PARTIAL_CPU; + return per_partial * static_cast(n_partials_obs) + HOST_INGEST_BYTES_PER_RUN * rawrun_start.size(); +} + +size_t RotationScaleMerge::HostMergeBytes() const { + // The device path's ingest counted the rocking events; without it every partial may be one. + const auto events = static_cast(rocking_events_at_ingest > 0 ? rocking_events_at_ingest : n_partials_obs); + return HOST_MERGE_BYTES_PER_EVENT * events + HOST_MERGE_BYTES_PER_HKL * rawrun_start.size() + + (OnGPU() ? 0 : HOST_MERGE_BYTES_PER_PARTIAL_CPU * static_cast(n_partials_obs)); +} + +size_t RotationScaleMerge::HostBytesBound(size_t n_reflections, bool on_gpu) { + constexpr size_t PER_PARTIAL_CPU = sizeof(PartialObs) + sizeof(float) + sizeof(uint8_t) + 4 * sizeof(int32_t) + + HOST_MERGE_BYTES_PER_PARTIAL_CPU; + const size_t per_reflection = (on_gpu ? HOST_INGEST_BYTES_PER_PARTIAL_GPU : PER_PARTIAL_CPU) + + HOST_INGEST_BYTES_PER_RUN + HOST_MERGE_BYTES_PER_EVENT + HOST_MERGE_BYTES_PER_HKL; + return per_reflection * n_reflections; +} + std::unique_ptr RotationScaleMerge::CloneIngest(int device, Logger &log) const { TimingMark mark(*logger, "rsm ingest copy"); std::unique_ptr c(new RotationScaleMerge(*this)); diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index fafd1c892..b5b90147c 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -250,6 +250,15 @@ public: [[nodiscard]] size_t DeviceBytes() const; // The device scaling's estimated peak (see the .cpp). Call after the ingest. [[nodiscard]] size_t GPUScalingPeakBytes() const; + + // The host memory an engine holds, for a caller deciding what may run beside it within the host budget + // (HostMemory.h); it decides nothing in the engine. HostIngestBytes: what the ingest holds from then on, + // between merges - call after IngestDevice(). HostMergeBytes: what one merge holds on top of it at its + // peak, from the counts the ingest knows (see the .cpp). HostBytesBound: the two together, before an + // ingest of `n_reflections`, at most. + [[nodiscard]] size_t HostIngestBytes() const; + [[nodiscard]] size_t HostMergeBytes() const; + [[nodiscard]] static size_t HostBytesBound(size_t n_reflections, bool on_gpu); // Above this fraction of the card's memory the scaling runs on the CPU. static constexpr double CPU_SCALING_ABOVE_CARD_FRACTION = 0.8; diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index e6e6eb076..6c2de20eb 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -32,6 +32,7 @@ #include "../reader/JFJochHDF5Reader.h" #include "../reader/JFJochStartMessageReader.h" +#include "../common/HostMemory.h" #include "../common/JFJochMath.h" #include "../common/ParallelFor.h" #include "../common/Logger.h" @@ -317,6 +318,20 @@ Rugnux::Rugnux(JFJochReader &reader, DiffractionExperiment experiment, LogHeaderBeamCenter(experiment_, pixel_mask_); + // What the run plans its host memory against, decided once from the machine's total (HostMemory.h): + // work that would not fit it side by side runs one part after the other, and the large per-worker + // buffers of the model's steps are held a share of it at a time. Only when things run changes. + const size_t machine = get_host_total_memory(); + host_budget_ = config_.host_memory_budget_bytes > 0 ? config_.host_memory_budget_bytes : DefaultHostMemoryBudget(); + HostWorkerPool::Instance().SetSize( + std::max(1, static_cast(HOST_WORKER_POOL_FRACTION * static_cast(host_budget_)))); + Logger("Rugnux").Info("Host memory budget {:.1f} GB ({}): what does not fit it side by side runs in series", + static_cast(host_budget_) / 1e9, + config_.host_memory_budget_bytes > 0 + ? std::string("set by the caller") + : fmt::format("{:.0f}% of the machine's {:.1f} GB", 100.0 * HOST_MEMORY_BUDGET_FRACTION, + static_cast(machine) / 1e9)); + // Every decision on a rotation run - per-frame scaling, the correction surfaces, the outlier test, // the error model, the CC1/2 half-sets and the resolution cut, the frame ledger, the space-group // search - is taken with Friedel mates merged, -A or not. Splitting them halves the multiplicity diff --git a/rugnux/Rugnux.h b/rugnux/Rugnux.h index e5cca61d9..92ca53d06 100644 --- a/rugnux/Rugnux.h +++ b/rugnux/Rugnux.h @@ -73,6 +73,11 @@ struct ProcessConfig { // The per-image loop is capped in that case (see RunPipeline); an explicit -N never is. bool nthreads_auto = false; + // The host memory the run plans against, in bytes; 0 (the default) takes the machine's + // (DefaultHostMemoryBudget). It decides only what runs side by side - never a number - so a caller + // setting it (the tests do) changes the run's time and its memory, and nothing it writes. + size_t host_memory_budget_bytes = 0; + // Output prefix for the _process.h5 (and scaled reflections). Empty => process without writing. std::string output_prefix; // What the reflection files record about where they came from (WriteReflections). @@ -800,6 +805,8 @@ class Rugnux { DiffractionExperiment experiment_; PixelMask pixel_mask_; ProcessConfig config_; + // The host memory budget (ProcessConfig::host_memory_budget_bytes, else the machine's), decided once. + size_t host_budget_ = 0; // The run's cancel flag, shared with the copies the run makes of itself to run passes side by // side (see RunAllPasses), so cancelling the run cancels those too. class CancelFlag { diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index 8706399d0..c50ccb5f6 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -948,13 +948,41 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { /*full_stats=*/false, /*measure_cc_before_corrections=*/true); return ScaleMergeResult{std::move(r.merged), std::move(r.statistics), r.cc_half_before_corrections}; }; + // What the host holds from here on (HostMemory.h): the reflections, and beside them the engines below. + // The counts are known now, before anything is ingested, so each concurrency below is decided on them: + // where two pieces of work would not fit the host budget together, one waits for the other. None of + // them reads anything another writes while it runs (see each), so only the order changes. + size_t n_reflections = 0, n_reflections_guard = 0; + { + const auto &outcomes = indexer->GetIntegrationOutcome(); + for (size_t o = 0; o < outcomes.size(); ++o) { + n_reflections += outcomes[o].reflections.size(); + if (static_cast(o) < prepass_images) + n_reflections_guard += outcomes[o].reflections.size(); + } + } + const size_t reflection_bytes = n_reflections * sizeof(Reflection); + const auto host_gb = [](size_t bytes) { return static_cast(bytes) / 1e9; }; std::future guard_ahead; - if (merge_first_part && !(postrefine_probe_ && !geometry_prepass && postrefine_probe_only_)) - guard_ahead = std::async(std::launch::async, [&] { - if (guard_on_second_gpu) - pin_gpu(1); // the card and its NUMA node, as an image-loop worker has them - return merge_guard(); - }); + if (merge_first_part && !(postrefine_probe_ && !geometry_prepass && postrefine_probe_only_)) { + // The guard's merge runs beside the main engine's ingest. Where an engine of the guard's frames and + // one of all the frames would not fit the budget together - counted at the most an engine of that + // many reflections can hold - the guard merges first, on its own, and the ingest starts after it. + const bool on_gpu = get_gpu_count() > 0; + const size_t beside = reflection_bytes + RotationScaleMerge::HostBytesBound(n_reflections_guard, on_gpu) + + RotationScaleMerge::HostBytesBound(n_reflections, on_gpu); + if (beside <= host_budget_ || guard_on_second_gpu) + guard_ahead = std::async(std::launch::async, [&] { + if (guard_on_second_gpu) + pin_gpu(1); // the card and its NUMA node, as an image-loop worker has them + return merge_guard(); + }); + else { + logger.Info("Host memory: the quality guard's merge and this pass's ingest need up to {:.1f} GB side by " + "side, over the {:.1f} GB budget - the guard merges first", host_gb(beside), host_gb(host_budget_)); + guard_merge = merge_guard(); + } + } // The scaling engine's ingest reads the outcomes and nothing else either, so its host half runs // beside the post-refinement too; with one GPU the device half waits for the guard to hand the // card back. @@ -1096,6 +1124,19 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { std::unique_ptr rsm_all; Logger all_arm_log = Logger::Buffered(); std::future all_arm_ahead; + const auto start_all_arm = [&rsm_all, &all_arm_log] { + return std::async(std::launch::async, [&rsm_all, &all_arm_log] { + pin_gpu(rsm_all->GetGpuDevice()); // the card and its NUMA node, as the guard's merge has them + TimingMark mark(all_arm_log, "merge P1, all observations"); + auto r = rsm_all->Run(/*for_search=*/true, /*full_stats=*/true, + /*measure_cc_before_corrections=*/false); + // This engine merges once: what is left for it to do is the write-back of its + // per-frame scales, so its observations and its merge's working arrays go now, + // rather than through the search the others are about to run. + rsm_all->ReleaseAllButPerFrameScale(); + return r; + }); + }; if (is_rotation) { if (experiment_.GetRefineRotationWedgeInScaling() || rot_ss.GetRotationWedgeForScaling().has_value()) @@ -1150,7 +1191,8 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { } // With one card the guard's merge has the card to itself first. if (merge_first_part && !guard_on_second_gpu) { - guard_merge = guard_ahead.valid() ? guard_ahead.get() : merge_guard(); + if (!guard_merge) + guard_merge = guard_ahead.valid() ? guard_ahead.get() : merge_guard(); guard_log.ReplayInto(logger); } rsm->IngestDevice(); @@ -1158,6 +1200,21 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { const int32_t cards = get_gpu_count(); engine_log.ReplayInto(logger); rsm->SetLogger(logger); + // The canonical pass's post-refinement gathers (from the moment the engine has ingested) and fits + // beside the merges that follow. Where its gather and the selection it ranks would not fit the host + // budget beside the reflections, this engine and its first merge, it is done before that merge. + const size_t host_engine = reflection_bytes + rsm->HostIngestBytes(); + const size_t host_merge = rsm->HostMergeBytes(); + constexpr size_t POSTREFINE_RANKED_BYTES = 16; // PostRefineRotationGeometry's ranking, per partial + const size_t host_postrefine = n_reflections * (sizeof(PostRefinePartial) + POSTREFINE_RANKED_BYTES); + bool postrefine_beside = p.postrefine_ahead.valid(); + if (postrefine_beside && host_engine + host_merge + host_postrefine > host_budget_) { + logger.Info("Host memory: the post-refinement ({:.1f} GB) does not fit beside the merge ({:.1f} GB, " + "{:.1f} GB held) within the {:.1f} GB budget - it is done before the merge", + host_gb(host_postrefine), host_gb(host_merge), host_gb(host_engine), host_gb(host_budget_)); + p.postrefine_ahead.wait(); + postrefine_beside = false; + } // On a card of its own from three cards up (card 0 has this engine, card 1 the quality guard's // and the speculative passes). Otherwise on the last card, beside the engine already there - // this one, or the guard's, which merges at most as many observations - and only where two @@ -1167,6 +1224,7 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // holds more, but a merge that does not fit beside it waits for it to end. On a 16 GB card that // is the data sets up to about 90 M observations at high multiplicity; the large ones merge // their arms one after the other. + bool all_arm_beside = false; if (!geometry_prepass && !experiment_.GetGemmiSpaceGroup().has_value() && rsm->GetSearchMinZeta() > 0.0 && cards > 0 && rsm->OnGPU()) { const int32_t device = std::min(cards - 1, 2); @@ -1177,23 +1235,23 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { rsm_all->SetSearchMinZeta(0.0); rsm_all->SetExportScaledFulls(false); // as scale_and_merge sets them for a search merge rsm_all->SetFrenchWilson(false); + // ... and beside this engine's merges only where two merges fit the host budget as well, + // with the post-refinement where it is still running; otherwise it merges where the search + // first reads it, after this engine's merges, on its own. + const size_t host_beside = host_engine + 2 * host_merge + (postrefine_beside ? host_postrefine : 0); + all_arm_beside = host_beside <= host_budget_; + if (!all_arm_beside) + logger.Info("Host memory: two P1 merges side by side need about {:.1f} GB, over the {:.1f} GB " + "budget - the all-observation merge is made after this engine's", + host_gb(host_beside), host_gb(host_budget_)); } } if (rsm_all) { // Its observations are copied from this engine's card arrays, not staged from the // reflections a second time. rsm_all->IngestDevice(&*rsm); - all_arm_ahead = std::async(std::launch::async, [&rsm_all, &all_arm_log] { - pin_gpu(rsm_all->GetGpuDevice()); // the card and its NUMA node, as the guard's merge has them - TimingMark mark(all_arm_log, "merge P1, all observations"); - auto r = rsm_all->Run(/*for_search=*/true, /*full_stats=*/true, - /*measure_cc_before_corrections=*/false); - // This engine merges once: what is left for it to do is the write-back of its - // per-frame scales, so its observations and its merge's working arrays go now, - // rather than through the search the others are about to run. - rsm_all->ReleaseAllButPerFrameScale(); - return r; - }); + if (all_arm_beside) + all_arm_ahead = start_all_arm(); } } @@ -1637,6 +1695,8 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // others is made again where it is wanted, in series, as it would have been without it. std::optional all_merge; std::future all_search_ahead; + if (rsm_all && !all_arm_ahead.valid()) + all_arm_ahead = start_all_arm(); // not beside the merges above (the host budget): now if (all_arm_ahead.valid()) { TimingMark wait_mark(logger, "wait for the all-observation merge"); try { diff --git a/tests/RugnuxLargeTest.cpp b/tests/RugnuxLargeTest.cpp index 19ba12005..38a6de228 100644 --- a/tests/RugnuxLargeTest.cpp +++ b/tests/RugnuxLargeTest.cpp @@ -9,11 +9,15 @@ #include #include +#include +#include #include +#include #include #include "TestData.h" #include "../common/DiffractionExperiment.h" +#include "../common/HostMemory.h" #include "../common/IndexingSettings.h" #include "../reader/JFJochHDF5Reader.h" #include "../rugnux/Rugnux.h" @@ -121,3 +125,63 @@ TEST_CASE("Rugnux_FirstPassMemoMatchesRecomputation", "[large]") { reader.Close(); } + +// The host memory budget decides only what runs side by side - the quality guard's merge beside the +// ingest, the post-refinement beside the merge chain, the two search merges beside each other, how many +// large per-worker buffers are held at once - and never what is computed. So a run planned against a +// budget nothing fits in (everything in series) writes what a run with room for everything writes, to +// the byte. +TEST_CASE("Rugnux_HostBudgetChangesOnlyTheSchedule", "[large]") { + const auto master = jfjoch_test::LargeDataFile("rotation_master.h5"); + if (!master) + SKIP("rotation_master.h5 not available (git-lfs data not pulled)"); + + RegisterHDF5Filter(); + const auto slurp = [](const std::string &path) { + std::ifstream f(path, std::ios::binary); + return std::string(std::istreambuf_iterator(f), std::istreambuf_iterator()); + }; + const auto run = [&](size_t budget, const std::string &prefix) { + JFJochHDF5Reader reader; + REQUIRE_NOTHROW(reader.ReadFile(*master)); + auto dataset = reader.GetDataset(); + REQUIRE(dataset); + DiffractionExperiment experiment(dataset->experiment); + IndexingSettings indexing; + indexing.Algorithm(IndexingAlgorithmEnum::Auto); + indexing.RotationIndexing(true); + experiment.ImportIndexingSettings(indexing); + + ProcessConfig config; + config.mode = ProcessMode::FullAnalysis; + config.nthreads = default_threads(); + config.spot_finding = DiffractionExperiment::DefaultDataProcessingSettings(); + config.spot_finding.indexing = true; + config.rotation_indexing = true; + config.two_pass_rotation = true; + config.rotation_postrefine_geometry = true; + config.run_scaling = true; + config.write_process_h5 = false; + config.output_prefix = prefix; + config.host_memory_budget_bytes = budget; + + Rugnux process(reader, experiment, *dataset->pixel_mask, config); + ProcessResult result; + REQUIRE_NOTHROW(result = process.Run()); + reader.Close(); + return result; + }; + const std::filesystem::path dir = std::filesystem::temp_directory_path() / "rugnux_host_budget_test"; + std::filesystem::create_directories(dir); + const ProcessResult roomy = run(size_t(1) << 40, (dir / "roomy").string()); + const ProcessResult tight = run(size_t(1) << 20, (dir / "tight").string()); + + REQUIRE(roomy.has_merge_statistics); + CHECK(tight.merge_statistics_text == roomy.merge_statistics_text); + CHECK(tight.merged_i_sigma == roomy.merged_i_sigma); + const std::string hkl = slurp((dir / "roomy.hkl").string()); + CHECK_FALSE(hkl.empty()); + CHECK(slurp((dir / "tight.hkl").string()) == hkl); + std::filesystem::remove_all(dir); + HostWorkerPool::Instance().SetSize(0); // process-wide: back to the machine's for the tests after this one +} From e2bc7a38707e9926e6075f233375d16ddfd2a747 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 17 Oct 2026 10:45:04 +0200 Subject: [PATCH 6/7] Rotation scaling: the fulls and the merge's working arrays go with each arm, and are really freed ReleaseMergeScratch and ReleaseAllButPerFrameScale assigned `= {}` to std::vector members, which is the initializer-list assignment: it empties a vector and keeps its memory, so only the struct-typed members were ever handed back. They now free (Free(v): v = T()). And every RunArm hands back the fulls, their staging and group CSR and the merge's working arrays when it returns - each is rebuilt before it is read - instead of keeping them beside the next arm's combine and through the analyses: 4nwv on the CPU at 12 threads, live peak in scaling 13.0 -> 11.9 GiB, in the analyses 11.7 -> 7.3. Output unchanged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../scale_merge/RotationScaleMerge.cpp | 11 +++- .../scale_merge/RotationScaleMerge.h | 64 ++++++++++--------- 2 files changed, 44 insertions(+), 31 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 120dc0753..8610dfc36 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -6912,7 +6912,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // to grow they are freed first: growing in place would copy the previous chain's records into // the new allocation only for the gather to overwrite them. if (static_cast(nf) > fulls.capacity()) { - fulls = {}; + Free(fulls); fulls_staging = {}; } fulls.resize(nf); @@ -7161,5 +7161,14 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // its CC1/2 is the answer without a second merge. A caller that did not ask keeps the NaN. if (measure_cc_before_corrections) r.cc_half_before_corrections = cc_half_before_corrections.value_or(r.statistics.overall.cc_half); + // The fulls and the merge's working arrays are rebuilt by every arm before they are read, so they go + // with it: kept, they sat beside the next arm's combine and through everything the caller did between + // merges - on a large cell in P1, the fulls and about 330 bytes per ASU group, gigabytes. Faulting them + // in again costs each arm a fraction of a second. + ReleaseMergeScratch(); + Free(fulls); + Free(fulls_staging); + Free(fulls_group_start); + Free(fulls_group_perm); return r; } diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index b5b90147c..fb3ad61cb 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -186,27 +186,31 @@ public: // beside others with it off, at the point where it would have run in series. void WriteBackPerFrameScale() { FinalizePerFrameScale(last_cc, last_cc_n, last_frame_scaled); } + // Hands an array (or a struct of them) back to the allocator. `v = {}` would not: on a std::vector it + // is the initializer-list assignment, which empties the vector and keeps its memory. + template static void Free(T &v) { v = T(); } + // Hand back the merge's working arrays - MergeAndStats' accumulators, R_meas sums, ISa scatter, // error-model pools and per-full fields. Each is resized and refilled before it is read, so the // next merge starts from them as an engine's first merge does, and gives the same answer. On a // large cell in P1 they are gigabytes, kept between merges only to save faulting them in again; // this is for a gap between merges where something else sits on the run's memory peak. void ReleaseMergeScratch() { - merge_fields = {}; - merge_acc = {}; - merge_accum = {}; - merge_merged_I = {}; - merge_rejected = {}; - merge_half_set = {}; - merge_obs_key = {}; - merge_wobs = {}; - merge_rmeas = {}; - merge_rmeas_staging = {}; - asymptote_scatter = {}; - em_samples = {}; - em_cut_pool = {}; - em_fit_pool = {}; - em_chi2 = {}; + Free(merge_fields); + Free(merge_acc); + Free(merge_accum); + Free(merge_merged_I); + Free(merge_rejected); + Free(merge_half_set); + Free(merge_obs_key); + Free(merge_wobs); + Free(merge_rmeas); + Free(merge_rmeas_staging); + Free(asymptote_scatter); + Free(em_samples); + Free(em_cut_pool); + Free(em_fit_pool); + Free(em_chi2); } // For an engine that will not merge again: everything on the host but the per-frame state @@ -214,21 +218,21 @@ public: // last_cc_n, last_frame_scaled). Its card memory goes with the engine. void ReleaseAllButPerFrameScale() { ReleaseMergeScratch(); - partials = {}; - fulls = {}; - fulls_staging = {}; - fulls_group_start = {}; - fulls_group_perm = {}; - perm = {}; - rawrun_start = {}; - rawrun_count = {}; - rawrun_h = {}; - rawrun_k = {}; - rawrun_l = {}; - rawrun_d = {}; - rawrun_group = {}; - finite_ok = {}; - corr_ingested = {}; + Free(partials); + Free(fulls); + Free(fulls_staging); + Free(fulls_group_start); + Free(fulls_group_perm); + Free(perm); + Free(rawrun_start); + Free(rawrun_count); + Free(rawrun_h); + Free(rawrun_k); + Free(rawrun_l); + Free(rawrun_d); + Free(rawrun_group); + Free(finite_ok); + Free(corr_ingested); obs_position.reset(); } From af1bad8f47e331602d7c105a0a33dc9e0e392ba1 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 17 Oct 2026 14:40:41 +0200 Subject: [PATCH 7/7] MR on the CPU: placement refinements reserve their grids against the host worker pool The first copies' placements are refined side by side, one thread each, and without a GPU each makes its structure factors and rigid body on the CPU - about a gigabyte of grids on a large cell, five at once at the 4nwv molecular-replacement peak (14.9 GiB live with a 16 GB machine's budget). Each now reserves 32 B per point of its grid before it starts, as a null replicate does. What each computes, and where its score goes, does not change. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../structure_refinement/MolecularReplacement.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 64d0ca0c1..aadd6eab8 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -2797,8 +2797,15 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec to_refine.push_back(&f); if (!to_refine.empty()) { std::vector> running; + // Without a GPU each placement's structure factors and rigid body are made on the CPU, a + // gigabyte or more of grids on a large cell, so each reserves them against the host's worker + // pool first (HostWorkerPool), as a null replicate does. + const size_t placement_bytes = get_gpu_count() == 0 + ? MODEL_FACTORS_BYTES_PER_GRID_POINT * ModelGridPoints(cell, &sg, d_min) + : 0; for (First *f : to_refine) running.push_back(std::async(std::launch::async, [&, f] { + const HostWorkerPool::Reservation reserved(HostWorkerPool::Instance(), placement_bytes); f->score = refine_placement(*f->placed, sg, *merge.merged, d_min, engines, nthreads, &scorers, &f->moved); f->clash = ClashFraction(*f->placed);