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