Jungfraujoch/image_analysis/structure_refinement/RigidBodyGPU.h

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// 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 <condition_variable>
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <stdexcept>
#include <string>
#include <tuple>
#include <vector>

#include "RigidBodyRefine.h"

class Logger;
class RigidBodyGPUEngine;
struct RigidBodyGPUZone;
struct RigidBodyGPUCapacity;

// A CUDA failure inside the rigid body, or in model validation's structure factors and maps
// (ModelStructureFactorsGPU). Model validation catches it and reports the validation as failed.
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 or search: each is a stream and the buffers for one
// fit at a time, sized for the finest zone of the ladder to d_min. Every fit takes one for its length, and
// waits for one when all are taken. Engines are interchangeable and every kernel deterministic, so which fit
// gets which engine does not change a number. A pool belongs to one cell and one d_min, and serves every
// model and group it was made for (Plan, Cover): what a zone needs of a model's atoms - their densities,
// radii and mask radii, which the placement does not change - is worked out once per group, atom count and
// zone and kept.
class RigidBodyGPUPool {
public:
    // A pool without engines yet, sized for the finest zone of the ladder to d_min: EngineBytes() is
    // what each engine will take, so a caller can decide how many first. Null where there is no GPU or
    // where the cell is too small for the gather (an atom's box wider than the cell), logged.
    // `max_observations`: the most working reflections any fit will be given in its finest zone.
    static std::unique_ptr<RigidBodyGPUPool> Plan(const gemmi::Model &model, const gemmi::UnitCell &cell,
                                                  const gemmi::SpaceGroup &sg, double d_min,
                                                  size_t max_observations, Logger &logger);
    ~RigidBodyGPUPool();

    // Sizes the engines, before Reserve(), for fits of `model` in `sg` as well - in the pool's cell, to its
    // d_min, with at most its observations. False, logged, where the cell is too small for the gather.
    bool Cover(const gemmi::Model &model, const gemmi::SpaceGroup &sg, Logger &logger);

    size_t EngineBytes() const;
    // Reserves `engines` engines, engine i on the i-th card from Device(), so a validation uses every
    // card there is. Engines that cannot be set up throw RigidBodyGPUFailure.
    void Reserve(size_t engines);
    size_t Engines() const { return engines_.size(); }
    // The card the pool was made from, which its first engine is on and the others count on from.
    int Device() const { return device_; }
    RigidBodyGPUEngine &Acquire();
    void Release(RigidBodyGPUEngine &engine);

    // The zone to d_min without its observations, for this pool's model.
    const RigidBodyGPUZone &Zone(const gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
                                 double d_min);
    // The zone with these observations and the composition of Fcalc at them. The null's replicates are all
    // fitted to the same reflections, so they share it.
    std::shared_ptr<const RigidBodyGPUZone> ObservedZone(const gemmi::Model &model, const gemmi::UnitCell &cell,
                                                         const gemmi::SpaceGroup &sg,
                                                         const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs,
                                                         double d_min);

private:
    RigidBodyGPUPool() = default;
    int device_ = 0;
    gemmi::UnitCell cell_;
    double d_min_ = 0;
    size_t max_observations_ = 0;
    std::unique_ptr<RigidBodyGPUCapacity> cap_;
    std::vector<std::unique_ptr<RigidBodyGPUEngine>> engines_;
    std::vector<RigidBodyGPUEngine *> idle_;
    std::mutex m_;
    std::condition_variable cv_;
    // By group, atom count and zone.
    std::map<std::tuple<std::string, size_t, double>, std::unique_ptr<RigidBodyGPUZone>> zones_;
    std::mutex zones_m_;
    struct ObservedEntry {
        std::string group;
        size_t atoms;
        double d_min;
        gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
        std::shared_ptr<const RigidBodyGPUZone> zone;
    };
    std::vector<ObservedEntry> observed_;
    std::mutex observed_m_;
};

// An engine taken from a pool for as long as this lives.
class RigidBodyGPULease {
public:
    explicit RigidBodyGPULease(RigidBodyGPUPool &pool) : pool_(pool), engine_(pool.Acquire()) {}
    ~RigidBodyGPULease() { pool_.Release(engine_); }
    RigidBodyGPULease(const RigidBodyGPULease &) = delete;
    RigidBodyGPULease &operator=(const RigidBodyGPULease &) = delete;
    RigidBodyGPUEngine &Engine() { return engine_; }

private:
    RigidBodyGPUPool &pool_;
    RigidBodyGPUEngine &engine_;
};

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 never moved by an evaluation, only 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<gemmi::ValueSigma<float>> &fobs, double d_min) override;
    size_t NumObservations() const override;

    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;
    void HostScale();

    RigidBodyGPUPool &pool_;
    RigidBodyGPULease lease_;
    RigidBodyGPUEngine &engine_;
    gemmi::Model &model_;
    const gemmi::UnitCell &cell_;
    const gemmi::SpaceGroup &sg_;
    size_t nthreads_;

    double d_min_ = 0;
    std::shared_ptr<const RigidBodyGPUZone> zone_;
    bool solvent_fitted_ = false;
    bool host_scale_ = false;             // the zone's scale is fitted on the host (HostScale)

    bool have_point_ = false;
    std::array<double, 6> q_{};
    double k_overall_ = 1;
    gemmi::SMat33<double> b_star_{0, 0, 0, 0, 0, 0};
};