Merge mr-gpumem (011b98838) into rc175: molecular replacement plans its GPU parts once per merge on model validation's budget (one rigid-body pool for every fit, the transform cropped to the indices it is read at, the pool trimmed at each stage), output unchanged

4nwv MR peak 13.77 -> 8.74 GB on the RTX 5080; myob, cytc and 4nwv outputs identical.

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

docs/CPU_DATA_ANALYSIS_DECISIONS.md

@@ -322,5 +322,10 @@ less: it may relabel the **enantiomorph** of the answer, past the anomalous chec
 reindexes; another group with the answer's rotations, or the twin hypothesis' subgroup, is written
 beside the answer as `<prefix>_mr.mtz`/`.cif` with the model as placed there, and the answer's own
 files stay what the run decided. The rotation and translation functions, the rigid body and the
-structure factors run on the GPU where there is one. The keys, thresholds and table are described in
+structure factors run on the GPU where there is one, decided for each merge before its search starts on
+the budget of §14.8 (80% of the card's total memory, less what a merge still at work holds there): the
+molecular transform, one rigid-body engine shared by every fit of the merge and the largest score engine
+take their place in that order where they still fit, more rigid-body engines are added only within half
+of the budget, and the score engines and the functions' working buffers share what is left, waiting for
+each other rather than spilling past it. The keys, thresholds and table are described in
 [the results report](RUGNUX_REPORT.md#molecular-replacement).

image_analysis/structure_refinement/ModelValidation.cpp

@@ -1883,10 +1883,14 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
 
 }  // namespace
 
+size_t ModelGPUBudget(size_t card, size_t resident) {
+    const size_t usable = card / 5 * 4;
+    return usable > resident ? usable - resident : 0;
+}
+
 ModelValidationGPUPlan PlanModelValidationGPU(const ModelValidationGPUNeed &need, size_t cards) {
     ModelValidationGPUPlan plan;
-    const size_t usable = need.card / 5 * 4;
-    plan.budget = usable > need.resident ? usable - need.resident : 0;
+    plan.budget = ModelGPUBudget(need.card, need.resident);
     auto take = [&](size_t bytes, size_t limit) {
         if (bytes == 0 || plan.bytes + bytes > limit)
             return false;
@@ -1999,7 +2003,8 @@ const gemmi::SpaceGroup *AdoptModelFrame(const ModelValidationResult &validation
 
 ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection> &merged,
                                             const gemmi::Structure &placed, const gemmi::SpaceGroup &data_space_group,
-                                            double k_sol, double b_sol, Logger &logger, size_t nthreads) {
+                                            double k_sol, double b_sol, Logger &logger, size_t nthreads,
+                                            size_t gpu_resident_bytes) {
     ModelValidationResult result;
     gemmi::Structure st = placed;
     st.setup_cell_images();
@@ -2018,13 +2023,16 @@ ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection>
     ModelStructureFactorsGPU *gpu = nullptr;
 #ifdef JFJOCH_USE_CUDA
     std::unique_ptr<ModelStructureFactorsGPU> engine;
-    // Read beside molecular replacement's engines, which keep to a quarter of the card
-    // (MolecularReplacement.cpp), so on the GPU where it fits a quarter as well.
-    if (get_gpu_count() > 0)
+    // After molecular replacement, whose engines are gone: alone on the card beside what is resident.
+    if (get_gpu_count() > 0) {
+        cuda_trim_allocation_pool();
         engine = StructureFactorsEngine(RigidBodyGPUEngine::CurrentDevice(), st.models[0], st.cell, *sg, d_min, logger);
-    if (engine && engine->DeviceBytes() > engine->DeviceTotalMemory() / 4) {
-        logger.Info("Model validation: structure factors to {:.2f} A on the CPU - {} of GPU memory, over a quarter "
-                    "of the card's {}", d_min, GBytes(engine->DeviceBytes()), GBytes(engine->DeviceTotalMemory()));
+    }
+    if (engine && engine->DeviceBytes() > ModelGPUBudget(engine->DeviceTotalMemory(), gpu_resident_bytes)) {
+        logger.Info("Model validation: structure factors to {:.2f} A on the CPU - {} of GPU memory, over the "
+                    "budget of 80% of the card's {}{}", d_min, GBytes(engine->DeviceBytes()),
+                    GBytes(engine->DeviceTotalMemory()),
+                    gpu_resident_bytes > 0 ? fmt::format(" less {} held by the merge", GBytes(gpu_resident_bytes)) : "");
         engine.reset();
     }
     if (engine)
@@ -2039,6 +2047,7 @@ ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection>
     }
 #else
     (void) logger;
+    (void) gpu_resident_bytes;
 #endif
     if (gpu == nullptr) {
         gemmi::DensityCalculator<Table, float> dc;

image_analysis/structure_refinement/ModelValidation.h

@@ -308,6 +308,11 @@ struct ModelValidationGPUPlan {
     bool may_run_beside = false;
 };
 
+// What model validation and molecular replacement may hold on a card: 80% of its TOTAL memory - never what
+// is free at the moment, so the same input on the same card always runs the same way - less what is
+// resident there already (a merge engine still at work beside them).
+size_t ModelGPUBudget(size_t card, size_t resident);
+
 // The one rule for what a validation puts on its GPU, decided before anything is reserved, from the
 // card's TOTAL memory and never from what is free at the moment, so the same input on the same card
 // always runs the same way. Everything the validation holds is alive at once - the engines are reserved
@@ -360,10 +365,12 @@ void KeepModelVerdict(ModelValidationResult &remade, const ModelValidationResult
 // F(+) - F(-) on the phases of F_calc + solvent, and its reading at the model's atoms - anomalous_pairs,
 // the scatterers' mean, the sites and anomalous_hands_disagree, by the same code and the same rule as the
 // validation's. No placement, no scale fit, no null, no file. For a caller that has already placed the
-// model and judged its fit (the --mr hand check). On the GPU where there is one.
+// model and judged its fit (the --mr hand check). On the GPU where the structure factors fit the budget
+// (ModelGPUBudget) beside `gpu_resident_bytes`.
 ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection> &merged,
                                             const gemmi::Structure &placed, const gemmi::SpaceGroup &data_space_group,
-                                            double k_sol, double b_sol, Logger &logger, size_t nthreads);
+                                            double k_sol, double b_sol, Logger &logger, size_t nthreads,
+                                            size_t gpu_resident_bytes = 0);
 
 // Reference intensities computed from an atomic model: |F_model|^2, with a flat bulk-solvent
 // contribution at the standard constants, for every reflection to `d_min`. This is the same

image_analysis/structure_refinement/MolecularReplacement.cpp

@@ -38,14 +38,15 @@
 #include "ModelScaling.h"
 #include "ModelValidation.h"
 #include "RigidBodyRefine.h"
-#ifdef JFJOCH_USE_CUDA
 #include "MolecularReplacementGPU.h"
+#ifdef JFJOCH_USE_CUDA
 #include "RigidBodyGPU.h"
 #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
 
 namespace {
@@ -171,6 +172,24 @@ int good_even_size(int n) {
     return n;
 }
 
+// The half-l grid `full` with only the indices |k|, l <= reach kept, on a grid of full.nu x (2 reach + 1) x
+// (reach + 1) that an index wraps into as it wraps into the whole one. The rows keep their length: cut to
+// 2 reach + 1 as well, the rotation function read them 15% slower on the GPU (measured on a large cell).
+gemmi::FPhiGrid<float> crop_half_l(const gemmi::FPhiGrid<float> &full, int reach) {
+    gemmi::FPhiGrid<float> out;
+    out.unit_cell = full.unit_cell;
+    out.spacegroup = full.spacegroup;
+    out.axis_order = full.axis_order;
+    out.half_l = true;
+    const int nu = full.nu, n = 2 * reach + 1;
+    out.set_size_without_checking(nu, n, reach + 1);
+    for (int l = 0; l <= reach; l++)
+        for (int k = -reach; k <= reach; k++)
+            std::copy_n(&full.data[static_cast<size_t>(nu) * (gemmi::modulo(k, full.nv) + static_cast<size_t>(full.nv) * l)],
+                        nu, &out.data[static_cast<size_t>(nu) * (gemmi::modulo(k, n) + static_cast<size_t>(n) * l)]);
+    return out;
+}
+
 // --- placements scored against the data: R-work / R-free after the bulk solvent and the scale ---
 
 struct PlacementScore {
@@ -282,16 +301,22 @@ PlacementScore r_factors(const gemmi::AsuData<std::complex<float>> &fmodel, cons
 
 // Model validation's structure-factor engines on the GPU, one per group, copy count and resolution, made
 // where the first placement of that kind is scored, each with the scale fit on the GPU beside it - or none,
-// where there is no GPU, the cell is too small for its gridding or the two would take over a quarter of the
-// card. Decided on the card's total memory, so the same input always takes the same path.
+// where there is no GPU, the cell is too small for its gridding or the two do not fit the room the search's
+// plan leaves them (PlanMolecularReplacementGPU). Decided on the engine's size, so the same input always takes
+// the same path.
 //
-// Together they keep within that quarter too. A search asks for one per group, copy count and resolution,
-// and kept to the end of the search they summed to well over the card on a large cell (seven of 1.4 GB
-// beside the rigid-body engines). So an engine that does not fit beside the ones kept takes the place of
-// the least recently used of those no score is using, and waits for a score to finish where every one is
-// in use; one taken out is made again where it is next asked for. Every engine computes the same numbers
-// whenever it is made, so this changes when the card holds them, not what a score gives.
+// Together they keep within that room, which the rotation and translation functions' working buffers share
+// (Hold). A search asks for one per group, copy count and resolution, and kept to the end of the search they
+// summed to well over the card on a large cell (seven of 1.4 GB beside the rigid-body engines). So an engine,
+// or a function's buffers, that do not fit beside what is held take the place of the least recently used
+// engine no score is using, and wait for a score or a function to finish where every one is in use; an engine
+// taken out is made again where it is next asked for. Every engine computes the same numbers whenever it is
+// made, so this changes when the card holds them, not what a score gives.
 #ifdef JFJOCH_USE_CUDA
+std::string in_gb(size_t bytes) {
+    return fmt::format("{:.2f} GB", bytes / 1e9);
+}
+
 class ScoreEngines {
 public:
     struct Engine {
@@ -304,50 +329,63 @@ public:
         CudaDevicePtr<std::complex<float>> fcalc, fmask;
     };
 
+    // The room of a merge's plan, set before its placements are scored; `scores`: whether they go on the GPU.
+    void SetRoom(size_t room, bool scores) {
+        std::lock_guard lock(m_);
+        room_ = room;
+        scores_ = scores;
+    }
+
+    bool Fits(size_t bytes) {
+        std::lock_guard lock(m_);
+        return bytes <= room_;
+    }
+
     // The engine for this kind of placement, kept on the card for as long as the returned pointer lives.
     std::shared_ptr<Engine> get(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, double d_min) {
         if (get_gpu_count() == 0)
             return nullptr;
-        size_t atoms = 0;
-        for (const gemmi::Chain &ch : st.models[0].chains)
-            for (const gemmi::Residue &r : ch.residues)
-                atoms += r.atoms.size();
         std::unique_lock lock(m_);
-        Entry &e = engines_[{sg.xhm(), atoms, d_min}];
+        if (!scores_)
+            return nullptr;
+        Entry &e = engines_[{sg.xhm(), gemmi::count_atom_sites(st.models[0]), d_min}];
         e.last_use = ++clock_;
         // Made by the first to ask for it, without holding up those asking for other ones.
         cv_.wait(lock, [&] { return !e.making; });
         if (e.engine || e.refused)
             return Lease(e);
         e.making = true;
+        const size_t room = room_;
         lock.unlock();
 
         std::unique_ptr<ModelStructureFactorsGPU> sf;
-        size_t points = 0, bytes = 0, quarter = 0;
+        size_t points = 0, bytes = 0;
         try {
             sf = std::make_unique<ModelStructureFactorsGPU>(RigidBodyGPUEngine::CurrentDevice(), st.models[0],
                                                             st.cell, sg, d_min);
             points = sf->Reflections();
-            bytes = sf->DeviceBytes() + ModelScaleGPU::DeviceBytes(points) + 2 * points * sizeof(std::complex<float>);
-            quarter = sf->DeviceTotalMemory() / 4;
+            bytes = Bytes(*sf);
         } catch (const std::exception &) {
             cuda_clear_error();
             sf.reset();
         }
         lock.lock();
-        if (!sf || !sf->Supported() || bytes > quarter) {
+        if (!sf || !sf->Supported() || bytes > room) {
             e.refused = true;
             e.making = false;
             cv_.notify_all();
             return nullptr;
         }
         const int card = sf->Device();
-        cv_.wait(lock, [&] { return MakeRoom(card, bytes, quarter); });
+        cv_.wait(lock, [&] { return MakeRoom(card, bytes); });
         resident_[card] += bytes;
         lock.unlock();
 
         std::shared_ptr<Engine> engine;
         try {
+            // The functions' buffers come from the stream-ordered pool, which keeps what they gave back; the
+            // structure factors are allocated outside it, and could not use it.
+            cuda_trim_allocation_pool();
             sf->Reserve();
             const int device = RigidBodyGPUEngine::CurrentDevice();
             set_gpu(card);
@@ -376,6 +414,39 @@ public:
         return Lease(e);
     }
 
+    // What a score engine for `sf` takes: its structure factors, the scale fit and the two arrays between them.
+    static size_t Bytes(const ModelStructureFactorsGPU &sf) {
+        const size_t points = sf.Reflections();
+        return sf.DeviceBytes() + ModelScaleGPU::DeviceBytes(points) + 2 * points * sizeof(std::complex<float>);
+    }
+
+    // `bytes` of the room on the current card for a function's working buffers, held for as long as the result
+    // lives.
+    class Held {
+    public:
+        Held(ScoreEngines &owner, int card, size_t bytes) : owner_(owner), card_(card), bytes_(bytes) {}
+        ~Held() {
+            {
+                std::lock_guard lock(owner_.m_);
+                owner_.resident_[card_] -= bytes_;
+            }
+            owner_.cv_.notify_all();
+        }
+        Held(const Held &) = delete;
+        Held &operator=(const Held &) = delete;
+    private:
+        ScoreEngines &owner_;
+        int card_;
+        size_t bytes_;
+    };
+    std::unique_ptr<Held> Hold(size_t bytes) {
+        const int card = RigidBodyGPUEngine::CurrentDevice();
+        std::unique_lock lock(m_);
+        cv_.wait(lock, [&] { return MakeRoom(card, bytes); });
+        resident_[card] += bytes;
+        return std::make_unique<Held>(*this, card, bytes);
+    }
+
 private:
     struct Entry {
         std::shared_ptr<Engine> engine;
@@ -386,7 +457,7 @@ private:
     };
 
     // The entry's engine for a score: the copy it holds is what marks the engine in use, and handing it back
-    // is what may let an engine waiting for room in.
+    // is what may let an engine or a function waiting for room in.
     std::shared_ptr<Engine> Lease(Entry &e) {
         if (!e.engine)
             return nullptr;
@@ -399,10 +470,10 @@ private:
         });
     }
 
-    // Whether `bytes` more fit the card's quarter beside the engines kept there, once the least recently
-    // used of those no score is using have been taken out. Under m_.
-    bool MakeRoom(int card, size_t bytes, size_t quarter) {
-        while (resident_[card] + bytes > quarter) {
+    // Whether `bytes` more fit the room on the card beside what is held there, once the least recently used
+    // engines no score is using have been taken out. Under m_.
+    bool MakeRoom(int card, size_t bytes) {
+        while (resident_[card] + bytes > room_) {
             Entry *oldest = nullptr;
             for (auto &[key, entry] : engines_)
                 if (entry.engine && entry.card == card && entry.engine.use_count() == 1
@@ -418,6 +489,8 @@ private:
 
     std::map<std::tuple<std::string, size_t, double>, Entry> engines_;
     std::map<int, size_t> resident_;
+    size_t room_ = 0;
+    bool scores_ = false;
     uint64_t clock_ = 0;
     std::mutex m_;
     std::condition_variable cv_;
@@ -582,37 +655,6 @@ PlacementScore score_placement(const gemmi::Structure &st, const gemmi::SpaceGro
     return out;
 }
 
-// The rigid-body engines for one model in one group, where there is a GPU.
-struct RigidBodyEngines {
-#ifdef JFJOCH_USE_CUDA
-    std::unique_ptr<RigidBodyGPUPool> pool;
-#endif
-    RigidBodyGPUPool *get() {
-#ifdef JFJOCH_USE_CUDA
-        return pool.get();
-#else
-        return nullptr;
-#endif
-    }
-};
-
-RigidBodyEngines make_engines(const gemmi::Structure &st, const gemmi::SpaceGroup &sg,
-                              const std::vector<MergedReflection> &merged, double d_min, size_t engines,
-                              Logger &logger) {
-    RigidBodyEngines out;
-#ifdef JFJOCH_USE_CUDA
-    const double finest = RigidBodyLadder(d_min).back();
-    size_t observations = 0;
-    for (const MergedReflection &r : merged)
-        if (!std::isnan(r.F) && r.d >= finest)
-            ++observations;
-    out.pool = RigidBodyGPUPool::Create(st.models[0], st.cell, sg, d_min, observations, engines, logger);
-#else
-    (void) st; (void) sg; (void) merged; (void) d_min; (void) engines; (void) logger;
-#endif
-    return out;
-}
-
 // The placement through the rigid body, and its R-factors where the rigid body left it. A search
 // placement is a few degrees and a grid step from its answer, well inside the rigid body's reach, so
 // the step is taken without the R-free check model validation makes before committing one.
@@ -642,21 +684,20 @@ PlacementScore refine_placement(gemmi::Structure &st, const gemmi::SpaceGroup &s
 // The copies of `assembly` (its chains split evenly, in order) as a crystal of group `sg`, each refined beside
 // the others, and scored.
 PlacementScore carry_copies(const gemmi::Structure &assembly, int copies, const gemmi::SpaceGroup &sg,
-                            const std::vector<MergedReflection> &merged, double d_min, size_t nthreads,
-                            ScoreEngines *scorers, Logger &logger) {
+                            const std::vector<MergedReflection> &merged, double d_min, RigidBodyGPUPool *pool,
+                            size_t nthreads, ScoreEngines *scorers) {
     std::vector<gemmi::Structure> body = SplitCopies(assembly, copies);
     for (gemmi::Structure &b : body) {
         b.spacegroup_hm = sg.xhm();
         b.setup_cell_images();
     }
-    RigidBodyEngines engines = make_engines(body[0], sg, merged, d_min, 1, logger);
     for (int k = 0; k < copies; k++) {
         std::vector<gemmi::Structure> rest;
         for (int o = 0; o < copies; o++)
             if (o != k)
                 rest.push_back(body[o]);
         const gemmi::Structure fixed = AssembleCopies(rest);
-        refine_rigid_body(body[k], sg, merged, d_min, engines.get(), nthreads, nullptr, &fixed.models[0]);
+        refine_rigid_body(body[k], sg, merged, d_min, pool, nthreads, nullptr, &fixed.models[0]);
     }
     return score_placement(AssembleCopies(body), sg, merged, d_min, nthreads, scorers);
 }
@@ -1154,12 +1195,18 @@ MolecularTransform::MolecularTransform(const SearchModel &model, double d_min, s
     dc.set_grid_cell_and_spacegroup(st);
     dc.set_refmac_compatible_blur(st.models[0]);
     PutModelDensityOnGrid(dc, st.models[0], {}, nthreads);
+    gemmi::FPhiGrid<float> full;
 #ifdef JFJOCH_USE_CUDA
-    if (get_gpu_count() > 0 && MapToFPhiGPUFits(dc.grid))
-        grid_ = MapToFPhiGPU(dc.grid);
-    else
+    if (get_gpu_count() > 0 && MapToFPhiGPUFits(dc.grid)) {
+        full = MapToFPhiGPU(dc.grid);
+        cuda_trim_allocation_pool();   // its buffers, which the engines allocated outside the pool cannot use
+    } else
 #endif
-        grid_ = MapToFPhi(dc.grid);
+        full = MapToFPhi(dc.grid);
+    // The transform is read at |s| <= 1 / d_min only - the search's reflections, turned - and interpolation
+    // reads one step further out: of the grid the FFT made, under half, which is all that is kept, here and
+    // on the GPU. The 2% are for a reflection's resolution taken in a cell slightly different from its search's.
+    grid_ = crop_half_l(full, static_cast<int>(std::ceil(1.02 * box_ / d_min)) + 1);
     d_min_ = d_min;
     blur_ = dc.blur;
 }
@@ -2177,9 +2224,64 @@ double ExpectedLogLikelihoodGain(const SearchModel &model, const MRMerge &merge,
     return ellg;
 }
 
+#ifdef JFJOCH_USE_CUDA
+// What one merge's placements need of the GPU, worked out on the CPU: the rigid body's pool - one for every fit
+// of the merge, sized for the largest, the most copies the cell holds (Matthews) in any of its groups, planned
+// but without engines; null where the GPU cannot grid them - and the largest score engine.
+struct MergeGPUSize {
+    std::unique_ptr<RigidBodyGPUPool> pool;
+    size_t rigid_body = 0;
+    size_t rigid_body_engines = 0;
+    size_t structure_factors = 0;
+};
+
+MergeGPUSize size_merge_gpu(const SearchModel &model, const MRMerge &merge, double d_min, int device) {
+    set_gpu(device);
+    MergeGPUSize out;
+    const gemmi::UnitCell cell(merge.cell);
+    const std::vector<gemmi::SpaceGroup> groups = SameRotationGroups(merge.space_group);
+    const std::vector<double> pm = MatthewsProbabilities(cell.volume / static_cast<double>(merge.space_group.operations().order()),
+                                                         model.mass_da);
+    int copies_max = 1;
+    for (size_t n = 0; n < pm.size(); n++)
+        if (pm[n] >= MATTHEWS_P_MIN)
+            copies_max = static_cast<int>(n) + 1;
+    // Every fit is given the reflections of its finest zone.
+    size_t observations = 0;
+    for (const MergedReflection &r : *merge.merged)
+        if (!std::isnan(r.F) && r.d >= RigidBodyLadder(d_min).back())
+            ++observations;
+    // As many fits as run at once: a group's first copies side by side, the groups side by side, or a null's
+    // replicates.
+    out.rigid_body_engines = std::max<size_t>(groups.size() * RESCORED_PER_GROUP, NULL_REPLICATES);
+    Logger quiet = Logger::Buffered();
+    bool rigid_body_gpu = true;   // every group's fits can be gridded
+    for (const gemmi::SpaceGroup &sg : groups) {
+        const gemmi::Structure one = PlaceSearchModel(model, gemmi::Mat33(), gemmi::Fractional(0, 0, 0), cell, sg);
+        const gemmi::Structure most = AssembleCopies(std::vector<gemmi::Structure>(copies_max, one));
+        for (const gemmi::Structure *st : {&one, &most}) {
+            if (!out.pool)
+                out.pool = RigidBodyGPUPool::Plan(st->models[0], cell, sg, d_min, observations, quiet);
+            rigid_body_gpu = rigid_body_gpu && out.pool && out.pool->Cover(st->models[0], sg, quiet);
+        }
+        try {   // as ScoreEngines::get() makes it, which leaves the scores on the CPU where it cannot
+            const ModelStructureFactorsGPU sf(device, most.models[0], cell, sg, d_min);
+            if (sf.Supported())
+                out.structure_factors = std::max(out.structure_factors, ScoreEngines::Bytes(sf));
+        } catch (const std::exception &) {
+            cuda_clear_error();
+        }
+    }
+    if (rigid_body_gpu)
+        out.rigid_body = out.pool->EngineBytes();
+    return out;
+}
+#endif
+
 // The search and the decision for one search model.
+// `gpu_budget`: what the search may hold on its card (PlanMolecularReplacementGPU).
 MolecularReplacementResult search_model(const SearchModel &model, const std::vector<MRMerge> &merges, Logger &logger,
-                                        size_t nthreads) {
+                                        size_t nthreads, size_t gpu_budget) {
     MolecularReplacementResult result;
     result.residues_read = model.residues_read;
     result.residues_kept = model.residues_kept;
@@ -2215,6 +2317,13 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
         return result;
     }
 
+#ifdef JFJOCH_USE_CUDA
+    // Each merge's GPU sizes; the answer merge's found on the CPU while the transform is made.
+    std::vector<std::future<MergeGPUSize>> sized(merges.size());
+    if (get_gpu_count() > 0)
+        sized[0] = std::async(std::launch::async, size_merge_gpu, std::cref(model), std::cref(merges[0]),
+                              result.d_min_score, RigidBodyGPUEngine::CurrentDevice());
+#endif
     t0 = std::chrono::steady_clock::now();
     std::optional<MolecularTransform> ft_made;
     try {
@@ -2224,11 +2333,12 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
         return result;
     }
     const MolecularTransform &ft = *ft_made;
+    [[maybe_unused]] const size_t transform_bytes = ft.Grid().data.size() * sizeof(std::complex<float>);
 #ifdef JFJOCH_USE_CUDA
-    // The rotation and translation functions on the GPU where there is one; the rest of the search is
-    // the same either way.
+    // The rotation and translation functions on the GPU where the transform fits the budget (the plan's first
+    // part, PlanMolecularReplacementGPU); the rest of the search is the same either way.
     std::unique_ptr<MolecularReplacementGPU> gpu;
-    if (get_gpu_count() > 0) {
+    if (get_gpu_count() > 0 && transform_bytes <= gpu_budget) {
         try {
             gpu = std::make_unique<MolecularReplacementGPU>(ft);
         } catch (const JFJochException &e) {
@@ -2245,6 +2355,85 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                 result.tncs.present ? " - tNCS" : "");
     const std::vector<gemmi::Mat33> sphere = UniformRotations(ROTATION_STEP_DEG);
     ScoreEngines scorers;
+#ifdef JFJOCH_USE_CUDA
+    // Each merge's rigid-body engines, reserved by its plan and kept until its choice is made (choose_in).
+    std::vector<std::unique_ptr<RigidBodyGPUPool>> rigid_body(merges.size());
+#endif
+    const auto rigid_body_pool = [&](size_t m) -> RigidBodyGPUPool * {
+#ifdef JFJOCH_USE_CUDA
+        return rigid_body[m].get();
+#else
+        (void) m;
+        return nullptr;
+#endif
+    };
+    // What merge m's search and placements put on the GPU, decided at once before its search starts
+    // (PlanMolecularReplacementGPU), on the sizes size_merge_gpu() found - for the answer's merge while the
+    // transform was made.
+    const auto plan_gpu = [&](size_t m) {
+#ifdef JFJOCH_USE_CUDA
+        if (get_gpu_count() == 0)
+            return;
+        for (auto &pool : rigid_body)
+            pool.reset();   // the other merge's choice is made
+        if (!sized[m].valid())
+            sized[m] = std::async(std::launch::deferred, size_merge_gpu, std::cref(model), std::cref(merges[m]),
+                                  result.d_min_score, RigidBodyGPUEngine::CurrentDevice());
+        MergeGPUSize size = sized[m].get();
+        MolecularReplacementGPUNeed need;
+        need.budget = gpu_budget;
+        need.transform = transform_bytes;
+        need.rigid_body = size.rigid_body;
+        need.rigid_body_engines = size.rigid_body_engines;
+        need.structure_factors = size.structure_factors;
+        const MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need);
+        scorers.SetRoom(plan.room, plan.structure_factors);
+        if (plan.rigid_body_engines > 0) {
+            size.pool->Reserve(plan.rigid_body_engines);
+            rigid_body[m] = std::move(size.pool);
+        }
+        logger.Info("Molecular replacement: {} on the GPU: transform {}, rigid body {}, scores {} - {} held of a {} "
+                    "budget, {} left for the score engines and the functions' buffers",
+                    merges[m].name, plan.search ? in_gb(need.transform) : "on the CPU",
+                    plan.rigid_body_engines > 0 ? fmt::format("{} x {}", plan.rigid_body_engines, in_gb(need.rigid_body))
+                                                : "on the CPU",
+                    plan.structure_factors ? "up to " + in_gb(need.structure_factors) + " each" : "on the CPU",
+                    in_gb(plan.bytes), in_gb(gpu_budget), in_gb(plan.room));
+#else
+        (void) m;
+#endif
+    };
+    // A rotation or a translation function on the GPU where the search is there and the call's working buffers
+    // fit the room (held for the call), on the CPU otherwise - the same function, equal to rounding.
+    const auto rotation_function = [&](const SearchReflections &refl, const gemmi::UnitCell &cell,
+                                       const gemmi::SpaceGroup &sg, const std::vector<gemmi::Mat33> &rotations) {
+#ifdef JFJOCH_USE_CUDA
+        const size_t bytes = MolecularReplacementGPU::RotationBytes(refl, sg, rotations.size());
+        if (gpu && scorers.Fits(bytes)) {
+            const auto held = scorers.Hold(bytes);
+            return gpu->RotationFunction(refl, cell, sg, rotations);
+        }
+#endif
+        return RotationFunction(ft, refl, cell, sg, rotations, nthreads);
+    };
+    const auto translation_function = [&](const TranslationSetup &setup, const std::vector<gemmi::Mat33> &rotations,
+                                          const std::vector<std::complex<double>> *known,
+                                          TranslationSetupGPU *prepared) {
+#ifdef JFJOCH_USE_CUDA
+        const size_t bytes = MolecularReplacementGPU::TranslationBytes(setup, rotations.size(), known != nullptr);
+        if (gpu && scorers.Fits(bytes)) {
+            const auto held = scorers.Hold(bytes);
+            return gpu->TranslationFunction(setup, rotations, known, prepared);
+        }
+#else
+        (void) prepared;
+#endif
+        std::vector<TranslationPeak> tf(rotations.size());
+        ParallelFor(static_cast<int>(rotations.size()), nthreads, [&](int i) {
+            tf[i] = TranslationFunction(ft, setup, rotations[i], known);
+        });
+        return tf;
+    };
 
     std::vector<gemmi::Mat33> orientations;   // of each candidate's placement, before the rigid body
     // One merge: the rotation function over its point group, then for each of its groups the
@@ -2271,12 +2460,8 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
             d_llg = std::max(d_llg, std::min(result.d_min_search, ds[LLG_TERMS - 1]));
         const SearchReflections lrefl = MakeSearchReflections(*merge.merged, cell, merge.space_group, d_llg,
                                                               result.d_max_search);
-#ifdef JFJOCH_USE_CUDA
-        const std::vector<double> rf = gpu ? gpu->RotationFunction(refl, cell, merge.space_group, rotations)
-                                           : RotationFunction(ft, refl, cell, merge.space_group, rotations, nthreads);
-#else
-        const std::vector<double> rf = RotationFunction(ft, refl, cell, merge.space_group, rotations, nthreads);
-#endif
+        plan_gpu(m);
+        const std::vector<double> rf = rotation_function(refl, cell, merge.space_group, rotations);
         double rf_mean = 0, rf_sd = 0;
         for (double x : rf)
             rf_mean += x;
@@ -2318,12 +2503,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                         for (int axis = 0; axis < 3; axis++)
                             for (double sign : {1.0, -1.0})
                                 trial.push_back(axis_rotation(axis, sign * step).multiply(peak_rotations[k]));
-#ifdef JFJOCH_USE_CUDA
-                const std::vector<double> v = gpu ? gpu->RotationFunction(refl, cell, merge.space_group, trial)
-                                                  : RotationFunction(ft, refl, cell, merge.space_group, trial, nthreads);
-#else
-                const std::vector<double> v = RotationFunction(ft, refl, cell, merge.space_group, trial, nthreads);
-#endif
+                const std::vector<double> v = rotation_function(refl, cell, merge.space_group, trial);
                 size_t next = 0;
                 for (size_t k = 0; k < peaks.size(); k++) {
                     if (!active[k])
@@ -2354,17 +2534,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
         std::vector<std::future<std::vector<TranslationPeak>>> tf_running;
         for (const gemmi::SpaceGroup &sg : groups)
             tf_running.push_back(std::async(std::launch::async, [&, sg] {
-                const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg);
-                std::vector<TranslationPeak> tf(peaks.size());
-#ifdef JFJOCH_USE_CUDA
-                if (gpu)
-                    tf = gpu->TranslationFunction(setup, peak_rotations);
-                else
-#endif
-                    ParallelFor(static_cast<int>(peaks.size()), nthreads, [&](int i) {
-                        tf[i] = TranslationFunction(ft, setup, peak_rotations[i]);
-                    });
-                return tf;
+                return translation_function(MakeTranslationSetup(refl, cell, sg), peak_rotations, nullptr, nullptr);
             }));
         std::vector<TranslationPeak> tf_all;
         for (auto &r : tf_running) {
@@ -2406,19 +2576,9 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
             std::iota(by_tf.begin(), by_tf.end(), 0);
             std::stable_sort(by_tf.begin(), by_tf.end(), [&](size_t a, size_t b) { return tfs[a].z > tfs[b].z; });
             if (!seeds.empty()) {
-                const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg);
-                std::vector<TranslationPeak> tf_seed(seeds.size());
-#ifdef JFJOCH_USE_CUDA
-                const std::vector<double> rf_seed = gpu ? gpu->RotationFunction(refl, cell, merge.space_group, seeds)
-                                                        : RotationFunction(ft, refl, cell, merge.space_group, seeds, nthreads);
-                if (gpu)
-                    tf_seed = gpu->TranslationFunction(setup, seeds);
-                else
-#else
-                const std::vector<double> rf_seed = RotationFunction(ft, refl, cell, merge.space_group, seeds, nthreads);
-#endif
-                    for (size_t k = 0; k < seeds.size(); k++)
-                        tf_seed[k] = TranslationFunction(ft, setup, seeds[k]);
+                const std::vector<double> rf_seed = rotation_function(refl, cell, merge.space_group, seeds);
+                const std::vector<TranslationPeak> tf_seed = translation_function(MakeTranslationSetup(refl, cell, sg),
+                                                                                  seeds, nullptr, nullptr);
                 std::vector<size_t> first;
                 for (size_t k = 0; k < seeds.size(); k++) {
                     first.push_back(rots.size());
@@ -2447,7 +2607,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
             };
 
             // --- the first copy: the best few through packing, the rigid body and R-free ---
-            RigidBodyEngines engines;
+            RigidBodyGPUPool *const engines = rigid_body_pool(m);
             bool have = false;
             gemmi::Mat33 orientation;
             // Where not even one copy fits the cell, only the best translation-function placement is
@@ -2497,11 +2657,10 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                 if (f.refine)
                     to_refine.push_back(&f);
             if (!to_refine.empty()) {
-                engines = make_engines(*to_refine[0]->placed, sg, *merge.merged, d_min, to_refine.size(), logger);
                 std::vector<std::future<void>> running;
                 for (First *f : to_refine)
                     running.push_back(std::async(std::launch::async, [&, f] {
-                        f->score = refine_placement(*f->placed, sg, *merge.merged, d_min, engines.get(), nthreads,
+                        f->score = refine_placement(*f->placed, sg, *merge.merged, d_min, engines, nthreads,
                                                     &scorers, &f->moved);
                         f->clash = ClashFraction(*f->placed);
                     }));
@@ -2555,9 +2714,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
             // --- further copies, one at a time, with the copies placed as a known partial structure ---
             // Each copy is refined on its own and all of them scored together.
             std::optional<TranslationSetup> setup2;
-#ifdef JFJOCH_USE_CUDA
             std::unique_ptr<TranslationSetupGPU, void (*)(TranslationSetupGPU *)> setup2_gpu{nullptr, nullptr};
-#endif
             // The first `count` copies refined in turn, each with all the others fixed, and the assembly scored.
             const auto refine_beside_others = [&](int count) {
                 for (int c = 0; c < count; c++) {
@@ -2567,7 +2724,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                             rest.push_back(copies[o]);
                     const gemmi::Structure fixed = AssembleCopies(rest);
                     gemmi::Mat33 moved;
-                    refine_rigid_body(copies[c], sg, *merge.merged, d_min, engines.get(), nthreads, &moved,
+                    refine_rigid_body(copies[c], sg, *merge.merged, d_min, engines, nthreads, &moved,
                                       &fixed.models[0]);
                     copy_rotation[c] = moved.multiply(copy_rotation[c]);
                     copy_position[c] = centroid(copies[c]);
@@ -2610,17 +2767,11 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                     }
                     rknown_llg = RefinedLogLikelihoodGain(rsetup(), rknown, model.error_A, &known_scale);
                 });
-                std::vector<TranslationPeak> tf2(rots2.size());
 #ifdef JFJOCH_USE_CUDA
-                if (gpu) {
-                    if (!setup2_gpu)
-                        setup2_gpu = gpu->Prepare(*setup2, true);
-                    tf2 = gpu->TranslationFunction(*setup2, rots2, &known, setup2_gpu.get());
-                } else
+                if (gpu && !setup2_gpu)
+                    setup2_gpu = gpu->Prepare(*setup2, true);
 #endif
-                    ParallelFor(static_cast<int>(rots2.size()), nthreads, [&](int i) {
-                        tf2[i] = TranslationFunction(ft, *setup2, rots2[i], &known);
-                    });
+                const std::vector<TranslationPeak> tf2 = translation_function(*setup2, rots2, &known, setup2_gpu.get());
                 std::vector<size_t> by_tf2(rots2.size());
                 std::iota(by_tf2.begin(), by_tf2.end(), 0);
                 std::stable_sort(by_tf2.begin(), by_tf2.end(), [&](size_t a, size_t b) { return tf2[a].z > tf2[b].z; });
@@ -2691,7 +2842,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                 for (Trial &trial : trials) {
                     std::vector<gemmi::Structure> all = copies;
                     gemmi::Mat33 moved;
-                    refine_rigid_body(trial.copy, sg, *merge.merged, d_min, engines.get(), nthreads, &moved,
+                    refine_rigid_body(trial.copy, sg, *merge.merged, d_min, engines, nthreads, &moved,
                                       &others.models[0]);
                     all.push_back(trial.copy);
                     trial.assembly = AssembleCopies(all);
@@ -2837,8 +2988,6 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
         // The null of candidate i.
         const auto run_null = [&](int i) {
             const MRCandidate &c = result.candidates[i];
-            RigidBodyEngines engines = make_engines(*c.placed, c.space_group, *merges[merge].merged,
-                                                    result.d_min_score, NULL_REPLICATES, logger);
             // One copy: its replicates placed by the translation function, as it was.
             std::function<std::vector<gemmi::Fractional>(const std::vector<gemmi::Mat33> &)> translate;
             if (c.copies == 1)
@@ -2847,24 +2996,16 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                                                                          gemmi::UnitCell(merges[merge].cell),
                                                                          merges[merge].space_group,
                                                                          result.d_min_search, result.d_max_search);
-                    const TranslationSetup setup = MakeTranslationSetup(refl, gemmi::UnitCell(merges[merge].cell),
-                                                                        c.space_group);
-                    std::vector<TranslationPeak> tf(turned.size());
-#ifdef JFJOCH_USE_CUDA
-                    if (gpu)
-                        tf = gpu->TranslationFunction(setup, turned);
-                    else
-#endif
-                        ParallelFor(static_cast<int>(turned.size()), nthreads, [&](int i) {
-                            tf[i] = TranslationFunction(ft, setup, turned[i]);
-                        });
+                    const std::vector<TranslationPeak> tf = translation_function(
+                            MakeTranslationSetup(refl, gemmi::UnitCell(merges[merge].cell), c.space_group), turned,
+                            nullptr, nullptr);
                     std::vector<gemmi::Fractional> out;
                     for (const TranslationPeak &p : tf)
                         out.push_back(p.t);
                     return out;
                 };
             return null_sigma(*c.placed, c.copies, orientations[i], c.space_group, *merges[merge].merged,
-                              result.d_min_score, engines.get(), nthreads, &scorers, translate);
+                              result.d_min_score, rigid_body_pool(merge), nthreads, &scorers, translate);
         };
         // The first candidate's null alone; where it fails, the nulls of all the other placed candidates side
         // by side, each taken where the choice comes to it.
@@ -2889,7 +3030,6 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
             std::vector<Option> options{{gemmi::Fractional(0, 0, 0), c.placed,
                                          score_placement(*c.placed, c.space_group, merged, d_min, nthreads, &scorers,
                                                          nullptr, nullptr, &v), c.clash_fraction}};
-            RigidBodyEngines engines;
             for (int o = 0; o < 8; o++) {
                 const gemmi::Fractional shift((v.x + 0.5 * (o & 1)) / 2, (v.y + 0.5 * ((o >> 1) & 1)) / 2,
                                               (v.z + 0.5 * ((o >> 2) & 1)) / 2);
@@ -2907,9 +3047,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                 SetModelPositions(moved->models[0], pos);
                 Option opt{shift, moved, {}, ClashFraction(*moved, c.copies)};
                 if (opt.clash <= 2 * CLASH_MAX) {
-                    if (engines.get() == nullptr)
-                        engines = make_engines(*moved, c.space_group, merged, d_min, 1, logger);
-                    refine_rigid_body(*moved, c.space_group, merged, d_min, engines.get(), nthreads);
+                    refine_rigid_body(*moved, c.space_group, merged, d_min, rigid_body_pool(merge), nthreads);
                     opt.clash = ClashFraction(*moved, c.copies);
                     opt.score = score_placement(*moved, c.space_group, merged, d_min, nthreads, &scorers, nullptr,
                                                 nullptr, &v);
@@ -3000,7 +3138,8 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                 // they build the same one.
                 const MRCandidate &b = result.candidates[best];
                 const PlacementScore carried = carry_copies(*b.placed, b.copies, c.space_group, *merges[merge].merged,
-                                                            result.d_min_score, nthreads, &scorers, logger);
+                                                            result.d_min_score, rigid_body_pool(merge), nthreads,
+                                                            &scorers);
                 logger.Info("Molecular replacement: {}'s {} copies carried into {}: R-free {:.3f} against {:.3f}",
                             b.space_group.xhm(), b.copies, c.space_group.xhm(), carried.r_free, b.r_free);
                 if (!(carried.r_free - b.r_free >= R_FREE_MARGIN) || !(b.r_free < CARRY_R_FREE_MAX)) {
@@ -3125,9 +3264,40 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
 
 }  // namespace
 
+MolecularReplacementGPUPlan PlanMolecularReplacementGPU(const MolecularReplacementGPUNeed &need) {
+    MolecularReplacementGPUPlan plan;
+    auto take = [&](size_t bytes, size_t limit) {
+        if (bytes == 0 || plan.bytes + bytes > limit)
+            return false;
+        plan.bytes += bytes;
+        return true;
+    };
+    plan.search = take(need.transform, need.budget);
+    plan.rigid_body_engines = need.rigid_body_engines > 0 && take(need.rigid_body, need.budget) ? 1 : 0;
+    plan.structure_factors = take(need.structure_factors, need.budget);
+    // More engines only within half of the budget (see the declaration).
+    while (plan.rigid_body_engines > 0 && plan.rigid_body_engines < need.rigid_body_engines
+           && take(need.rigid_body, need.budget / 2))
+        plan.rigid_body_engines++;
+    plan.room = need.budget - (plan.search ? need.transform : 0) - plan.rigid_body_engines * need.rigid_body;
+    return plan;
+}
+
 MolecularReplacementResult RunMolecularReplacement(const std::string &model_path, const std::vector<MRMerge> &merges,
-                                                   Logger &logger, size_t nthreads) {
+                                                   Logger &logger, size_t nthreads, size_t gpu_resident_bytes) {
     const auto t0 = std::chrono::steady_clock::now();
+    // What the search may hold on the card: the budget model validation keeps to. What the stream-ordered pool
+    // holds with no buffer in it - freed by the steps before - goes back to the driver first, as before a
+    // validation: the engines are allocated outside the pool and cannot use it.
+    size_t gpu_budget = 0;
+#ifdef JFJOCH_USE_CUDA
+    if (get_gpu_count() > 0) {
+        cuda_trim_allocation_pool();
+        gpu_budget = ModelGPUBudget(get_gpu_total_memory(RigidBodyGPUEngine::CurrentDevice()), gpu_resident_bytes);
+    }
+#else
+    (void) gpu_resident_bytes;
+#endif
     SearchModel model = PrepareSearchModel(model_path);
     if (!model.failure_reason.empty()) {
         MolecularReplacementResult result;
@@ -3145,7 +3315,7 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
     if (!pae_path.empty())
         logger.Info("Molecular replacement: {} rigid domain(s) in the predicted aligned error {}", domains.size(), pae_path);
 
-    MolecularReplacementResult result = search_model(model, merges, logger, nthreads);
+    MolecularReplacementResult result = search_model(model, merges, logger, nthreads, gpu_budget);
     result.model_path = model_path;
     result.pae_domains = static_cast<int>(domains.size());
     result.t_model += t_model;
@@ -3155,14 +3325,16 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
     // own.
     if (result.ok && result.chosen < 0 && result.decision.starts_with("NOT_PLACED") && domains.size() >= 2) {
         const auto t_domains = std::chrono::steady_clock::now();
-        // Each domain on a thread of its own, its log kept and replayed in order.
+        // Each domain on a thread of its own, its log kept and replayed in order, and with an equal share of the
+        // budget.
         std::vector<Logger> logs;
         for (size_t i = 0; i < domains.size(); i++)
             logs.push_back(Logger::Buffered());
         std::vector<std::future<MolecularReplacementResult>> running;
         for (size_t i = 0; i < domains.size(); i++)
             running.push_back(std::async(std::launch::async, [&, i] {
-                return search_model(DomainModel(model, domains[i]), merges, logs[i], nthreads);
+                return search_model(DomainModel(model, domains[i]), merges, logs[i], nthreads,
+                                    gpu_budget / domains.size());
             }));
         std::vector<MolecularReplacementResult> found;
         for (size_t i = 0; i < domains.size(); i++) {
@@ -3214,13 +3386,15 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
 }
 
 void CheckEnantiomorphHand(MolecularReplacementResult &result, const std::vector<MergedReflection> &merged,
-                           const gemmi::SpaceGroup &data_space_group, Logger &logger, size_t nthreads) {
+                           const gemmi::SpaceGroup &data_space_group, Logger &logger, size_t nthreads,
+                           size_t gpu_resident_bytes) {
     const MRCandidate &c = result.candidates[result.chosen];
     const auto t0 = std::chrono::steady_clock::now();
     Logger quiet = Logger::Buffered();
     ModelValidationResult v;
     try {
-        v = ReadModelAnomalousMap(merged, *c.placed, data_space_group, c.k_sol, c.b_sol, quiet, nthreads);
+        v = ReadModelAnomalousMap(merged, *c.placed, data_space_group, c.k_sol, c.b_sol, quiet, nthreads,
+                                  gpu_resident_bytes);
     } catch (const std::runtime_error &e) {   // a GPU failure: no reading, so no relabelling
         logger.Warning("Molecular replacement: the anomalous map could not be read ({})", e.what());
     }

image_analysis/structure_refinement/MolecularReplacement.h

@@ -370,6 +370,36 @@ struct MolecularReplacementResult {
     double t_domains = 0;   // the PAE domains' own searches, wall clock (they run side by side)
 };
 
+// What a molecular-replacement search puts on its GPU, for one merge: the molecular transform the rotation and
+// translation functions read, the rigid body's engines and the score engines, and what the functions' working
+// buffers may take beside them. Bytes of one of each; 0 where that part cannot run on the GPU.
+struct MolecularReplacementGPUNeed {
+    size_t budget = 0;               // ModelGPUBudget (ModelValidation.h), or a PAE domain's share of it
+    size_t transform = 0;
+    size_t rigid_body = 0;
+    size_t rigid_body_engines = 0;   // how many rigid-body fits can run at once
+    size_t structure_factors = 0;    // the largest score engine of the merge's placements
+};
+
+struct MolecularReplacementGPUPlan {
+    bool search = false;             // the transform held on the GPU, the functions run there
+    size_t rigid_body_engines = 0;   // 0: the rigid body on the CPU
+    bool structure_factors = false;  // the scores on the GPU
+    size_t room = 0;                 // for the score engines and the functions' working buffers
+    size_t bytes = 0;                // the transform, the rigid-body engines and one score engine
+};
+
+// The one rule for what a search puts on its GPU, decided before anything is reserved, from the budget model
+// validation keeps to (80% of the card's TOTAL memory less what is resident). The transform and the rigid
+// body's engines are held for the whole merge, so it is their sum, with the largest score engine, that is held
+// against the budget. Each part takes one where it still fits, in the order of what it saves: the transform
+// (the rotation and translation functions), the rigid body, the scores; a part that does not fit runs on the
+// CPU. More rigid-body engines only within half of the budget, as model validation adds them. What is left of
+// the budget beside the transform and the engines is the room the score engines - kept for as long as they fit,
+// the least recently used given up first - and the functions' working buffers share, each waiting for the other
+// to be done rather than spilling past it.
+MolecularReplacementGPUPlan PlanMolecularReplacementGPU(const MolecularReplacementGPUNeed &need);
+
 // Search `merges` (the answer's first) for the predicted model at `model_path`, in every space group with
 // the rotations of each merge's group - one copy, then further copies one at a time while each new one
 // stands out, packs and lowers R-free - and decide. A candidate is PLACED where a copy's translation
@@ -377,8 +407,9 @@ struct MolecularReplacementResult {
 // null. The model chooses among the answer merge's groups where one placed candidate leads
 // every other by an R-free margin; it chooses the twin hypothesis' merge only where none of the
 // answer's groups can hold the model at all (Matthews, packing) and a group of the subgroup can.
+// `gpu_resident_bytes`: what is held on the card meanwhile by what came before (a merge engine still at work).
 MolecularReplacementResult RunMolecularReplacement(const std::string &model_path, const std::vector<MRMerge> &merges,
-                                                   Logger &logger, size_t nthreads);
+                                                   Logger &logger, size_t nthreads, size_t gpu_resident_bytes = 0);
 
 // The hand check of an ENANTIOMORPH decision, filled into result.hand_check: taken into the answer by
 // --model's rule for a model's enantiomorph - the model has to fit against its own null, which the choice
@@ -386,4 +417,5 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
 // `data_space_group`) must not come out inverted, read by model validation's own code on the placement
 // and the bulk solvent MR scored. Where it is not taken, the decision becomes ENANTIOMORPH_VETOED.
 void CheckEnantiomorphHand(MolecularReplacementResult &result, const std::vector<MergedReflection> &merged,
-                           const gemmi::SpaceGroup &data_space_group, Logger &logger, size_t nthreads);
+                           const gemmi::SpaceGroup &data_space_group, Logger &logger, size_t nthreads,
+                           size_t gpu_resident_bytes = 0);

image_analysis/structure_refinement/MolecularReplacementGPU.cu

@@ -409,11 +409,65 @@ std::vector<double> MolecularReplacementGPU::RotationFunction(const SearchReflec
     return out;
 }
 
-// What the translation function needs of one setup, on the device: the gather's lists - for every point of the
-// half grid in cuFFT's layout ((u nv + v) nwh + l), the (term, pair) entries landing on it, in entry order, and
-// the same for the cross terms with a known partial structure - the terms' vectors, phases, weights and bins,
-// and the transform's plan. Made once where a setup is used for several functions in a row.
+// What the translation function needs of one setup, worked out on the host: the gather's lists - for every
+// point of the half grid in cuFFT's layout ((u nv + v) nwh + l), the (term, pair) entries landing on it, in
+// entry order, and the same for the cross terms with a known partial structure - and the terms' vectors,
+// phases and weights as the device takes them. Made once where a setup is used for several functions in a row;
+// uploaded by each, so that the device holds it only while a function runs.
 struct TranslationSetupGPU {
+    bool cross = false;
+    std::vector<int> offset, entry, cross_offset, cross_entry, bin_n;
+    std::vector<float3> s;
+    std::vector<float2> phase;
+    std::vector<float> weight;
+};
+
+namespace {
+
+TranslationSetupGPU prepare_translation(const TranslationSetup &setup, bool cross) {
+    TranslationSetupGPU out;
+    out.cross = cross;
+    const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1;
+    const size_t cells = static_cast<size_t>(nu) * nv * nwh;
+    std::vector<int> &offset = out.offset;
+    offset.assign(cells + 1, 0);
+    for (size_t e = 0; e < setup.target_u.size(); e++)
+        if (setup.target_u[e] >= 0)
+            offset[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh + setup.target_l[e] + 1]++;
+    for (size_t c = 0; c < cells; c++)
+        offset[c + 1] += offset[c];
+    out.entry.resize(offset[cells]);
+    {
+        std::vector<int> fill(offset.begin(), offset.end() - 1);
+        for (size_t e = 0; e < setup.target_u.size(); e++)
+            if (setup.target_u[e] >= 0)
+                out.entry[fill[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh
+                               + setup.target_l[e]]++] = static_cast<int>(e);
+    }
+    if (cross) {
+        std::vector<int> &cross_offset = out.cross_offset;
+        cross_offset.assign(cells + 1, 0);
+        for (size_t e = 0; e < setup.cross_term.size(); e++)
+            cross_offset[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh + setup.cross_l[e] + 1]++;
+        for (size_t c = 0; c < cells; c++)
+            cross_offset[c + 1] += cross_offset[c];
+        out.cross_entry.resize(cross_offset[cells]);
+        std::vector<int> fill(cross_offset.begin(), cross_offset.end() - 1);
+        for (size_t e = 0; e < setup.cross_term.size(); e++)
+            out.cross_entry[fill[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh
+                                 + setup.cross_l[e]]++] = static_cast<int>(e);
+    }
+    out.phase.resize(setup.phase.size());
+    for (size_t i = 0; i < out.phase.size(); i++)
+        out.phase[i] = make_float2(static_cast<float>(setup.phase[i].real()), static_cast<float>(setup.phase[i].imag()));
+    out.weight.assign(setup.weight.begin(), setup.weight.end());
+    out.s = float3s(setup.s);
+    out.bin_n = bin_counts(setup.bin, setup.bins);
+    return out;
+}
+
+// A prepared setup on the device, with the transform's plan, for one translation function.
+struct TranslationDevice {
     CudaDevicePtr<int> cross_offset, cross_entry, cross_term, cross_op, bin, bin_n, pair_j, pair_k, offset, entry;
     CudaDevicePtr<char> cross_conj;
     CudaDevicePtr<float3> s;
@@ -422,67 +476,74 @@ struct TranslationSetupGPU {
     CudaFFTPlan plan;
 };
 
-namespace {
-
-std::unique_ptr<TranslationSetupGPU> prepare_translation(const TranslationSetup &setup, bool cross, cudaStream_t stream) {
-    auto out = std::make_unique<TranslationSetupGPU>();
-    const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1;
-    const size_t cells = static_cast<size_t>(nu) * nv * nwh;
-    std::vector<int> offset(cells + 1, 0);
-    for (size_t e = 0; e < setup.target_u.size(); e++)
-        if (setup.target_u[e] >= 0)
-            offset[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh + setup.target_l[e] + 1]++;
-    for (size_t c = 0; c < cells; c++)
-        offset[c + 1] += offset[c];
-    std::vector<int> entry(offset[cells]);
-    {
-        std::vector<int> fill(offset.begin(), offset.end() - 1);
-        for (size_t e = 0; e < setup.target_u.size(); e++)
-            if (setup.target_u[e] >= 0)
-                entry[fill[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh
-                           + setup.target_l[e]]++] = static_cast<int>(e);
-    }
-    std::vector<int> cross_offset(cells + 1, 0), cross_entry;
-    if (cross) {
-        for (size_t e = 0; e < setup.cross_term.size(); e++)
-            cross_offset[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh + setup.cross_l[e] + 1]++;
-        for (size_t c = 0; c < cells; c++)
-            cross_offset[c + 1] += cross_offset[c];
-        cross_entry.resize(cross_offset[cells]);
-        std::vector<int> fill(cross_offset.begin(), cross_offset.end() - 1);
-        for (size_t e = 0; e < setup.cross_term.size(); e++)
-            cross_entry[fill[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh
-                             + setup.cross_l[e]]++] = static_cast<int>(e);
-    }
-    out->cross_offset = upload(cross_offset, stream);
-    out->cross_entry = upload(cross_entry, stream);
-    out->cross_term = upload(setup.cross_term, stream);
-    out->cross_op = upload(setup.cross_op, stream);
-    out->cross_conj = upload(setup.cross_conj, stream);
-    std::vector<float2> phase(setup.phase.size());
-    for (size_t i = 0; i < phase.size(); i++)
-        phase[i] = make_float2(static_cast<float>(setup.phase[i].real()), static_cast<float>(setup.phase[i].imag()));
-    const std::vector<float> weight(setup.weight.begin(), setup.weight.end());
-    out->s = upload(float3s(setup.s), stream);
-    out->phase = upload(phase, stream);
-    out->weight = upload(weight, stream);
-    out->bin = upload(setup.bin, stream);
-    out->bin_n = upload(bin_counts(setup.bin, setup.bins), stream);
-    out->pair_j = upload(setup.pair_j, stream);
-    out->pair_k = upload(setup.pair_k, stream);
-    out->offset = upload(offset, stream);
-    out->entry = upload(entry, stream);
-    out->plan = CudaFFTPlan(CudaFFTShape{.n = {nu, nv, nw}, .type = CUFFT_C2R}, stream);
-    cuda_err(cudaStreamSynchronize(stream));   // the host vectors go out of scope
-    return out;
+void upload_translation(const TranslationSetup &setup, const TranslationSetupGPU &p, cudaStream_t stream,
+                        TranslationDevice &out) {
+    out.cross_offset = upload(p.cross_offset, stream);
+    out.cross_entry = upload(p.cross_entry, stream);
+    out.cross_term = upload(setup.cross_term, stream);
+    out.cross_op = upload(setup.cross_op, stream);
+    out.cross_conj = upload(setup.cross_conj, stream);
+    out.s = upload(p.s, stream);
+    out.phase = upload(p.phase, stream);
+    out.weight = upload(p.weight, stream);
+    out.bin = upload(setup.bin, stream);
+    out.bin_n = upload(p.bin_n, stream);
+    out.pair_j = upload(setup.pair_j, stream);
+    out.pair_k = upload(setup.pair_k, stream);
+    out.offset = upload(p.offset, stream);
+    out.entry = upload(p.entry, stream);
+    out.plan = CudaFFTPlan(CudaFFTShape{.n = {setup.nu, setup.nv, setup.nw}, .type = CUFFT_C2R}, stream);
 }
 
 }  // namespace
 
 std::unique_ptr<TranslationSetupGPU, void (*)(TranslationSetupGPU *)> MolecularReplacementGPU::Prepare(
         const TranslationSetup &setup, bool known_part) {
-    return {prepare_translation(setup, known_part, state_->stream).release(),
-            [](TranslationSetupGPU *p) { delete p; }};
+    return {new TranslationSetupGPU(prepare_translation(setup, known_part)), [](TranslationSetupGPU *p) { delete p; }};
+}
+
+namespace {
+
+// What upload() of `count` values of `bytes` each takes.
+size_t uploaded(size_t count, size_t bytes) {
+    return std::max<size_t>(count, 1) * bytes;
+}
+
+int stat_block_count(size_t points) {
+    return static_cast<int>(std::min<size_t>(1024, (points + THREADS - 1) / THREADS));
+}
+
+}  // namespace
+
+size_t MolecularReplacementGPU::RotationBytes(const SearchReflections &refl, const gemmi::SpaceGroup &sg,
+                                              size_t rotations) {
+    const size_t n = refl.hkl.size(), nops = sg.operations().sym_ops.size();
+    const size_t batch = std::min(ROTATION_BATCH, std::max<size_t>(rotations, 1));
+    return uploaded(n * nops, sizeof(float3)) + 2 * uploaded(n, sizeof(float)) + uploaded(refl.bins, sizeof(int))
+           + uploaded(rotations, sizeof(Rot)) + batch * std::max<size_t>(n, 1) * sizeof(float)
+           + uploaded(rotations, sizeof(double));
+}
+
+size_t MolecularReplacementGPU::TranslationBytes(const TranslationSetup &setup, size_t rotations, bool known) {
+    const size_t cells = static_cast<size_t>(setup.nu) * setup.nv * (setup.nw / 2 + 1);
+    const size_t points = static_cast<size_t>(setup.nu) * setup.nv * setup.nw;
+    const size_t terms = setup.Terms(), nops = setup.ops;
+    const size_t entries = std::count_if(setup.target_u.begin(), setup.target_u.end(), [](int u) { return u >= 0; });
+    const size_t nstat = static_cast<size_t>(stat_block_count(points)) * rotations;
+    // The setup, as upload_translation() puts it on the device
+    size_t bytes = uploaded(known ? cells + 1 : 0, sizeof(int)) + uploaded(known ? setup.cross_term.size() : 0, sizeof(int))
+                   + 2 * uploaded(setup.cross_term.size(), sizeof(int)) + uploaded(setup.cross_conj.size(), sizeof(char))
+                   + uploaded(terms * nops, sizeof(float3)) + uploaded(terms * nops, sizeof(float2))
+                   + 2 * uploaded(terms, sizeof(int)) + uploaded(setup.bins, sizeof(int))
+                   + 2 * uploaded(setup.pair_j.size(), sizeof(int)) + uploaded(cells + 1, sizeof(int))
+                   + uploaded(entries, sizeof(int))
+                   + std::max<size_t>(CudaFFTPlan::WorkBytes(CudaFFTShape{.n = {setup.nu, setup.nv, setup.nw},
+                                                                         .type = CUFFT_C2R}), 1);
+    // and the function's own buffers
+    bytes += uploaded(known ? terms : 0, sizeof(float2)) + terms * nops * sizeof(float2) + 2 * terms * sizeof(float)
+             + cells * sizeof(float2) + points * sizeof(float)
+             + nstat * (2 * sizeof(double) + sizeof(float) + sizeof(long long));
+    return bytes;
 }
 
 std::vector<TranslationPeak> MolecularReplacementGPU::TranslationFunction(
@@ -496,12 +557,13 @@ std::vector<TranslationPeak> MolecularReplacementGPU::TranslationFunction(
     cudaStream_t stream = state_->stream;
     const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1;
     const size_t cells = static_cast<size_t>(nu) * nv * nwh, points = static_cast<size_t>(nu) * nv * nw;
-    std::unique_ptr<TranslationSetupGPU> made;
+    TranslationSetupGPU made;
     if (prepared == nullptr) {
-        made = prepare_translation(setup, known != nullptr, stream);
-        prepared = made.get();
+        made = prepare_translation(setup, known != nullptr);
+        prepared = &made;
     }
-    TranslationSetupGPU &p = *prepared;
+    TranslationDevice p;
+    upload_translation(setup, *prepared, stream, p);
     std::vector<float2> known_f;
     if (known != nullptr)
         for (const std::complex<double> &k : *known)
@@ -511,7 +573,7 @@ std::vector<TranslationPeak> MolecularReplacementGPU::TranslationFunction(
     CudaDevicePtr<float> term_sum(terms), scaled(terms);
     CudaDevicePtr<float2> coef(cells);
     CudaDevicePtr<float> map(points);
-    const int stat_blocks = static_cast<int>(std::min<size_t>(1024, (points + THREADS - 1) / THREADS));
+    const int stat_blocks = stat_block_count(points);
     const size_t nstat = static_cast<size_t>(stat_blocks) * rotations.size();
     CudaDevicePtr<double> d_sum(nstat), d_sum2(nstat);
     CudaDevicePtr<float> d_best(nstat);

image_analysis/structure_refinement/MolecularReplacementGPU.h

@@ -44,11 +44,17 @@ public:
                                                      const std::vector<gemmi::Mat33> &rotations,
                                                      const std::vector<std::complex<double>> *known = nullptr,
                                                      TranslationSetupGPU *prepared = nullptr);
-    // What the translation function needs of `setup` on the device, made once; with `known_part`, for the
+    // What the translation function needs of `setup`, worked out once on the host; with `known_part`, for the
     // function beside a known partial structure.
     std::unique_ptr<TranslationSetupGPU, void (*)(TranslationSetupGPU *)> Prepare(const TranslationSetup &setup,
                                                                                 bool known_part);
 
+    // The device memory a call takes for as long as it runs, and gives back when it returns: RotationFunction()
+    // of `rotations` against `refl`'s reflections in `sg`, and TranslationFunction() of `rotations` (with
+    // `known`, beside a known partial structure).
+    static size_t RotationBytes(const SearchReflections &refl, const gemmi::SpaceGroup &sg, size_t rotations);
+    static size_t TranslationBytes(const TranslationSetup &setup, size_t rotations, bool known);
+
 private:
     std::unique_ptr<MolecularReplacementGPUState> state_;
 };

image_analysis/structure_refinement/RigidBodyGPU.cpp

@@ -10,6 +10,7 @@
 #include <ceres/rotation.h>
 
 #include "gemmi/dencalc.hpp"   // DensityCalculator
+#include "gemmi/calculate.hpp"  // count_atom_sites
 #include "gemmi/it92.hpp"      // IT92 x-ray form factors
 #include "gemmi/scaling.hpp"   // Scaling
 
@@ -132,17 +133,33 @@ std::unique_ptr<RigidBodyGPUPool> RigidBodyGPUPool::Plan(const gemmi::Model &mod
                                                          size_t max_observations, Logger &logger) {
     if (get_gpu_count() == 0)
         return nullptr;
+    std::unique_ptr<RigidBodyGPUPool> pool(new RigidBodyGPUPool);
+    pool->cell_ = cell;
+    pool->d_min_ = d_min;
+    pool->max_observations_ = max_observations;
+    pool->cap_ = std::make_unique<RigidBodyGPUCapacity>();
+    pool->cap_->observations = max_observations;
+    pool->cap_->rows = max_observations;
+    if (!pool->Cover(model, sg, logger))
+        return nullptr;
     try {
-        std::unique_ptr<RigidBodyGPUPool> pool(new RigidBodyGPUPool);
-        // Sized for the largest zone of the ladder, which every fit in the validation walks a part of.
-        RigidBodyGPUCapacity cap;
-        const size_t ops = sg.operations().sym_ops.size();
-        for (double zone_d : RigidBodyLadder(d_min)) {
-            const RigidBodyGPUZone &zone = pool->Zone(model, cell, sg, zone_d);
+        pool->device_ = RigidBodyGPUEngine::CurrentDevice();
+    } catch (const JFJochException &e) {
+        throw RigidBodyGPUFailure(e.what());
+    }
+    return pool;
+}
+
+bool RigidBodyGPUPool::Cover(const gemmi::Model &model, const gemmi::SpaceGroup &sg, Logger &logger) {
+    try {
+        // Sized for the largest zone of the ladder, which every fit walks a part of.
+        RigidBodyGPUCapacity &cap = *cap_;
+        for (double zone_d : RigidBodyLadder(d_min_)) {
+            const RigidBodyGPUZone &zone = Zone(model, cell_, sg, zone_d);
             if (!RigidBodyGPUEngine::Supports(zone)) {
                 logger.Info("Model validation: the rigid body runs on the CPU - the cell is too small for the "
                             "GPU's gridding at {:.1f} A", zone_d);
-                return nullptr;
+                return false;
             }
             cap.atoms = std::max(cap.atoms, zone.atoms.size());
             cap.grid_points = std::max(cap.grid_points, static_cast<size_t>(zone.nu) * zone.nv * zone.nw);
@@ -151,12 +168,8 @@ std::unique_ptr<RigidBodyGPUPool> RigidBodyGPUPool::Plan(const gemmi::Model &mod
             cap.pairs = std::max(cap.pairs, RigidBodyGPUEngine::PairBound(zone));
             cap.fft_work_bytes = std::max(cap.fft_work_bytes, RigidBodyGPUEngine::FFTWorkBytes(zone.nu, zone.nv, zone.nw));
         }
-        cap.observations = max_observations;
-        cap.rows = max_observations;
-        cap.terms = max_observations * ops;
-        pool->cap_ = std::make_unique<RigidBodyGPUCapacity>(cap);
-        pool->device_ = RigidBodyGPUEngine::CurrentDevice();
-        return pool;
+        cap.terms = std::max(cap.terms, max_observations_ * sg.operations().sym_ops.size());
+        return true;
     } catch (const JFJochException &e) {
         throw RigidBodyGPUFailure(e.what());
     }
@@ -181,36 +194,16 @@ void RigidBodyGPUPool::Reserve(size_t engines) {
         idle_.push_back(e.get());
 }
 
-std::unique_ptr<RigidBodyGPUPool> RigidBodyGPUPool::Create(const gemmi::Model &model, const gemmi::UnitCell &cell,
-                                                           const gemmi::SpaceGroup &sg, double d_min,
-                                                           size_t max_observations, size_t max_engines,
-                                                           Logger &logger) {
-    std::unique_ptr<RigidBodyGPUPool> pool = Plan(model, cell, sg, d_min, max_observations, logger);
-    if (!pool)
-        return nullptr;
-    const size_t bytes = pool->EngineBytes();
-    const size_t total = get_gpu_total_memory(pool->device_);
-    const size_t fit = bytes > 0 ? total / 4 / bytes : 0;
-    if (fit == 0) {
-        logger.Info("Model validation: the rigid body runs on the CPU - one GPU engine needs {:.0f} MB, over a "
-                    "quarter of the card's {:.0f} MB", bytes / 1e6, total / 1e6);
-        return nullptr;
-    }
-    pool->Reserve(std::min<size_t>(fit, std::max<size_t>(max_engines, 1)));
-    logger.Info("Model validation: the rigid body runs on the GPU, {} engine(s) of {:.0f} MB", pool->Engines(),
-                bytes / 1e6);
-    return pool;
-}
-
 RigidBodyGPUPool::~RigidBodyGPUPool() = default;
 
 std::shared_ptr<const RigidBodyGPUZone> RigidBodyGPUPool::ObservedZone(
         const gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
         const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) {
     const RigidBodyGPUZone &model_zone = Zone(model, cell, sg, d_min);
+    const size_t atoms = gemmi::count_atom_sites(model);
     std::lock_guard lock(observed_m_);
     for (const ObservedEntry &o : observed_)
-        if (o.d_min == d_min && SameObservations(o.fobs, fobs))
+        if (o.group == sg.xhm() && o.atoms == atoms && o.d_min == d_min && SameObservations(o.fobs, fobs))
             return o.zone;
     const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min);
     std::vector<gemmi::Miller> hkl;
@@ -220,16 +213,17 @@ std::shared_ptr<const RigidBodyGPUZone> RigidBodyGPUPool::ObservedZone(
     const SymmetryComposition composition(grid, d_min, hkl);
     auto zone = std::make_shared<RigidBodyGPUZone>(model_zone);
     AddObservations(*zone, cell, composition, fobs);
-    observed_.push_back({d_min, fobs, zone});
+    observed_.push_back({sg.xhm(), atoms, d_min, fobs, zone});
     return zone;
 }
 
 const RigidBodyGPUZone &RigidBodyGPUPool::Zone(const gemmi::Model &model, const gemmi::UnitCell &cell,
                                               const gemmi::SpaceGroup &sg, double d_min) {
+    const std::tuple<std::string, size_t, double> key{sg.xhm(), gemmi::count_atom_sites(model), d_min};
     std::lock_guard lock(zones_m_);
-    auto it = zones_.find(d_min);
+    auto it = zones_.find(key);
     if (it == zones_.end())
-        it = zones_.emplace(d_min, std::make_unique<RigidBodyGPUZone>(ModelZone(model, cell, sg, d_min))).first;
+        it = zones_.emplace(key, std::make_unique<RigidBodyGPUZone>(ModelZone(model, cell, sg, d_min))).first;
     return *it->second;
 }
 

image_analysis/structure_refinement/RigidBodyGPU.h

@@ -16,6 +16,7 @@
 #include <optional>
 #include <stdexcept>
 #include <string>
+#include <tuple>
 #include <vector>
 
 #include "RigidBodyRefine.h"
@@ -32,12 +33,13 @@ public:
     explicit RigidBodyGPUFailure(const std::string &what) : std::runtime_error(what) {}
 };
 
-// A few engines, reserved once for a whole validation: each is a stream and the buffers for one fit at
-// a time, sized for the finest zone of the ladder to d_min. The real fit and the null's replicates each
-// take one for the length of their fit, and wait for one when all are taken. Engines are
-// interchangeable and every kernel deterministic, so which replicate gets which engine does not change
-// a number. A pool belongs to one model: what a zone needs of its atoms - their densities, radii and mask
-// radii, which the placement does not change - is worked out once per zone and kept.
+// A few engines, reserved once for a whole validation or search: each is a stream and the buffers for one
+// fit at a time, sized for the finest zone of the ladder to d_min. Every fit takes one for its length, and
+// waits for one when all are taken. Engines are interchangeable and every kernel deterministic, so which fit
+// gets which engine does not change a number. A pool belongs to one cell and one d_min, and serves every
+// model and group it was made for (Plan, Cover): what a zone needs of a model's atoms - their densities,
+// radii and mask radii, which the placement does not change - is worked out once per group, atom count and
+// zone and kept.
 class RigidBodyGPUPool {
 public:
     // A pool without engines yet, sized for the finest zone of the ladder to d_min: EngineBytes() is
@@ -47,13 +49,12 @@ public:
     static std::unique_ptr<RigidBodyGPUPool> Plan(const gemmi::Model &model, const gemmi::UnitCell &cell,
                                                   const gemmi::SpaceGroup &sg, double d_min,
                                                   size_t max_observations, Logger &logger);
-    // Plan(), then as many engines as fit a quarter of the card's total memory, at most `max_engines`
-    // (molecular replacement's rule). Null, logged, where not even one fits.
-    static std::unique_ptr<RigidBodyGPUPool> Create(const gemmi::Model &model, const gemmi::UnitCell &cell,
-                                                    const gemmi::SpaceGroup &sg, double d_min,
-                                                    size_t max_observations, size_t max_engines, Logger &logger);
     ~RigidBodyGPUPool();
 
+    // Sizes the engines, before Reserve(), for fits of `model` in `sg` as well - in the pool's cell, to its
+    // d_min, with at most its observations. False, logged, where the cell is too small for the gather.
+    bool Cover(const gemmi::Model &model, const gemmi::SpaceGroup &sg, Logger &logger);
+
     size_t EngineBytes() const;
     // Reserves `engines` engines, engine i on the i-th card from Device(), so a validation uses every
     // card there is. Engines that cannot be set up throw RigidBodyGPUFailure.
@@ -77,14 +78,20 @@ public:
 private:
     RigidBodyGPUPool() = default;
     int device_ = 0;
+    gemmi::UnitCell cell_;
+    double d_min_ = 0;
+    size_t max_observations_ = 0;
     std::unique_ptr<RigidBodyGPUCapacity> cap_;
     std::vector<std::unique_ptr<RigidBodyGPUEngine>> engines_;
     std::vector<RigidBodyGPUEngine *> idle_;
     std::mutex m_;
     std::condition_variable cv_;
-    std::map<double, std::unique_ptr<RigidBodyGPUZone>> zones_;
+    // By group, atom count and zone.
+    std::map<std::tuple<std::string, size_t, double>, std::unique_ptr<RigidBodyGPUZone>> zones_;
     std::mutex zones_m_;
     struct ObservedEntry {
+        std::string group;
+        size_t atoms;
         double d_min;
         gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
         std::shared_ptr<const RigidBodyGPUZone> zone;

rugnux/RugnuxScaleMerge.cpp

@@ -4302,17 +4302,18 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
         // memory goes back before the search reserves its rigid-body engines on the card.
         if (!ledger_ahead.valid())
             rsm.reset();
+        const size_t merge_on_card = rsm && rsm->OnGPU() ? rsm->DeviceBytes() : 0;
         std::vector<MRMerge> merges{{"MERGE", &sm.merged, *result.consensus_cell, experiment_.GetSpaceGroupOrP1()}};
         if (model_twin)
             merges.push_back({"TWIN_SUBGROUP_MERGE", &model_twin->merged, model_twin->cell, model_twin->space_group,
                               model_twin->law});
         MolecularReplacementResult mr = RunMolecularReplacement(config_.model_path, merges, logger,
-                                                                static_cast<size_t>(config_.nthreads));
+                                                                static_cast<size_t>(config_.nthreads), merge_on_card);
         if (mr.chosen >= 0) {
             const MRCandidate &c = mr.candidates[mr.chosen];
             if (mr.decision == "ENANTIOMORPH")
                 CheckEnantiomorphHand(mr, sm.merged, experiment_.GetSpaceGroupOrP1(), logger,
-                                      static_cast<size_t>(config_.nthreads));
+                                      static_cast<size_t>(config_.nthreads), merge_on_card);
             if (mr.decision == "ENANTIOMORPH") {
                 // A change of label only (see ValidateAgainstModel): the two groups index identically.
                 experiment_.SetSpaceGroup(c.space_group);
@@ -4346,6 +4347,7 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
             const UnitCell cell = own ? own->cell : *result.consensus_cell;
             ModelValidationSchedule placed_elsewhere;
             placed_elsewhere.fit_test = false;
+            placed_elsewhere.gpu_resident_bytes = merge_on_card;
             phase("Validating against model (molecular replacement)");
             ModelValidationResult v = ValidateAgainstModel(
                     own ? own->merged : sm.merged, cell, config_.model_path,

tests/ModelValidationTest.cpp

@@ -1276,12 +1276,13 @@ TEST_CASE("ModelValidation_RigidBodyTranslationDerivativeIsExact", "[ModelValida
 
 #ifdef JFJOCH_USE_CUDA
 namespace {
-    // A pool of `engines`: decided on the card's total memory, so there is one wherever there is a GPU.
+    // A pool of `engines`, wherever there is a GPU.
     std::unique_ptr<RigidBodyGPUPool> TestPool(const gemmi::Model &model, const gemmi::UnitCell &cell,
                                                const gemmi::SpaceGroup &sg, double d_min, size_t observations,
                                                size_t engines, Logger &logger) {
-        auto pool = RigidBodyGPUPool::Create(model, cell, sg, d_min, observations, engines, logger);
+        auto pool = RigidBodyGPUPool::Plan(model, cell, sg, d_min, observations, logger);
         REQUIRE(pool);
+        pool->Reserve(engines);
         return pool;
     }
 }

tests/MolecularReplacementTest.cpp

@@ -93,6 +93,72 @@ namespace {
     }
 }
 
+namespace {
+    constexpr size_t GB = 1000000000;
+
+    // A search's parts, sized as a large cell measures: the transform 0.84 GB, rigid-body engines of 2.10 GB
+    // (nine fits at once), score engines of up to 1.86 GB; the budget 80% of the card.
+    MolecularReplacementGPUNeed LargeCellNeed(size_t card) {
+        MolecularReplacementGPUNeed need;
+        need.budget = card / 5 * 4;
+        need.transform = 835500000;
+        need.rigid_body = 2101000000;
+        need.rigid_body_engines = 9;
+        need.structure_factors = 1860000000;
+        return need;
+    }
+}
+
+TEST_CASE("MolecularReplacement_GPUPlanOnCardSizes", "[MolecularReplacement]") {
+    size_t last_rb = 0;
+    for (size_t card : {6 * GB, 8 * GB, 16 * GB, 24 * GB, 64 * GB}) {
+        const MolecularReplacementGPUNeed need = LargeCellNeed(card);
+        const MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need);
+        CAPTURE(card, plan.rigid_body_engines, plan.bytes, plan.room);
+        CHECK(plan.bytes <= need.budget);
+        // Every part on the GPU on every one of these cards; only the engine count grows with it.
+        CHECK(plan.search);
+        CHECK(plan.structure_factors);
+        CHECK(plan.rigid_body_engines >= std::max<size_t>(last_rb, 1));
+        // What the transform and the engines leave is the room, and it holds the largest score engine.
+        CHECK(plan.room == need.budget - need.transform - plan.rigid_body_engines * need.rigid_body);
+        CHECK(plan.room >= need.structure_factors);
+        last_rb = plan.rigid_body_engines;
+    }
+    // 16 GB, a 12.8 GB budget: one of each, 4.80 GB, and a second engine would pass half of it.
+    CHECK(PlanMolecularReplacementGPU(LargeCellNeed(16 * GB)).rigid_body_engines == 1);
+    CHECK(PlanMolecularReplacementGPU(LargeCellNeed(24 * GB)).rigid_body_engines == 3);
+    CHECK(PlanMolecularReplacementGPU(LargeCellNeed(64 * GB)).rigid_body_engines == 9);
+}
+
+TEST_CASE("MolecularReplacement_GPUPlanPartThatDoesNotFitGoesToTheCPU", "[MolecularReplacement]") {
+    // 6 GB, a 4.8 GB budget, which the three parts fill to 4.80 GB: with a larger score engine, the scores
+    // go to the CPU and the transform and the rigid body stay.
+    MolecularReplacementGPUNeed need = LargeCellNeed(6 * GB);
+    need.structure_factors = 2 * GB;
+    MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need);
+    CHECK(plan.search);
+    CHECK(plan.rigid_body_engines == 1);
+    CHECK(!plan.structure_factors);
+    CHECK(plan.bytes <= need.budget);
+
+    // A part the GPU cannot grid (0 bytes) is never placed; the room is then what the transform leaves.
+    need = LargeCellNeed(16 * GB);
+    need.rigid_body = 0;
+    plan = PlanMolecularReplacementGPU(need);
+    CHECK(plan.rigid_body_engines == 0);
+    CHECK(plan.room == need.budget - need.transform);
+
+    // No budget (no GPU, or all of it held beside): everything on the CPU.
+    need = LargeCellNeed(16 * GB);
+    need.budget = 0;
+    plan = PlanMolecularReplacementGPU(need);
+    CHECK(!plan.search);
+    CHECK(plan.rigid_body_engines == 0);
+    CHECK(!plan.structure_factors);
+    CHECK(plan.room == 0);
+}
+
 TEST_CASE("MolecularReplacement_UniformRotationsCoverSO3", "[MolecularReplacement]") {
     const auto rotations = UniformRotations(6.0);
     const double step = 6.0 * M_PI / 180.0;
@@ -443,8 +509,9 @@ TEST_CASE("MolecularReplacement_RigidBodyBesideFixedCopy", "[MolecularReplacemen
 #ifdef JFJOCH_USE_CUDA
     if (get_gpu_count() > 0) {
         gemmi::Structure gpu = start;
-        auto pool = RigidBodyGPUPool::Create(gpu.models[0], cell, sg, 3.0, fobs.v.size(), 1, logger);
+        auto pool = RigidBodyGPUPool::Plan(gpu.models[0], cell, sg, 3.0, fobs.v.size(), logger);
         REQUIRE(pool);
+        pool->Reserve(1);
         RefineRigidBody(gpu.models[0], cell, sg, fobs, 3.0, logger, 4, pool.get(), &first.models[0]);
         CHECK(worst(gpu) < 0.1);
     }