// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only #pragma once // The device half of restrained refinement on a GPU (RefinementGPU.h, CUDA builds only). The host half works // out with gemmi everything that depends on the cell, the group and the reflections; this half grids the atoms, // transforms, composes the symmetry, evaluates the likelihood target, takes its derivative back to every atom // by the FFT method, and evaluates the geometry restraints. Every kernel is deterministic - gathers and fixed // reduction trees, no floating-point atomics - so the same input gives the same bits on every run. #include #include #include #include #include #include "ModelDensityGPU.h" #include "ModelMaskGPU.h" #include "RefinementMath.h" // Everything an engine needs that does not depend on the atoms. struct RefinementGPUSetup { ModelDensityGrid grid; // the structure-factor grid, u fastest ModelDensityGrid mask_grid; // the bulk-solvent mask's (BulkSolventMask::MaskGrid), u fastest double volume = 0; std::vector> sym_ops; // gemmi's Op: rot row-major and tran, in units of 1 / den int den = 24; double centring = 1; // the number of centring vectors std::vector mask_ops; // every operator with every centring vector std::vector> rows; // the reflections, in the order their numbers come back std::vector inv_d2; std::vector fo; std::vector work, centric; // With a twin law: per row the row of its twin mate (RefinementReflections::mate), an involution; empty // without one. std::vector mate; // The adjoint of the composition, as a gather: the half-u spectrum points the gradient's coefficients // land on (ascending), and for each the terms (row * ops + op) * 2 + conj adding to it, in order. std::vector scatter_point, scatter_start, scatter_term; size_t max_atoms = 0, max_pairs = 0; }; struct RefinementGPUEngineImpl; // What the engines of one setup on one card share (RefinementGPU.cu): the reflections' tables, the bulk-solvent mask's // buffers, and lanes of the transforms' working buffers, which the engines of a lane use one at a time. struct RefinementGPUCard; class RefinementGPUEngine { public: // What a card holds for the engines on it together (MakeCard), and each lane of working buffers beside. static size_t CardBytes(const RefinementGPUSetup &setup); static size_t LaneBytes(const RefinementGPUSetup &setup, size_t fft_work_bytes); // What each engine holds beside, with restraint sets of up to `records` records and plane atoms and `slots` // gradient slots together. static size_t DeviceBytes(const RefinementGPUSetup &setup, size_t records, size_t slots); static size_t FFTWorkBytes(int device, const RefinementGPUSetup &setup); static std::shared_ptr MakeCard(int device, const RefinementGPUSetup &setup); // `setup` is read for as long as the engine lives; `card` is MakeCard(device, setup)'s. `own_lane`: on a lane of // working buffers of its own, else on the card's lane with the fewest engines. RefinementGPUEngine(int device, const RefinementGPUSetup &setup, std::shared_ptr card, size_t fft_work_bytes, bool own_lane); ~RefinementGPUEngine(); // n_cen exp(blur s^2 / 4) per reflection: prepare_asu_data()'s unblur. void SetBlur(double blur); // Per macro-cycle: F_model = k_total (F + fmask), the likelihood's D and variance, and with a twin law the twin // fraction (XrayRowModel). void SetRowModel(const std::vector &k_total, const std::vector> &fmask, const std::vector &d, const std::vector &var, float twin_fraction = 0); // The atoms' densities, as ModelStructureFactorsGPU takes them (the boxes of ModelDensityGPU::Boxes() made // from them): called again only where some atom's density changed. Read until the call returns; from pinned // memory they go up by DMA. void SetDensities(const ModelDensityAtom *atoms, size_t n); // The atoms' grid positions (fractional times the grid size, not wrapped), in the order of SetDensities(): // read until the next evaluation returns, so from a pinned buffer the copy goes up by DMA. void SetPositions(const std::array *pos, size_t n); // F of the atoms at the rows. void Fcalc(std::vector> &f); // The bulk-solvent mask's F at the rows (the mask atoms fractional, wrapped), on the mask's own grid. void Fmask(const std::vector &mask_atoms, std::vector> &f); // -ln L over the working rows for the atoms last set and, with `grad`, dT/d(x, y, z, B) per atom into // HostGradient(). double Evaluate(bool grad); // The gradients of the last evaluation, in the engine's pinned buffers: valid until the next one. const std::array *HostGradient() const; const std::array *HostRestraintGradient() const; // dT/d(U11 U22 U33 U12 U13 U23) of the anisotropic atoms (six per atom, zero for the isotropic ones) of the last // evaluation with a gradient where some atom was anisotropic. const float *HostUGradient() const; // Restraint set `which` (0 or 1), as RestraintSet holds it. void SetRestraints(int which, const std::vector &records, const std::vector &plane_atoms, const std::vector &slot_start, const std::vector &slot_index, int n_slots); // The sum of E of set `which` at `xyzb` (x, y, z, B per atom) and dE/d(x, y, z, B) per atom. double EvaluateRestraints(int which, const std::vector> &xyzb, std::vector> *grad); // Evaluate() and EvaluateRestraints() in one submission with one wait: the X-ray target is queued first, and // `xyzb` writes the restraints' x, y, z, B of every atom while the device works on it. The gradients into // HostGradient() and HostRestraintGradient(). double EvaluateWithRestraints(bool grad, int which, const std::function *)> &xyzb, bool restraint_grad, double &restraints); // The most atoms of a plane restraint the kernel takes. static constexpr int MAX_PLANE_ATOMS = 32; private: std::unique_ptr impl_; };