From 43e836573d4a73601eac00dcfae4df42127c1e2d Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 17 Oct 2026 09:31:05 +0200 Subject: [PATCH] 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);