diff --git a/rugnux/CMakeLists.txt b/rugnux/CMakeLists.txt index ad785e24a..02bc5c336 100644 --- a/rugnux/CMakeLists.txt +++ b/rugnux/CMakeLists.txt @@ -18,6 +18,10 @@ ADD_LIBRARY(JFJochRugnux STATIC ModelValidation.h RigidBodyRefine.cpp RigidBodyRefine.h + $<$:RigidBodyGPU.cpp> + $<$:RigidBodyGPU.cu> + RigidBodyGPU.h + RigidBodyGPUEngine.h SigmaA.cpp SigmaA.h ResultReport.cpp diff --git a/rugnux/ModelGrid.cpp b/rugnux/ModelGrid.cpp index 125c987f2..65eb64c39 100644 --- a/rugnux/ModelGrid.cpp +++ b/rugnux/ModelGrid.cpp @@ -222,7 +222,10 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator &dc, const gem // SolventMasker(AtomicRadiiSet::Refmac).put_mask_on_grid() (gemmi solmask.hpp) with the atoms and the // symmetry spread over the threads as above - the same mask, bit for bit. The island removal and the -// shrink that follow are gemmi's own. +// shrink that follow are gemmi's own, except that the shrink is skipped where it cannot change a point: +// it looks at the grid offsets within rshrink of each point (set_margin_around), and on a grid whose +// spacing is coarser than rshrink along all three axes there are none - which is every rigid-body zone - +// yet it still walks the whole grid. void PutMaskOnGrid(gemmi::Grid &grid, const gemmi::Model &model, const std::vector &orbit_leaders, size_t nthreads) { const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac); @@ -241,5 +244,9 @@ void PutMaskOnGrid(gemmi::Grid &grid, const gemmi::Model &model, const st }); SymmetrizeOrbits(grid, orbit_leaders, nthreads, [](float a, float b) { return a < b ? a : b; }); masker.remove_islands(grid); - masker.shrink(grid); + bool shrink_has_offsets = false; + for (int i = 0; i < 3; i++) + shrink_has_offsets = shrink_has_offsets || static_cast(std::floor(masker.rshrink / grid.spacing[i])) > 0; + if (shrink_has_offsets) + masker.shrink(grid); } diff --git a/rugnux/ModelValidation.cpp b/rugnux/ModelValidation.cpp index d16a53328..ca67c89dd 100644 --- a/rugnux/ModelValidation.cpp +++ b/rugnux/ModelValidation.cpp @@ -37,6 +37,11 @@ #include "ModelFFT.h" #include "RigidBodyRefine.h" #include "SigmaA.h" +#ifdef JFJOCH_USE_CUDA +#include +#include "RigidBodyGPU.h" +#include "../common/CUDAWrapper.h" +#endif namespace { @@ -330,16 +335,21 @@ double frame_probe_r(const gemmi::Structure &st, const gemmi::SpaceGroup *sg, } // namespace -ModelValidationResult ValidateAgainstModel(const std::vector &merged, - const UnitCell &cell, - const std::string &model_path, - const std::string &output_prefix, - Logger &logger, - const gemmi::SpaceGroup *data_space_group, - bool probe_indexing_ambiguity, - size_t nthreads, - double wavelength_A, - const std::vector &report_shell_d_min) { +namespace { + +// ValidateAgainstModel() with the rigid body on the GPU where `rigid_body_gpu` allows it and a card is +// there to take it, on the CPU otherwise - for the whole validation either way. +ModelValidationResult Validate(const std::vector &merged, + const UnitCell &cell, + const std::string &model_path, + const std::string &output_prefix, + Logger &logger, + const gemmi::SpaceGroup *data_space_group, + bool probe_indexing_ambiguity, + size_t nthreads, + double wavelength_A, + const std::vector &report_shell_d_min, + bool rigid_body_gpu) { ModelValidationResult result; result.model_path = model_path; @@ -539,6 +549,26 @@ ModelValidationResult ValidateAgainstModel(const std::vector & }; compute_model_factors(mdl, d_min); + // The rigid body's device, decided once for the whole validation - the real fit and every replicate + // of the null on the same one, so that a verdict never depends on how full the card was when a + // replicate started. The engines are reserved here, up front, and never grow. + RigidBodyGPUPool *rigid_body_pool = nullptr; +#ifdef JFJOCH_USE_CUDA + std::unique_ptr rigid_body_engines; + if (rigid_body_gpu && std::getenv("JFJOCH_RIGID_BODY_CPU") == nullptr) { + const double finest_zone = RigidBodyLadder(d_min).back(); + size_t zone_observations = 0; + for (const MergedReflection &r : obs) + if (!std::isnan(r.F) && r.d >= finest_zone) + ++zone_observations; + rigid_body_engines = RigidBodyGPUPool::Create(st.models[0], ucell, *sg, d_min, zone_observations, + std::min(std::max(nthreads, 1), 4), logger); + rigid_body_pool = rigid_body_engines.get(); + } +#else + (void) rigid_body_gpu; +#endif + gemmi::GroupOps gops = sg->operations(); gemmi::ReciprocalAsu asu(sg); @@ -658,7 +688,7 @@ ModelValidationResult ValidateAgainstModel(const std::vector & // exactly where it was read. const std::vector before = ModelPositions(ms.st.models[0]); const RigidBodyRefineResult rb = - RefineRigidBody(ms.st.models[0], ucell, *sg, out.fit.fobs_work, to_d, logger, nthreads); + RefineRigidBody(ms.st.models[0], ucell, *sg, out.fit.fobs_work, to_d, logger, nthreads, rigid_body_pool); if (!rb.converged) { // RefineRigidBody leaves the model wherever the solver left it, usable answer or not, so // the restore cannot be conditional on the same flag the re-fit is. @@ -982,6 +1012,10 @@ ModelValidationResult ValidateAgainstModel(const std::vector & "reflections and no null was run; R-work {:.4f}, R-free {:.4f}", real.fit.r_work, real.fit.r_free); } +#ifdef JFJOCH_USE_CUDA + } catch (const RigidBodyGPUFailure &) { + throw; // not the null's failure but the device's: the whole validation runs again on the CPU +#endif } catch (const std::exception &e) { // The null could not be built, so no decision is licensed and none is taken - which is the // NOT_TESTED state, not a failure of the run. If the indexing probe had already won on a @@ -1399,6 +1433,35 @@ ModelValidationResult ValidateAgainstModel(const std::vector & return result; } +} // namespace + +ModelValidationResult ValidateAgainstModel(const std::vector &merged, + const UnitCell &cell, + const std::string &model_path, + const std::string &output_prefix, + Logger &logger, + const gemmi::SpaceGroup *data_space_group, + bool probe_indexing_ambiguity, + size_t nthreads, + double wavelength_A, + const std::vector &report_shell_d_min) { +#ifdef JFJOCH_USE_CUDA + // A CUDA failure in the rigid body does not end the run: the validation is started again from the + // model as read, on the CPU throughout. It is re-runnable, and a dead model check should not take a + // finished merge with it. + try { + return Validate(merged, cell, model_path, output_prefix, logger, data_space_group, + probe_indexing_ambiguity, nthreads, wavelength_A, report_shell_d_min, true); + } catch (const RigidBodyGPUFailure &e) { + cuda_clear_error(); + logger.Warning("Model validation: the rigid body failed on the GPU ({}); validating again on the CPU", + e.what()); + } +#endif + return Validate(merged, cell, model_path, output_prefix, logger, data_space_group, probe_indexing_ambiguity, + nthreads, wavelength_A, report_shell_d_min, false); +} + namespace { // The reindexing operator as it reads on Miller indices ("k,h,-l" rather than "y,x,-z"). diff --git a/rugnux/RigidBodyGPU.cpp b/rugnux/RigidBodyGPU.cpp new file mode 100644 index 000000000..5206ea917 --- /dev/null +++ b/rugnux/RigidBodyGPU.cpp @@ -0,0 +1,375 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "RigidBodyGPU.h" + +#include +#include + +#include +#include + +#include "gemmi/dencalc.hpp" // DensityCalculator +#include "gemmi/it92.hpp" // IT92 x-ray form factors +#include "gemmi/scaling.hpp" // Scaling +#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent + +#include "ModelGrid.h" // PutMaskOnGrid +#include "ModelScaling.h" // FitModelScale +#include "RigidBodyGPUEngine.h" +#include "../common/CUDAWrapper.h" +#include "../common/JFJochException.h" +#include "../common/Logger.h" + +namespace { + +using Table = gemmi::IT92; + +// Everything about a zone that does not depend on the placement. The observations are optional: without +// them the zone only says how big an engine it needs. +RigidBodyGPUZone MakeZone(const gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + double d_min, const SymmetryComposition *composition, + const gemmi::AsuData> *fobs) { + RigidBodyGPUZone zone; + const gemmi::Grid grid = RigidBodyZoneGrid(cell, sg, d_min); + zone.nu = grid.nu; + zone.nv = grid.nv; + zone.nw = grid.nw; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) { + zone.orth[3 * i + j] = cell.orth.mat[i][j]; + zone.frac[3 * i + j] = cell.frac.mat[i][j]; + } + zone.volume = cell.volume; + + // The density of each atom as PutModelDensityOnGrid() (ModelGrid.cpp) precalculates it. + gemmi::DensityCalculator dc; + dc.d_min = d_min; + dc.rate = 1.5; + dc.grid.unit_cell = cell; + dc.grid.spacegroup = &sg; + dc.set_refmac_compatible_blur(model); + const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac); + int index = 0; + for (const gemmi::Chain &ch : model.chains) + for (const gemmi::Residue &r : ch.residues) + for (const gemmi::Atom &atom : r.atoms) { + using CReal = Table::Coef::coef_type; + const auto &coef = Table::get(atom.element, atom.charge, atom.serial); + const float addend = dc.addends.get(atom.element); + RigidBodyGPUAtom a{}; + a.occ = atom.occ; + a.aniso = atom.aniso.nonzero(); + if (!a.aniso) { + const CReal b = static_cast(atom.b_iso + dc.blur); + const auto precal = coef.precalculate_density_iso(b, addend); + a.radius = dc.estimate_radius(precal, b); + for (int k = 0; k < 5; k++) { + a.a[k] = precal.a[k]; + a.b[k][0] = precal.b[k]; + } + } else { + const auto aniso_b = atom.aniso.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(dc.blur)); + const CReal b_max = std::max(std::max(aniso_b.u11, aniso_b.u22), aniso_b.u33); + a.radius = static_cast(dc.estimate_radius(coef.precalculate_density_iso(b_max, addend), b_max)); + const auto precal = coef.precalculate_density_aniso_b(aniso_b, addend); + for (int k = 0; k < 5; k++) { + a.a[k] = precal.a[k]; + const gemmi::SMat33 &m = precal.b[k]; + const float e[6] = {m.u11, m.u22, m.u33, m.u12, m.u13, m.u23}; + std::copy(e, e + 6, a.b[k]); + } + } + zone.atoms.push_back(a); + // The bulk-solvent mask's atoms, as PutMaskOnGrid() takes them. + if (!((masker.ignore_hydrogen && atom.is_hydrogen()) || + (masker.ignore_zero_occupancy_atoms && atom.occ <= 0))) { + zone.mask_atom.push_back(index); + zone.mask_radius.push_back(static_cast( + masker.constant_r + masker.rprobe + gemmi::refmac_radius_for_bulk_solvent(atom.element.elem))); + } + ++index; + } + + const gemmi::GroupOps gops = sg.operations(); + for (const gemmi::Op::Tran &cen : gops.cen_ops) + for (const gemmi::Op &op : gops.sym_ops) { + std::array image{}; + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) + image[3 * i + j] = static_cast(op.rot[i][j]) / gemmi::Op::DEN; + image[9 + i] = static_cast(op.tran[i] + cen[i]) / gemmi::Op::DEN; + } + zone.images.push_back(image); + } + for (const gemmi::Vec6 &c : gemmi::adp_symmetry_constraints(&sg)) + zone.constraints.push_back(c); + + if (composition == nullptr || fobs == nullptr) + return zone; + + zone.ops = composition->Ops(); + for (size_t m = 0; m < composition->Hkl().size(); m++) { + const gemmi::Miller &h = composition->Hkl()[m]; + zone.row_hkl.push_back({h[0], h[1], h[2]}); + // prepare_asu_data()'s unblur, exp(B_blur |s|^2 / 4), as Compose() applies it + zone.row_scale.push_back(composition->Centring() * std::exp(dc.blur * 0.25 * composition->InvD2()[m])); + zone.row_stol2.push_back(cell.calculate_stol_sq(h)); + } + for (const SymmetryComposition::Term &t : composition->Terms()) { + RigidBodyGPUTerm g{}; + for (int k = 0; k < 3; k++) + g.k[k] = t.k[k]; + g.phase[0] = t.phase.real(); + g.phase[1] = t.phase.imag(); + g.s[0] = t.s.x; + g.s[1] = t.s.y; + g.s[2] = t.s.z; + zone.terms.push_back(g); + } + double sum = 0; + for (size_t i = 0; i < fobs->v.size(); i++) { + const int m = composition->Row()[i]; + const gemmi::ValueSigma &vs = fobs->v[i].value; + zone.obs_row.push_back(m); + zone.obs_fobs.push_back(vs.value); + sum += vs.value; + // the points gemmi's prepare_points() takes + if (m >= 0 && !std::isnan(vs.value) && !std::isnan(vs.sigma)) + zone.point_obs.push_back(static_cast(i)); + } + zone.f_mean = fobs->v.empty() ? 1.0 : sum / static_cast(fobs->v.size()); + return zone; +} + +} // namespace + +std::unique_ptr RigidBodyGPUPool::Create(const gemmi::Model &model, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, double d_min, + size_t max_observations, size_t max_engines, + Logger &logger) { + if (get_gpu_count() == 0) + return nullptr; + try { + // Sized for the largest zone of the ladder, which every fit in the validation walks a part of. + RigidBodyGPUCapacity cap; + const size_t ops = sg.operations().sym_ops.size(); + for (double zone_d : RigidBodyLadder(d_min)) { + const RigidBodyGPUZone zone = MakeZone(model, cell, sg, zone_d, nullptr, nullptr); + if (!RigidBodyGPUEngine::Supports(zone)) { + logger.Info("Model validation: the rigid body runs on the CPU - the cell is too small for the " + "GPU's gridding at {:.1f} A", zone_d); + return nullptr; + } + cap.atoms = std::max(cap.atoms, zone.atoms.size()); + cap.grid_points = std::max(cap.grid_points, static_cast(zone.nu) * zone.nv * zone.nw); + cap.complex_points = std::max(cap.complex_points, static_cast(zone.nu / 2 + 1) * zone.nv * zone.nw); + cap.bricks = std::max(cap.bricks, RigidBodyGPUEngine::Bricks(zone.nu, zone.nv, zone.nw)); + cap.pairs = std::max(cap.pairs, RigidBodyGPUEngine::PairBound(zone)); + cap.fft_work_bytes = std::max(cap.fft_work_bytes, RigidBodyGPUEngine::FFTWorkBytes(zone.nu, zone.nv, zone.nw)); + } + cap.observations = max_observations; + cap.rows = max_observations; + cap.terms = max_observations * ops; + const size_t bytes = RigidBodyGPUEngine::DeviceBytes(cap); + + size_t free = 0, total = 0; + RigidBodyGPUEngine::MemoryInfo(free, total); + constexpr size_t HEADROOM = 1ull << 30; + const size_t budget = std::min(total / 4, free > HEADROOM ? free - HEADROOM : 0); + const size_t fit = bytes > 0 ? budget / bytes : 0; + const size_t want = std::min({fit, std::max(max_engines, 1), 4}); + std::unique_ptr pool(new RigidBodyGPUPool); + const int device = RigidBodyGPUEngine::CurrentDevice(); + for (size_t i = 0; i < want; i++) + pool->engines_.push_back(std::make_unique(cap, device)); + if (pool->engines_.empty()) { + logger.Info("Model validation: the rigid body runs on the CPU - one GPU engine needs {:.0f} MB and the " + "budget is {:.0f} MB", bytes / 1e6, budget / 1e6); + return nullptr; + } + for (auto &e : pool->engines_) + pool->idle_.push_back(e.get()); + logger.Info("Model validation: the rigid body runs on the GPU, {} engine(s) of {:.0f} MB", pool->Engines(), + bytes / 1e6); + return pool; + } catch (const JFJochException &e) { + cuda_clear_error(); + logger.Warning("Model validation: the rigid body runs on the CPU - the GPU engines could not be set up ({})", + e.what()); + return nullptr; + } +} + +RigidBodyGPUPool::~RigidBodyGPUPool() = default; + +RigidBodyGPUEngine &RigidBodyGPUPool::Acquire() { + std::unique_lock lock(m_); + cv_.wait(lock, [this] { return !idle_.empty(); }); + RigidBodyGPUEngine *e = idle_.back(); + idle_.pop_back(); + return *e; +} + +void RigidBodyGPUPool::Release(RigidBodyGPUEngine &engine) { + { + std::lock_guard lock(m_); + idle_.push_back(&engine); + } + cv_.notify_one(); +} + +RigidBodyTargetGPU::RigidBodyTargetGPU(RigidBodyGPUPool &pool, gemmi::Model &model, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, size_t nthreads) + : RigidBodyTargetBase(model), pool_(pool), engine_(pool.Acquire()), model_(model), cell_(cell), sg_(sg), + nthreads_(nthreads) { + std::vector> relative; + for (const gemmi::Position &p : base_) + relative.push_back({p.x - centre_.x, p.y - centre_.y, p.z - centre_.z}); + try { + engine_.SetBody(relative); + } catch (const JFJochException &e) { + pool_.Release(engine_); + throw RigidBodyGPUFailure(e.what()); + } +} + +RigidBodyTargetGPU::~RigidBodyTargetGPU() { + pool_.Release(engine_); +} + +// Place()'s placement as a matrix: the columns are the rotated axes, rotated as Place() rotates a point. +void RigidBodyTargetGPU::Placement(const double q[6], double rotation[9], double translation[3]) const { + const double aa[3] = {q[0] / rms_radius_, q[1] / rms_radius_, q[2] / rms_radius_}; + for (int j = 0; j < 3; j++) { + const double e[3] = {j == 0 ? 1.0 : 0.0, j == 1 ? 1.0 : 0.0, j == 2 ? 1.0 : 0.0}; + double column[3]; + ceres::AngleAxisRotatePoint(aa, e, column); + for (int i = 0; i < 3; i++) + rotation[3 * i + j] = column[i]; + } + translation[0] = centre_.x + q[3]; + translation[1] = centre_.y + q[4]; + translation[2] = centre_.z + q[5]; +} + +void RigidBodyTargetGPU::SetZone(const gemmi::AsuData> &fobs, double d_min) { + fobs_ = fobs; + d_min_ = d_min; + solvent_fitted_ = false; + have_point_ = false; + const gemmi::Grid grid = RigidBodyZoneGrid(cell_, sg_, d_min); + orbit_leaders_ = OrbitLeaders(grid, nthreads_); + std::vector hkl; + for (const auto &hv : fobs_.v) + hkl.push_back(hv.hkl); + const SymmetryComposition composition(grid, d_min, hkl); + zone_ = std::make_unique(MakeZone(model_, cell_, sg_, d_min, &composition, &fobs_)); + try { + engine_.SetZone(*zone_); + } catch (const JFJochException &e) { + throw RigidBodyGPUFailure(e.what()); + } +} + +// RigidBodyTarget::Residuals() with the density, the composition, the transforms and the residuals on +// the device. +bool RigidBodyTargetGPU::Residuals(const double q[6], double *residuals) { + try { + ++evaluations; + double rotation[9], translation[3]; + Placement(q, rotation, translation); + engine_.Fcalc(rotation, translation); + if (!(hold_mask && have_point_)) { + Place(q, model_); + gemmi::Grid mask_grid = RigidBodyZoneGrid(cell_, sg_, d_min_); + PutMaskOnGrid(mask_grid, model_, orbit_leaders_, nthreads_); + engine_.Fmask(mask_grid.data.data()); + } + if (zone_->row_hkl.empty()) + return false; + + std::vector> fc, fm; + engine_.DownloadPoints(fc, fm); + gemmi::Scaling scaling(cell_, &sg_); + scaling.use_solvent = true; + for (size_t p = 0; p < zone_->point_obs.size(); p++) { + const auto &o = fobs_.v[zone_->point_obs[p]]; + const int m = zone_->obs_row[zone_->point_obs[p]]; + scaling.points.push_back({o.hkl, zone_->row_stol2[m], std::complex(fc[p][0], fc[p][1]), + std::complex(fm[p][0], fm[p][1]), o.value.value, o.value.sigma}); + } + if (scaling.points.empty()) + return false; + // The scale as RigidBodyTarget::Residuals() fits it: the solvent once per zone, then the + // overall scale and the anisotropic B at every evaluation. + if (!solvent_fitted_) { + FitModelScale(scaling, {}, nthreads_); + k_sol = scaling.k_sol; + b_sol = scaling.b_sol; + solvent_fitted_ = true; + } + scaling.k_sol = k_sol; + scaling.b_sol = b_sol; + scaling.fix_k_sol = true; + scaling.fix_b_sol = true; + scaling.fit_isotropic_b_approximately(); + scaling.fit_parameters(); + + const gemmi::SMat33 &b = scaling.b_star; + const double b_star[6] = {b.u11, b.u22, b.u33, b.u12, b.u13, b.u23}; + engine_.Residuals(scaling.k_overall, b_star, k_sol, b_sol, residuals); + unmatched = static_cast(std::count(zone_->obs_row.begin(), zone_->obs_row.end(), -1)); + + have_point_ = true; + std::copy(q, q + 6, q_.begin()); + k_overall_ = scaling.k_overall; + b_star_ = scaling.b_star; + return true; + } catch (const JFJochException &e) { + throw RigidBodyGPUFailure(e.what()); + } +} + +// RigidBodyTarget::Jacobian(): the rotation columns by forward difference, three copies of the body +// gridded and transformed together, the translation columns exact, the scale re-fit projected out. +bool RigidBodyTargetGPU::Jacobian(const double q[6], double *jacobian) { + if (!have_point_ || !std::equal(q, q + 6, q_.begin())) { + std::vector residuals(NumObservations()); + if (!Residuals(q, residuals.data())) + return false; + } + try { + ++jacobians; + const double step = RigidBodyJacobianStep(d_min_); + double rotation[3][9], translation[3][3]; + for (int j = 0; j < 3; j++) { + double qj[6]; + std::copy(q_.begin(), q_.end(), qj); + qj[j] += step; + Placement(qj, rotation[j], translation[j]); + } + const double b_star[6] = {b_star_.u11, b_star_.u22, b_star_.u33, b_star_.u12, b_star_.u13, b_star_.u23}; + std::vector jtj, jtq; + engine_.Jacobian(rotation, translation, step, k_overall_, b_star, k_sol, b_sol, jtj, jtq); + + // Following Golub & Pereyra (1973) SIAM J. Numer. Anal. 10, 413-432, in the form of Kaufman (1975) BIT 15, 49-57 + const int p = static_cast(zone_->constraints.size()) + 1; + Eigen::MatrixXd a(p, p), c(p, 6); + for (int i = 0; i < p; i++) { + for (int j = 0; j < p; j++) + a(i, j) = jtj[i * p + j]; + for (int j = 0; j < 6; j++) + c(i, j) = jtq[i * 6 + j]; + } + const Eigen::MatrixXd x = a.ldlt().solve(c); + std::vector xv(p * 6); + for (int i = 0; i < p; i++) + for (int j = 0; j < 6; j++) + xv[i * 6 + j] = x(i, j); + engine_.ProjectJacobian(xv, jacobian); + return true; + } catch (const JFJochException &e) { + throw RigidBodyGPUFailure(e.what()); + } +} diff --git a/rugnux/RigidBodyGPU.cu b/rugnux/RigidBodyGPU.cu new file mode 100644 index 000000000..c90c372b7 --- /dev/null +++ b/rugnux/RigidBodyGPU.cu @@ -0,0 +1,893 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "RigidBodyGPUEngine.h" + +#include +#include +#include +#include + +#include +#include + +#include "../common/JFJochException.h" +#include "../image_analysis/indexing/CUDAMemHelpers.h" + +// Every kernel below is deterministic: nothing is summed with a floating-point atomic, each grid point +// adds its atoms in model order, and every reduction runs in a fixed order. The same placement gives the +// same residuals, bit for bit, on every run and whichever engine runs it. + +namespace { + +void cuda_err(cudaError_t val) { + if (val != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); +} + +void cufft_err(cufftResult val) { + if (val != CUFFT_SUCCESS) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, + "cuFFT error " + std::to_string(static_cast(val))); +} + +constexpr int BRICK = 8; // the gather's bricks are BRICK^3 grid points, one block each +constexpr int GATHER_TILE = 64; // atoms staged in shared memory at a time +constexpr int MAX_BRICKS_PER_AXIS = 8; // an atom's box touches at most this many bricks along an axis +constexpr int THREADS = 256; +constexpr int MAX_SCALE_PARAMS = 7; // k_overall + up to six B* constraints + +// The grid of a zone as the kernels see it. +struct GridGeom { + int nu, nv, nw; + float orth_n[9]; // orth * diag(1/nu, 1/nv, 1/nw), row-major: grid offset -> Cartesian +}; + +// Where the body is: x -> R x_rel + t, and the fractionalization. +struct Placement { + double r[9]; + double t[3]; + double frac[9]; +}; + +struct DeviceTerm { + int index; // of F1(k), or of F1(-k), in the half-u transform + int conj; // read as the Friedel mate + double phase[2]; + double s[3]; +}; + +// The scale a residual or a Jacobian row is taken at. +struct Scale { + double k_overall; + double b_star[6]; // u11 u22 u33 u12 u13 u23 + double k_sol, b_sol; + int n_constraints; + double constraints[6][6]; +}; + +__device__ __forceinline__ int imod(int a, int n) { + const int r = a % n; + return r < 0 ? r + n : r; +} + +// gemmi's unsafe_expapprox() (formfact.hpp), the exponential the density is computed with. +__device__ __forceinline__ float expapprox(float x) { + const float val = 12102203.1615614f * x + 1065353216.f; + const int vali = static_cast(val); + const float a = __int_as_float(vali & 0x7F800000); + const float b = __int_as_float((vali & 0x7FFFFF) | 0x3F800000); + return a * (0.509871020f + b * (0.312146713f + b * (0.166617139f + b * (-2.190619930e-3f + b * 1.3555747234e-2f)))); +} + +// b_star.r_u_r(h), as gemmi's SMat33. +__device__ __forceinline__ double r_u_r(const double *u, const int *h) { + const double x = h[0], y = h[1], z = h[2]; + return x * x * u[0] + y * y * u[1] + z * z * u[2] + 2 * (x * y * u[3] + x * z * u[4] + y * z * u[5]); +} + +// Each atom placed: its grid coordinates (fractional times the grid size, wrapped into the cell) for +// the gather and, for a mask atom, its wrapped fractional coordinates and mask radius for the mask. +__global__ void place_kernel(const double *rel, int n, Placement p, int nu, int nv, int nw, const int *mask_slot, + const float *mask_radius, float4 *pos, double *mask_atoms) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) + return; + const double x = rel[3 * i], y = rel[3 * i + 1], z = rel[3 * i + 2]; + const double c[3] = {p.r[0] * x + p.r[1] * y + p.r[2] * z + p.t[0], + p.r[3] * x + p.r[4] * y + p.r[5] * z + p.t[1], + p.r[6] * x + p.r[7] * y + p.r[8] * z + p.t[2]}; + double f[3]; + for (int k = 0; k < 3; k++) { + f[k] = p.frac[3 * k] * c[0] + p.frac[3 * k + 1] * c[1] + p.frac[3 * k + 2] * c[2]; + f[k] -= floor(f[k]); + if (f[k] >= 1.0) // a tiny negative coordinate wraps to 1 - epsilon, which rounds to 1 + f[k] = 0.0; + } + const int n3[3] = {nu, nv, nw}; + float g[3]; + for (int k = 0; k < 3; k++) { + g[k] = static_cast(f[k] * n3[k]); + if (g[k] >= n3[k]) + g[k] -= n3[k]; + } + pos[i] = make_float4(g[0], g[1], g[2], 0.0f); + const int slot = mask_slot[i]; + if (slot >= 0) { + for (int k = 0; k < 3; k++) + mask_atoms[4 * slot + k] = f[k]; + mask_atoms[4 * slot + 3] = mask_radius[i]; + } +} + +// The distinct bricks the points c - d ... c + d of one axis fall in, wrapped into the cell. A grid size +// that is not a multiple of BRICK leaves the last brick partial, which is why this is done on wrapped +// points and not in brick coordinates. +__device__ int axis_bricks(int c, int d, int n, int *out) { + int k = 0; + for (int p = c - d; p <= c + d; p++) { + const int b = imod(p, n) / BRICK; + bool seen = false; + for (int j = 0; j < k; j++) + seen = seen || out[j] == b; + if (!seen && k < MAX_BRICKS_PER_AXIS) + out[k++] = b; + } + return k; +} + +// The bricks atom i's box touches, per axis. The box is gemmi's: `box` points either side of the +// point nearest the atom. +__device__ void atom_bricks(const float4 &pos, const int *box, int nu, int nv, int nw, int *bu, int &nbu, int *bv, + int &nbv, int *bw, int &nbw) { + nbu = axis_bricks(__float2int_rn(pos.x), box[0], nu, bu); + nbv = axis_bricks(__float2int_rn(pos.y), box[1], nv, bv); + nbw = axis_bricks(__float2int_rn(pos.z), box[2], nw, bw); +} + +__global__ void count_pairs_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw, int *count) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i > n) + return; + if (i == n) { + count[n] = 0; // so the exclusive scan's last entry is the total + return; + } + int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c; + atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c); + count[i] = a * b * c; +} + +// The (brick, atom) pairs, atom-major: sorted stably by brick, each brick's atoms stay in model order. +__global__ void fill_pairs_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw, + const int *offset, int *key, int *value) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) + return; + int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c; + atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c); + const int nbu = (nu + BRICK - 1) / BRICK, nbv = (nv + BRICK - 1) / BRICK; + int o = offset[i]; + for (int k = 0; k < c; k++) + for (int j = 0; j < b; j++) + for (int l = 0; l < a; l++) { + key[o] = bu[l] + nbu * (bv[j] + nbv * bw[k]); + value[o] = i; + o++; + } +} + +__global__ void brick_ranges_kernel(const int *key, int npairs, int *start, int *end) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= npairs) + return; + const int k = key[i]; + if (i == 0 || key[i - 1] != k) + start[k] = i; + if (i == npairs - 1 || key[i + 1] != k) + end[k] = i + 1; +} + +// The density of the model on the grid, as gemmi's do_add_atom_density_to_grid() puts it there: every +// point adds, in model order, each atom whose sphere it is inside. One block per brick; the atoms of the +// brick are staged in shared memory. +__global__ void gather_kernel(const RigidBodyGPUAtom *atoms, const float4 *pos, const int *pair_atom, + const int *brick_start, const int *brick_end, GridGeom g, float *grid) { + const int nbu = (g.nu + BRICK - 1) / BRICK, nbv = (g.nv + BRICK - 1) / BRICK; + const int brick = blockIdx.x; + const int s = brick_start[brick], e = brick_end[brick]; // s == e: an empty brick, written as zeros + const int u = (brick % nbu) * BRICK + threadIdx.x; + const int v = (brick / nbu % nbv) * BRICK + threadIdx.y; + const int w = (brick / (nbu * nbv)) * BRICK + threadIdx.z; + const bool inside = u < g.nu && v < g.nv && w < g.nw; + const int tid = threadIdx.x + BRICK * (threadIdx.y + BRICK * threadIdx.z); + + __shared__ RigidBodyGPUAtom sh_atom[GATHER_TILE]; + __shared__ float4 sh_pos[GATHER_TILE]; + float acc = 0.0f; + for (int c = s; c < e; c += GATHER_TILE) { + const int m = min(GATHER_TILE, e - c); + __syncthreads(); + if (tid < m) { + const int a = pair_atom[c + tid]; + sh_atom[tid] = atoms[a]; + sh_pos[tid] = pos[a]; + } + __syncthreads(); + if (!inside) + continue; + for (int k = 0; k < m; k++) { + const RigidBodyGPUAtom &at = sh_atom[k]; + // The atom's offset from the point, in grid units, to the nearest periodic image. + float fx = sh_pos[k].x - u, fy = sh_pos[k].y - v, fz = sh_pos[k].z - w; + fx -= g.nu * rintf(fx / g.nu); + fy -= g.nv * rintf(fy / g.nv); + fz -= g.nw * rintf(fz / g.nw); + const float x = g.orth_n[0] * fx + g.orth_n[1] * fy + g.orth_n[2] * fz; + const float y = g.orth_n[3] * fx + g.orth_n[4] * fy + g.orth_n[5] * fz; + const float z = g.orth_n[6] * fx + g.orth_n[7] * fy + g.orth_n[8] * fz; + const float r2 = x * x + y * y + z * z; + if (r2 > at.radius * at.radius) + continue; + float density = 0.0f; + if (!at.aniso) { + for (int q = 0; q < 5; q++) + density += at.a[q] * expapprox(fmaxf(at.b[q][0] * r2, -88.f)); + } else { + for (int q = 0; q < 5; q++) { + const float *b = at.b[q]; + const float rur = x * x * b[0] + y * y * b[1] + z * z * b[2] + 2 * (x * y * b[3] + x * z * b[4] + y * z * b[5]); + density += at.a[q] * expapprox(fmaxf(rur, -88.f)); + } + } + acc += at.occ * density; + } + } + if (inside) + grid[u + static_cast(g.nu) * (v + static_cast(g.nv) * w)] = acc; +} + +// F at the zone's rows composed from the transform of one copy (SymmetryComposition::Compose), with +// dF/dt where `df_dt` is not null. `f1` is cuFFT's forward transform, which carries exp(-2 pi i h.x): +// conjugated and scaled by V/N it is gemmi's F1. +__global__ void compose_kernel(const float2 *f1, float norm, const DeviceTerm *terms, const double *row_scale, + int rows, int ops, double2 *f, double2 *df_dt) { + const int m = blockIdx.x * blockDim.x + threadIdx.x; + if (m >= rows) + return; + double sr = 0, si = 0, d[3][2] = {}; + for (int o = 0; o < ops; o++) { + const DeviceTerm t = terms[static_cast(m) * ops + o]; + const float2 out = f1[t.index]; + double vr = out.x * norm, vi = -out.y * norm; // F1 at the index read + if (t.conj) + vi = -vi; + const double tr = t.phase[0] * vr - t.phase[1] * vi, ti = t.phase[0] * vi + t.phase[1] * vr; + sr += tr; + si += ti; + for (int k = 0; k < 3; k++) { + d[k][0] += tr * t.s[k]; + d[k][1] += ti * t.s[k]; + } + } + const double scale = row_scale[m]; + f[m] = make_double2(scale * sr, scale * si); + if (df_dt != nullptr) { + // times 2 pi i + const double two_pi = 2 * 3.14159265358979323846; + for (int k = 0; k < 3; k++) + df_dt[3 * static_cast(m) + k] = make_double2(-scale * two_pi * d[k][1], scale * two_pi * d[k][0]); + } +} + +// Fmask at each row, read at h directly (the mask is symmetric): gemmi's get_value_by_hkl() on the +// transform of the mask, taking the Friedel mate where h is outside the stored half. +__global__ void fmask_kernel(const float2 *fm, float norm, const int *row_hkl, int rows, int nu, int nv, int nw, + float2 *fmask) { + const int m = blockIdx.x * blockDim.x + threadIdx.x; + if (m >= rows) + return; + int h = row_hkl[3 * m], k = row_hkl[3 * m + 1], l = row_hkl[3 * m + 2]; + const bool conj = h < 0; + if (conj) { + h = -h; + k = -k; + l = -l; + } + const float2 out = fm[h + static_cast(nu / 2 + 1) * (imod(k, nv) + static_cast(nv) * imod(l, nw))]; + const float re = out.x * norm, im = -out.y * norm; + fmask[m] = make_float2(re, conj ? -im : im); +} + +// Scaling::Point's fcmol and fmask of each point. +__global__ void pack_points_kernel(const int *point_obs, const int *obs_row, int n, const double2 *f, + const float2 *fmask, float2 *fcmol_p, float2 *fmask_p) { + const int p = blockIdx.x * blockDim.x + threadIdx.x; + if (p >= n) + return; + const int m = obs_row[point_obs[p]]; + fcmol_p[p] = make_float2(static_cast(f[m].x), static_cast(f[m].y)); + fmask_p[p] = fmask[m]; +} + +// gemmi's Scaling::get_fcalc() of a row: fcmol + (float) solvent scale * fmask, in float. +__device__ float2 solvent_fcalc(const double2 &f, const float2 &fm, double stol2, const Scale &sc) { + const float solvent = static_cast(sc.k_sol * exp(-sc.b_sol * stol2)); + return make_float2(static_cast(f.x) + solvent * fm.x, static_cast(f.y) + solvent * fm.y); +} + +// The residual of each observation: scale_data() applied to the row, then (Fobs - |F|) / . +__global__ void residual_kernel(const int *obs_row, const float *obs_fobs, int n, const double2 *f, + const float2 *fmask, const int *row_hkl, const double *row_stol2, Scale sc, + double f_mean, double *residual) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) + return; + const int m = obs_row[i]; + if (m < 0) { + residual[i] = 0.0; + return; + } + float2 value = solvent_fcalc(f[m], fmask[m], row_stol2[m], sc); + const float k = static_cast(sc.k_overall * exp(-0.25 * r_u_r(sc.b_star, row_hkl + 3 * m))); + value.x *= k; + value.y *= k; + residual[i] = (obs_fobs[i] - hypotf(value.x, value.y)) / f_mean; +} + +// One row of the fixed-scale Jacobian and of the scale parameters' own columns (RigidBodyTarget::Jacobian). +__global__ void jacobian_kernel(const int *obs_row, int n, const double2 *f, const double2 *df_dt, + const double2 *shifted, size_t rows, const float2 *fmask, const int *row_hkl, + const double *row_stol2, Scale sc, double f_mean, double step, double *jacobian, + double *jk) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) + return; + const int p = 1 + sc.n_constraints; + const int m = obs_row[i]; + if (m < 0) { + for (int j = 0; j < 6; j++) + jacobian[6 * static_cast(i) + j] = 0.0; + for (int a = 0; a < p; a++) + jk[static_cast(i) * p + a] = 0.0; + return; + } + const int *h = row_hkl + 3 * m; + // compute_value_and_derivatives() with k_sol and b_sol fixed + const float2 fcalc = solvent_fcalc(f[m], fmask[m], row_stol2[m], sc); + const double fcalc_abs = hypotf(fcalc.x, fcalc.y); + const double kaniso = exp(-0.25 * r_u_r(sc.b_star, h)); + const double fe = fcalc_abs * kaniso; + const double y = sc.k_overall * fe; + const double x = h[0], yy = h[1], z = h[2]; + const double du[6] = {-0.25 * y * (x * x), -0.25 * y * (yy * yy), -0.25 * y * (z * z), + -0.5 * y * (x * yy), -0.5 * y * (x * z), -0.5 * y * (yy * z)}; + jk[static_cast(i) * p] = -fe / f_mean; + for (int c = 0; c < sc.n_constraints; c++) { + double dot = 0; + for (int q = 0; q < 6; q++) + dot += sc.constraints[c][q] * du[q]; + jk[static_cast(i) * p + 1 + c] = -dot / f_mean; + } + // d|F_scaled|/dq = K Re(conj(F_total) dFcalc/dq) / |F_total| + const double k = sc.k_overall * kaniso / (fcalc_abs * f_mean); + for (int j = 0; j < 6; j++) { + double2 d; + if (j < 3) { + const double2 s = shifted[j * rows + m]; + d = make_double2((s.x - f[m].x) / step, (s.y - f[m].y) / step); + } else { + d = df_dt[3 * static_cast(m) + j - 3]; + } + jacobian[6 * static_cast(i) + j] = -k * (fcalc.x * d.x + fcalc.y * d.y); + } +} + +// out[b] = sum over i of a[i * lda + ia] * c[i * ldc + ic] for the pair b = (ia, ic): one block per pair, +// a fixed stride per thread and a fixed shared-memory tree, so the sum is the same on every run. +__global__ void products_kernel(const double *a, int lda, const double *c, int ldc, int n, int nc, double *out) { + const int ia = blockIdx.x / nc, ic = blockIdx.x % nc; + double sum = 0; + for (int i = threadIdx.x; i < n; i += THREADS) + sum += a[static_cast(i) * lda + ia] * c[static_cast(i) * ldc + ic]; + __shared__ double sh[THREADS]; + sh[threadIdx.x] = sum; + __syncthreads(); + for (int s = THREADS / 2; s > 0; s >>= 1) { + if (threadIdx.x < s) + sh[threadIdx.x] += sh[threadIdx.x + s]; + __syncthreads(); + } + if (threadIdx.x == 0) + out[blockIdx.x] = sh[0]; +} + +// J - J_k x, the scale re-fit projected out. +__global__ void project_kernel(double *jacobian, const double *jk, int p, const double *x, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) + return; + for (int j = 0; j < 6; j++) { + double v = jacobian[6 * static_cast(i) + j]; + for (int a = 0; a < p; a++) + v -= jk[static_cast(i) * p + a] * x[a * 6 + j]; + jacobian[6 * static_cast(i) + j] = v; + } +} + +int blocks(size_t n) { + return static_cast((n + THREADS - 1) / THREADS); +} + +// cub's temporary storage for the pair scan and sort at the capacity. +size_t CubBytes(size_t atoms, size_t pairs) { + size_t scan = 0, sort = 0; + cuda_err(cub::DeviceScan::ExclusiveSum(nullptr, scan, static_cast(nullptr), static_cast(nullptr), + static_cast(atoms + 1))); + cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, sort, static_cast(nullptr), static_cast(nullptr), + static_cast(nullptr), static_cast(nullptr), + static_cast(pairs))); + return std::max(scan, sort); +} + +// A plan for an (nu, nv, nw) R2C transform, u fastest, `batch` grids back to back, whose work area the +// caller provides; its size in `work`. +cufftHandle MakePlan(int nu, int nv, int nw, int batch, size_t &work) { + cufftHandle plan; + cufft_err(cufftCreate(&plan)); + cufft_err(cufftSetAutoAllocation(plan, 0)); + int n[3] = {nw, nv, nu}; + const int idist = nu * nv * nw, odist = (nu / 2 + 1) * nv * nw; + const cufftResult r = cufftMakePlanMany(plan, 3, n, nullptr, 1, idist, nullptr, 1, odist, CUFFT_R2C, batch, &work); + if (r != CUFFT_SUCCESS) + cufftDestroy(plan); + cufft_err(r); + return plan; +} + +// gemmi's box around each atom (MakeAtomBox, ModelGrid.cpp): the points within ceil(radius / spacing) +// of the nearest one along each axis, spacing being the distance between the grid's lattice planes. +std::vector AtomBoxes(const RigidBodyGPUZone &zone) { + const int n3[3] = {zone.nu, zone.nv, zone.nw}; + double spacing[3]; + for (int k = 0; k < 3; k++) + spacing[k] = 1.0 / (n3[k] * std::sqrt(zone.frac[3 * k] * zone.frac[3 * k] + zone.frac[3 * k + 1] * zone.frac[3 * k + 1] + + zone.frac[3 * k + 2] * zone.frac[3 * k + 2])); + std::vector box(3 * zone.atoms.size()); + for (size_t i = 0; i < zone.atoms.size(); i++) + for (int k = 0; k < 3; k++) + box[3 * i + k] = static_cast(std::ceil(zone.atoms[i].radius / spacing[k])); + return box; +} + +// The most bricks the 2 d + 1 wrapped points of a box can fall in along an axis of n points. +size_t AxisBrickBound(int d, int n) { + const int nb = (n + BRICK - 1) / BRICK; + return std::min(nb, (2 * d + 1 + BRICK - 1) / BRICK + 1); +} + +} // namespace + +struct RigidBodyGPUEngineImpl { + int device = 0; + RigidBodyGPUCapacity cap; + CudaStream stream; + + CudaDevicePtr atoms; + CudaDevicePtr box; + CudaDevicePtr rel; + CudaDevicePtr pos; + CudaDevicePtr mask_slot; + CudaDevicePtr mask_radius; + CudaDevicePtr mask_atoms; // 4 per mask atom: fractional x, y, z in [0,1), radius + + CudaDevicePtr grid; // 3 grids back to back + CudaDevicePtr spectrum; // their 3 transforms + CudaDevicePtr fft_work; + + CudaDevicePtr count, offset, key, value, key_sorted, value_sorted, brick_start, brick_end; + CudaDevicePtr cub_temp; + size_t cub_bytes = 0; + + CudaDevicePtr terms; + CudaDevicePtr row_hkl; + CudaDevicePtr row_scale, row_stol2; + CudaDevicePtr f, df_dt, shifted; + CudaDevicePtr fmask; + + CudaDevicePtr obs_row, point_obs; + CudaDevicePtr obs_fobs; + CudaDevicePtr residual, jacobian, jk; + CudaDevicePtr point_fcmol, point_fmask; + CudaDevicePtr products, x; + + std::map, cufftHandle> plans; + + // The zone. + int n_atoms = 0; + GridGeom geom{}; + double frac[9] = {}; + double volume = 0; + size_t rows = 0, ops = 1, n_obs = 0, n_points = 0, n_constraints = 0; + std::vector> constraints; + double f_mean = 1; + + explicit RigidBodyGPUEngineImpl(const RigidBodyGPUCapacity &c, int dev) + : device(dev), cap(c) { + const auto sync = CudaAlloc::Synchronous; + const size_t na = std::max(cap.atoms, 1); + atoms = CudaDevicePtr(na, sync); + box = CudaDevicePtr(3 * na, sync); + rel = CudaDevicePtr(3 * na, sync); + pos = CudaDevicePtr(na, sync); + mask_slot = CudaDevicePtr(na, sync); + mask_radius = CudaDevicePtr(na, sync); + mask_atoms = CudaDevicePtr(4 * na, sync); + grid = CudaDevicePtr(3 * cap.grid_points, sync); + spectrum = CudaDevicePtr(3 * cap.complex_points, sync); + fft_work = CudaDevicePtr(std::max(cap.fft_work_bytes, 1), sync); + count = CudaDevicePtr(na + 1, sync); + offset = CudaDevicePtr(na + 1, sync); + const size_t np = std::max(cap.pairs, 1); + key = CudaDevicePtr(np, sync); + value = CudaDevicePtr(np, sync); + key_sorted = CudaDevicePtr(np, sync); + value_sorted = CudaDevicePtr(np, sync); + brick_start = CudaDevicePtr(std::max(cap.bricks, 1), sync); + brick_end = CudaDevicePtr(std::max(cap.bricks, 1), sync); + cub_bytes = CubBytes(na, np); + cub_temp = CudaDevicePtr(std::max(cub_bytes, 1), sync); + const size_t nr = std::max(cap.rows, 1), no = std::max(cap.observations, 1); + terms = CudaDevicePtr(std::max(cap.terms, 1), sync); + row_hkl = CudaDevicePtr(3 * nr, sync); + row_scale = CudaDevicePtr(nr, sync); + row_stol2 = CudaDevicePtr(nr, sync); + f = CudaDevicePtr(nr, sync); + df_dt = CudaDevicePtr(3 * nr, sync); + shifted = CudaDevicePtr(3 * nr, sync); + fmask = CudaDevicePtr(nr, sync); + obs_row = CudaDevicePtr(no, sync); + point_obs = CudaDevicePtr(no, sync); + obs_fobs = CudaDevicePtr(no, sync); + residual = CudaDevicePtr(no, sync); + jacobian = CudaDevicePtr(6 * no, sync); + jk = CudaDevicePtr(MAX_SCALE_PARAMS * no, sync); + point_fcmol = CudaDevicePtr(no, sync); + point_fmask = CudaDevicePtr(no, sync); + products = CudaDevicePtr(MAX_SCALE_PARAMS * (MAX_SCALE_PARAMS + 6), sync); + x = CudaDevicePtr(MAX_SCALE_PARAMS * 6, sync); + } + + ~RigidBodyGPUEngineImpl() { + cudaSetDevice(device); + cudaStreamSynchronize(stream); + for (auto &[k, plan] : plans) + cufftDestroy(plan); + } + + cufftHandle Plan(int batch) { + const std::array k{geom.nu, geom.nv, geom.nw, batch}; + auto it = plans.find(k); + if (it != plans.end()) + return it->second; + size_t work = 0; + const cufftHandle plan = MakePlan(geom.nu, geom.nv, geom.nw, batch, work); + if (work > cap.fft_work_bytes) { + cufftDestroy(plan); + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: cuFFT work area over the reserve"); + } + cufft_err(cufftSetWorkArea(plan, fft_work)); + cufft_err(cufftSetStream(plan, stream)); + plans.emplace(k, plan); + return plan; + } + + size_t GridPoints() const { return static_cast(geom.nu) * geom.nv * geom.nw; } + size_t ComplexPoints() const { return static_cast(geom.nu / 2 + 1) * geom.nv * geom.nw; } + float Norm() const { return static_cast(volume / static_cast(GridPoints())); } + + // The body placed at (R, t) and its density gathered into grid `slot`. + void Density(const double r[9], const double t[3], int slot) { + Placement p{}; + std::copy(r, r + 9, p.r); + std::copy(t, t + 3, p.t); + std::copy(frac, frac + 9, p.frac); + place_kernel<<>>(rel, n_atoms, p, geom.nu, geom.nv, geom.nw, mask_slot, + mask_radius, pos, mask_atoms); + cuda_err(cudaGetLastError()); + count_pairs_kernel<<>>(pos, box, n_atoms, geom.nu, geom.nv, geom.nw, + count); + cuda_err(cudaGetLastError()); + size_t temp = cub_bytes; + cuda_err(cub::DeviceScan::ExclusiveSum(cub_temp.get(), temp, count.get(), offset.get(), n_atoms + 1, stream)); + int npairs = 0; + cuda_err(cudaMemcpyAsync(&npairs, offset.get() + n_atoms, sizeof(int), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); + if (static_cast(npairs) > cap.pairs) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: gather pairs over the reserve"); + fill_pairs_kernel<<>>(pos, box, n_atoms, geom.nu, geom.nv, geom.nw, + offset, key, value); + cuda_err(cudaGetLastError()); + const int nbu = (geom.nu + BRICK - 1) / BRICK, nbv = (geom.nv + BRICK - 1) / BRICK, + nbw = (geom.nw + BRICK - 1) / BRICK; + const int nb = nbu * nbv * nbw; + int bits = 1; + while ((1 << bits) < nb) + bits++; + temp = cub_bytes; + cuda_err(cub::DeviceRadixSort::SortPairs(cub_temp.get(), temp, key.get(), key_sorted.get(), value.get(), + value_sorted.get(), npairs, 0, bits, stream)); + cuda_err(cudaMemsetAsync(brick_start, 0, nb * sizeof(int), stream)); + cuda_err(cudaMemsetAsync(brick_end, 0, nb * sizeof(int), stream)); + if (npairs > 0) { + brick_ranges_kernel<<>>(key_sorted, npairs, brick_start, brick_end); + cuda_err(cudaGetLastError()); + } + gather_kernel<<>>(atoms, pos, value_sorted, brick_start, brick_end, + geom, grid.get() + slot * GridPoints()); + cuda_err(cudaGetLastError()); + } + + void Compose(int slot, double2 *out, double2 *derivative) { + compose_kernel<<>>(spectrum.get() + slot * ComplexPoints(), Norm(), terms, + row_scale, static_cast(rows), static_cast(ops), + out, derivative); + cuda_err(cudaGetLastError()); + } + + Scale MakeScale(double k_overall, const double b_star[6], double k_sol, double b_sol) const { + Scale s{}; + s.k_overall = k_overall; + std::copy(b_star, b_star + 6, s.b_star); + s.k_sol = k_sol; + s.b_sol = b_sol; + s.n_constraints = static_cast(n_constraints); + for (size_t c = 0; c < n_constraints; c++) + for (int q = 0; q < 6; q++) + s.constraints[c][q] = constraints[c][q]; + return s; + } +}; + +size_t RigidBodyGPUEngine::Bricks(int nu, int nv, int nw) { + return static_cast((nu + BRICK - 1) / BRICK) * ((nv + BRICK - 1) / BRICK) * ((nw + BRICK - 1) / BRICK); +} + +size_t RigidBodyGPUEngine::PairBound(const RigidBodyGPUZone &zone) { + const std::vector box = AtomBoxes(zone); + const int n3[3] = {zone.nu, zone.nv, zone.nw}; + size_t pairs = 0; + for (size_t i = 0; i < zone.atoms.size(); i++) { + size_t product = 1; + for (int k = 0; k < 3; k++) + product *= AxisBrickBound(box[3 * i + k], n3[k]); + pairs += product; + } + return pairs; +} + +bool RigidBodyGPUEngine::Supports(const RigidBodyGPUZone &zone) { + const std::vector box = AtomBoxes(zone); + const int n3[3] = {zone.nu, zone.nv, zone.nw}; + for (size_t i = 0; i < box.size(); i++) + if (2 * box[i] + 1 > n3[i % 3] || AxisBrickBound(box[i], n3[i % 3]) > MAX_BRICKS_PER_AXIS) + return false; + return true; +} + +void RigidBodyGPUEngine::MemoryInfo(size_t &free, size_t &total) { + cuda_err(cudaMemGetInfo(&free, &total)); +} + +int RigidBodyGPUEngine::CurrentDevice() { + int device = 0; + cuda_err(cudaGetDevice(&device)); + return device; +} + +size_t RigidBodyGPUEngine::FFTWorkBytes(int nu, int nv, int nw) { + size_t largest = 0; + for (int batch : {1, 3}) { + size_t work = 0; + const cufftHandle plan = MakePlan(nu, nv, nw, batch, work); + cufftDestroy(plan); + largest = std::max(largest, work); + } + return largest; +} + +size_t RigidBodyGPUEngine::DeviceBytes(const RigidBodyGPUCapacity &c) { + const size_t na = c.atoms + 1, no = c.observations + 1, nr = c.rows + 1; + size_t bytes = na * (sizeof(RigidBodyGPUAtom) + 3 * sizeof(int) + 7 * sizeof(double) + sizeof(float4) + + sizeof(int) + sizeof(float) + 2 * sizeof(int)); + bytes += 3 * c.grid_points * sizeof(float) + 3 * c.complex_points * sizeof(float2) + c.fft_work_bytes; + bytes += 4 * c.pairs * sizeof(int) + 2 * c.bricks * sizeof(int); + bytes += c.terms * sizeof(DeviceTerm); + bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 7 * sizeof(double2) + sizeof(float2)); + bytes += no * (2 * sizeof(int) + sizeof(float) + (1 + 6 + MAX_SCALE_PARAMS) * sizeof(double) + 2 * sizeof(float2)); + bytes += CubBytes(na, c.pairs + 1); + return bytes; +} + +RigidBodyGPUEngine::RigidBodyGPUEngine(const RigidBodyGPUCapacity &capacity, int device) { + cuda_err(cudaSetDevice(device)); + impl_ = std::make_unique(capacity, device); +} + +RigidBodyGPUEngine::~RigidBodyGPUEngine() = default; + +void RigidBodyGPUEngine::SetBody(const std::vector> &relative) { + RigidBodyGPUEngineImpl &e = *impl_; + cuda_err(cudaSetDevice(e.device)); + if (relative.size() > e.cap.atoms) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: more atoms than the reserve"); + e.n_atoms = static_cast(relative.size()); + cuda_err(cudaMemcpyAsync(e.rel, relative.data(), relative.size() * 3 * sizeof(double), cudaMemcpyHostToDevice, + e.stream)); + cuda_err(cudaStreamSynchronize(e.stream)); +} + +void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) { + RigidBodyGPUEngineImpl &e = *impl_; + cuda_err(cudaSetDevice(e.device)); + const size_t np = static_cast(zone.nu) * zone.nv * zone.nw; + const size_t nc = static_cast(zone.nu / 2 + 1) * zone.nv * zone.nw; + if (static_cast(zone.atoms.size()) != e.n_atoms || np > e.cap.grid_points || nc > e.cap.complex_points || + Bricks(zone.nu, zone.nv, zone.nw) > e.cap.bricks || zone.row_hkl.size() > e.cap.rows || + zone.terms.size() > e.cap.terms || zone.obs_row.size() > e.cap.observations || + zone.constraints.size() + 1 > MAX_SCALE_PARAMS) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: zone over the engine's reserve"); + + e.geom.nu = zone.nu; + e.geom.nv = zone.nv; + e.geom.nw = zone.nw; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) + e.geom.orth_n[3 * i + j] = static_cast(zone.orth[3 * i + j] / (j == 0 ? zone.nu : j == 1 ? zone.nv : zone.nw)); + std::copy(zone.frac, zone.frac + 9, e.frac); + e.volume = zone.volume; + e.rows = zone.row_hkl.size(); + e.ops = zone.ops; + e.n_obs = zone.obs_row.size(); + e.n_points = zone.point_obs.size(); + e.n_constraints = zone.constraints.size(); + e.constraints = zone.constraints; + e.f_mean = zone.f_mean; + + // An atom whose box is wider than the cell would reach a point through more than one image, which + // gemmi's box walk adds and the gather's nearest image does not; such a cell is left to the CPU. + if (!Supports(zone)) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: an atom wider than the cell"); + const std::vector box = AtomBoxes(zone); + + std::vector slot(zone.atoms.size(), -1); + std::vector radius(zone.atoms.size(), 0.0f); + for (size_t j = 0; j < zone.mask_atom.size(); j++) { + slot[zone.mask_atom[j]] = static_cast(j); + radius[zone.mask_atom[j]] = zone.mask_radius[j]; + } + + std::vector terms(zone.terms.size()); + for (size_t i = 0; i < zone.terms.size(); i++) { + const RigidBodyGPUTerm &t = zone.terms[i]; + DeviceTerm &d = terms[i]; + d.conj = t.k[0] < 0; // the half-u transform holds h >= 0, and F1(-k) = conj F1(k) for a real map + const int ku = d.conj ? -t.k[0] : t.k[0], kv = d.conj ? -t.k[1] : t.k[1], kw = d.conj ? -t.k[2] : t.k[2]; + d.index = ku + (zone.nu / 2 + 1) * (((kv % zone.nv) + zone.nv) % zone.nv + + zone.nv * (((kw % zone.nw) + zone.nw) % zone.nw)); + d.phase[0] = t.phase[0]; + d.phase[1] = t.phase[1]; + std::copy(t.s, t.s + 3, d.s); + } + + auto up = [&](void *dst, const void *src, size_t bytes) { + if (bytes > 0) + cuda_err(cudaMemcpyAsync(dst, src, bytes, cudaMemcpyHostToDevice, e.stream)); + }; + up(e.atoms, zone.atoms.data(), zone.atoms.size() * sizeof(RigidBodyGPUAtom)); + up(e.box, box.data(), box.size() * sizeof(int)); + up(e.mask_slot, slot.data(), slot.size() * sizeof(int)); + up(e.mask_radius, radius.data(), radius.size() * sizeof(float)); + up(e.terms, terms.data(), terms.size() * sizeof(DeviceTerm)); + up(e.row_hkl, zone.row_hkl.data(), zone.row_hkl.size() * 3 * sizeof(int)); + up(e.row_scale, zone.row_scale.data(), zone.row_scale.size() * sizeof(double)); + up(e.row_stol2, zone.row_stol2.data(), zone.row_stol2.size() * sizeof(double)); + up(e.obs_row, zone.obs_row.data(), zone.obs_row.size() * sizeof(int)); + up(e.obs_fobs, zone.obs_fobs.data(), zone.obs_fobs.size() * sizeof(float)); + up(e.point_obs, zone.point_obs.data(), zone.point_obs.size() * sizeof(int)); + e.Plan(1); + e.Plan(3); + cuda_err(cudaStreamSynchronize(e.stream)); +} + +void RigidBodyGPUEngine::Fcalc(const double rotation[9], const double translation[3]) { + RigidBodyGPUEngineImpl &e = *impl_; + cuda_err(cudaSetDevice(e.device)); + e.Density(rotation, translation, 0); + cufft_err(cufftExecR2C(e.Plan(1), e.grid, reinterpret_cast(e.spectrum.get()))); + e.Compose(0, e.f, e.df_dt); +} + +void RigidBodyGPUEngine::Fmask(const float *host_mask) { + RigidBodyGPUEngineImpl &e = *impl_; + cuda_err(cudaSetDevice(e.device)); + float *mask_grid = e.grid.get() + e.GridPoints(); + cufftComplex *mask_spectrum = reinterpret_cast(e.spectrum.get() + e.ComplexPoints()); + cuda_err(cudaMemcpyAsync(mask_grid, host_mask, e.GridPoints() * sizeof(float), cudaMemcpyHostToDevice, e.stream)); + cufft_err(cufftExecR2C(e.Plan(1), mask_grid, mask_spectrum)); + fmask_kernel<<>>(reinterpret_cast(mask_spectrum), e.Norm(), + e.row_hkl, static_cast(e.rows), e.geom.nu, e.geom.nv, + e.geom.nw, e.fmask); + cuda_err(cudaGetLastError()); + cuda_err(cudaStreamSynchronize(e.stream)); // the host mask may be freed once this returns +} + +void RigidBodyGPUEngine::DownloadPoints(std::vector> &fcmol, + std::vector> &fmask) { + RigidBodyGPUEngineImpl &e = *impl_; + cuda_err(cudaSetDevice(e.device)); + fcmol.resize(e.n_points); + fmask.resize(e.n_points); + if (e.n_points == 0) + return; + pack_points_kernel<<>>(e.point_obs, e.obs_row, + static_cast(e.n_points), e.f, e.fmask, + e.point_fcmol, e.point_fmask); + cuda_err(cudaGetLastError()); + cuda_err(cudaMemcpyAsync(fcmol.data(), e.point_fcmol, e.n_points * sizeof(float2), cudaMemcpyDeviceToHost, e.stream)); + cuda_err(cudaMemcpyAsync(fmask.data(), e.point_fmask, e.n_points * sizeof(float2), cudaMemcpyDeviceToHost, e.stream)); + cuda_err(cudaStreamSynchronize(e.stream)); +} + +void RigidBodyGPUEngine::Residuals(double k_overall, const double b_star[6], double k_sol, double b_sol, + double *residuals) { + RigidBodyGPUEngineImpl &e = *impl_; + cuda_err(cudaSetDevice(e.device)); + const Scale sc = e.MakeScale(k_overall, b_star, k_sol, b_sol); + residual_kernel<<>>(e.obs_row, e.obs_fobs, static_cast(e.n_obs), e.f, + e.fmask, e.row_hkl, e.row_stol2, sc, e.f_mean, + e.residual); + cuda_err(cudaGetLastError()); + cuda_err(cudaMemcpyAsync(residuals, e.residual, e.n_obs * sizeof(double), cudaMemcpyDeviceToHost, e.stream)); + cuda_err(cudaStreamSynchronize(e.stream)); +} + +void RigidBodyGPUEngine::Jacobian(const double rotation[3][9], const double translation[3][3], double step, + double k_overall, const double b_star[6], double k_sol, double b_sol, + std::vector &jtj, std::vector &jtq) { + RigidBodyGPUEngineImpl &e = *impl_; + cuda_err(cudaSetDevice(e.device)); + for (int j = 0; j < 3; j++) + e.Density(rotation[j], translation[j], j); + cufft_err(cufftExecR2C(e.Plan(3), e.grid, reinterpret_cast(e.spectrum.get()))); + for (int j = 0; j < 3; j++) + e.Compose(j, e.shifted.get() + j * e.rows, nullptr); + + const Scale sc = e.MakeScale(k_overall, b_star, k_sol, b_sol); + const int p = 1 + static_cast(e.n_constraints); + const int n = static_cast(e.n_obs); + jacobian_kernel<<>>(e.obs_row, n, e.f, e.df_dt, e.shifted, e.rows, e.fmask, + e.row_hkl, e.row_stol2, sc, e.f_mean, step, e.jacobian, + e.jk); + cuda_err(cudaGetLastError()); + products_kernel<<

>>(e.jk, p, e.jk, p, n, p, e.products); + cuda_err(cudaGetLastError()); + products_kernel<<

>>(e.jk, p, e.jacobian, 6, n, 6, e.products.get() + p * p); + cuda_err(cudaGetLastError()); + std::vector out(p * p + p * 6); + cuda_err(cudaMemcpyAsync(out.data(), e.products, out.size() * sizeof(double), cudaMemcpyDeviceToHost, e.stream)); + cuda_err(cudaStreamSynchronize(e.stream)); + jtj.assign(out.begin(), out.begin() + p * p); + jtq.assign(out.begin() + p * p, out.end()); +} + +void RigidBodyGPUEngine::ProjectJacobian(const std::vector &x, double *jacobian) { + RigidBodyGPUEngineImpl &e = *impl_; + cuda_err(cudaSetDevice(e.device)); + const int p = 1 + static_cast(e.n_constraints); + cuda_err(cudaMemcpyAsync(e.x, x.data(), x.size() * sizeof(double), cudaMemcpyHostToDevice, e.stream)); + project_kernel<<>>(e.jacobian, e.jk, p, e.x, static_cast(e.n_obs)); + cuda_err(cudaGetLastError()); + cuda_err(cudaMemcpyAsync(jacobian, e.jacobian, e.n_obs * 6 * sizeof(double), cudaMemcpyDeviceToHost, e.stream)); + cuda_err(cudaStreamSynchronize(e.stream)); +} diff --git a/rugnux/RigidBodyGPU.h b/rugnux/RigidBodyGPU.h new file mode 100644 index 000000000..a1de80d0d --- /dev/null +++ b/rugnux/RigidBodyGPU.h @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// The rigid body's target evaluated on a GPU (CUDA builds only; the header itself needs no CUDA). +// It is the function RigidBodyTarget computes on the CPU - the same density, the same composition of +// Fcalc from one copy, the same bulk-solvent mask, scale fit, residuals and Jacobian - moved to the +// device so that a fit takes a fraction of a second instead of several seconds. The two agree to +// rounding, not bit for bit: float distances, cuFFT for FFTW. The GPU is deterministic on its own. + +#include +#include +#include +#include +#include +#include +#include + +#include "RigidBodyRefine.h" + +class Logger; +class RigidBodyGPUEngine; +struct RigidBodyGPUZone; + +// A CUDA failure inside the rigid body. Model validation catches it and starts again on the CPU. +class RigidBodyGPUFailure : public std::runtime_error { +public: + explicit RigidBodyGPUFailure(const std::string &what) : std::runtime_error(what) {} +}; + +// A few engines, reserved once for a whole validation: each is a stream and the buffers for one fit at +// a time, sized for the finest zone of the ladder to d_min. The real fit and the null's replicates each +// take one for the length of their fit, and wait for one when all are taken. Engines are +// interchangeable and every kernel deterministic, so which replicate gets which engine does not change +// a number. +class RigidBodyGPUPool { +public: + // Null where there is no GPU, where not even one engine fits the budget - a quarter of the card, + // and never the last gigabyte of what is free, since the merge may be running beside it - or where + // the cell is too small for the gather (an atom's box wider than the cell). Logged either way. + // `max_observations`: the most working reflections any fit will be given in its finest zone. + // `max_engines`: at most this many, however much memory there is (1 to 4). + static std::unique_ptr Create(const gemmi::Model &model, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, double d_min, + size_t max_observations, size_t max_engines, Logger &logger); + ~RigidBodyGPUPool(); + + size_t Engines() const { return engines_.size(); } + RigidBodyGPUEngine &Acquire(); + void Release(RigidBodyGPUEngine &engine); + +private: + RigidBodyGPUPool() = default; + std::vector> engines_; + std::vector idle_; + std::mutex m_; + std::condition_variable cv_; +}; + +class RigidBodyTargetGPU : public RigidBodyTargetBase { +public: + // Takes an engine from `pool` for its lifetime. q = 0 is the placement `model` has now; unlike the + // CPU target the model is only moved by Place(). + RigidBodyTargetGPU(RigidBodyGPUPool &pool, gemmi::Model &model, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, size_t nthreads); + ~RigidBodyTargetGPU() override; + + void SetZone(const gemmi::AsuData> &fobs, double d_min) override; + size_t NumObservations() const override { return fobs_.v.size(); } + + bool Residuals(const double q[6], double *residuals) override; + bool Jacobian(const double q[6], double *jacobian) override; + +private: + // The placement at q as x -> R (x - centroid) + centroid + t. + void Placement(const double q[6], double rotation[9], double translation[3]) const; + + RigidBodyGPUPool &pool_; + RigidBodyGPUEngine &engine_; + gemmi::Model &model_; + const gemmi::UnitCell &cell_; + const gemmi::SpaceGroup &sg_; + size_t nthreads_; + + gemmi::AsuData> fobs_; + double d_min_ = 0; + std::unique_ptr zone_; + bool solvent_fitted_ = false; + std::vector orbit_leaders_; // for the host mask + + bool have_point_ = false; + std::array q_{}; + double k_overall_ = 1; + gemmi::SMat33 b_star_{0, 0, 0, 0, 0, 0}; +}; diff --git a/rugnux/RigidBodyGPUEngine.h b/rugnux/RigidBodyGPUEngine.h new file mode 100644 index 000000000..35ca075b1 --- /dev/null +++ b/rugnux/RigidBodyGPUEngine.h @@ -0,0 +1,120 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// The device half of RigidBodyTargetGPU (RigidBodyGPU.h): one engine is one CUDA stream and the buffers +// for one resolution zone of one fit. The host half works out everything that depends only on the +// model, the cell and the zone - with gemmi, which stays out of nvcc - and hands it over in the plain +// structures below; the engine then does, per evaluation, what depends on the placement. CUDA builds +// only. + +#include +#include +#include +#include + +// One atom's density on one zone's grid, as PutModelDensityOnGrid() (ModelGrid.cpp) sets it up: gemmi's +// precalculated five-Gaussian sum and the radius it cuts the sum at. +struct RigidBodyGPUAtom { + float a[5]; + float b[5][6]; // isotropic: b[k][0] multiplies r^2; anisotropic: the matrix, u11 u22 u33 u12 u13 u23 + float occ; + float radius; + int aniso; +}; + +// One term of SymmetryComposition: F1 at k = hR, times the phase of the operator's translation. +struct RigidBodyGPUTerm { + int k[3]; + double phase[2]; // exp(+2 pi i h.t), real and imaginary + double s[3]; // k as a Cartesian reciprocal vector +}; + +// Everything a zone needs that does not depend on the placement. +struct RigidBodyGPUZone { + int nu = 0, nv = 0, nw = 0; // the zone's grid, u fastest + double orth[9] = {}, frac[9] = {}; // row-major + double volume = 0; + std::vector atoms; // model order + + // The bulk-solvent mask's atoms: the model's index of each, and its radius (probe included). + std::vector mask_atom; + std::vector mask_radius; + // Every image: the group's operators, each with each centring vector, fractional. + std::vector> images; // rot[9] row-major, tran[3] + + // The composition: rows (composed indices) and Ops() terms per row, row-major. + std::vector> row_hkl; + std::vector row_scale; // n_cen * unblur + std::vector row_stol2; + size_t ops = 1; + std::vector terms; + + // The observations the residuals are over: each one's row (-1 without one) and amplitude. + std::vector obs_row; + std::vector obs_fobs; + double f_mean = 1; + // The scale's points: the observations gemmi's prepare_points() would take, in order. + std::vector point_obs; + std::vector> constraints; // adp_symmetry_constraints() +}; + +// Upper bounds an engine is sized for. +struct RigidBodyGPUCapacity { + size_t atoms = 0; + size_t grid_points = 0; // nu * nv * nw + size_t complex_points = 0; // (nu / 2 + 1) * nv * nw + size_t bricks = 0; + size_t pairs = 0; // (brick, atom) pairs of the gather + size_t rows = 0, terms = 0; + size_t observations = 0; + size_t fft_work_bytes = 0; +}; + +struct RigidBodyGPUEngineImpl; + +class RigidBodyGPUEngine { +public: + // The bytes an engine of this capacity reserves on the device. + static size_t DeviceBytes(const RigidBodyGPUCapacity &capacity); + // The largest cuFFT work area a (nu, nv, nw) grid needs, batch 1 or 3. + static size_t FFTWorkBytes(int nu, int nv, int nw); + // (brick, atom) pairs of the gather over at most, for a zone's atoms on its grid. + static size_t PairBound(const RigidBodyGPUZone &zone); + static size_t Bricks(int nu, int nv, int nw); + // Whether the gather reproduces gemmi's box walk on this zone: every atom's box narrower than the cell. + static bool Supports(const RigidBodyGPUZone &zone); + static void MemoryInfo(size_t &free, size_t &total); + static int CurrentDevice(); + + RigidBodyGPUEngine(const RigidBodyGPUCapacity &capacity, int device); + ~RigidBodyGPUEngine(); + + // Per fit: each atom's position relative to the model centroid, which the placements rotate. + void SetBody(const std::vector> &relative); + // Per zone. Throws if the zone does not fit the capacity. + void SetZone(const RigidBodyGPUZone &zone); + + // Per evaluation, at the placement x -> R x_rel + t: Fcalc and dF/dt at the rows, then the bulk- + // solvent mask and Fmask. `host_mask`, where not null, is the mask grid computed on the host + // instead (the stub until ModelMaskGPU is in). + void Fcalc(const double rotation[9], const double translation[3]); + void Fmask(const float *host_mask); + // The scale's points: fcmol (Fcalc as complex) and fmask, downloaded (the stub). + void DownloadPoints(std::vector> &fcmol, std::vector> &fmask); + // The residuals at the scale given, NumObservations() of them, into `residuals` (host). + void Residuals(double k_overall, const double b_star[6], double k_sol, double b_sol, double *residuals); + + // The Jacobian at the last evaluation: `rotation[j]` and `translation[j]` place the body one step + // along rotation axis j. Writes J (fixed scale) and J_k (the scale parameters' columns) on the + // device and returns J_k^T J_k (p x p) and J_k^T J (p x 6), p = 1 + constraints, row-major. + void Jacobian(const double rotation[3][9], const double translation[3][3], double step, double k_overall, + const double b_star[6], double k_sol, double b_sol, std::vector &jtj, + std::vector &jtq); + // J - J_k x, x p x 6 row-major, into `jacobian` (host, NumObservations() x 6). + void ProjectJacobian(const std::vector &x, double *jacobian); + +private: + std::unique_ptr impl_; +}; diff --git a/rugnux/RigidBodyRefine.cpp b/rugnux/RigidBodyRefine.cpp index 3bd814177..16535e169 100644 --- a/rugnux/RigidBodyRefine.cpp +++ b/rugnux/RigidBodyRefine.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -22,6 +23,9 @@ #include "ModelFFT.h" // MapToFPhi #include "ModelGrid.h" // PutModelDensityOnGrid, PutMaskOnGrid #include "ModelScaling.h" // FitModelScale +#ifdef JFJOCH_USE_CUDA +#include "RigidBodyGPU.h" // RigidBodyTargetGPU +#endif #include "../common/JFJochMath.h" // PI #include "../common/ParallelFor.h" #include "../common/Logger.h" @@ -49,20 +53,11 @@ constexpr double JACOBIAN_STEP_FRACTION = 0.01; // How finely a zone's maps are sampled: DensityCalculator's rate, for a spacing of d_min / (2 * rate). constexpr double GRID_RATE = 1.5; -// An empty grid of the zone's size, the size DensityCalculator gives its own at GRID_RATE. -gemmi::Grid ZoneGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) { - gemmi::Grid grid; - grid.unit_cell = cell; - grid.spacegroup = &sg; - grid.set_size_from_spacing(d_min / (2 * GRID_RATE), gemmi::GridSizeRounding::Up); - return grid; -} - // Ceres' own numeric differentiation steps by |x| * relative_step_size, which is zero at the start of // every zone (the placement begins at no shift), so the Jacobian is supplied by RigidBodyTarget. class RigidBodyCost : public ceres::CostFunction { public: - explicit RigidBodyCost(RigidBodyTarget &target) : target_(target) { + explicit RigidBodyCost(RigidBodyTargetBase &target) : target_(target) { set_num_residuals(static_cast(target.NumObservations())); mutable_parameter_block_sizes()->push_back(6); } @@ -85,7 +80,7 @@ public: } private: - RigidBodyTarget &target_; + RigidBodyTargetBase &target_; mutable std::array last_q_{}; mutable std::vector last_residuals_; // at last_q_; empty until an evaluation succeeds }; @@ -97,7 +92,9 @@ private: // sees only grid noise and commits or not by coin flip, while the other five parameters carry the // noise in with them. The free directions are the common fixed subspace of the group's rotation // parts, and the projector onto it is simply their average. -gemmi::Mat33 GaugeProjector(const gemmi::SpaceGroup &sg, const gemmi::UnitCell &cell) { +} // namespace + +gemmi::Mat33 RigidBodyGaugeProjector(const gemmi::SpaceGroup &sg, const gemmi::UnitCell &cell) { const gemmi::GroupOps gops = sg.operations(); double m[3][3] = {}; const double n = static_cast(gops.sym_ops.size()) * gemmi::Op::DEN; @@ -112,7 +109,28 @@ gemmi::Mat33 GaugeProjector(const gemmi::SpaceGroup &sg, const gemmi::UnitCell & return cell.orth.mat.multiply(mean).multiply(cell.frac.mat); } -} // namespace +// An empty grid of the zone's size, the size DensityCalculator gives its own at GRID_RATE. +gemmi::Grid RigidBodyZoneGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min) { + gemmi::Grid grid; + grid.unit_cell = cell; + grid.spacegroup = &sg; + grid.set_size_from_spacing(d_min / (2 * GRID_RATE), gemmi::GridSizeRounding::Up); + return grid; +} + +double RigidBodyJacobianStep(double d_min) { + return JACOBIAN_STEP_FRACTION * d_min; +} + +std::vector RigidBodyLadder(double d_min) { + std::vector ladder; + for (double zone : LADDER) + if (zone >= d_min) + ladder.push_back(zone); + if (ladder.empty()) + ladder.push_back(d_min); + return ladder; +} std::vector ModelPositions(const gemmi::Model &model) { std::vector pos; @@ -165,6 +183,7 @@ SymmetryComposition::SymmetryComposition(const gemmi::Grid &grid, double static_cast(grid.nu) * (gemmi::modulo(kv, grid.nv) + static_cast(grid.nv) * kw); t.phase = std::polar(1.0, -op.phase_shift(h)); // gemmi's phase_shift is -2 pi h.t t.s = cell.frac.mat.left_multiply(gemmi::Vec3(k[0], k[1], k[2])); + t.k = k; terms_.push_back(t); } } @@ -202,10 +221,7 @@ void SymmetryComposition::Compose(const gemmi::FPhiGrid &f1, double unblu }); } -RigidBodyTarget::RigidBodyTarget(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, - size_t nthreads) - : model_(model), cell_(cell), sg_(sg), nthreads_(nthreads), base_(ModelPositions(model)), - column_models_(3, model) { +RigidBodyTargetBase::RigidBodyTargetBase(const gemmi::Model &model) : base_(ModelPositions(model)) { if (base_.empty()) return; for (const gemmi::Position &p : base_) @@ -217,11 +233,16 @@ RigidBodyTarget::RigidBodyTarget(gemmi::Model &model, const gemmi::UnitCell &cel rms_radius_ = std::sqrt(r2 / static_cast(base_.size())); } +RigidBodyTarget::RigidBodyTarget(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + size_t nthreads) + : RigidBodyTargetBase(model), model_(model), cell_(cell), sg_(sg), nthreads_(nthreads), + column_models_(3, model) {} + // Parameters are carried as six lengths in angstroms - the first three are the angle-axis rotation // vector multiplied by the model's rms radius, so a unit of each of the six moves a typical atom by // the same amount. The rotation is about the model centroid, which decorrelates it from the // translation, and is applied to the positions only: an anisotropic U is not turned with the body. -void RigidBodyTarget::Place(const double q[6], gemmi::Model &model) const { +void RigidBodyTargetBase::Place(const double q[6], gemmi::Model &model) const { const double aa[3] = {q[0] / rms_radius_, q[1] / rms_radius_, q[2] / rms_radius_}; size_t i = 0; for (gemmi::Chain &ch : model.chains) @@ -246,7 +267,7 @@ void RigidBodyTarget::SetZone(const gemmi::AsuData> &fo f_mean_ = fobs_.v.empty() ? 1.0 : sum / static_cast(fobs_.v.size()); solvent_fitted_ = false; have_point_ = false; - const gemmi::Grid grid = ZoneGrid(cell_, sg_, d_min); + const gemmi::Grid grid = RigidBodyZoneGrid(cell_, sg_, d_min); orbit_leaders_ = OrbitLeaders(grid, nthreads_); std::vector hkl; for (const auto &hv : fobs_.v) @@ -282,7 +303,7 @@ bool RigidBodyTarget::Residuals(const double q[6], double *residuals) { fmask = fmask_; } else { // The mask is symmetric as gridded, so it is read at h directly. - gemmi::Grid mask_grid = ZoneGrid(cell_, sg_, d_min_); + gemmi::Grid mask_grid = RigidBodyZoneGrid(cell_, sg_, d_min_); PutMaskOnGrid(mask_grid, model_, orbit_leaders_, nthreads_); const gemmi::FPhiGrid fm = MapToFPhi(mask_grid); for (const gemmi::Miller &h : hkl) @@ -367,7 +388,7 @@ bool RigidBodyTarget::Jacobian(const double q[6], double *jacobian) { return false; } ++jacobians; - const double step = JACOBIAN_STEP_FRACTION * d_min_; + const double step = RigidBodyJacobianStep(d_min_); std::array>, 3> shifted; const std::launch policy = nthreads_ > 1 ? std::launch::async : std::launch::deferred; std::vector> running; @@ -475,23 +496,25 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, const gemmi::AsuData> &fobs, double d_min, Logger &logger, - size_t nthreads) { + size_t nthreads, + RigidBodyGPUPool *gpu) { const auto t0 = std::chrono::steady_clock::now(); RigidBodyRefineResult result; if (fobs.v.empty()) return result; - RigidBodyTarget target(model, cell, sg, nthreads); + std::unique_ptr target_backend; +#ifdef JFJOCH_USE_CUDA + if (gpu != nullptr) + target_backend = std::make_unique(*gpu, model, cell, sg, nthreads); +#endif + if (!target_backend) + target_backend = std::make_unique(model, cell, sg, nthreads); + RigidBodyTargetBase &target = *target_backend; if (!target.Usable()) return result; - std::vector ladder; - for (double zone : LADDER) - if (zone >= d_min) - ladder.push_back(zone); - if (ladder.empty()) - ladder.push_back(d_min); - - const gemmi::Mat33 gauge = GaugeProjector(sg, cell); + const std::vector ladder = RigidBodyLadder(d_min); + const gemmi::Mat33 gauge = RigidBodyGaugeProjector(sg, cell); double q[6] = {0, 0, 0, 0, 0, 0}; bool any_zone_solved = false; for (double zone : ladder) { @@ -554,9 +577,9 @@ RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, if (!result.zones.empty()) logger.Info("Model validation: rigid body over {} resolution zone(s) down to {:.1f} A, " - "{} evaluations and {} Jacobians in {:.2f} s: rotation {:.3f} deg, translation {:.3f} A", + "{} evaluations and {} Jacobians in {:.3f} s: rotation {:.3f} deg, translation {:.3f} A{}", result.zones.size(), result.zones.back(), result.evaluations, result.jacobians, - result.seconds, result.angle_deg, result.shift_A); + result.seconds, result.angle_deg, result.shift_A, gpu != nullptr ? " (GPU)" : ""); else logger.Info("Model validation: rigid body had no resolution zone with enough reflections to " "run in; the model is left where it arrived"); diff --git a/rugnux/RigidBodyRefine.h b/rugnux/RigidBodyRefine.h index 13317df97..bbccc538f 100644 --- a/rugnux/RigidBodyRefine.h +++ b/rugnux/RigidBodyRefine.h @@ -50,15 +50,33 @@ double RigidBodyReachDeg(const gemmi::Model &model, double d_min); // `fobs` should be the working set only - the free reflections are what the caller decides on. The // model is left MOVED (whether or not the refinement helped): the caller scores it and puts it back // with SetModelPositions() if it did not. +// +// `gpu`: evaluate the target on an engine of this pool (RigidBodyGPU.h, CUDA builds only) rather than +// on the CPU. The same target either way; null is the CPU. +class RigidBodyGPUPool; RigidBodyRefineResult RefineRigidBody(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, const gemmi::AsuData> &fobs, double d_min, Logger &logger, - size_t nthreads = 1); + size_t nthreads = 1, + RigidBodyGPUPool *gpu = nullptr); -// --- The pieces RefineRigidBody is built from, exposed for the tests --- +// --- The pieces RefineRigidBody is built from, exposed for the tests and the GPU backend --- + +// The resolution zones RefineRigidBody walks down to d_min, coarsest first, before a zone too thin to +// run in is dropped. +std::vector RigidBodyLadder(double d_min); + +// A zone's empty grid: its size, cell and group, as DensityCalculator sizes its own for the zone. +gemmi::Grid RigidBodyZoneGrid(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min); + +// The step of the forward-difference rotation columns of a zone's Jacobian, in angstroms of q. +double RigidBodyJacobianStep(double d_min); + +// The projector onto the directions in which the origin of `sg` is free (see RefineRigidBody). +gemmi::Mat33 RigidBodyGaugeProjector(const gemmi::SpaceGroup &sg, const gemmi::UnitCell &cell); // Fcalc of the whole crystal composed from the transform of ONE copy of its content, gridded on the // zone's grid without symmetrization: F(h) = n_cen * sum over the group's operators x -> Rx + t of @@ -74,9 +92,22 @@ public: // and not 000 - so an observation it does not match stays unmatched, as it was. SymmetryComposition(const gemmi::Grid &grid, double d_min, const std::vector &hkl); + struct Term { + size_t index; // of F1(hR) in the half-l grid, or of its Friedel mate + bool conj; // read as the Friedel mate: conj F1(-hR) + std::complex phase; // exp(+2 pi i h.t) + gemmi::Vec3 s; // hR as a Cartesian reciprocal vector + gemmi::Miller k; // hR + }; + // For each of `hkl`, its position among the composed indices, or -1. const std::vector &Row() const { return row_; } const std::vector &Hkl() const { return hkl_; } + // Ops() terms per composed index, index-major, and each index's 1/d^2 and n_cen. + const std::vector &Terms() const { return terms_; } + size_t Ops() const { return ops_; } + const std::vector &InvD2() const { return inv_d2_; } + double Centring() const { return centring_; } // F at the composed indices from `f1` = MapToFPhi() of the unsymmetrized grid, times the unblur // factor prepare_asu_data() applies; with `df_dt`, also dF/dt for a Cartesian translation t. @@ -84,12 +115,6 @@ public: std::vector, 3>> *df_dt, size_t nthreads) const; private: - struct Term { - size_t index; // of F1(hR) in the half-l grid, or of its Friedel mate - bool conj; // read as the Friedel mate: conj F1(-hR) - std::complex phase; // exp(+2 pi i h.t) - gemmi::Vec3 s; // hR as a Cartesian reciprocal vector - }; std::vector row_; std::vector hkl_; std::vector inv_d2_; @@ -100,24 +125,25 @@ private: // The target RefineRigidBody minimises over one resolution zone: the amplitude residuals of the model // placed at q (the angle-axis rotation about the centroid times the rms radius, then the translation, -// all in angstroms), with the scale re-fitted at every evaluation. -class RigidBodyTarget { +// all in angstroms), with the scale re-fitted at every evaluation. The placement and the counters are +// common; what evaluates the residuals and the Jacobian is the backend: RigidBodyTarget below on the +// CPU, RigidBodyTargetGPU (RigidBodyGPU.h) on a GPU. Both compute the same function. +class RigidBodyTargetBase { public: - // q = 0 is the placement `model` has now. Every evaluation moves `model`. - RigidBodyTarget(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, - size_t nthreads); + explicit RigidBodyTargetBase(const gemmi::Model &model); + virtual ~RigidBodyTargetBase() = default; bool Usable() const { return !base_.empty() && rms_radius_ > 0; } double RmsRadius() const { return rms_radius_; } void Place(const double q[6], gemmi::Model &model) const; - void SetZone(const gemmi::AsuData> &fobs, double d_min); - size_t NumObservations() const { return fobs_.v.size(); } + virtual void SetZone(const gemmi::AsuData> &fobs, double d_min) = 0; + virtual size_t NumObservations() const = 0; - bool Residuals(const double q[6], double *residuals); + virtual bool Residuals(const double q[6], double *residuals) = 0; // d residuals / d q at q, NumObservations() x 6 row-major, with the scale re-fit folded in by // projection and the bulk-solvent mask held. - bool Jacobian(const double q[6], double *jacobian); + virtual bool Jacobian(const double q[6], double *jacobian) = 0; bool hold_mask = false; // for the tests: keep the last evaluation's bulk-solvent mask int evaluations = 0; @@ -125,6 +151,25 @@ public: int unmatched = 0; // zone observations with no calculated amplitude to compare against double k_sol = 0, b_sol = 0; // the bulk solvent the zone's target was evaluated with +protected: + std::vector base_; + gemmi::Position centre_; + double rms_radius_ = 0; +}; + +// The CPU backend. +class RigidBodyTarget : public RigidBodyTargetBase { +public: + // q = 0 is the placement `model` has now. Every evaluation moves `model`. + RigidBodyTarget(gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + size_t nthreads); + + void SetZone(const gemmi::AsuData> &fobs, double d_min) override; + size_t NumObservations() const override { return fobs_.v.size(); } + + bool Residuals(const double q[6], double *residuals) override; + bool Jacobian(const double q[6], double *jacobian) override; + private: void CopyFcalc(const gemmi::Model &model, std::vector> &f, std::vector, 3>> *df_dt) const; @@ -133,9 +178,6 @@ private: const gemmi::UnitCell &cell_; const gemmi::SpaceGroup &sg_; size_t nthreads_; - std::vector base_; - gemmi::Position centre_; - double rms_radius_ = 0; std::vector column_models_; // the rotation columns' own copies gemmi::AsuData> fobs_; diff --git a/tests/ModelValidationTest.cpp b/tests/ModelValidationTest.cpp index 531d55dac..535b63aec 100644 --- a/tests/ModelValidationTest.cpp +++ b/tests/ModelValidationTest.cpp @@ -21,6 +21,10 @@ #include "../rugnux/ModelGrid.h" #include "../rugnux/ModelValidation.h" #include "../rugnux/RigidBodyRefine.h" +#ifdef JFJOCH_USE_CUDA +#include "../rugnux/RigidBodyGPU.h" +#include "../common/CUDAWrapper.h" +#endif #include "../rugnux/SigmaA.h" #include "../rugnux/WriteModel.h" #include "../image_analysis/scale_merge/ReindexAmbiguity.h" @@ -1027,8 +1031,34 @@ namespace { // it. Per column (those in `columns`): the cosine between the two and the ratio of their norms. // Also logged, not checked: how far the difference moves once the mask is let move with the body, // which is what holding it costs. + // The model's own amplitudes to `d_min`: "observed" data whose minimum is where the model is. + gemmi::AsuData> OwnAmplitudes(const char *cryst, const gemmi::Structure &st, + double d_min, Logger &logger) { + const auto path = WriteTemp("rigid_body_own_amplitudes.pdb", ClusterPdb(cryst).c_str()); + const auto ref = ModelReferenceIntensities(path, {}, {}, 3.0, logger); + std::filesystem::remove(path); + gemmi::AsuData> fobs; + fobs.unit_cell_ = st.cell; + fobs.spacegroup_ = st.find_spacegroup(); + for (const auto &r : ref) + if (st.cell.calculate_d({{r.h, r.k, r.l}}) >= d_min) + fobs.v.push_back({{{r.h, r.k, r.l}}, {std::sqrt(r.I), 1.0f}}); + fobs.ensure_sorted(); + return fobs; + } + + // The rigid body's target on the CPU, or on an engine of `gpu` where it is given. + std::unique_ptr MakeTarget(gemmi::Model &model, const gemmi::Structure &st, + RigidBodyGPUPool *gpu) { +#ifdef JFJOCH_USE_CUDA + if (gpu != nullptr) + return std::make_unique(*gpu, model, st.cell, *st.find_spacegroup(), 4); +#endif + return std::make_unique(model, st.cell, *st.find_spacegroup(), 4); + } + void CheckRigidBodyJacobian(const char *cryst, const double q0[6], const std::vector &columns, - double min_cosine, double max_norm_error) { + double min_cosine, double max_norm_error, bool gpu = false) { Logger logger("CheckRigidBodyJacobian"); const gemmi::Structure st = AnisoCluster(cryst); const gemmi::SpaceGroup *sg = st.find_spacegroup(); @@ -1048,7 +1078,15 @@ namespace { fobs.ensure_sorted(); gemmi::Model model = st.models[0]; - RigidBodyTarget target(model, st.cell, *sg, 4); + std::unique_ptr pool; +#ifdef JFJOCH_USE_CUDA + if (gpu) { + pool = RigidBodyGPUPool::Create(model, st.cell, *sg, zone, fobs.v.size(), 1, logger); + REQUIRE(pool); + } +#endif + const std::unique_ptr target_backend = MakeTarget(model, st, pool.get()); + RigidBodyTargetBase &target = *target_backend; target.SetZone(fobs, zone); const size_t n = target.NumObservations(); std::vector r(n), jacobian(n * 6); @@ -1190,3 +1228,150 @@ TEST_CASE("ModelValidation_RigidBodyReportsTheRotationItApplied", "[ModelValidat CHECK(std::acos(std::clamp((trace - 1) / 2, -1.0, 1.0)) * 180 / PI == Catch::Approx(result.angle_deg).margin(1e-6)); } + +#ifdef JFJOCH_USE_CUDA + +// The GPU target is the CPU target's function, computed on the device: at the same placement the two +// give the same residuals and the same Jacobian to rounding (float distances on the device, cuFFT for +// FFTW), on every group of the composition tests, both zones, anisotropic atoms included. +TEST_CASE("RigidBodyGPU_MatchesCPU", "[ModelValidation][gpu]") { + if (get_gpu_count() == 0) { + SUCCEED("no GPU"); + return; + } + Logger logger("RigidBodyGPU_MatchesCPU"); + const double q0[6] = {0.12, -0.09, 0.07, 0.15, -0.10, 0.05}; + for (const char *cryst : kRigidBodyCrysts) { + const gemmi::Structure st = AnisoCluster(cryst); + const gemmi::SpaceGroup *sg = st.find_spacegroup(); + for (double zone : {6.0, 3.5}) { + const auto fobs = OwnAmplitudes(cryst, st, zone, logger); + gemmi::Model cpu_model = st.models[0], gpu_model = st.models[0]; + auto pool = RigidBodyGPUPool::Create(gpu_model, st.cell, *sg, zone, fobs.v.size(), 1, logger); + REQUIRE(pool); + RigidBodyTarget cpu(cpu_model, st.cell, *sg, 4); + RigidBodyTargetGPU gpu(*pool, gpu_model, st.cell, *sg, 4); + cpu.SetZone(fobs, zone); + gpu.SetZone(fobs, zone); + const size_t n = cpu.NumObservations(); + REQUIRE(gpu.NumObservations() == n); + std::vector rc(n), rg(n), jc(6 * n), jg(6 * n); + REQUIRE(cpu.Residuals(q0, rc.data())); + REQUIRE(gpu.Residuals(q0, rg.data())); + REQUIRE(cpu.Jacobian(q0, jc.data())); + REQUIRE(gpu.Jacobian(q0, jg.data())); + CHECK(gpu.unmatched == cpu.unmatched); + CHECK(gpu.k_sol == cpu.k_sol); + CHECK(gpu.b_sol == cpu.b_sol); + + double worst = 0, rms = 0; + for (size_t i = 0; i < n; i++) { + worst = std::max(worst, std::fabs(rc[i] - rg[i])); + rms += rc[i] * rc[i] / static_cast(n); + } + INFO(cryst << " at " << zone << " A: residuals differ by at most " << worst << ", rms residual " + << std::sqrt(rms)); + CHECK(worst <= 1e-4); + for (int j = 0; j < 6; j++) { + double diff = 0, norm = 0; + for (size_t i = 0; i < n; i++) { + diff += gemmi::sq(jc[6 * i + j] - jg[6 * i + j]); + norm += gemmi::sq(jc[6 * i + j]); + } + INFO(cryst << " at " << zone << " A, column " << j << ": relative difference " + << std::sqrt(diff / norm)); + CHECK(std::sqrt(diff) <= 2e-3 * std::sqrt(norm) + 1e-9); + } + } + } +} + +// The whole Jacobian on the GPU against the central difference of the GPU's own residuals, with the +// bounds of ModelValidation_RigidBodyJacobianMatchesNumericScaleRefit. +TEST_CASE("RigidBodyGPU_JacobianMatchesNumericScaleRefit", "[ModelValidation][gpu]") { + if (get_gpu_count() == 0) { + SUCCEED("no GPU"); + return; + } + const double close[6] = {0.0, 0.0, 0.0, 0.05, -0.04, 0.03}; + const double rotated[6] = {0.12, -0.09, 0.07, 0.15, -0.10, 0.05}; + for (const char *cryst : kRigidBodyCrysts) { + CheckRigidBodyJacobian(cryst, close, {0, 1, 2, 3, 4, 5}, 0.99, 0.02, true); + CheckRigidBodyJacobian(cryst, rotated, {0, 1, 2}, 0.99, 0.03, true); + } +} + +namespace { + // A whole placement on the GPU from the model displaced by 0.54 A and turned 2 deg, against the + // model's own amplitudes; the model is left where the fit put it. + RigidBodyRefineResult DisplacedFit(const char *cryst, gemmi::Structure &st, RigidBodyGPUPool *pool, + Logger &logger) { + const gemmi::SpaceGroup *sg = st.find_spacegroup(); + const auto fobs = OwnAmplitudes(cryst, st, 3.0, logger); + std::vector moved = ModelPositions(st.models[0]); + gemmi::Vec3 centre; + for (const gemmi::Position &p : moved) + centre += p; + centre *= 1.0 / static_cast(moved.size()); + const double a = 2.0 * PI / 180.0; + const gemmi::Mat33 rz(std::cos(a), -std::sin(a), 0, std::sin(a), std::cos(a), 0, 0, 0, 1); + for (gemmi::Position &p : moved) + p = gemmi::Position(rz.multiply(gemmi::Vec3(p) - centre) + centre + gemmi::Vec3(0.40, -0.30, 0.20)); + SetModelPositions(st.models[0], moved); + return RefineRigidBody(st.models[0], st.cell, *sg, fobs, 3.0, logger, 4, pool); + } +} + +// The GPU fit walks where the CPU fit walks: from the same displaced start, to the same placement +// within a few thousandths of an angstrom. +TEST_CASE("RigidBodyGPU_FitAgreesWithCPU", "[ModelValidation][gpu]") { + if (get_gpu_count() == 0) { + SUCCEED("no GPU"); + return; + } + Logger logger("RigidBodyGPU_FitAgreesWithCPU"); + for (const char *cryst : {kCryst, kPolarCryst, kRigidBodyCrysts[3]}) { + gemmi::Structure cpu_st = AnisoCluster(cryst), gpu_st = AnisoCluster(cryst); + auto pool = RigidBodyGPUPool::Create(gpu_st.models[0], gpu_st.cell, *gpu_st.find_spacegroup(), 3.0, 100000, 1, + logger); + REQUIRE(pool); + const RigidBodyRefineResult cpu = DisplacedFit(cryst, cpu_st, nullptr, logger); + const RigidBodyRefineResult gpu = DisplacedFit(cryst, gpu_st, pool.get(), logger); + CHECK(gpu.converged == cpu.converged); + const std::vector pc = ModelPositions(cpu_st.models[0]), pg = ModelPositions(gpu_st.models[0]); + double rmsd = 0; + for (size_t i = 0; i < pc.size(); i++) + rmsd += pc[i].dist_sq(pg[i]) / static_cast(pc.size()); + INFO(cryst << ": CPU " << cpu.angle_deg << " deg " << cpu.shift_A << " A, GPU " << gpu.angle_deg << " deg " + << gpu.shift_A << " A, " << std::sqrt(rmsd) << " A apart"); + CHECK(std::sqrt(rmsd) < 2e-3); + } +} + +// Deterministic: the same fit twice, and on a pool of one engine and of four, gives the same placement +// bit for bit. +TEST_CASE("RigidBodyGPU_Deterministic", "[ModelValidation][gpu]") { + if (get_gpu_count() == 0) { + SUCCEED("no GPU"); + return; + } + Logger logger("RigidBodyGPU_Deterministic"); + const char *cryst = kRigidBodyCrysts[4]; + std::vector> placed; + for (size_t engines : {1, 1, 4}) { + gemmi::Structure st = AnisoCluster(cryst); + auto pool = RigidBodyGPUPool::Create(st.models[0], st.cell, *st.find_spacegroup(), 3.0, 100000, engines, logger); + REQUIRE(pool); + const RigidBodyRefineResult r = DisplacedFit(cryst, st, pool.get(), logger); + CHECK(r.evaluations > 0); + placed.push_back(ModelPositions(st.models[0])); + } + for (size_t k = 1; k < placed.size(); k++) + for (size_t i = 0; i < placed[0].size(); i++) { + CHECK(placed[k][i].x == placed[0][i].x); + CHECK(placed[k][i].y == placed[0][i].y); + CHECK(placed[k][i].z == placed[0][i].z); + } +} + +#endif