From bfc6ca8487e1539b0027394d234f0d59d789e414 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Fri, 16 Oct 2026 17:16:28 +0200 Subject: [PATCH] MR and model validation: use every card's memory, split MR's functions over the cards On a multi-GPU host the structure stage ran essentially on one card: the rotation function's launches all went to the search's card, and the rigid-body engines of model validation and MR were admitted only within half of that one card's budget, although RigidBodyGPUPool already placed engines on other cards. On a 4x T4 server 4nwv's MR rotation function took 75.8 s on one card, validation queued its 10 fits on one engine and MR's rescoring and null got one engine. - PlanModelValidationGPU / PlanMolecularReplacementGPU: this card is planned exactly as on a machine with it alone (GPU-vs-CPU placement unchanged: 80% of the total, less what is resident). The rigid-body engines it has no room for go on the other cards in turn, each against its own budget (80% of its total, less the resident bytes, counted on every card) within half of it. MR also puts a copy of the molecular transform on each other card within half of its budget. - RigidBodyGPUPool::Reserve takes per-card engine counts. - Validation forwards the most it holds on any card as gpu_bytes to a validation beside, and may run one beside only where twice that fits every card's budget. - MR splits every rotation and translation function into contiguous slices of the orientation list over the cards holding a transform copy (and room for the call's buffers), one thread per card, gathered back in order. Each orientation is scored by itself (one block / one map), nothing is summed across orientations, so N cards give the bytes of one. - ScoreEngines keeps a room per card; score engines always go on the search's own card; function buffers are held on the card they run on. Byte-identical at N=1 (workstation, one RTX 5080): timing benchmark (myob, cytc, 4nwv) and 4nwv, 8tyy, 6oel, 3inp with --model give the same p.hkl, MTZ data, p_model.pdb, maps and report against rc176. New tests: per-card plans on four 15.6 GB cards (validation 1/3/3/3 engines, MR T+1/T+2/T+2/T+2), and the split driven over 1-4 virtual cards on one device, bitwise equal to one card. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../structure_refinement/ModelValidation.cpp | 56 ++++- .../structure_refinement/ModelValidation.h | 24 +- .../MolecularReplacement.cpp | 218 +++++++++++++++--- .../MolecularReplacement.h | 20 +- .../MolecularReplacementGPU.h | 28 +++ .../structure_refinement/RigidBodyGPU.cpp | 9 +- .../structure_refinement/RigidBodyGPU.h | 6 +- tests/ModelValidationTest.cpp | 39 +++- tests/MolecularReplacementTest.cpp | 108 ++++++++- 9 files changed, 454 insertions(+), 54 deletions(-) diff --git a/image_analysis/structure_refinement/ModelValidation.cpp b/image_analysis/structure_refinement/ModelValidation.cpp index a79dd085d..7be413fe1 100644 --- a/image_analysis/structure_refinement/ModelValidation.cpp +++ b/image_analysis/structure_refinement/ModelValidation.cpp @@ -799,6 +799,7 @@ ModelValidationResult Validate(const std::vector &merged, cuda_trim_allocation_pool(); ModelValidationGPUNeed need; need.card = get_gpu_total_memory(device); + need.other_cards = OtherCardsTotalMemory(device); need.resident = schedule.gpu_resident_bytes; need.beside = schedule.gpu_beside_bytes; sf_engine = StructureFactorsEngine(device, st.models[0], ucell, *sg, d_min, logger); @@ -837,8 +838,8 @@ ModelValidationResult Validate(const std::vector &merged, need.scale_fit = ModelScaleGPUPool::DeviceBytes(obs.size(), 1); need.scale_fit_engines = SCALE_GPU_ENGINES; - gpu_plan = PlanModelValidationGPU(need, static_cast(get_gpu_count())); - gpu_bytes = gpu_plan.bytes; + gpu_plan = PlanModelValidationGPU(need); + gpu_bytes = gpu_plan.most_on_a_card; may_run_beside = gpu_plan.may_run_beside; std::string parts; @@ -858,6 +859,12 @@ ModelValidationResult Validate(const std::vector &merged, part(fmt::format("the null's to {:.2f} A", null_d_min), gpu_plan.null_structure_factors, need.null_structure_factors, 0); part("rigid body", gpu_plan.rigid_body_engines > 0, need.rigid_body, gpu_plan.rigid_body_engines); + if (gpu_plan.rigid_body_engines > gpu_plan.rigid_body_per_card[0]) { + std::string per_card; + for (size_t n : gpu_plan.rigid_body_per_card) + per_card += fmt::format("{}{}", per_card.empty() ? "" : "/", n); + parts += fmt::format(" - {} on the cards from this one", per_card); + } part("scale fits", gpu_plan.scale_fit_engines > 0, need.scale_fit, gpu_plan.scale_fit_engines); logger.Info("Model validation: {} - {} held of a {} budget, 80% of the card's {}{}{}", parts, GBytes(gpu_plan.bytes), GBytes(gpu_plan.budget), GBytes(need.card), @@ -927,7 +934,7 @@ ModelValidationResult Validate(const std::vector &merged, #ifdef JFJOCH_USE_CUDA if (rigid_body_engines && gpu_plan.rigid_body_engines > 0) { - rigid_body_engines->Reserve(gpu_plan.rigid_body_engines); + rigid_body_engines->Reserve(gpu_plan.rigid_body_per_card); rigid_body_pool = rigid_body_engines.get(); } if (gpu_plan.scale_fit_engines > 0) { @@ -1888,7 +1895,7 @@ size_t ModelGPUBudget(size_t card, size_t resident) { return usable > resident ? usable - resident : 0; } -ModelValidationGPUPlan PlanModelValidationGPU(const ModelValidationGPUNeed &need, size_t cards) { +ModelValidationGPUPlan PlanModelValidationGPU(const ModelValidationGPUNeed &need) { ModelValidationGPUPlan plan; plan.budget = ModelGPUBudget(need.card, need.resident); auto take = [&](size_t bytes, size_t limit) { @@ -1903,22 +1910,55 @@ ModelValidationGPUPlan PlanModelValidationGPU(const ModelValidationGPUNeed &need plan.scale_fit_engines = need.scale_fit_engines > 0 && take(need.scale_fit, plan.budget) ? 1 : 0; // More engines only within half of the budget (see the declaration). - const size_t room = std::min(plan.budget / 2, plan.budget > need.beside ? plan.budget - need.beside : 0); + const auto room = [&](size_t budget) { return std::min(budget / 2, budget > need.beside ? budget - need.beside : 0); }; while (plan.rigid_body_engines > 0 && plan.rigid_body_engines < need.rigid_body_engines - && take(need.rigid_body, room)) + && take(need.rigid_body, room(plan.budget))) plan.rigid_body_engines++; while (plan.scale_fit_engines > 0 && plan.scale_fit_engines < need.scale_fit_engines - && take(need.scale_fit, room)) + && take(need.scale_fit, room(plan.budget))) plan.scale_fit_engines++; + // The rigid-body engines this card has no room for, on the other cards in turn. + plan.rigid_body_per_card.assign(1 + need.other_cards.size(), 0); + plan.rigid_body_per_card[0] = plan.rigid_body_engines; + std::vector held(plan.rigid_body_per_card.size(), 0); + held[0] = plan.bytes; + for (bool added = plan.rigid_body_engines > 0; added;) { + added = false; + for (size_t k = 1; k < held.size() && plan.rigid_body_engines < need.rigid_body_engines; k++) + if (held[k] + need.rigid_body <= room(ModelGPUBudget(need.other_cards[k - 1], need.resident))) { + held[k] += need.rigid_body; + plan.rigid_body_per_card[k]++; + plan.rigid_body_engines++; + added = true; + } + } + plan.most_on_a_card = *std::max_element(held.begin(), held.end()); + + const size_t cards = held.size(); + size_t smallest_budget = plan.budget; + for (size_t total : need.other_cards) + smallest_budget = std::min(smallest_budget, ModelGPUBudget(total, need.resident)); const size_t whole = (plan.structure_factors ? need.structure_factors : 0) + (plan.null_structure_factors ? need.null_structure_factors : 0) + (plan.rigid_body_engines > 0 ? need.rigid_body * need.rigid_body_engines : 0) + (plan.scale_fit_engines > 0 ? need.scale_fit * need.scale_fit_engines : 0); - plan.may_run_beside = 2 * whole <= cards * plan.budget; + plan.may_run_beside = 2 * whole <= cards * plan.budget && 2 * plan.most_on_a_card <= smallest_budget; return plan; } +std::vector OtherCardsTotalMemory(int device) { + std::vector out; +#ifdef JFJOCH_USE_CUDA + const int cards = get_gpu_count(); + for (int k = 1; k < cards; k++) + out.push_back(get_gpu_total_memory((device + k) % cards)); +#else + (void) device; +#endif + return out; +} + ModelValidationResult ValidateAgainstModel(const std::vector &merged, const UnitCell &cell, const std::string &model_path, diff --git a/image_analysis/structure_refinement/ModelValidation.h b/image_analysis/structure_refinement/ModelValidation.h index 4bc05dff7..eee350c03 100644 --- a/image_analysis/structure_refinement/ModelValidation.h +++ b/image_analysis/structure_refinement/ModelValidation.h @@ -288,6 +288,9 @@ struct ModelValidationSchedule { // where the null is scored to the data's own resolution and shares the engine to d_min). struct ModelValidationGPUNeed { size_t card = 0; // the card's total memory + // The total memory of every other card, in the order the rigid body's engines count on to them from this + // one (RigidBodyGPUPool::Reserve, OtherCardsTotalMemory). + std::vector other_cards; size_t resident = 0; // ModelValidationSchedule::gpu_resident_bytes size_t beside = 0; // ModelValidationSchedule::gpu_beside_bytes size_t structure_factors = 0; // to the data's resolution, with the maps @@ -302,9 +305,12 @@ struct ModelValidationGPUPlan { size_t budget = 0; // 80% of the card, less what is resident there bool structure_factors = false; bool null_structure_factors = false; - size_t rigid_body_engines = 0; // 0: the rigid body on the CPU + size_t rigid_body_engines = 0; // 0: the rigid body on the CPU; on every card together + // The rigid body's engines on this card and on each of need.other_cards, in that order. + std::vector rigid_body_per_card; size_t scale_fit_engines = 0; // 0: the scale fits on the CPU - size_t bytes = 0; // what the plan holds on the card + size_t bytes = 0; // what the plan holds on this card + size_t most_on_a_card = 0; // what it holds on the card it holds most on bool may_run_beside = false; }; @@ -323,9 +329,17 @@ size_t ModelGPUBudget(size_t card, size_t resident); // only within half of the budget, less what a validation beside holds: they change only how many fits // run at once, and buy little - on a large cell a fit already fills the card (measured on a 16 GB card: // 13.6 s with one rigid-body engine, 13.1 s with two, 12.6 s with five) - so the card keeps headroom -// rather than filling with them. A second validation may run beside where twice the whole of this one's GPU parts, at every -// engine it can use, fits the budgets of the cards together. -ModelValidationGPUPlan PlanModelValidationGPU(const ModelValidationGPUNeed &need, size_t cards = 1); +// rather than filling with them. Where this card has no room for all the rigid-body engines the +// validation can use, the rest go on the other cards, one card after the other in turn, each by the same +// rule against its own budget (80% of its total, less what is resident): this card's parts are decided as +// on a machine with this card alone, and how many cards there are changes only where fits run. A second +// validation may run beside where twice the whole of this one's GPU parts, at every engine it can use, fits +// the budgets of the cards together, and twice what this one holds on any card fits every card's budget. +ModelValidationGPUPlan PlanModelValidationGPU(const ModelValidationGPUNeed &need); + +// The total memory of every card but `device`, in the order from it that the rigid body's engines are put on +// them (RigidBodyGPUPool::Reserve): device + 1, device + 2, ... wrapping round. Empty without CUDA. +std::vector OtherCardsTotalMemory(int device); ModelValidationResult ValidateAgainstModel(const std::vector &merged, const UnitCell &cell, diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 47d165261..e18add1a2 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -312,16 +312,18 @@ 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) { + // The room of a merge's plan on each card, set before its placements are scored; the score engines go on + // `card`, the search's own, and `scores`: whether they go on the GPU at all. + void SetRoom(const std::map &room, int card, bool scores) { std::lock_guard lock(m_); room_ = room; + card_ = card; scores_ = scores; } - bool Fits(size_t bytes) { + bool Fits(int card, size_t bytes) { std::lock_guard lock(m_); - return bytes <= room_; + return bytes <= Room(card); } // The engine for this kind of placement, kept on the card for as long as the returned pointer lives. @@ -338,14 +340,14 @@ public: if (e.engine || e.refused) return Lease(e); e.making = true; - const size_t room = room_; + const int card = card_; + const size_t room = Room(card); lock.unlock(); std::unique_ptr sf; size_t points = 0, bytes = 0; try { - sf = std::make_unique(RigidBodyGPUEngine::CurrentDevice(), st.models[0], - st.cell, sg, d_min); + sf = std::make_unique(card, st.models[0], st.cell, sg, d_min); points = sf->Reflections(); bytes = Bytes(*sf); } catch (const std::exception &) { @@ -359,7 +361,6 @@ public: cv_.notify_all(); return nullptr; } - const int card = sf->Device(); cv_.wait(lock, [&] { return MakeRoom(card, bytes); }); resident_[card] += bytes; lock.unlock(); @@ -403,8 +404,7 @@ public: 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. + // `bytes` of the room on `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) {} @@ -422,8 +422,7 @@ public: int card_; size_t bytes_; }; - std::unique_ptr Hold(size_t bytes) { - const int card = RigidBodyGPUEngine::CurrentDevice(); + std::unique_ptr Hold(int card, size_t bytes) { std::unique_lock lock(m_); cv_.wait(lock, [&] { return MakeRoom(card, bytes); }); resident_[card] += bytes; @@ -456,7 +455,7 @@ private: // 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_) { + while (resident_[card] + bytes > Room(card)) { Entry *oldest = nullptr; for (auto &[key, entry] : engines_) if (entry.engine && entry.card == card && entry.engine.use_count() == 1 @@ -470,9 +469,16 @@ private: return true; } + // Under m_. + size_t Room(int card) const { + const auto it = room_.find(card); + return it == room_.end() ? 0 : it->second; + } + std::map, Entry> engines_; std::map resident_; - size_t room_ = 0; + std::map room_; + int card_ = 0; bool scores_ = false; uint64_t clock_ = 0; std::mutex m_; @@ -2323,9 +2329,10 @@ MergeGPUSize size_merge_gpu(const SearchModel &model, const MRMerge &merge, doub #endif // The search and the decision for one search model. -// `gpu_budget`: what the search may hold on its card (PlanMolecularReplacementGPU). +// `gpu_budget`: what the search may hold on its card (PlanMolecularReplacementGPU); `other_budgets`: on each other +// card, in the order from it. MolecularReplacementResult search_model(const SearchModel &model, const std::vector &merges, Logger &logger, - size_t nthreads, size_t gpu_budget) { + size_t nthreads, size_t gpu_budget, const std::vector &other_budgets) { MolecularReplacementResult result; result.residues_read = model.residues_read; result.residues_kept = model.residues_kept; @@ -2382,6 +2389,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec // 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; + const int device = get_gpu_count() > 0 ? RigidBodyGPUEngine::CurrentDevice() : 0; if (get_gpu_count() > 0 && transform_bytes <= gpu_budget) { try { gpu = std::make_unique(ft); @@ -2391,6 +2399,12 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec } } result.gpu_search = gpu != nullptr; + // The transform's copies on the other cards the plan gives one (made by the first merge's plan, the same for + // every merge), and the cards that score a function's orientations: this one first. + std::vector> copies; + std::vector search_cards; + if (gpu) + search_cards.push_back({device, gpu.get()}); #endif result.t_transform = seconds_since(t0); result.tncs = DetectTranslationalNcs(*merges[0].merged, gemmi::UnitCell(merges[0].cell), merges[0].space_group); @@ -2426,16 +2440,41 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec MergeGPUSize size = sized[m].get(); MolecularReplacementGPUNeed need; need.budget = gpu_budget; + need.other_budgets = other_budgets; 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); + const int cards = get_gpu_count(); + std::map room; + std::vector engines; + for (size_t k = 0; k < plan.cards.size(); k++) { + room[(device + static_cast(k)) % cards] = plan.cards[k].room; + engines.push_back(plan.cards[k].rigid_body_engines); + } + scorers.SetRoom(room, device, plan.structure_factors); if (plan.rigid_body_engines > 0) { - size.pool->Reserve(plan.rigid_body_engines); + size.pool->Reserve(engines); rigid_body[m] = std::move(size.pool); } + if (gpu && search_cards.size() == 1) { + std::vector on; + for (size_t k = 1; k < plan.cards.size(); k++) + if (plan.cards[k].transform) + on.push_back((device + static_cast(k)) % cards); + copies.resize(on.size()); + std::vector> uploading; + for (size_t i = 0; i < on.size(); i++) + uploading.push_back(std::async(std::launch::async, [&, i] { + const ScopedGpuPin pin(on[i]); + copies[i] = std::make_unique(ft); + })); + for (auto &u : uploading) + u.get(); + for (size_t i = 0; i < on.size(); i++) + search_cards.push_back({on[i], copies[i].get()}); + } 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", @@ -2443,20 +2482,44 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec : "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)); + if (plan.cards.size() > 1) { + std::string per_card; + for (const MolecularReplacementGPUCard &c : plan.cards) + per_card += fmt::format("{}{}{}", per_card.empty() ? "" : "/", c.transform ? "T+" : "", + c.rigid_body_engines); + logger.Info("Molecular replacement: {} on the cards from this one: {} (T: a copy of the transform, then " + "the rigid-body engines); the functions split over {} card(s)", + merges[m].name, per_card, search_cards.size()); + } #else (void) m; #endif }; +#ifdef JFJOCH_USE_CUDA + // The cards whose room holds the buffers of a whole call of `bytes`, which a slice of it never needs more + // than - this one always, as the call is made on the GPU only where it fits there - and the room a slice + // holds on its card while it runs. + const auto fitting = [&](size_t bytes) { + std::vector out{search_cards[0]}; + for (size_t k = 1; k < search_cards.size(); k++) + if (scorers.Fits(search_cards[k].device, bytes)) + out.push_back(search_cards[k]); + return out; + }; + const MolecularReplacementGPUHold hold = [&](int card, size_t bytes) -> std::shared_ptr { + return scorers.Hold(card, bytes); + }; +#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. + // fit the room on this card, on the CPU otherwise - the same function, equal to rounding. On the GPU its + // orientations are split over the cards with a copy of the transform and room for them, each slice holding + // its buffers' room on its card while it runs; the result is the same on any number of cards. 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); - } + if (gpu && scorers.Fits(device, bytes)) + return SplitRotationFunction(fitting(bytes), refl, cell, sg, rotations, hold); #endif return RotationFunction(ft, refl, cell, sg, rotations, nthreads); }; @@ -2465,10 +2528,8 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec 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); - } + if (gpu && scorers.Fits(device, bytes)) + return SplitTranslationFunction(fitting(bytes), setup, rotations, known, prepared, hold); #else (void) prepared; #endif @@ -3320,20 +3381,112 @@ MolecularReplacementGPUPlan PlanMolecularReplacementGPU(const MolecularReplaceme && 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; + plan.cards.push_back({plan.search, plan.rigid_body_engines, plan.room}); + + // The other cards: a copy of the transform each, then the engines this card has no room for, in turn. + std::vector held; + for (size_t budget : need.other_budgets) { + MolecularReplacementGPUCard card; + card.transform = plan.search && need.transform <= budget / 2; + held.push_back(card.transform ? need.transform : 0); + plan.cards.push_back(card); + } + for (bool added = plan.rigid_body_engines > 0; added;) { + added = false; + for (size_t k = 0; k < held.size() && plan.rigid_body_engines < need.rigid_body_engines; k++) + if (held[k] + need.rigid_body <= need.other_budgets[k] / 2) { + held[k] += need.rigid_body; + plan.cards[k + 1].rigid_body_engines++; + plan.rigid_body_engines++; + added = true; + } + } + for (size_t k = 0; k < held.size(); k++) + plan.cards[k + 1].room = need.other_budgets[k] - held[k]; return plan; } +#ifdef JFJOCH_USE_CUDA +namespace { + +// slice(k, first, last) for the k-th of `parts` contiguous slices [first, last) of `n` items, each on a thread of +// its own; returns when every slice is done. +void on_slices(size_t n, size_t parts, const std::function &slice) { + const size_t per = (n + parts - 1) / parts; + std::vector> running; + for (size_t k = 0; k < parts; k++) + running.push_back(std::async(std::launch::async, [&, k] { + const size_t first = std::min(k * per, n); + slice(k, first, std::min(first + per, n)); + })); + for (auto &r : running) + r.get(); +} + +} // namespace + +std::vector SplitRotationFunction(const std::vector &cards, + const SearchReflections &refl, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, const std::vector &rotations, + const MolecularReplacementGPUHold &hold) { + std::vector out(rotations.size()); + on_slices(rotations.size(), cards.size(), [&](size_t k, size_t first, size_t last) { + if (first == last) + return; + const std::vector slice(rotations.begin() + first, rotations.begin() + last); + const ScopedGpuPin pin(cards[k].device); + const std::shared_ptr held = + hold ? hold(cards[k].device, MolecularReplacementGPU::RotationBytes(refl, sg, slice.size())) : nullptr; + const std::vector v = cards[k].gpu->RotationFunction(refl, cell, sg, slice); + std::copy(v.begin(), v.end(), out.begin() + first); + }); + return out; +} + +std::vector SplitTranslationFunction(const std::vector &cards, + const TranslationSetup &setup, + const std::vector &rotations, + const std::vector> *known, + TranslationSetupGPU *prepared, + const MolecularReplacementGPUHold &hold) { + // What every card's slice needs of the setup, worked out once. + std::unique_ptr made{nullptr, nullptr}; + if (prepared == nullptr && cards.size() > 1 && rotations.size() > 1) { + made = cards[0].gpu->Prepare(setup, known != nullptr); + prepared = made.get(); + } + std::vector out(rotations.size()); + on_slices(rotations.size(), cards.size(), [&](size_t k, size_t first, size_t last) { + if (first == last) + return; + const std::vector slice(rotations.begin() + first, rotations.begin() + last); + const ScopedGpuPin pin(cards[k].device); + const std::shared_ptr held = + hold ? hold(cards[k].device, MolecularReplacementGPU::TranslationBytes(setup, slice.size(), known != nullptr)) + : nullptr; + const std::vector v = cards[k].gpu->TranslationFunction(setup, slice, known, prepared); + std::copy(v.begin(), v.end(), out.begin() + first); + }); + return out; +} +#endif + MolecularReplacementResult RunMolecularReplacement(const std::string &model_path, const std::vector &merges, 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. + // The same on every other card, counting what is resident as held on each. size_t gpu_budget = 0; + std::vector other_budgets; #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); + const int device = RigidBodyGPUEngine::CurrentDevice(); + gpu_budget = ModelGPUBudget(get_gpu_total_memory(device), gpu_resident_bytes); + for (size_t total : OtherCardsTotalMemory(device)) + other_budgets.push_back(ModelGPUBudget(total, gpu_resident_bytes)); } #else (void) gpu_resident_bytes; @@ -3355,7 +3508,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, gpu_budget); + MolecularReplacementResult result = search_model(model, merges, logger, nthreads, gpu_budget, other_budgets); result.model_path = model_path; result.pae_domains = static_cast(domains.size()); result.t_model += t_model; @@ -3366,7 +3519,10 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path 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, and with an equal share of the - // budget. + // budget on every card. + std::vector other_shares; + for (size_t b : other_budgets) + other_shares.push_back(b / domains.size()); std::vector logs; for (size_t i = 0; i < domains.size(); i++) logs.push_back(Logger::Buffered()); @@ -3374,7 +3530,7 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path 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, - gpu_budget / domains.size()); + gpu_budget / domains.size(), other_shares); })); std::vector found; for (size_t i = 0; i < domains.size(); i++) { diff --git a/image_analysis/structure_refinement/MolecularReplacement.h b/image_analysis/structure_refinement/MolecularReplacement.h index ea163349a..c3820666c 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.h +++ b/image_analysis/structure_refinement/MolecularReplacement.h @@ -380,18 +380,30 @@ struct MolecularReplacementResult { // 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 + // The same for every other card, in the order the rigid body's engines count on to them from this one + // (RigidBodyGPUPool::Reserve, OtherCardsTotalMemory). + std::vector other_budgets; 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 }; +// What a search puts on one card. +struct MolecularReplacementGPUCard { + bool transform = false; // a copy of the transform: the card scores its share of every function + size_t rigid_body_engines = 0; + size_t room = 0; // for the functions' working buffers (and on this card, the score engines) +}; + 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 + size_t rigid_body_engines = 0; // 0: the rigid body on the CPU; on every card together 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 + // This card and each card of need.other_budgets, in that order. + std::vector cards; }; // The one rule for what a search puts on its GPU, decided before anything is reserved, from the budget model @@ -402,7 +414,11 @@ struct MolecularReplacementGPUPlan { // 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. +// to be done rather than spilling past it. This card is planned as on a machine with it alone; each other card +// then takes, within half of its own budget, a copy of the transform - the rotation and translation functions' +// orientations are split over the cards that hold one, each scored as on one card - and, one card after the +// other in turn, the rigid-body engines this card has no room for. How many cards there are changes only +// where the work runs, never a number. MolecularReplacementGPUPlan PlanMolecularReplacementGPU(const MolecularReplacementGPUNeed &need); // Search `merges` (the answer's first) for the predicted model at `model_path`, in every space group with diff --git a/image_analysis/structure_refinement/MolecularReplacementGPU.h b/image_analysis/structure_refinement/MolecularReplacementGPU.h index 1d68e31e0..c8f3ad538 100644 --- a/image_analysis/structure_refinement/MolecularReplacementGPU.h +++ b/image_analysis/structure_refinement/MolecularReplacementGPU.h @@ -10,6 +10,8 @@ // rounding, not bit for bit. Every reduction runs in a fixed order and nothing is summed with an // atomic, so the GPU is deterministic on its own. +#include +#include #include #include @@ -57,3 +59,29 @@ public: private: std::unique_ptr state_; }; + +// The search on one card: the transform uploaded on `device`. +struct MolecularReplacementGPUOnCard { + int device = 0; + MolecularReplacementGPU *gpu = nullptr; +}; + +// What a function's slice holds on its card while it runs: called on the slice's thread, with the card and the +// bytes its buffers take, and kept until the slice is done. Empty: nothing is held. +using MolecularReplacementGPUHold = std::function(int device, size_t bytes)>; + +// The rotation and translation functions with the orientations split over `cards`: the k-th card scores the +// k-th of cards.size() contiguous slices of `rotations`, on a thread of its own beside the others, and the +// results are put back in the order of `rotations`. Every orientation is scored by itself - one block of the +// rotation function, one map of the translation function - and nothing is summed across orientations, so the +// result is that of one card, bit for bit, on any number of cards. +std::vector SplitRotationFunction(const std::vector &cards, + const SearchReflections &refl, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, const std::vector &rotations, + const MolecularReplacementGPUHold &hold = {}); +std::vector SplitTranslationFunction(const std::vector &cards, + const TranslationSetup &setup, + const std::vector &rotations, + const std::vector> *known, + TranslationSetupGPU *prepared, + const MolecularReplacementGPUHold &hold = {}); diff --git a/image_analysis/structure_refinement/RigidBodyGPU.cpp b/image_analysis/structure_refinement/RigidBodyGPU.cpp index 184015750..65d075b2c 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.cpp +++ b/image_analysis/structure_refinement/RigidBodyGPU.cpp @@ -179,13 +179,16 @@ size_t RigidBodyGPUPool::EngineBytes() const { return RigidBodyGPUEngine::DeviceBytes(*cap_); } -void RigidBodyGPUPool::Reserve(size_t engines) { +void RigidBodyGPUPool::Reserve(const std::vector &per_card) { // The engines give the same bits on any card of one model, so where they are changes only where they run. // A card that is merely busy is waited for by the engines' allocations, never left for the CPU. try { const int cards = get_gpu_count(); - for (size_t i = 0; i < engines; i++) - engines_.push_back(std::make_unique(*cap_, (device_ + static_cast(i)) % cards)); + const size_t rounds = per_card.empty() ? 0 : *std::max_element(per_card.begin(), per_card.end()); + for (size_t round = 0; round < rounds; round++) + for (size_t k = 0; k < per_card.size(); k++) + if (round < per_card[k]) + engines_.push_back(std::make_unique(*cap_, (device_ + static_cast(k)) % cards)); set_gpu(device_); // each engine was made on its own card } catch (const JFJochException &e) { throw RigidBodyGPUFailure(e.what()); diff --git a/image_analysis/structure_refinement/RigidBodyGPU.h b/image_analysis/structure_refinement/RigidBodyGPU.h index 29732ec70..d99ad72e7 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.h +++ b/image_analysis/structure_refinement/RigidBodyGPU.h @@ -56,9 +56,9 @@ public: 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. - void Reserve(size_t engines); + // Reserves per_card[k] engines on the k-th card from Device() (k = 0 is Device() itself), one card after + // the other in turn. Engines that cannot be set up throw RigidBodyGPUFailure. + void Reserve(const std::vector &per_card); size_t Engines() const { return engines_.size(); } // The card the pool was made from, which its first engine is on and the others count on from. int Device() const { return device_; } diff --git a/tests/ModelValidationTest.cpp b/tests/ModelValidationTest.cpp index cec04d31d..cd6852bab 100644 --- a/tests/ModelValidationTest.cpp +++ b/tests/ModelValidationTest.cpp @@ -162,6 +162,43 @@ TEST_CASE("ModelValidation_GPUPlanOnCardSizes", "[ModelValidation]") { CHECK(large.may_run_beside); } +TEST_CASE("ModelValidation_GPUPlanEnginesOnOtherCards", "[ModelValidation]") { + // Four cards of 15.6 GB each: this card cannot take a second rigid-body engine within half of its 12.48 GB + // budget (4.42 + 1.88 GB > 6.24 GB), and the other nine go three to a card on the other three. + const size_t card = 15600000000; + const ModelValidationGPUPlan alone = PlanModelValidationGPU(LargeCellNeed(card)); + ModelValidationGPUNeed need = LargeCellNeed(card); + need.other_cards = {card, card, card}; + const ModelValidationGPUPlan plan = PlanModelValidationGPU(need); + // This card's parts are those of the card alone. + CHECK(plan.budget == alone.budget); + CHECK(plan.bytes == alone.bytes); + CHECK(plan.structure_factors == alone.structure_factors); + CHECK(plan.null_structure_factors == alone.null_structure_factors); + CHECK(plan.scale_fit_engines == alone.scale_fit_engines); + REQUIRE(plan.rigid_body_per_card.size() == 4); + CHECK(plan.rigid_body_per_card[0] == alone.rigid_body_engines); + CHECK(alone.rigid_body_engines == 1); + CHECK(plan.rigid_body_per_card == std::vector{1, 3, 3, 3}); + CHECK(plan.rigid_body_engines == 10); + CHECK(plan.most_on_a_card == 3 * need.rigid_body); + CHECK(plan.may_run_beside); + // Each other card within half of its own budget, less what a validation beside holds there. + need.beside = plan.most_on_a_card; + const ModelValidationGPUPlan beside = PlanModelValidationGPU(need); + for (size_t k = 1; k < 4; k++) + CHECK(beside.rigid_body_per_card[k] * need.rigid_body + <= std::min(plan.budget / 2, plan.budget - need.beside)); + // What is resident is counted on every card: a merge holding most of each leaves no room for extra engines. + need.beside = 0; + need.resident = 8 * GB; + const ModelValidationGPUPlan full = PlanModelValidationGPU(need); + CHECK(full.rigid_body_per_card == std::vector{1, 1, 1, 1}); + // One card alone: the plan as it was. + CHECK(alone.rigid_body_per_card == std::vector{1}); + CHECK(alone.most_on_a_card == alone.bytes); +} + TEST_CASE("ModelValidation_GPUPlanPartThatDoesNotFitGoesToTheCPU", "[ModelValidation]") { ModelValidationGPUNeed need = LargeCellNeed(8 * GB); need.structure_factors = 7 * GB; // over the 6.4 GB budget on its own @@ -1282,7 +1319,7 @@ namespace { size_t engines, Logger &logger) { auto pool = RigidBodyGPUPool::Plan(model, cell, sg, d_min, observations, logger); REQUIRE(pool); - pool->Reserve(engines); + pool->Reserve({engines}); return pool; } } diff --git a/tests/MolecularReplacementTest.cpp b/tests/MolecularReplacementTest.cpp index e03353a9f..cf8f3054e 100644 --- a/tests/MolecularReplacementTest.cpp +++ b/tests/MolecularReplacementTest.cpp @@ -3,8 +3,10 @@ #include +#include #include #include +#include #include #include @@ -131,6 +133,46 @@ TEST_CASE("MolecularReplacement_GPUPlanOnCardSizes", "[MolecularReplacement]") { CHECK(PlanMolecularReplacementGPU(LargeCellNeed(64 * GB)).rigid_body_engines == 9); } +TEST_CASE("MolecularReplacement_GPUPlanOnOtherCards", "[MolecularReplacement]") { + // Four cards of 15.6 GB, a 12.48 GB budget each: this card holds what it holds alone (one engine), each other + // card a copy of the transform and, within half of its budget, two engines. + const size_t card = 15600000000; + const MolecularReplacementGPUPlan alone = PlanMolecularReplacementGPU(LargeCellNeed(card)); + MolecularReplacementGPUNeed need = LargeCellNeed(card); + need.other_budgets = {need.budget, need.budget, need.budget}; + const MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need); + CHECK(plan.search == alone.search); + CHECK(plan.structure_factors == alone.structure_factors); + CHECK(plan.room == alone.room); + CHECK(plan.bytes == alone.bytes); + REQUIRE(plan.cards.size() == 4); + CHECK(plan.cards[0].transform); + CHECK(plan.cards[0].rigid_body_engines == alone.rigid_body_engines); + CHECK(plan.cards[0].room == alone.room); + CHECK(alone.rigid_body_engines == 1); + for (size_t k = 1; k < 4; k++) { + CHECK(plan.cards[k].transform); + CHECK(plan.cards[k].rigid_body_engines == 2); + CHECK(plan.cards[k].room == need.budget - need.transform - 2 * need.rigid_body); + } + CHECK(plan.rigid_body_engines == 7); + // No more engines than the search can use: with 64 GB cards this card alone takes all nine. + need = LargeCellNeed(64 * GB); + need.other_budgets = {need.budget}; + const MolecularReplacementGPUPlan large = PlanMolecularReplacementGPU(need); + CHECK(large.rigid_body_engines == 9); + CHECK(large.cards[0].rigid_body_engines == 9); + CHECK(large.cards[1].rigid_body_engines == 0); + CHECK(large.cards[1].transform); + // A search on the CPU puts nothing on the other cards either. + need = LargeCellNeed(card); + need.transform = need.budget + 1; + need.other_budgets = {need.budget}; + const MolecularReplacementGPUPlan cpu = PlanMolecularReplacementGPU(need); + CHECK(!cpu.search); + CHECK(!cpu.cards[1].transform); +} + 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. @@ -511,7 +553,7 @@ TEST_CASE("MolecularReplacement_RigidBodyBesideFixedCopy", "[MolecularReplacemen gemmi::Structure gpu = start; auto pool = RigidBodyGPUPool::Plan(gpu.models[0], cell, sg, 3.0, fobs.v.size(), logger); REQUIRE(pool); - pool->Reserve(1); + pool->Reserve({1}); RefineRigidBody(gpu.models[0], cell, sg, fobs, 3.0, logger, 4, pool.get(), &first.models[0]); CHECK(worst(gpu) < 0.1); } @@ -577,3 +619,67 @@ TEST_CASE("MolecularReplacement_GPUMatchesCPU", "[MolecularReplacement]") { } } #endif + +#ifdef JFJOCH_USE_CUDA +TEST_CASE("MolecularReplacement_SplitOverCardsMatchesOneCard", "[MolecularReplacement]") { + if (get_gpu_count() == 0) + SKIP("no GPU"); + const SearchModel model = synthetic_model(); + const gemmi::UnitCell cell(42, 50, 58, 90, 90, 90); + const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21"); + const gemmi::Mat33 truth = axis_angle({1, 2, 3}, 50); + const auto merged = observed(PlaceSearchModel(model, truth, gemmi::Fractional(0.17, 0.31, 0.09), cell, sg)); + const SearchReflections refl = MakeSearchReflections(merged, cell, sg, 4.0, 12.0); + const MolecularTransform ft(model, 4.0, 4); + // Four "cards", each a search of its own on this one: the split's slices, gather and order are those of four + // cards, so this is the multi-card path, run where there is one card. + MolecularReplacementGPU one(ft); + std::vector> copies; + std::vector cards; + for (int k = 0; k < 4; k++) { + copies.push_back(std::make_unique(ft)); + cards.push_back({0, copies.back().get()}); // the test thread's device + } + std::atomic holds{0}; + const MolecularReplacementGPUHold hold = [&](int, size_t bytes) -> std::shared_ptr { + CHECK(bytes > 0); + holds++; + return nullptr; + }; + const auto all = RotationAsu(UniformRotations(15.0), CartesianPointGroup(sg, cell)); + REQUIRE(all.size() % 4 != 0); // the last slice shorter than the others + for (size_t n : {all.size(), size_t(3), size_t(1)}) { + const std::vector rotations(all.begin(), all.begin() + n); + const std::vector want = one.RotationFunction(refl, cell, sg, rotations); + for (size_t parts : {1, 2, 3, 4}) { + const std::vector some(cards.begin(), cards.begin() + parts); + const std::vector got = SplitRotationFunction(some, refl, cell, sg, rotations, hold); + CAPTURE(n, parts); + REQUIRE(got.size() == want.size()); + CHECK(std::memcmp(got.data(), want.data(), want.size() * sizeof(double)) == 0); + } + } + CHECK(holds > 0); + + const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg); + const std::vector orientations{truth, all[3], all[17], all[20], all[31]}; + const TranslationSetup setup2 = MakeTranslationSetup(refl, cell, sg, true); + const auto known = KnownStructureFactors(ft, setup2, cell, {all[5]}, {gemmi::Fractional(0.3, 0.1, 0.4)}); + for (const std::vector> *k : {(const std::vector> *) nullptr, &known}) { + const TranslationSetup &s = k == nullptr ? setup : setup2; + const std::vector want = one.TranslationFunction(s, orientations, k); + for (size_t parts : {2, 4}) { + const std::vector some(cards.begin(), cards.begin() + parts); + const std::vector got = SplitTranslationFunction(some, s, orientations, k, nullptr, hold); + CAPTURE(parts, k != nullptr); + REQUIRE(got.size() == want.size()); + for (size_t i = 0; i < want.size(); i++) { + CHECK(std::memcmp(&got[i].z, &want[i].z, sizeof(double)) == 0); + CHECK(got[i].t.x == want[i].t.x); + CHECK(got[i].t.y == want[i].t.y); + CHECK(got[i].t.z == want[i].t.z); + } + } + } +} +#endif