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