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