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

MR and model validation: use every card's memory, split MR's functions over the cards

Browse Source 
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-10-16 17:16:28 +02:00
1 parent 19b135186d
commit bfc6ca8487
9 files changed
+454 -54
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/structure_refinement/ModelValidation.cpp
+48 -8
View File
@@ -799,6 +799,7 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &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<MergedReflection> &merged,
need.scale_fit = ModelScaleGPUPool::DeviceBytes(obs.size(), 1);
need.scale_fit_engines = SCALE_GPU_ENGINES;
gpu_plan = PlanModelValidationGPU(need, static_cast<size_t>(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<MergedReflection> &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<MergedReflection> &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<size_t> 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<size_t> OtherCardsTotalMemory(int device) {
std::vector<size_t> 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<MergedReflection> &merged,
const UnitCell &cell,
const std::string &model_path,