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

Merge branch 'modelpar' into gpusf

Browse Source 
Both sides kept: modelpar's parallel basis scoring, the second validation
started on a forecast beside the first (write gate, schedule parameter),
the multi-GPU placement and delete-before-rewrite; gpusf's GPU structure
factors, maps and null engine, failure-instead-of-restart, and the merge
engine released before the validation (now just before modelpar's
ValidateAgainstModel call, after the forecast lambda is set up).

Placement: each validation's structure-factor engines (d_min and the
null's) are made on the card of the thread that runs it - the main
thread's for the first validation, card 1 % count for the speculative
second, which pins itself there - so two validations on two cards use
both, as the rigid-body pools do.

One memory rule for both, per card, from total memory, up front: a
validation plans at most half of its card - its structure-factor engines
a quarter together (was half for the d_min engine alone), its rigid-body
engines a quarter (RigidBodyGPUPool) - and the second validation runs
beside the first only where twice the first's plan (rigid-body planned
bytes + the d_min engine, twice it where a null is coming, the null's
engine being no larger) fits half of all cards' memory together (was:
twice the rigid-body plan within a quarter). The card's total is read
once when the engine is made; a CUDA error there fails the validation
like any other.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
This commit is contained in:
leonarski_f
2026-10-09 13:51:41 +02:00
co-authored by Claude Opus 5.5
parent 0f707caab9 89f52fef95
commit 278c488c37
9 changed files with 229 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/ModelValidation.cpp
+79 -21
View File
@@ -7,6 +7,7 @@
#include <algorithm>
#include <cmath>
#include <complex>
#include <cstdio>
#include <array>
#include <atomic>
#include <functional>
@@ -33,6 +34,7 @@
#include "../../common/CorrelationCoefficient.h"
#include "../../common/JFJochMath.h" // PI (M_PI is not standard, and MSVC does not define it)
#include "../../common/Logger.h"
#include "../../common/ParallelFor.h"
#include "../scale_merge/ReindexAmbiguity.h" // ReindexReflections
#include "../scale_merge/CrystalSetting.h" // CellMappingOperators
#include "ModelFFT.h"
@@ -75,11 +77,20 @@ gemmi::Grid<float> map_from_coefficients(gemmi::AsuData<std::complex<float>> &co
return MapFromFPhi(gemmi::get_f_phi_on_grid<float>(coef, size, true));
}
// A file this run may already have written - the validation in the model's setting writes over the
// first one's maps - is removed before it is written again, not truncated: XFS (and ext4) flush a file
// truncated and rewritten when it is closed, and three 150 MB maps forced out to a disk that way cost
// seconds where a fresh file costs nothing.
void remove_before_rewriting(const std::string &path) {
std::remove(path.c_str());
}
// Write a map as CCP4; return its RMS (the sigma the map is read in).
double write_ccp4(const gemmi::Grid<float> &map, const std::string &path) {
gemmi::Ccp4<float> ccp4;
ccp4.grid = map;
ccp4.update_ccp4_header(2);
remove_before_rewriting(path);
ccp4.write_ccp4_map(path);
return ccp4.hstats.rms;
}
@@ -357,28 +368,29 @@ double frame_probe_r(const gemmi::Structure &st, const gemmi::SpaceGroup *sg,
// once, before any of them is computed, and never revisited: nothing is moved to the CPU part way through.
// Decided on what the GPU path needs against the card's TOTAL memory and not against what happens to be
// free at the moment, so the same input on the same machine always takes the same path, whatever else is
// running beside it. The structure factors may take half the card - `reserved` is what an engine of the
// same validation already holds there; the rigid body's engines take at most a quarter beside them
// (RigidBodyGPUPool::Create), and the last quarter is left to everything else.
// running beside it. A validation plans at most half of its card: its structure-factor engines a quarter
// together - `reserved` is what one of them already holds there - and its rigid-body engines a quarter
// (RigidBodyGPUPool::Create). A second validation runs beside it only where twice that plan fits half of
// the cards together (see on_forecast below), so validations never plan more than half of the memory.
std::unique_ptr<ModelStructureFactorsGPU> StructureFactorsGPU(int device, const gemmi::Model &model,
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
double d_min, size_t reserved, Logger &logger) {
auto sf = std::make_unique<ModelStructureFactorsGPU>(device, model, cell, sg, d_min);
std::unique_ptr<ModelStructureFactorsGPU> sf;
try {
sf = std::make_unique<ModelStructureFactorsGPU>(device, model, cell, sg, d_min);
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
const std::array<int, 3> n = sf->GridSize();
if (!sf->Supported()) {
logger.Info("Model validation: structure factors to {:.2f} A on the CPU - the cell is too small for the "
"GPU's gridding", d_min);
return nullptr;
}
size_t total = 0;
try {
total = sf->DeviceTotalMemory();
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
if (reserved + sf->DeviceBytes() > total / 2) {
const size_t total = sf->DeviceTotalMemory();
if (reserved + sf->DeviceBytes() > total / 4) {
logger.Info("Model validation: structure factors to {:.2f} A on the CPU - on a {}x{}x{} grid they need "
"{:.2f} GB of GPU memory beside {:.2f} GB already reserved, over half of the card's {:.2f} GB",
"{:.2f} GB of GPU memory beside {:.2f} GB already reserved, over a quarter of the card's {:.2f} GB",
d_min, n[0], n[1], n[2], sf->DeviceBytes() / 1e9, reserved / 1e9, total / 1e9);
return nullptr;
}
@@ -410,6 +422,7 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
size_t nthreads,
double wavelength_A,
const std::vector<float> &report_shell_d_min,
const ModelValidationSchedule &schedule,
bool rigid_body_gpu) {
ModelValidationResult result;
result.model_path = model_path;
@@ -539,13 +552,15 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
"is not how the data describe this lattice; scoring {} change(s) of basis",
st.cell.a, st.cell.b, st.cell.c, st.cell.alpha, st.cell.beta, st.cell.gamma,
sg->xhm(), frames.size() - 1);
size_t best = 0;
double best_r = 0;
for (size_t i = 0; i < frames.size(); i++) {
// The candidates are scored at the same time, each on its own copy of the model, and read
// in their own order below, so the ranking and its tie-break are those of one after another.
std::vector<double> frame_r(frames.size(), 0.0);
std::vector<const gemmi::SpaceGroup *> frame_sg(frames.size(), nullptr);
ParallelFor(static_cast<int>(frames.size()), nthreads, [&](int i) {
gemmi::Structure trial = st;
const gemmi::SpaceGroup *trial_sg = sg;
if (i > 0 && !change_model_basis(trial, trial_sg, frames[i]))
continue; // no origin shift names the transformed group; not a basis we can take
return; // no origin shift names the transformed group; not a basis we can take
refractionalize_into(trial, data_cell);
trial.setup_cell_images();
// Each frame at the data's best indexing: the alternative indexings are only probed
@@ -557,8 +572,17 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
if (same_point_group)
for (const gemmi::Op &law : reindex_ops)
r = std::min(r, frame_probe_r(trial, trial_sg, ReindexReflections(obs, law), d_min));
frame_r[i] = r;
frame_sg[i] = trial_sg;
});
size_t best = 0;
double best_r = 0;
for (size_t i = 0; i < frames.size(); i++) {
if (frame_sg[i] == nullptr)
continue;
const double r = frame_r[i];
logger.Info("Model validation: {:<12} -> {:<12} R {:.4f} on the coarse shell",
frames[i].triplet(), trial_sg->xhm(), r);
frames[i].triplet(), frame_sg[i]->xhm(), r);
if (i == 0 || r < best_r) {
best_r = r;
best = i;
@@ -880,6 +904,27 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
const bool decision_pending = result.model_enantiomorph_candidate
|| !(indexing.op == gemmi::Op::identity())
|| !(result.change_of_basis_op == gemmi::Op::identity());
if (schedule.on_forecast) {
// Two validations at once take twice the engines, and are allowed where twice what this one plans -
// its rigid-body engines, its structure-factor engine and, where a null is coming, the null's,
// which is no larger - fits half of the cards' memory taken together (see StructureFactorsGPU).
bool may_run_beside = true;
size_t planned = 0, card = 0;
if (rigid_body_pool != nullptr) {
planned += rigid_body_pool->PlannedBytes();
card = rigid_body_pool->CardBytes();
}
#ifdef JFJOCH_USE_CUDA
if (sf_gpu != nullptr) {
planned += (decision_pending ? 2 : 1) * sf_gpu->DeviceBytes();
card = sf_gpu->DeviceTotalMemory();
}
#endif
if (planned > 0)
may_run_beside = 2 * planned <= static_cast<size_t>(get_gpu_count()) * card / 2;
schedule.on_forecast({result.change_of_basis_op, indexing.op, result.model_enantiomorph_candidate,
may_run_beside});
}
std::vector<gemmi::Mat33> equivalent; // the orientations the crystal cannot tell from the model's
double reach_deg = 0;
std::vector<gemmi::Mat33> null_rotation;
@@ -1024,7 +1069,11 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
std::vector<std::future<void>> replicate_threads;
if (decision_pending)
for (size_t t = 0; t < std::min<size_t>(std::max<size_t>(nthreads, 1), NULL_REPLICATES); t++)
replicate_threads.push_back(std::async(std::launch::async, [&] {
replicate_threads.push_back(std::async(std::launch::async, [&, t] {
// Each on a card in turn from the one after the real fit's, which its engines are
// spread the same way over (RigidBodyGPUPool::Create).
if (rigid_body_pool != nullptr)
pin_gpu((rigid_body_pool->Device() + 1 + static_cast<int>(t)) % get_gpu_count());
for (int i = next_replicate++; i < NULL_REPLICATES; i = next_replicate++)
run_replicate(i);
}));
@@ -1354,6 +1403,13 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
mapfofc.v.push_back({terms[i].hkl, delfwt[i] * ph});
}
// A validation started before the one it follows had decided writes nothing until that one has,
// and nothing at all where it decided otherwise.
if (schedule.write_gate.valid() && !schedule.write_gate.get()) {
result.failure_reason = "superseded before its files were written";
return result;
}
// --- write the maps and score the 2mFo-DFc map at atom centres (a real map peaks there) ---
// The three maps - this one, the difference map and the anomalous map - share nothing but what
// they read, so the other two are made and written on threads of their own beside this one.
@@ -1525,6 +1581,7 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
}
mtz.nreflections = static_cast<int>(terms.size());
mtz.data = std::move(data);
remove_before_rewriting(output_prefix + "_maps.mtz");
mtz.write_to_file(output_prefix + "_maps.mtz");
} catch (const std::exception &e) {
logger.Warning("Model validation: could not write map MTZ: {}", e.what());
@@ -1581,14 +1638,15 @@ ModelValidationResult ValidateAgainstModel(const std::vector<MergedReflection> &
bool probe_indexing_ambiguity,
size_t nthreads,
double wavelength_A,
const std::vector<float> &report_shell_d_min) {
const std::vector<float> &report_shell_d_min,
const ModelValidationSchedule &schedule) {
#ifdef JFJOCH_USE_CUDA
// A CUDA failure ends the validation, not the run, and nothing is moved to the CPU part way through:
// the validation reports why it did not finish, and a validation that did not finish decides nothing,
// so the reflection files are those of a run without a model.
try {
return Validate(merged, cell, model_path, output_prefix, logger, data_space_group,
probe_indexing_ambiguity, nthreads, wavelength_A, report_shell_d_min, true);
probe_indexing_ambiguity, nthreads, wavelength_A, report_shell_d_min, schedule, true);
} catch (const RigidBodyGPUFailure &e) {
cuda_clear_error();
ModelValidationResult failed;
@@ -1599,7 +1657,7 @@ ModelValidationResult ValidateAgainstModel(const std::vector<MergedReflection> &
}
#else
return Validate(merged, cell, model_path, output_prefix, logger, data_space_group, probe_indexing_ambiguity,
nthreads, wavelength_A, report_shell_d_min, false);
nthreads, wavelength_A, report_shell_d_min, schedule, false);
#endif
}