Jungfraujoch/image_analysis/structure_refinement/RefinementGPUEngine.h

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// 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 <array>
#include <cstddef>
#include <functional>
#include <memory>
#include <vector>

#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<std::array<int, 12>> 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<ModelMaskOp> mask_ops;         // every operator with every centring vector
    std::vector<std::array<int, 3>> rows;      // the reflections, in the order their numbers come back
    std::vector<double> inv_d2;
    std::vector<float> fo;
    std::vector<char> work, centric;
    // With twin laws: per law and row (law * rows + row) the row of its twin mate under the law, and the observed
    // row whose mate it is (-1: none) - RefinementReflections::mates and mates_from; empty without a law.
    int laws = 0;
    std::vector<int> mate, mate_from;
    // 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<int> 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<RefinementGPUCard> 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<RefinementGPUCard> 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 twin laws their
    // fractions (XrayRowModel).
    void SetRowModel(const std::vector<float> &k_total, const std::vector<std::array<float, 2>> &fmask,
                     const std::vector<float> &d, const std::vector<float> &var,
                     const std::vector<float> &twin_fractions = {});

    // 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<float, 4> *pos, size_t n);
    // F of the atoms at the rows.
    void Fcalc(std::vector<std::array<float, 2>> &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<ModelMaskAtom> &mask_atoms, std::vector<std::array<float, 2>> &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<float, 4> *HostGradient() const;
    const std::array<float, 4> *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<RestraintRecord> &records, const std::vector<int> &plane_atoms,
                       const std::vector<int> &slot_start, const std::vector<int> &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<std::array<float, 4>> &xyzb,
                              std::vector<std::array<float, 4>> *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<void(std::array<float, 4> *)> &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<RefinementGPUEngineImpl> impl_;
};