From 1fe73bab4130de808528859b5ff7a65314787825 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 12 Oct 2026 21:40:54 +0200 Subject: [PATCH] Model validation: scale fits on the GPU, change-of-basis probe gridded once per frame on all threads The scale fit of every model-validation fit (fit_model: the real model, the indexing probe's relabellings and every null replicate) was FitModelScale on the host - a long-double Levenberg-Marquardt per solvent grid point, ~80% of the validation's CPU on a 400 A cell. It now runs on the GPU through ModelScaleGPU::FitSolvent, the engine the rigid body and molecular replacement already use, from a small pool (ModelScaleGPUPool: up to four engines with their own streams, sized on the card's total memory, at most an eighth of it). The grid, the fits and the winner rule are FitModelScale's; the sums are taken in double on the device, so R-factors agree with the host's to rounding. The CPU build, and a GPU build without a card, keep FitModelScale. The change-of-basis probe gridded each candidate frame with gemmi's serial put_model_density_on_grid, once for the data as merged and once more for each alternative indexing. It now grids each frame once, with ModelGrid's parallel PutModelDensityOnGrid (the same grid, bit for bit), from threads outside the pool so the gridding spreads over it, and scores every labelling on that F_calc. Exact. 4nwv (122k atoms, I222, 400 A cell): model validation 33.2 s -> 14.0 s; every logged R, sigma and decision the same. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- .../structure_refinement/ModelScaling.cpp | 106 +++++++++++++ .../structure_refinement/ModelScaling.h | 36 +++++ .../structure_refinement/ModelValidation.cpp | 140 +++++++++++++----- tests/ModelScaleGPUTest.cpp | 35 +++++ 4 files changed, 282 insertions(+), 35 deletions(-) diff --git a/image_analysis/structure_refinement/ModelScaling.cpp b/image_analysis/structure_refinement/ModelScaling.cpp index 84575bd22..d7b941c5a 100644 --- a/image_analysis/structure_refinement/ModelScaling.cpp +++ b/image_analysis/structure_refinement/ModelScaling.cpp @@ -16,6 +16,13 @@ #include "../../common/ParallelFor.h" +#ifdef JFJOCH_USE_CUDA +#include "../../common/CUDAWrapper.h" +#include "../indexing/CUDAMemHelpers.h" +#include "ModelScaleGPU.h" +#include "RigidBodyGPUEngine.h" // ModelScaleGPUTooFewReflections, RigidBodyGPUEngine::CurrentDevice +#endif + namespace { // gemmi::Scaling as a grid point fits it: the same starting point, parameters, model values, @@ -415,3 +422,102 @@ double SolventBMax(const gemmi::Model &model) { const double b_max = ModelScaleBox{}.b_hi; return n > 0 ? std::max(b_max, 2.0 * sum / static_cast(n)) : b_max; } + +#ifdef JFJOCH_USE_CUDA +struct ModelScaleGPUPool::Engine { + CudaStream stream; + ModelScaleGPU scale; + CudaDevicePtr fcalc, fmask; + explicit Engine(size_t max_points) : scale(stream, max_points), fcalc(max_points), fmask(max_points) {} +}; + +size_t ModelScaleGPUPool::DeviceBytes(size_t max_points, size_t engines) { + return engines * (ModelScaleGPU::DeviceBytes(max_points) + 2 * max_points * sizeof(float2)); +} + +ModelScaleGPUPool::ModelScaleGPUPool(int device, size_t max_points, size_t engines) + : device_(device), max_points_(max_points) { + const int current = RigidBodyGPUEngine::CurrentDevice(); + set_gpu(device_); + for (size_t i = 0; i < engines; i++) { + engines_.push_back(std::make_unique(max_points)); + idle_.push_back(engines_.back().get()); + } + set_gpu(current); +} + +ModelScaleGPUPool::~ModelScaleGPUPool() { + const int current = RigidBodyGPUEngine::CurrentDevice(); + set_gpu(device_); + engines_.clear(); + set_gpu(current); +} + +ModelScaleReport ModelScaleGPUPool::Fit(gemmi::Scaling &scaling, double b_hi, size_t nthreads) { + const size_t n = scaling.points.size(); + std::vector> hkl(n); + std::vector stol2(n); + std::vector fobs(n), sigma(n); + std::vector fc(n), fm(n); + for (size_t i = 0; i < n; i++) { + const auto &p = scaling.points[i]; + hkl[i] = {p.hkl[0], p.hkl[1], p.hkl[2]}; + stol2[i] = p.stol2; + fobs[i] = p.fobs; + sigma[i] = p.sigma; + fc[i] = {p.fcmol.real(), p.fcmol.imag()}; + fm[i] = {p.fmask.real(), p.fmask.imag()}; + } + std::vector> constraints; + for (const gemmi::Vec6 &c : scaling.constraint_matrix) + constraints.push_back(c); + double frac[9]; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) + frac[3 * i + j] = scaling.cell.frac.mat[i][j]; + + Engine *e; + { + std::unique_lock lock(m_); + cv_.wait(lock, [&] { return !idle_.empty(); }); + e = idle_.back(); + idle_.pop_back(); + } + auto release = [&] { + std::lock_guard lock(m_); + idle_.push_back(e); + cv_.notify_one(); + }; + const int current = RigidBodyGPUEngine::CurrentDevice(); + ModelSolventFit fit; + try { + set_gpu(device_); + e->scale.SetPoints(hkl, stol2, fobs, sigma, constraints, frac); + if (cudaMemcpyAsync(e->fcalc, fc.data(), n * sizeof(float2), cudaMemcpyHostToDevice, e->stream) != cudaSuccess + || cudaMemcpyAsync(e->fmask, fm.data(), n * sizeof(float2), cudaMemcpyHostToDevice, e->stream) != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "ModelScaleGPUPool: the upload failed"); + fit = e->scale.FitSolvent(e->fcalc, e->fmask, b_hi); + } catch (const ModelScaleGPUTooFewReflections &) { + set_gpu(current); + release(); + return FitModelScale(scaling, {.b_hi = b_hi}, nthreads); + } catch (...) { + set_gpu(current); + release(); + throw; + } + set_gpu(current); + release(); + + scaling.k_sol = fit.k_sol; + scaling.b_sol = fit.b_sol; + scaling.k_overall = fit.scale.k_overall; + const double *b = fit.scale.b_star; + scaling.b_star = {b[0], b[1], b[2], b[3], b[4], b[5]}; + ModelScaleReport report; + report.n_points = static_cast(n); + report.n_grid = fit.n_grid; + report.r_work_fit = fit.r; + return report; +} +#endif diff --git a/image_analysis/structure_refinement/ModelScaling.h b/image_analysis/structure_refinement/ModelScaling.h index 039fcb64e..a3d465e85 100644 --- a/image_analysis/structure_refinement/ModelScaling.h +++ b/image_analysis/structure_refinement/ModelScaling.h @@ -3,7 +3,11 @@ #pragma once +#include #include +#include +#include +#include #include "gemmi/model.hpp" #include "gemmi/scaling.hpp" @@ -46,3 +50,35 @@ struct ModelScaleReport { // nthreads: the grid points are fitted in parallel on up to this many threads; the answer does not // depend on it. ModelScaleReport FitModelScale(gemmi::Scaling &scaling, ModelScaleBox box = {}, size_t nthreads = 1); + +#ifdef JFJOCH_USE_CUDA +// FitModelScale on a GPU (ModelScaleGPU::FitSolvent): the same grid over the same box, the same fits and +// the same winner, with the sums over the reflections taken in double on the device where the host takes +// them in long double - so the answer is FitModelScale's to rounding, not to the bit. A few engines, +// reserved once, each a stream and room for one fit at a time; a fit takes a free one and waits for one +// where all are taken. Every kernel is deterministic, so which engine a fit gets changes no number. +class ModelScaleGPUPool { +public: + // What `engines` engines for fits of up to `max_points` reflections take on the card. + static size_t DeviceBytes(size_t max_points, size_t engines); + // On `device`. Throws JFJochException on a CUDA failure. + ModelScaleGPUPool(int device, size_t max_points, size_t engines); + ~ModelScaleGPUPool(); + + // FitModelScale(scaling, {.b_hi = b_hi}, nthreads). Where FitModelScale would walk the grid point to + // point (five or fewer reflections for the isotropic start), it is FitModelScale on the host. Safe to + // call from several threads. Throws JFJochException on a CUDA failure. + ModelScaleReport Fit(gemmi::Scaling &scaling, double b_hi, size_t nthreads); + // What the pool holds on the card. + size_t Bytes() const { return DeviceBytes(max_points_, engines_.size()); } + +private: + struct Engine; + int device_; + size_t max_points_; + std::vector> engines_; + std::vector idle_; + std::mutex m_; + std::condition_variable cv_; +}; +#endif diff --git a/image_analysis/structure_refinement/ModelValidation.cpp b/image_analysis/structure_refinement/ModelValidation.cpp index 0b971faa2..6792f8336 100644 --- a/image_analysis/structure_refinement/ModelValidation.cpp +++ b/image_analysis/structure_refinement/ModelValidation.cpp @@ -344,49 +344,63 @@ constexpr double FRAME_PROBE_D_MAX = 8.0; constexpr size_t FRAME_PROBE_MIN_REFLECTIONS = 200; // R of the model against the observed amplitudes after an overall scale and an isotropic B, over -// that coarse shell. Only the ranking is ever used, never the value. +// that coarse shell - the lowest over `labellings`, the labellings of the data it is scored against, +// each on the same F_calc. Only the ranking is ever used, never the value. double frame_probe_r(const gemmi::Structure &st, const gemmi::SpaceGroup *sg, - const std::vector &obs, double d_min) { + const std::vector *> &labellings, double d_min, + size_t nthreads) { const double probe_d_min = std::max(d_min, FRAME_PROBE_D_MIN); gemmi::DensityCalculator dc; dc.d_min = probe_d_min; dc.rate = 1.5; dc.set_grid_cell_and_spacegroup(st); dc.set_refmac_compatible_blur(st.models[0]); - dc.put_model_density_on_grid(st.models[0]); + // put_model_density_on_grid() spread over the threads, the same grid bit for bit (ModelGrid.h); the + // orbits are those of the grid initialize_grid() sizes. + gemmi::Grid sized; + sized.unit_cell = dc.grid.unit_cell; + sized.spacegroup = dc.grid.spacegroup; + sized.set_size_from_spacing(dc.requested_grid_spacing(), gemmi::GridSizeRounding::Up); + PutModelDensityOnGrid(dc, st.models[0], OrbitLeaders(sized, nthreads), nthreads); gemmi::AsuData> fcalc = MapToFPhi(dc.grid) .prepare_asu_data(probe_d_min, dc.blur, false, false, false); gemmi::GroupOps gops = sg->operations(); gemmi::ReciprocalAsu asu(sg); - gemmi::AsuData> fobs; - fobs.unit_cell_ = st.cell; - fobs.spacegroup_ = sg; - // The low-resolution cap is dropped, rather than the probe abandoned, when a small cell does not - // put enough reflections in the shell. - for (double d_max : {FRAME_PROBE_D_MAX, 1e9}) { - fobs.v.clear(); - for (const MergedReflection &r : obs) { - if (std::isnan(r.F) || r.d < probe_d_min || r.d > d_max) - continue; - gemmi::Miller h{{r.h, r.k, r.l}}; - if (!asu.is_in(h)) - h = asu.to_asu(h, gops).first; - fobs.v.push_back({h, {r.F, 1.0f}}); + double best = 0; + for (size_t i = 0; i < labellings.size(); i++) { + gemmi::AsuData> fobs; + fobs.unit_cell_ = st.cell; + fobs.spacegroup_ = sg; + // The low-resolution cap is dropped, rather than the probe abandoned, when a small cell does not + // put enough reflections in the shell. + for (double d_max : {FRAME_PROBE_D_MAX, 1e9}) { + fobs.v.clear(); + for (const MergedReflection &r : *labellings[i]) { + if (std::isnan(r.F) || r.d < probe_d_min || r.d > d_max) + continue; + gemmi::Miller h{{r.h, r.k, r.l}}; + if (!asu.is_in(h)) + h = asu.to_asu(h, gops).first; + fobs.v.push_back({h, {r.F, 1.0f}}); + } + if (fobs.v.size() >= FRAME_PROBE_MIN_REFLECTIONS) + break; } - if (fobs.v.size() >= FRAME_PROBE_MIN_REFLECTIONS) - break; + double r = 1.0; + if (!fobs.v.empty()) { + fobs.ensure_asu(); + fobs.ensure_sorted(); + gemmi::Scaling scaling(st.cell, sg); + scaling.use_solvent = false; + scaling.prepare_points(fcalc, fobs, nullptr); + scaling.fit_isotropic_b_approximately(); + r = scaling.calculate_r_factor(); + } + best = i == 0 ? r : std::min(best, r); } - if (fobs.v.empty()) - return 1.0; - fobs.ensure_asu(); - fobs.ensure_sorted(); - gemmi::Scaling scaling(st.cell, sg); - scaling.use_solvent = false; - scaling.prepare_points(fcalc, fobs, nullptr); - scaling.fit_isotropic_b_approximately(); - return scaling.calculate_r_factor(); + return best; } // The anomalous difference map's coefficients on the reflections of `fmodel`: F(+) - F(-) carried on the @@ -516,6 +530,33 @@ std::unique_ptr StructureFactorsGPU(int device, const "card's {:.2f} GB", d_min, n[0], n[1], n[2], sf->DeviceBytes() / 1e9, total / 1e9); return sf; } + +// The scale fits on the GPU (ModelScaleGPUPool) - or not, decided here, once, on the card's total memory as +// the structure factors are: SCALE_GPU_ENGINES engines, fewer where they would take over an eighth of the +// card, none where not even one fits. +constexpr size_t SCALE_GPU_ENGINES = 4; + +std::unique_ptr ScaleFitsGPU(int device, size_t max_points, Logger &logger) { + const size_t total = get_gpu_total_memory(device); + size_t engines = SCALE_GPU_ENGINES; + while (engines > 0 && ModelScaleGPUPool::DeviceBytes(max_points, engines) > total / 8) + engines--; + if (engines == 0) { + logger.Info("Model validation: scale fits on the CPU - one fit of {} reflections needs {:.2f} GB of GPU " + "memory, over an eighth of the card's {:.2f} GB", max_points, + ModelScaleGPUPool::DeviceBytes(max_points, 1) / 1e9, total / 1e9); + return nullptr; + } + std::unique_ptr pool; + try { + pool = std::make_unique(device, max_points, engines); + } catch (const JFJochException &e) { + throw RigidBodyGPUFailure(e.what()); + } + logger.Info("Model validation: scale fits on the GPU, {} engine(s), {:.2f} GB of the card's {:.2f} GB", + engines, pool->Bytes() / 1e9, total / 1e9); + return pool; +} #endif } // namespace @@ -669,9 +710,20 @@ ModelValidationResult Validate(const std::vector &merged, sg->xhm(), frames.size() - 1); // 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. + // On threads of their own, not on ParallelFor's pool, so that each one's gridding spreads over + // the pool (a parallel pass reached from a pool worker runs inline, ParallelFor.h). std::vector frame_r(frames.size(), 0.0); std::vector frame_sg(frames.size(), nullptr); - ParallelFor(static_cast(frames.size()), nthreads, [&](int i) { + // The data as merged and, where the alternative indexings are left to the probe below, in each + // of them: the lowest R over these is the frame's. + std::vector> relabelled; + if (same_point_group) + for (const gemmi::Op &law : reindex_ops) + relabelled.push_back(ReindexReflections(obs, law)); + std::vector *> labellings{&obs}; + for (const auto &v : relabelled) + labellings.push_back(&v); + auto score_frame = [&](size_t i) { gemmi::Structure trial = st; const gemmi::SpaceGroup *trial_sg = sg; if (i > 0 && !change_model_basis(trial, trial_sg, frames[i])) @@ -683,13 +735,14 @@ ModelValidationResult Validate(const std::vector &merged, // as every other and the ranking is noise. Measured on a pseudo-orthorhombic P2 // crystal: every frame read R 0.77-0.81 in the merged indexing, and the frame chosen // on that noise swapped two axes the model shares with the data. - double r = frame_probe_r(trial, trial_sg, obs, d_min); - 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_r[i] = frame_probe_r(trial, trial_sg, labellings, d_min, nthreads); frame_sg[i] = trial_sg; - }); + }; + std::vector> scored; + for (size_t i = 0; i < frames.size(); i++) + scored.push_back(std::async(std::launch::async, score_frame, i)); + for (std::future &f : scored) + f.get(); size_t best = 0; double best_r = 0; for (size_t i = 0; i < frames.size(); i++) { @@ -737,6 +790,10 @@ ModelValidationResult Validate(const std::vector &merged, sf_engine = StructureFactorsGPU(RigidBodyGPUEngine::CurrentDevice(), st.models[0], ucell, *sg, d_min, 0, logger); sf_gpu = sf_engine.get(); + // The scale fits beside them, on the same card, for as many reflections as the data have. + std::unique_ptr scale_gpu; + if (rigid_body_gpu && get_gpu_count() > 0) + scale_gpu = ScaleFitsGPU(RigidBodyGPUEngine::CurrentDevice(), obs.size(), logger); #endif // --- Fcalc (atomic) via electron density on a grid + FFT, plus a flat bulk-solvent mask -> Fmask. @@ -874,6 +931,15 @@ ModelValidationResult Validate(const std::vector &merged, gemmi::Scaling scaling(ucell, sg); scaling.use_solvent = true; scaling.prepare_points(out.fmodel, fobs_work, &ms.fmask); +#ifdef JFJOCH_USE_CUDA + if (scale_gpu) { + try { + scale_gpu->Fit(scaling, SolventBMax(ms.st.models[0]), nthreads); + } catch (const JFJochException &e) { + throw RigidBodyGPUFailure(e.what()); + } + } else +#endif FitModelScale(scaling, {.b_hi = SolventBMax(ms.st.models[0])}, nthreads); scaling.scale_data(out.fmodel, &ms.fmask); // out.fmodel now holds the scaled, solvent-corrected Fmodel out.k_sol = scaling.k_sol; @@ -1046,6 +1112,10 @@ ModelValidationResult Validate(const std::vector &merged, planned += (decision_pending ? 2 : 1) * sf_gpu->DeviceBytes(); card = sf_gpu->DeviceTotalMemory(); } + if (scale_gpu) { + planned += scale_gpu->Bytes(); + card = get_gpu_total_memory(RigidBodyGPUEngine::CurrentDevice()); + } if (planned > 0) may_run_beside = 2 * planned <= static_cast(get_gpu_count()) * card / 2; #endif diff --git a/tests/ModelScaleGPUTest.cpp b/tests/ModelScaleGPUTest.cpp index 7406565c2..86c167f83 100644 --- a/tests/ModelScaleGPUTest.cpp +++ b/tests/ModelScaleGPUTest.cpp @@ -15,6 +15,7 @@ #include #include +#include #include "gemmi/asumask.hpp" #include "gemmi/mmread_gz.hpp" @@ -347,4 +348,38 @@ TEST_CASE("ModelScaleGPU_Deterministic", "[ModelValidation][gpu]") { } } +// What model validation's fits take: FitModelScale's answer, to rounding, from a pool shared by more threads +// than it has engines - and the same answer, to the bit, whichever engine a fit got. +TEST_CASE("ModelScaleGPUPool_MatchesFitModelScale", "[ModelValidation][gpu]") { + if (get_gpu_count() == 0) + SKIP("No GPU"); + for (const Case &c : CASES) { + INFO(c.name); + gemmi::UnitCell cell(c.a, c.b, c.c, c.alpha, c.beta, c.gamma); + const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(c.hm); + const gemmi::Scaling points = MakeScaling(cell, sg, 2.8); + gemmi::Scaling cpu = points; + const ModelScaleReport report = FitModelScale(cpu, {}, 8); + + ModelScaleGPUPool pool(0, points.points.size(), 2); + std::vector> gpu(6, points); + std::vector gpu_report(gpu.size()); + std::vector> fits; + for (size_t i = 0; i < gpu.size(); i++) + fits.push_back(std::async(std::launch::async, [&, i] { gpu_report[i] = pool.Fit(gpu[i], 80.0, 2); })); + for (auto &f : fits) + f.get(); + + CHECK(gpu_report[0].n_grid == report.n_grid); + CHECK(std::fabs(gpu_report[0].r_work_fit - report.r_work_fit) < R_TOLERANCE); + CHECK(std::fabs(gpu[0].calculate_r_factor() - cpu.calculate_r_factor()) < R_TOLERANCE); + for (size_t i = 1; i < gpu.size(); i++) { + CHECK(gpu[i].k_sol == gpu[0].k_sol); + CHECK(gpu[i].b_sol == gpu[0].b_sol); + CHECK(gpu[i].k_overall == gpu[0].k_overall); + CHECK(std::memcmp(&gpu[i].b_star, &gpu[0].b_star, sizeof(gpu[0].b_star)) == 0); + } + } +} + #endif