Jungfraujoch/image_analysis/structure_refinement/RefinementGPU.cu

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "RefinementGPUEngine.h"

#include <algorithm>
#include <string>

#include <cub/cub.cuh>
#include <cufft.h>

#include "../../common/JFJochException.h"
#include "../indexing/CUDAMemHelpers.h"

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<int>(val)));
}

constexpr int THREADS = 256;
constexpr int WARP = 32;
constexpr double FOUR_PI2 = 4 * refine_math::PI * refine_math::PI;

int blocks(size_t n) {
    return static_cast<int>((n + THREADS - 1) / THREADS);
}

__device__ __forceinline__ int imod(int a, int n) {
    const int r = a % n;
    return r < 0 ? r + n : r;
}

// gemmi's unsafe_expapprox(), as ModelDensityGPU.cu has it.
__device__ __forceinline__ float expapprox(float x) {
    const float val = 12102203.1615614f * x + 1065353216.f;
    const int vali = static_cast<int>(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))));
}

// exp(+2 pi i h.t) of row h and operator o, exact for the fractions a translation is made of.
__device__ __forceinline__ void op_phase(const int *h, const int *tran, int den, double &c, double &s) {
    sincospi(2.0 * (h[0] * tran[0] + h[1] * tran[1] + h[2] * tran[2]) / den, &s, &c);
}

// F of the atoms at each row, composed from the half-u transform of one copy (ModelStructureFactorsGPU.cu's
// compose_kernel), and with `evaluate` the target's term of a working row and the coefficient
// conj(dT/dF) n_cen unblur the gradient's map is made of.
__global__ void compose_target_kernel(const float2 *f1, float norm, int nu, int nv, int nw, const int *ops, int n_ops,
                                      int den, const int *row_hkl, const double *row_scale, int rows, bool evaluate,
                                      const float *k_total, const float2 *fmask, const float *ml_d, const float *ml_var,
                                      const float *fo, const char *work, const char *centric, float2 *f, double *nll,
                                      float2 *coef) {
    const int m = blockIdx.x * blockDim.x + threadIdx.x;
    if (m >= rows)
        return;
    const int h[3] = {row_hkl[3 * m], row_hkl[3 * m + 1], row_hkl[3 * m + 2]};
    double sr = 0, si = 0;
    for (int o = 0; o < n_ops; o++) {
        const int *rot = ops + 12 * o, *tran = rot + 9;
        int k[3];
        for (int i = 0; i < 3; i++)
            k[i] = (rot[i] * h[0] + rot[3 + i] * h[1] + rot[6 + i] * h[2]) / den;
        const bool conj = k[0] < 0;
        if (conj)
            for (int i = 0; i < 3; i++)
                k[i] = -k[i];
        const float2 out = f1[k[0] + static_cast<size_t>(nu / 2 + 1) * (imod(k[1], nv) + static_cast<size_t>(nv) * imod(k[2], nw))];
        const double vr = out.x * norm, vi = (conj ? out.y : -out.y) * norm;
        double s, c;
        op_phase(h, tran, den, c, s);
        sr += c * vr - s * vi;
        si += c * vi + s * vr;
    }
    const double scale = row_scale[m];
    sr *= scale;
    si *= scale;
    f[m] = make_float2(static_cast<float>(sr), static_cast<float>(si));
    if (!evaluate)
        return;
    double term = 0, cr = 0, ci = 0;
    if (work[m]) {
        const double k = k_total[m];
        const double mr = k * (sr + fmask[m].x), mi = k * (si + fmask[m].y);
        const double fc = sqrt(mr * mr + mi * mi);
        double dfc = 0;
        term = refine_math::AmplitudeNegLogLikelihood<double>(fo[m], fc, ml_d[m], ml_var[m], centric[m], dfc);
        if (fc > 0) {
            // dT/dF = k dT/d|F_model| F_model / |F_model|; the coefficient is its conjugate times the scale
            const double g = k * dfc / fc * scale;
            cr = g * mr;
            ci = -g * mi;
        }
    }
    nll[m] = term;
    coef[m] = make_float2(static_cast<float>(cr), static_cast<float>(ci));
}

// The gradient's coefficients on the half-u spectrum: each point the sum, in order, of the terms
// (coef[row] exp(2 pi i h.t_op)) / 2 that land on it, conjugated where they land as a Friedel mate.
__global__ void scatter_kernel(const float2 *coef, const int *row_hkl, const int *ops, int n_ops, int den,
                               const int *point, const int *start, const int *term, int points, float2 *spectrum) {
    const int s = blockIdx.x * blockDim.x + threadIdx.x;
    if (s >= points)
        return;
    double sr = 0, si = 0;
    for (int e = start[s]; e < start[s + 1]; e++) {
        const int t = term[e];
        const bool conj = t & 1;
        const int m = (t >> 1) / n_ops, o = (t >> 1) % n_ops;
        double c, sn;
        op_phase(row_hkl + 3 * m, ops + 12 * o + 9, den, c, sn);
        const double ar = 0.5 * (coef[m].x * c - coef[m].y * sn), ai = 0.5 * (coef[m].x * sn + coef[m].y * c);
        sr += ar;
        si += conj ? -ai : ai;
    }
    spectrum[point[s]] = make_float2(static_cast<float>(sr), static_cast<float>(si));
}

// F_mask at each row, read at h directly (ModelStructureFactorsGPU.cu's fmask_kernel).
__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<size_t>(nu / 2 + 1) * (imod(k, nv) + static_cast<size_t>(nv) * imod(l, nw))];
    const float re = out.x * norm, im = -out.y * norm;
    fmask[m] = make_float2(re, conj ? -im : im);
}

// dT/d(x, y, z, B) of every atom: the derivative map summed against the derivatives of the atom's density over
// its box, one warp per atom, the lanes' sums added in a fixed tree. The box, the points within the radius and
// their offsets from the centre are the density gather's (ModelDensityGPU.cu).
__global__ void atom_gradient_kernel(const ModelDensityAtom *atoms, const float4 *pos, const int *box, int n,
                                     ModelDensityGPU::Geometry g, const float *map, float map_scale, float4 *grad) {
    const int a = (blockIdx.x * blockDim.x + threadIdx.x) / WARP;
    const int lane = threadIdx.x % WARP;
    if (a >= n)
        return;
    const ModelDensityAtom at = atoms[a];
    const float4 p = pos[a];
    const float nctr[3] = {p.x, p.y, p.z};
    int lo[3], span[3];
    for (int k = 0; k < 3; k++) {
        lo[k] = __float2int_rn(nctr[k]) - box[3 * a + k];
        span[k] = 2 * box[3 * a + k] + 1;
    }
    const int total = span[0] * span[1] * span[2];
    float gx = 0, gy = 0, gz = 0, gb = 0;
    for (int t = lane; t < total; t += WARP) {
        const int u = lo[0] + t % span[0], v = lo[1] + t / span[0] % span[1], w = lo[2] + t / (span[0] * span[1]);
        const float fx = nctr[0] - static_cast<float>(u);
        const float fy = nctr[1] - static_cast<float>(v);
        const float fz = nctr[2] - static_cast<float>(w);
        const float z = g.m22 * fz;
        const float y = g.m11 * fy + g.m12 * fz;
        const float x = g.m00 * fx + (g.m01 * fy + g.m02 * fz);
        const float r2 = (y * y + z * z) + x * x;
        if (r2 > at.max_r2)
            continue;
        const float value = map[imod(u, g.nu) + static_cast<size_t>(g.nu) * (imod(v, g.nv) + static_cast<size_t>(g.nv) * imod(w, g.nw))];
        float dr2 = 0, db = 0;
        for (int q = 0; q < 5; q++) {
            const float b = at.b[q][0];
            const float e = at.a[q] * expapprox(fmaxf(b * r2, -88.f));
            dr2 += e * b;
            db += e * b * (1.5f + b * r2);
        }
        gx += value * dr2 * x;
        gy += value * dr2 * y;
        gz += value * dr2 * z;
        gb += value * db;
    }
    for (int off = WARP / 2; off > 0; off /= 2) {
        gx += __shfl_down_sync(0xffffffff, gx, off);
        gy += __shfl_down_sync(0xffffffff, gy, off);
        gz += __shfl_down_sync(0xffffffff, gz, off);
        gb += __shfl_down_sync(0xffffffff, gb, off);
    }
    if (lane == 0) {
        const float s = 2.0f * at.occ * map_scale;
        grad[a] = make_float4(s * gx, s * gy, s * gz, at.occ * map_scale * gb / static_cast<float>(FOUR_PI2));
    }
}

using V3f = refine_math::V3<float>;

__device__ V3f xyz(const float4 *p, int i) {
    const float4 q = p[i];
    return {q.x, q.y, q.z};
}

// Every restraint's E, and its gradient into its slots, one thread per restraint.
__global__ void restraint_kernel(const RestraintRecord *records, int n, const int *plane_atoms, const float4 *xyzb,
                                 float4 *contrib, double *energy) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n)
        return;
    const RestraintRecord &r = records[i];
    V3f g[4] = {};
    int na = 0;
    float e = 0;
    switch (static_cast<RestraintKind>(r.kind)) {
        case RestraintKind::Bond:
            e = refine_math::BondTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), r.ideal, r.sigma, g);
            na = 2;
            break;
        case RestraintKind::Angle:
            e = refine_math::AngleTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), xyz(xyzb, r.atom[2]), r.ideal,
                                       r.sigma, g);
            na = 3;
            break;
        case RestraintKind::Torsion:
            e = refine_math::TorsionTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), xyz(xyzb, r.atom[2]),
                                         xyz(xyzb, r.atom[3]), r.ideal, r.sigma, r.period, g);
            na = 4;
            break;
        case RestraintKind::Chiral:
            e = refine_math::ChiralTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), xyz(xyzb, r.atom[2]),
                                        xyz(xyzb, r.atom[3]), r.ideal, r.sigma, g);
            na = 4;
            break;
        case RestraintKind::Repulsion: {
            const V3f q = xyz(xyzb, r.atom[1]);
            const V3f image{r.rot[0] * q.x + r.rot[1] * q.y + r.rot[2] * q.z + r.tran[0],
                            r.rot[3] * q.x + r.rot[4] * q.y + r.rot[5] * q.z + r.tran[1],
                            r.rot[6] * q.x + r.rot[7] * q.y + r.rot[8] * q.z + r.tran[2]};
            e = refine_math::RepulsionTerm(xyz(xyzb, r.atom[0]), image, r.ideal, r.sigma, g);
            const V3f gi = g[1];
            g[1] = {r.rot[0] * gi.x + r.rot[3] * gi.y + r.rot[6] * gi.z,
                    r.rot[1] * gi.x + r.rot[4] * gi.y + r.rot[7] * gi.z,
                    r.rot[2] * gi.x + r.rot[5] * gi.y + r.rot[8] * gi.z};
            na = 2;
            break;
        }
        case RestraintKind::Plane: {
            V3f px[RefinementGPUEngine::MAX_PLANE_ATOMS], pg[RefinementGPUEngine::MAX_PLANE_ATOMS];
            const int m = r.atom[1];
            for (int k = 0; k < m; k++)
                px[k] = xyz(xyzb, plane_atoms[r.atom[0] + k]);
            e = refine_math::PlaneTerm(px, m, r.sigma, pg);
            for (int k = 0; k < m; k++)
                contrib[r.slot + k] = make_float4(pg[k].x, pg[k].y, pg[k].z, 0.0f);
            break;
        }
        case RestraintKind::BSimilarity: {
            const float rr = (xyzb[r.atom[0]].w - xyzb[r.atom[1]].w) / r.sigma;
            e = rr * rr;
            contrib[r.slot] = make_float4(0, 0, 0, 2 * rr / r.sigma);
            contrib[r.slot + 1] = make_float4(0, 0, 0, -2 * rr / r.sigma);
            break;
        }
    }
    for (int k = 0; k < na; k++)
        contrib[r.slot + k] = make_float4(g[k].x, g[k].y, g[k].z, 0.0f);
    energy[i] = e;
}

// Each atom's restraint gradient: its slots summed in order.
__global__ void slot_gather_kernel(const float4 *contrib, const int *start, const int *index, int n, float4 *grad) {
    const int a = blockIdx.x * blockDim.x + threadIdx.x;
    if (a >= n)
        return;
    float4 s = make_float4(0, 0, 0, 0);
    for (int k = start[a]; k < start[a + 1]; k++) {
        const float4 c = contrib[index[k]];
        s.x += c.x;
        s.y += c.y;
        s.z += c.z;
        s.w += c.w;
    }
    grad[a] = s;
}

CudaFFTShape ForwardShape(const ModelDensityGrid &g) {
    return {.n = {g.nw, g.nv, g.nu}, .idist = g.nu * g.nv * g.nw, .odist = (g.nu / 2 + 1) * g.nv * g.nw,
            .type = CUFFT_R2C};
}

CudaFFTShape InverseShape(const ModelDensityGrid &g) {
    return {.n = {g.nw, g.nv, g.nu}, .idist = (g.nu / 2 + 1) * g.nv * g.nw, .odist = g.nu * g.nv * g.nw,
            .type = CUFFT_C2R};
}

class DeviceGuard {
public:
    explicit DeviceGuard(int device) {
        cuda_err(cudaGetDevice(&previous_));
        cuda_err(cudaSetDevice(device));
    }
    ~DeviceGuard() { cudaSetDevice(previous_); }
    DeviceGuard(const DeviceGuard &) = delete;
    DeviceGuard &operator=(const DeviceGuard &) = delete;

private:
    int previous_ = 0;
};

size_t GridPoints(const ModelDensityGrid &g) {
    return static_cast<size_t>(g.nu) * g.nv * g.nw;
}

size_t ComplexPoints(const ModelDensityGrid &g) {
    return static_cast<size_t>(g.nu / 2 + 1) * g.nv * g.nw;
}

template <typename T>
void Upload(CudaDevicePtr<T> &dst, const std::vector<T> &src, cudaStream_t stream) {
    if (!src.empty())
        cuda_err(cudaMemcpyAsync(dst, src.data(), src.size() * sizeof(T), cudaMemcpyHostToDevice, stream));
}

// A device buffer of at least n elements, kept between calls and grown when it is too small.
template <typename T>
void Ensure(CudaDevicePtr<T> &p, size_t &capacity, size_t n) {
    if (n <= capacity && p)
        return;
    capacity = std::max<size_t>(n, 1);
    p = CudaDevicePtr<T>(capacity, CudaAlloc::Synchronous);
}

// cub's temporary storage to reduce n doubles.
size_t ReduceBytes(size_t n) {
    size_t bytes = 0;
    cuda_err(cub::DeviceReduce::Sum(nullptr, bytes, static_cast<double *>(nullptr), static_cast<double *>(nullptr),
                                    static_cast<int>(std::max<size_t>(n, 1))));
    return bytes;
}

}  // namespace

struct RefinementGPUEngineImpl {
    struct Restraints {
        int records = 0, atoms = 0;
        size_t cap_records = 0, cap_planes = 0, cap_start = 0, cap_index = 0, cap_slots = 0, cap_reduce = 0,
               cap_atoms = 0, cap_xyzb = 0, reduce_bytes = 0;
        CudaDevicePtr<RestraintRecord> rec;
        CudaDevicePtr<int> plane_atoms, slot_start, slot_index;
        CudaDevicePtr<float4> contrib, xyzb, grad;
        CudaDevicePtr<double> energy, sum;
        CudaDevicePtr<char> reduce_temp;
    };

    RefinementGPUSetup setup;
    size_t work_bytes;
    int device = 0;
    CudaStream stream;
    int n_atoms = 0;

    CudaDevicePtr<float4> pos;
    CudaDevicePtr<int> box;
    CudaDevicePtr<ModelMaskAtom> mask_atoms;
    CudaDevicePtr<float> grid;
    CudaDevicePtr<float2> spectrum;
    CudaDevicePtr<char> fft_work;
    CudaDevicePtr<int> ops, row_hkl;
    CudaDevicePtr<double> row_scale;
    CudaDevicePtr<float> k_total, ml_d, ml_var, fo;
    CudaDevicePtr<float2> fmask_scaled, f, fm, coef;
    CudaDevicePtr<char> work, centric;
    CudaDevicePtr<double> nll, nll_sum;
    CudaDevicePtr<char> reduce_temp;
    size_t reduce_bytes = 0;
    CudaDevicePtr<int> scatter_point, scatter_start, scatter_term;
    CudaDevicePtr<float4> grad;
    CudaDevicePtr<ModelDensityAtom> atoms_copy;   // the density gather keeps its own; this one is the gradient's
    std::unique_ptr<ModelDensityGPU> density;
    std::unique_ptr<ModelMaskGPU> mask;
    CudaFFTPlan forward, inverse;
    ModelDensityGPU::Geometry geom{};
    Restraints restraints[2];

    RefinementGPUEngineImpl(int dev, const RefinementGPUSetup &s, size_t fft_work_bytes) : setup(s), work_bytes(fft_work_bytes), device(dev) {
        const auto sync = CudaAlloc::Synchronous;
        const size_t na = std::max<size_t>(setup.max_atoms, 1), nr = std::max<size_t>(setup.rows.size(), 1);
        pos = CudaDevicePtr<float4>(na, sync);
        box = CudaDevicePtr<int>(3 * na, sync);
        mask_atoms = CudaDevicePtr<ModelMaskAtom>(na, sync);
        atoms_copy = CudaDevicePtr<ModelDensityAtom>(na, sync);
        grad = CudaDevicePtr<float4>(na, sync);
        grid = CudaDevicePtr<float>(GridPoints(setup.grid), sync);
        spectrum = CudaDevicePtr<float2>(ComplexPoints(setup.grid), sync);
        fft_work = CudaDevicePtr<char>(std::max<size_t>(work_bytes, 1), sync);
        ops = CudaDevicePtr<int>(std::max<size_t>(12 * setup.sym_ops.size(), 1), sync);
        row_hkl = CudaDevicePtr<int>(3 * nr, sync);
        row_scale = CudaDevicePtr<double>(nr, sync);
        k_total = CudaDevicePtr<float>(nr, sync);
        ml_d = CudaDevicePtr<float>(nr, sync);
        ml_var = CudaDevicePtr<float>(nr, sync);
        fo = CudaDevicePtr<float>(nr, sync);
        fmask_scaled = CudaDevicePtr<float2>(nr, sync);
        f = CudaDevicePtr<float2>(nr, sync);
        fm = CudaDevicePtr<float2>(nr, sync);
        coef = CudaDevicePtr<float2>(nr, sync);
        work = CudaDevicePtr<char>(nr, sync);
        centric = CudaDevicePtr<char>(nr, sync);
        nll = CudaDevicePtr<double>(nr, sync);
        nll_sum = CudaDevicePtr<double>(1, sync);
        reduce_bytes = ReduceBytes(nr);
        reduce_temp = CudaDevicePtr<char>(std::max<size_t>(reduce_bytes, 1), sync);
        scatter_point = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_point.size(), 1), sync);
        scatter_start = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_start.size(), 1), sync);
        scatter_term = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_term.size(), 1), sync);
        density = std::make_unique<ModelDensityGPU>(stream, setup.max_atoms, setup.max_pairs,
                                                    ModelDensityGPU::Bricks(setup.grid.nu, setup.grid.nv, setup.grid.nw));
        mask = std::make_unique<ModelMaskGPU>(stream, GridPoints(setup.grid));
        ModelMaskGrid mask_grid;
        mask_grid.nu = setup.grid.nu;
        mask_grid.nv = setup.grid.nv;
        mask_grid.nw = setup.grid.nw;
        std::copy(setup.grid.orth, setup.grid.orth + 9, mask_grid.orth);
        mask_grid.volume = setup.volume;
        mask->SetGrid(mask_grid, setup.mask_ops);

        std::vector<int> flat_ops;
        for (const auto &o : setup.sym_ops)
            flat_ops.insert(flat_ops.end(), o.begin(), o.end());
        std::vector<int> hkl;
        for (const auto &h : setup.rows)
            hkl.insert(hkl.end(), h.begin(), h.end());
        Upload(ops, flat_ops, stream);
        Upload(row_hkl, hkl, stream);
        Upload(fo, setup.fo, stream);
        Upload(work, setup.work, stream);
        Upload(centric, setup.centric, stream);
        Upload(scatter_point, setup.scatter_point, stream);
        Upload(scatter_start, setup.scatter_start, stream);
        Upload(scatter_term, setup.scatter_term, stream);

        const double inv_n[3] = {1.0 / setup.grid.nu, 1.0 / setup.grid.nv, 1.0 / setup.grid.nw};
        geom.nu = setup.grid.nu;
        geom.nv = setup.grid.nv;
        geom.nw = setup.grid.nw;
        geom.m00 = static_cast<float>(setup.grid.orth[0] * inv_n[0]);
        geom.m01 = static_cast<float>(setup.grid.orth[1] * inv_n[1]);
        geom.m02 = static_cast<float>(setup.grid.orth[2] * inv_n[2]);
        geom.m11 = static_cast<float>(setup.grid.orth[4] * inv_n[1]);
        geom.m12 = static_cast<float>(setup.grid.orth[5] * inv_n[2]);
        geom.m22 = static_cast<float>(setup.grid.orth[8] * inv_n[2]);

        forward = CudaFFTPlan(ForwardShape(setup.grid), stream, fft_work);
        inverse = CudaFFTPlan(InverseShape(setup.grid), stream, fft_work);
        if (forward.WorkBytes() > work_bytes || inverse.WorkBytes() > work_bytes)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: cuFFT work area over the reserve");
        cuda_err(cudaStreamSynchronize(stream));
    }

    ~RefinementGPUEngineImpl() { cudaStreamSynchronize(stream); }

    float Norm() const { return static_cast<float>(setup.volume / static_cast<double>(GridPoints(setup.grid))); }

    // The atoms' density on the grid and its transform into the spectrum.
    void Transform() {
        density->Compute(pos, grid);
        cufft_err(forward.R2C(grid, reinterpret_cast<cufftComplex *>(spectrum.get())));
    }

    void Compose(bool evaluate) {
        const int rows = static_cast<int>(setup.rows.size());
        if (rows == 0)
            return;
        compose_target_kernel<<<blocks(rows), THREADS, 0, stream>>>(
            spectrum, Norm(), setup.grid.nu, setup.grid.nv, setup.grid.nw, ops, static_cast<int>(setup.sym_ops.size()),
            setup.den, row_hkl, row_scale, rows, evaluate, k_total, fmask_scaled, ml_d, ml_var, fo, work, centric, f,
            nll, coef);
        cuda_err(cudaGetLastError());
    }
};

size_t RefinementGPUEngine::FFTWorkBytes(int device, const ModelDensityGrid &grid) {
    DeviceGuard guard(device);
    return std::max(CudaFFTPlan::WorkBytes(ForwardShape(grid)), CudaFFTPlan::WorkBytes(InverseShape(grid)));
}

size_t RefinementGPUEngine::DeviceBytes(const RefinementGPUSetup &s, size_t fft_work_bytes, size_t max_records,
                                        size_t max_slots) {
    const size_t na = s.max_atoms + 1, nr = s.rows.size() + 1;
    size_t bytes = na * (2 * sizeof(float4) + 3 * sizeof(int) + sizeof(ModelMaskAtom) + sizeof(ModelDensityAtom));
    bytes += GridPoints(s.grid) * sizeof(float) + ComplexPoints(s.grid) * sizeof(float2) + fft_work_bytes;
    bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 4 * sizeof(float) + 4 * sizeof(float2) + 2);
    bytes += (s.scatter_point.size() + s.scatter_start.size() + s.scatter_term.size()) * sizeof(int);
    bytes += ModelDensityGPU::DeviceBytes(s.max_atoms, s.max_pairs, ModelDensityGPU::Bricks(s.grid.nu, s.grid.nv, s.grid.nw));
    bytes += ModelMaskGPU::DeviceBytes(GridPoints(s.grid));
    bytes += max_records * (sizeof(RestraintRecord) + sizeof(double)) + max_slots * (sizeof(float4) + sizeof(int)) +
             na * (2 * sizeof(float4) + sizeof(int));
    return bytes;
}

RefinementGPUEngine::RefinementGPUEngine(int device, const RefinementGPUSetup &setup, size_t fft_work_bytes) {
    DeviceGuard guard(device);
    impl_ = std::make_unique<RefinementGPUEngineImpl>(device, setup, fft_work_bytes);
}

RefinementGPUEngine::~RefinementGPUEngine() {
    if (impl_) {
        DeviceGuard guard(impl_->device);
        impl_.reset();
    }
}

void RefinementGPUEngine::SetBlur(double blur) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    std::vector<double> scale(e.setup.rows.size());
    for (size_t m = 0; m < scale.size(); m++)
        scale[m] = e.setup.centring * std::exp(blur * 0.25 * e.setup.inv_d2[m]);
    Upload(e.row_scale, scale, e.stream);
    cuda_err(cudaStreamSynchronize(e.stream));
}

void RefinementGPUEngine::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) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    Upload(e.k_total, k_total, e.stream);
    if (!fmask.empty())
        cuda_err(cudaMemcpyAsync(e.fmask_scaled, fmask.data(), fmask.size() * sizeof(float2), cudaMemcpyHostToDevice,
                                 e.stream));
    Upload(e.ml_d, d, e.stream);
    Upload(e.ml_var, var, e.stream);
    cuda_err(cudaStreamSynchronize(e.stream));
}

void RefinementGPUEngine::SetAtoms(const std::vector<ModelDensityAtom> &atoms, const std::vector<std::array<float, 4>> &pos) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    if (atoms.size() > e.setup.max_atoms || pos.size() != atoms.size())
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: atoms over the reserve");
    e.density->SetAtoms(e.setup.grid, atoms);
    e.n_atoms = static_cast<int>(atoms.size());
    const std::vector<int> b = ModelDensityGPU::Boxes(e.setup.grid, atoms);
    if (!atoms.empty()) {
        cuda_err(cudaMemcpyAsync(e.pos, pos.data(), pos.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream));
        cuda_err(cudaMemcpyAsync(e.atoms_copy, atoms.data(), atoms.size() * sizeof(ModelDensityAtom),
                                 cudaMemcpyHostToDevice, e.stream));
        Upload(e.box, b, e.stream);
    }
    cuda_err(cudaStreamSynchronize(e.stream));
}

void RefinementGPUEngine::Fcalc(std::vector<std::array<float, 2>> &f) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    e.Transform();
    e.Compose(false);
    f.resize(e.setup.rows.size());
    if (!f.empty())
        cuda_err(cudaMemcpyAsync(f.data(), e.f, f.size() * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
    cuda_err(cudaStreamSynchronize(e.stream));
}

void RefinementGPUEngine::Fmask(const std::vector<ModelMaskAtom> &mask_atoms, std::vector<std::array<float, 2>> &f) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    if (mask_atoms.size() > e.setup.max_atoms)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: mask atoms over the reserve");
    if (!mask_atoms.empty())
        cuda_err(cudaMemcpyAsync(e.mask_atoms, mask_atoms.data(), mask_atoms.size() * sizeof(ModelMaskAtom),
                                 cudaMemcpyHostToDevice, e.stream));
    e.mask->Compute(e.mask_atoms, static_cast<int>(mask_atoms.size()), e.grid);
    cufft_err(e.forward.R2C(e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
    const int rows = static_cast<int>(e.setup.rows.size());
    f.resize(rows);
    if (rows > 0) {
        fmask_kernel<<<blocks(rows), THREADS, 0, e.stream>>>(e.spectrum, e.Norm(), e.row_hkl, rows, e.setup.grid.nu,
                                                             e.setup.grid.nv, e.setup.grid.nw, e.fm);
        cuda_err(cudaGetLastError());
        cuda_err(cudaMemcpyAsync(f.data(), e.fm, rows * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
    }
    cuda_err(cudaStreamSynchronize(e.stream));
}

double RefinementGPUEngine::Evaluate(std::vector<std::array<float, 4>> *grad) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    const int rows = static_cast<int>(e.setup.rows.size());
    e.Transform();
    e.Compose(true);
    double sum = 0;
    if (rows > 0) {
        size_t temp = e.reduce_bytes;
        cuda_err(cub::DeviceReduce::Sum(e.reduce_temp.get(), temp, e.nll.get(), e.nll_sum.get(), rows, e.stream));
        cuda_err(cudaMemcpyAsync(&sum, e.nll_sum, sizeof(double), cudaMemcpyDeviceToHost, e.stream));
    }
    if (grad != nullptr) {
        const ModelDensityGrid &g = e.setup.grid;
        cuda_err(cudaMemsetAsync(e.spectrum, 0, ComplexPoints(g) * sizeof(float2), e.stream));
        const int points = static_cast<int>(e.setup.scatter_point.size());
        if (points > 0) {
            scatter_kernel<<<blocks(points), THREADS, 0, e.stream>>>(e.coef, e.row_hkl, e.ops,
                                                                     static_cast<int>(e.setup.sym_ops.size()), e.setup.den,
                                                                     e.scatter_point, e.scatter_start, e.scatter_term,
                                                                     points, e.spectrum);
            cuda_err(cudaGetLastError());
        }
        cufft_err(e.inverse.C2R(reinterpret_cast<cufftComplex *>(e.spectrum.get()), e.grid));
        if (e.n_atoms > 0) {
            const size_t threads = static_cast<size_t>(e.n_atoms) * WARP;
            atom_gradient_kernel<<<blocks(threads), THREADS, 0, e.stream>>>(e.atoms_copy, e.pos, e.box, e.n_atoms, e.geom,
                                                                            e.grid, e.Norm(), e.grad);
            cuda_err(cudaGetLastError());
        }
        grad->resize(e.n_atoms);
        if (e.n_atoms > 0)
            cuda_err(cudaMemcpyAsync(grad->data(), e.grad, e.n_atoms * sizeof(float4), cudaMemcpyDeviceToHost, e.stream));
    }
    cuda_err(cudaStreamSynchronize(e.stream));
    return sum;
}

void RefinementGPUEngine::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) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    RefinementGPUEngineImpl::Restraints &r = e.restraints[which];
    for (const RestraintRecord &rec : records)
        if (rec.kind == static_cast<int>(RestraintKind::Plane) && rec.atom[1] > MAX_PLANE_ATOMS)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: a plane of too many atoms");
    r.records = static_cast<int>(records.size());
    r.atoms = static_cast<int>(slot_start.empty() ? 0 : slot_start.size() - 1);
    Ensure(r.rec, r.cap_records, records.size());
    Ensure(r.energy, r.cap_reduce, records.size());
    Ensure(r.plane_atoms, r.cap_planes, plane_atoms.size());
    Ensure(r.slot_start, r.cap_start, slot_start.size());
    Ensure(r.slot_index, r.cap_index, slot_index.size());
    Ensure(r.contrib, r.cap_slots, static_cast<size_t>(n_slots));
    Ensure(r.xyzb, r.cap_xyzb, r.atoms);
    Ensure(r.grad, r.cap_atoms, r.atoms);
    if (!r.sum)
        r.sum = CudaDevicePtr<double>(1, CudaAlloc::Synchronous);
    const size_t bytes = ReduceBytes(r.cap_reduce);
    if (!r.reduce_temp || bytes > r.reduce_bytes) {
        r.reduce_bytes = bytes;
        r.reduce_temp = CudaDevicePtr<char>(std::max<size_t>(bytes, 1), CudaAlloc::Synchronous);
    }
    Upload(r.rec, records, e.stream);
    Upload(r.plane_atoms, plane_atoms, e.stream);
    Upload(r.slot_start, slot_start, e.stream);
    Upload(r.slot_index, slot_index, e.stream);
    cuda_err(cudaStreamSynchronize(e.stream));
}

double RefinementGPUEngine::EvaluateRestraints(int which, const std::vector<std::array<float, 4>> &xyzb,
                                               std::vector<std::array<float, 4>> *grad) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    RefinementGPUEngineImpl::Restraints &r = e.restraints[which];
    if (static_cast<int>(xyzb.size()) != r.atoms)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: restraint atoms do not match");
    if (r.atoms > 0)
        cuda_err(cudaMemcpyAsync(r.xyzb, xyzb.data(), xyzb.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream));
    double sum = 0;
    if (r.records > 0) {
        restraint_kernel<<<blocks(r.records), THREADS, 0, e.stream>>>(r.rec, r.records, r.plane_atoms, r.xyzb, r.contrib,
                                                                      r.energy);
        cuda_err(cudaGetLastError());
        size_t temp = r.reduce_bytes;
        cuda_err(cub::DeviceReduce::Sum(r.reduce_temp.get(), temp, r.energy.get(), r.sum.get(), r.records, e.stream));
        cuda_err(cudaMemcpyAsync(&sum, r.sum, sizeof(double), cudaMemcpyDeviceToHost, e.stream));
    }
    if (grad != nullptr) {
        grad->resize(r.atoms);
        if (r.atoms > 0) {
            if (r.records > 0) {
                slot_gather_kernel<<<blocks(r.atoms), THREADS, 0, e.stream>>>(r.contrib, r.slot_start, r.slot_index,
                                                                              r.atoms, r.grad);
                cuda_err(cudaGetLastError());
            } else
                cuda_err(cudaMemsetAsync(r.grad, 0, r.atoms * sizeof(float4), e.stream));
            cuda_err(cudaMemcpyAsync(grad->data(), r.grad, r.atoms * sizeof(float4), cudaMemcpyDeviceToHost, e.stream));
        }
    }
    cuda_err(cudaStreamSynchronize(e.stream));
    return sum;
}