mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 82 Wiki Activity

Molecular replacement: one GPU plan per merge on model validation's budget, the transform cropped, output unchanged

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-10-15 13:33:39 +02:00
co-authored by Claude Opus 5.5
parent 1ebb28e3ab
commit 011b98838e
12 changed files with 644 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/ModelValidation.cpp
+18 -9
View File
@@ -1883,10 +1883,14 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
} // namespace
size_t ModelGPUBudget(size_t card, size_t resident) {
const size_t usable = card / 5 * 4;
return usable > resident ? usable - resident : 0;
}
ModelValidationGPUPlan PlanModelValidationGPU(const ModelValidationGPUNeed &need, size_t cards) {
ModelValidationGPUPlan plan;
const size_t usable = need.card / 5 * 4;
plan.budget = usable > need.resident ? usable - need.resident : 0;
plan.budget = ModelGPUBudget(need.card, need.resident);
auto take = [&](size_t bytes, size_t limit) {
if (bytes == 0 || plan.bytes + bytes > limit)
return false;
@@ -1999,7 +2003,8 @@ const gemmi::SpaceGroup *AdoptModelFrame(const ModelValidationResult &validation
ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection> &merged,
const gemmi::Structure &placed, const gemmi::SpaceGroup &data_space_group,
double k_sol, double b_sol, Logger &logger, size_t nthreads) {
double k_sol, double b_sol, Logger &logger, size_t nthreads,
size_t gpu_resident_bytes) {
ModelValidationResult result;
gemmi::Structure st = placed;
st.setup_cell_images();
@@ -2018,13 +2023,16 @@ ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection>
ModelStructureFactorsGPU *gpu = nullptr;
#ifdef JFJOCH_USE_CUDA
std::unique_ptr<ModelStructureFactorsGPU> engine;
// Read beside molecular replacement's engines, which keep to a quarter of the card
// (MolecularReplacement.cpp), so on the GPU where it fits a quarter as well.
if (get_gpu_count() > 0)
// After molecular replacement, whose engines are gone: alone on the card beside what is resident.
if (get_gpu_count() > 0) {
cuda_trim_allocation_pool();
engine = StructureFactorsEngine(RigidBodyGPUEngine::CurrentDevice(), st.models[0], st.cell, *sg, d_min, logger);
if (engine && engine->DeviceBytes() > engine->DeviceTotalMemory() / 4) {
logger.Info("Model validation: structure factors to {:.2f} A on the CPU - {} of GPU memory, over a quarter "
"of the card's {}", d_min, GBytes(engine->DeviceBytes()), GBytes(engine->DeviceTotalMemory()));
}
if (engine && engine->DeviceBytes() > ModelGPUBudget(engine->DeviceTotalMemory(), gpu_resident_bytes)) {
logger.Info("Model validation: structure factors to {:.2f} A on the CPU - {} of GPU memory, over the "
"budget of 80% of the card's {}{}", d_min, GBytes(engine->DeviceBytes()),
GBytes(engine->DeviceTotalMemory()),
gpu_resident_bytes > 0 ? fmt::format(" less {} held by the merge", GBytes(gpu_resident_bytes)) : "");
engine.reset();
}
if (engine)
@@ -2039,6 +2047,7 @@ ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection>
}
#else
(void) logger;
(void) gpu_resident_bytes;
#endif
if (gpu == nullptr) {
gemmi::DensityCalculator<Table, float> dc;