From 011b98838e402d2eb901256b65e0a397893bc6a7 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 13:33:39 +0200 Subject: [PATCH] Molecular replacement: one GPU plan per merge on model validation's budget, the transform cropped, output unchanged What molecular replacement put on its GPU was decided part by part on a quarter of the card each: every group's rigid-body pool (up to three engines), every null's (up to nine), the score engines' cache, the --mr hand check's structure factors. The groups and the nulls run side by side, so the quarters added up: on 4nwv on a 16 GB card the search held 4 rigid-body engines of 1.88 GB, two score engines of 1.64 GB, the 1.82 GB molecular transform, two translation functions' 1.2 GB and the 2.4 GB the allocation pool kept from the validation before - 13.8 GB at its peak, 83% of the card. On an 8.55 GB card that pages. Now each merge's search is planned once, before it starts, on the budget model validation uses (ModelGPUBudget: 80% of the card's total, less what a merge engine still at work holds): the transform, one rigid-body engine and the largest score engine take their place in that order where they still fit, more rigid-body engines only within half of the budget, and what is left is a room the score engines (least recently used given up first, as before) and the rotation and translation functions' working buffers share, each waiting for the other to finish rather than spilling past it (PlanMolecularReplacementGPU). A part that does not fit runs on the CPU, decided on its size alone. - One rigid-body pool per merge serves every fit of the merge - the groups' first and further copies, the nulls, the tNCS shifts, the carried copies - sized for the largest (the most copies Matthews allows, in any group) and reserved once, instead of a pool made and freed for each. RigidBodyGPUPool::Cover widens a planned pool to another model or group; its zones are kept by group, atom count and resolution. RigidBodyGPUPool::Create (the quarter rule) is gone. - The sizes are worked out on the CPU while the transform is computed. - The molecular transform keeps only the indices it is read at (|k|, l <= box / d_min + 1): 0.84 GB instead of 1.82 on 4nwv, host and device. Its rows keep their full length: cut short too, the rotation function read them 15% slower. - The translation function's prepared setup is kept on the host and uploaded by each call, so the device holds it only while a function runs; the cross-term offsets are uploaded only where there are cross terms. - The allocation pool is trimmed before the search, after the transform's GPU FFT, before a score engine is made and before the hand check, whose structure factors are now judged against the budget. - rugnux passes the merge engine's bytes, where the finalist ledger keeps it, to the search, the hand check and the validation of the placed model. Measured on the RTX 5080 (16.6 GB), timing benchmark: peak device memory during MR 13.8 -> 8.7 GB on 4nwv (start 3.0 -> 0.66 GB after the trim), 1.67 -> 1.14 GB on cytc, 6.3 -> 0.76 GB on myob. On 4nwv the plan holds the transform 0.84 GB, one rigid-body engine of 2.10 GB and leaves 10.35 GB of a 13.29 GB budget to the score engines and the functions; a second engine would pass half of the budget (one engine against two: +1.2 s on the rescoring and the null). MR 31.7 -> 31.6 s on 4nwv. Reflection files, maps, placed models, reports and the MR candidate lines byte-identical to rc175-metal-merge on myob, cytc and 4nwv; [MolecularReplacement], [ModelValidation], RigidBody*, ModelScale*, [CUDAMemHelpers] pass. Co-Authored-By: Claude Opus 5.5 --- docs/CPU_DATA_ANALYSIS_DECISIONS.md | 7 +- .../structure_refinement/ModelValidation.cpp | 27 +- .../structure_refinement/ModelValidation.h | 11 +- .../MolecularReplacement.cpp | 462 ++++++++++++------ .../MolecularReplacement.h | 36 +- .../MolecularReplacementGPU.cu | 190 ++++--- .../MolecularReplacementGPU.h | 8 +- .../structure_refinement/RigidBodyGPU.cpp | 70 ++- .../structure_refinement/RigidBodyGPU.h | 31 +- rugnux/RugnuxScaleMerge.cpp | 6 +- tests/ModelValidationTest.cpp | 5 +- tests/MolecularReplacementTest.cpp | 69 ++- 12 files changed, 644 insertions(+), 278 deletions(-) diff --git a/docs/CPU_DATA_ANALYSIS_DECISIONS.md b/docs/CPU_DATA_ANALYSIS_DECISIONS.md index 24003342e..fc3387e05 100644 --- a/docs/CPU_DATA_ANALYSIS_DECISIONS.md +++ b/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 `_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). diff --git a/image_analysis/structure_refinement/ModelValidation.cpp b/image_analysis/structure_refinement/ModelValidation.cpp index 46b5d7b9b..a79dd085d 100644 --- a/image_analysis/structure_refinement/ModelValidation.cpp +++ b/image_analysis/structure_refinement/ModelValidation.cpp @@ -1883,10 +1883,14 @@ ModelValidationResult Validate(const std::vector &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 &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 ModelStructureFactorsGPU *gpu = nullptr; #ifdef JFJOCH_USE_CUDA std::unique_ptr 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 } #else (void) logger; + (void) gpu_resident_bytes; #endif if (gpu == nullptr) { gemmi::DensityCalculator dc; diff --git a/image_analysis/structure_refinement/ModelValidation.h b/image_analysis/structure_refinement/ModelValidation.h index 78cfeb972..4bc05dff7 100644 --- a/image_analysis/structure_refinement/ModelValidation.h +++ b/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 &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 diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 52e3ffd80..395d4c232 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/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 crop_half_l(const gemmi::FPhiGrid &full, int reach) { + gemmi::FPhiGrid 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(nu) * (gemmi::modulo(k, full.nv) + static_cast(full.nv) * l)], + nu, &out.data[static_cast(nu) * (gemmi::modulo(k, n) + static_cast(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> &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> 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 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 sf; - size_t points = 0, bytes = 0, quarter = 0; + size_t points = 0, bytes = 0; try { sf = std::make_unique(RigidBodyGPUEngine::CurrentDevice(), st.models[0], st.cell, sg, d_min); points = sf->Reflections(); - bytes = sf->DeviceBytes() + ModelScaleGPU::DeviceBytes(points) + 2 * points * sizeof(std::complex); - 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; 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); + } + + // `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 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(*this, card, bytes); + } + private: struct Entry { std::shared_ptr 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 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, Entry> engines_; std::map 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 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 &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 &merged, double d_min, size_t nthreads, - ScoreEngines *scorers, Logger &logger) { + const std::vector &merged, double d_min, RigidBodyGPUPool *pool, + size_t nthreads, ScoreEngines *scorers) { std::vector 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 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 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(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 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 groups = SameRotationGroups(merge.space_group); + const std::vector pm = MatthewsProbabilities(cell.volume / static_cast(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(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(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(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 &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> 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 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); #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 gpu; - if (get_gpu_count() > 0) { + if (get_gpu_count() > 0 && transform_bytes <= gpu_budget) { try { gpu = std::make_unique(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 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> 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 &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 &rotations, + const std::vector> *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 tf(rotations.size()); + ParallelFor(static_cast(rotations.size()), nthreads, [&](int i) { + tf[i] = TranslationFunction(ft, setup, rotations[i], known); + }); + return tf; + }; std::vector 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 rf = gpu ? gpu->RotationFunction(refl, cell, merge.space_group, rotations) - : RotationFunction(ft, refl, cell, merge.space_group, rotations, nthreads); -#else - const std::vector rf = RotationFunction(ft, refl, cell, merge.space_group, rotations, nthreads); -#endif + plan_gpu(m); + const std::vector 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 v = gpu ? gpu->RotationFunction(refl, cell, merge.space_group, trial) - : RotationFunction(ft, refl, cell, merge.space_group, trial, nthreads); -#else - const std::vector v = RotationFunction(ft, refl, cell, merge.space_group, trial, nthreads); -#endif + const std::vector 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>> 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 tf(peaks.size()); -#ifdef JFJOCH_USE_CUDA - if (gpu) - tf = gpu->TranslationFunction(setup, peak_rotations); - else -#endif - ParallelFor(static_cast(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 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 tf_seed(seeds.size()); -#ifdef JFJOCH_USE_CUDA - const std::vector 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 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 rf_seed = rotation_function(refl, cell, merge.space_group, seeds); + const std::vector tf_seed = translation_function(MakeTranslationSetup(refl, cell, sg), + seeds, nullptr, nullptr); std::vector 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> 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 setup2; -#ifdef JFJOCH_USE_CUDA std::unique_ptr 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 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(rots2.size()), nthreads, [&](int i) { - tf2[i] = TranslationFunction(ft, *setup2, rots2[i], &known); - }); + const std::vector tf2 = translation_function(*setup2, rots2, &known, setup2_gpu.get()); std::vector 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 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(const std::vector &)> 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 tf(turned.size()); -#ifdef JFJOCH_USE_CUDA - if (gpu) - tf = gpu->TranslationFunction(setup, turned); - else -#endif - ParallelFor(static_cast(turned.size()), nthreads, [&](int i) { - tf[i] = TranslationFunction(ft, setup, turned[i]); - }); + const std::vector tf = translation_function( + MakeTranslationSetup(refl, gemmi::UnitCell(merges[merge].cell), c.space_group), turned, + nullptr, nullptr); std::vector 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