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 <mutex>
#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) for h.t = ht / den, exact for the fractions a translation is made of.
__device__ __forceinline__ void op_phase(int ht, int den, double &c, double &s) {
    sincospi(2.0 * ht / den, &s, &c);
}

// The phases of every h.t the rows and operators make, as a table over the integer h.t from `min`: the hot
// kernels read it in place of op_phase() - the same function at the same argument, so the same bits.
__global__ void phase_table_kernel(int min, int count, int den, double2 *phase) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= count)
        return;
    double s, c;
    op_phase(min + i, den, c, s);
    phase[i] = make_double2(c, s);
}

// 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. `ops` holds each operator's rotation in whole
// units and its translation in units of 1 / den.
__global__ void compose_target_kernel(const float2 *f1, float norm, int nu, int nv, int nw, const int *ops, int n_ops,
                                      const double2 *phase, int phase_min, 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];
        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;
        const double2 p = phase[h[0] * tran[0] + h[1] * tran[1] + h[2] * tran[2] - phase_min];
        sr += p.x * vr - p.y * vi;
        si += p.x * vi + p.y * 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));
}

// With a twin law, the target's term of a working row against the twinned model,
// |F_model| = k sqrt((1-a) |U(h)|^2 + a |U(Th)|^2) with U = F + fmask, from F of every row (compose_target_kernel
// without `evaluate`); and the conjugates of its derivatives with respect to U(h) and to U(Th), unscaled, in `own`
// and `other` - zero for a row that is not a working observation.
__global__ void twin_target_kernel(const float2 *f, const float2 *fmask, const int *mate, const float *twin_fraction,
                                   int rows, const float *k_total, const float *ml_d, const float *ml_var,
                                   const float *fo, const char *work, const char *centric, double *nll, float2 *own,
                                   float2 *other) {
    const int m = blockIdx.x * blockDim.x + threadIdx.x;
    if (m >= rows)
        return;
    double term = 0, g = 0;
    const int j = mate[m];
    const double u1r = f[m].x + fmask[m].x, u1i = f[m].y + fmask[m].y;
    const double u2r = f[j].x + fmask[j].x, u2i = f[j].y + fmask[j].y;
    const double a = twin_fraction[0];
    if (work[m]) {
        const double k = k_total[m];
        const double fc = k * sqrt((1.0 - a) * (u1r * u1r + u1i * u1i) + a * (u2r * u2r + u2i * u2i));
        double dfc = 0;
        term = refine_math::AmplitudeNegLogLikelihood<double>(fo[m], fc, ml_d[m], ml_var[m], centric[m], dfc);
        if (fc > 0)
            g = dfc * k * k / fc;
    }
    nll[m] = term;
    own[m] = make_float2(static_cast<float>(g * (1.0 - a) * u1r), static_cast<float>(-g * (1.0 - a) * u1i));
    other[m] = make_float2(static_cast<float>(g * a * u2r), static_cast<float>(-g * a * u2i));
}

// The gradient's coefficient of every row under a twin law: its own term's derivative and that of the row whose
// mate it is - its own mate, the law being an involution - times the row's scale.
__global__ void twin_gather_kernel(const float2 *own, const float2 *other, const int *mate, const double *row_scale,
                                   int rows, float2 *coef) {
    const int n = blockIdx.x * blockDim.x + threadIdx.x;
    if (n >= rows)
        return;
    const double scale = row_scale[n];
    const float2 a = own[n], b = other[mate[n]];
    coef[n] = make_float2(static_cast<float>(scale * (static_cast<double>(a.x) + b.x)),
                          static_cast<float>(scale * (static_cast<double>(a.y) + b.y)));
}

// 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. A term is
// its row * 2 + conj, and its phase the table's entry at term_phase.
__global__ void scatter_kernel(const float2 *coef, const double2 *phase, const int *point, const int *start,
                               const int *term, const int *term_phase, 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;
        const double2 p = phase[term_phase[e]];
        const double ar = 0.5 * (coef[m].x * p.x - coef[m].y * p.y), ai = 0.5 * (coef[m].x * p.y + coef[m].y * p.x);
        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).
// The lanes are dealt the points of the box's rows (v, w) that can be within the radius, not the whole box: a warp's
// width of rows at a time, each cut to the chord the radius leaves it with a point to spare at either end (the radius
// test below decides, as the gather's), and the points of those rows numbered one after another.
// An anisotropic atom (rho = occ sum_q a_q exp(r^T M_q r), RefinementXray.cpp's AnisoGradient()) gives its B derivative
// along U + dB / (8 pi^2) I and, into grad_u, its six U derivatives; an isotropic one writes zeros there.
__global__ void __launch_bounds__(THREADS, 4)
atom_gradient_kernel(const ModelDensityAtom *__restrict__ atoms, const float4 *__restrict__ pos,
                     const int *__restrict__ box, int n, ModelDensityGPU::Geometry g, const float *__restrict__ map,
                     float map_scale, float4 *grad, float *grad_u) {
    constexpr unsigned FULL = 0xffffffff;
    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};
    // the box's first point, unwrapped and wrapped into the grid; a box is narrower than the cell (Supports())
    int lo[3], first[3], span[3];
    const int n3[3] = {g.nu, g.nv, g.nw};
    for (int k = 0; k < 3; k++) {
        lo[k] = __float2int_rn(nctr[k]) - box[3 * a + k];
        first[k] = imod(lo[k], n3[k]);
        span[k] = 2 * box[3 * a + k] + 1;
    }
    const int rows = span[1] * span[2];
    // an isotropic atom's Gaussians, in registers
    float qa[5], qb[5];
    for (int q = 0; q < 5; q++) {
        qa[q] = at.a[q];
        qb[q] = at.b[q][0];
    }
    const float max_r2 = at.max_r2;
    float gx = 0, gy = 0, gz = 0, gb = 0;
    float gu[6] = {0, 0, 0, 0, 0, 0};
    const bool aniso = at.aniso != 0;
    for (int r0 = 0; r0 < rows; r0 += WARP) {
        // The lane's row: its (iv, iw) in the box, its first point iu0 and its number of points. With
        // x = m00 fx + c and fx = centre - u, x^2 <= max_r2 - y^2 - z^2 bounds u.
        const int r = r0 + lane;
        int iv = 0, iw = 0, iu0 = 0, count = 0;
        if (r < rows) {
            iv = r % span[1];
            iw = r / span[1];
            const float fy = nctr[1] - static_cast<float>(lo[1] + iv);
            const float fz = nctr[2] - static_cast<float>(lo[2] + iw);
            const float z = g.m22 * fz;
            const float y = g.m11 * fy + g.m12 * fz;
            const float left = max_r2 - (y * y + z * z);
            if (left >= 0) {
                const float c = g.m01 * fy + g.m02 * fz;
                const float half = sqrtf(left);
                iu0 = max(0, static_cast<int>(floorf(nctr[0] - (half - c) / g.m00)) - 1 - lo[0]);
                const int iu1 = min(span[0] - 1, static_cast<int>(ceilf(nctr[0] + (half + c) / g.m00)) + 1 - lo[0]);
                count = max(0, iu1 - iu0 + 1);
            }
        }
        // the rows' points numbered in lane order: `end` is one past the lane's last
        int end = count;
        for (int off = 1; off < WARP; off *= 2) {
            const int below = __shfl_up_sync(FULL, end, off);
            if (lane >= off)
                end += below;
        }
        const int total = __shfl_sync(FULL, end, WARP - 1);
        const int skip = end - count - iu0;   // point number t of this row is at iu = t - skip
        for (int t0 = 0; t0 < total; t0 += WARP) {
            const int t = t0 + lane;
            // the lane j whose row holds point t: the number of lanes whose rows end at or before it
            int j = 0;
            for (int step = WARP / 2; step > 0; step /= 2)
                if (__shfl_sync(FULL, end, j + step - 1) <= t)
                    j += step;
            const int jv = __shfl_sync(FULL, iv, j), jw = __shfl_sync(FULL, iw, j);
            const int iu = t - __shfl_sync(FULL, skip, j);
            if (t >= total)
                continue;
            const int u = lo[0] + iu, v = lo[1] + jv, w = lo[2] + jw;
            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 > max_r2)
                continue;
            const int mu = first[0] + iu, mv = first[1] + jv, mw = first[2] + jw;
            const float value = map[(mu < g.nu ? mu : mu - g.nu) +
                                    static_cast<size_t>(g.nu) * ((mv < g.nv ? mv : mv - g.nv) +
                                                                 static_cast<size_t>(g.nv) * (mw < g.nw ? mw : mw - g.nw))];
            if (!aniso) {
                float dr2 = 0, db = 0;
                for (int q = 0; q < 5; q++) {
                    const float b = qb[q];
                    const float e = qa[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;
            } else {
                for (int q = 0; q < 5; q++) {
                    const float *m = at.b[q];
                    const float rur = x * x * m[0] + y * y * m[1] + z * z * m[2] + 2 * (x * y * m[3] + x * z * m[4] + y * z * m[5]);
                    const float e = value * at.a[q] * expapprox(fmaxf(rur, -88.f));
                    const float mr0 = m[0] * x + m[3] * y + m[4] * z, mr1 = m[3] * x + m[1] * y + m[5] * z,
                                mr2 = m[4] * x + m[5] * y + m[2] * z;
                    gx += e * mr0;
                    gy += e * mr1;
                    gz += e * mr2;
                    gu[0] += e * (m[0] + 2 * mr0 * mr0);
                    gu[1] += e * (m[1] + 2 * mr1 * mr1);
                    gu[2] += e * (m[2] + 2 * mr2 * mr2);
                    gu[3] += e * (m[3] + 2 * mr0 * mr1);
                    gu[4] += e * (m[4] + 2 * mr0 * mr2);
                    gu[5] += e * (m[5] + 2 * mr1 * mr2);
                }
            }
        }
    }
    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);
        for (int c = 0; c < 6; c++)
            gu[c] += __shfl_down_sync(0xffffffff, gu[c], off);
    }
    if (lane == 0) {
        const float s = 2.0f * at.occ * map_scale;
        const float o = at.occ * map_scale;
        if (!aniso)
            grad[a] = make_float4(s * gx, s * gy, s * gz, o * gb / static_cast<float>(FOUR_PI2));
        else
            grad[a] = make_float4(s * gx, s * gy, s * gz, o * (gu[0] + gu[1] + gu[2]) / static_cast<float>(2 * FOUR_PI2));
        if (grad_u != nullptr) {
            const float f[6] = {o * gu[0], o * gu[1], o * gu[2], 2 * o * gu[3], 2 * o * gu[4], 2 * o * gu[5]};
            for (int c = 0; c < 6; c++)
                grad_u[6 * a + c] = f[c];
        }
    }
}

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::NcsDistance:
            e = refine_math::NcsDistanceTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), xyz(xyzb, r.atom[2]),
                                             xyz(xyzb, r.atom[3]), r.sigma, r.ideal, 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: {
            float gi = 0, gj = 0;
            e = refine_math::BSimilarityTerm(xyzb[r.atom[0]].w, xyzb[r.atom[1]].w, r.sigma, gi, gj);
            contrib[r.slot] = make_float4(0, 0, 0, gi);
            contrib[r.slot + 1] = make_float4(0, 0, 0, gj);
            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};
}

bool SameGrid(const ModelDensityGrid &a, const ModelDensityGrid &b) {
    return a.nu == b.nu && a.nv == b.nv && a.nw == b.nw;
}

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);
}

// h.t of a row and an operator, in units of 1 / den.
int HT(const std::array<int, 3> &h, const std::array<int, 12> &op) {
    return h[0] * op[9] + h[1] * op[10] + h[2] * op[11];
}

// The range of h.t over the rows and operators: the phase table's.
void PhaseRange(const RefinementGPUSetup &setup, int &min, int &max) {
    min = max = 0;
    for (const auto &h : setup.rows)
        for (const auto &op : setup.sym_ops) {
            const int ht = HT(h, op);
            min = std::min(min, ht);
            max = std::max(max, ht);
        }
}

// 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;
}

// The engine's wait on its stream, polling: an evaluation is a millisecond or a few, and where a slot has a
// card to itself the wake-up of a sleeping thread is on the critical path of every one of them.
class StreamPoller {
public:
    cudaError_t Wait(cudaStream_t stream) const {
        cudaError_t err;
        while ((err = cudaStreamQuery(stream)) == cudaErrorNotReady) {}
        return err;
    }
};

// What a function queues on a stream, recorded once as a CUDA graph and replayed from then on: the same
// launches in the same order, so the same bits, without the host between them. Recorded again where `key`
// differs from the recording's (another shape: the atoms, a gradient asked or not, another restraint set).
// Where the stream cannot be captured the function is run as it is, every time.
class RecordedLaunches {
public:
    RecordedLaunches() = default;
    RecordedLaunches(const RecordedLaunches &) = delete;
    RecordedLaunches &operator=(const RecordedLaunches &) = delete;
    ~RecordedLaunches() {
        if (exec_ != nullptr)
            cudaGraphExecDestroy(exec_);
    }

    void Launch(cudaStream_t stream, long key, const std::function<void()> &queue) {
        if (direct_) {
            queue();
            return;
        }
        if (exec_ == nullptr || key != key_) {
            if (exec_ != nullptr)
                cuda_err(cudaGraphExecDestroy(exec_));
            exec_ = nullptr;
            cudaGraph_t graph = nullptr;
            cuda_err(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal));
            try {
                queue();
            } catch (const JFJochException &) {
                cudaStreamEndCapture(stream, &graph);
                if (graph != nullptr)
                    cudaGraphDestroy(graph);
                cudaGetLastError();
                direct_ = true;
                queue();
                return;
            }
            cuda_err(cudaStreamEndCapture(stream, &graph));
            cuda_err(cudaGraphInstantiate(&exec_, graph, 0));
            cuda_err(cudaGraphDestroy(graph));
            key_ = key;
        }
        cuda_err(cudaGraphLaunch(exec_, stream));
    }

private:
    cudaGraphExec_t exec_ = nullptr;
    long key_ = 0;
    bool direct_ = false;
};

}  // namespace

// The working buffers of a transform - the grid, its spectrum, cuFFT's work area, the mask's grid and spectrum where
// it has its own - which the engines of a lane use one at a time, under `lock`, from their launches to their wait:
// while one engine's evaluation is on the card the others of its lane go on with their host work.
struct RefinementGPULane {
    CudaDevicePtr<float> grid, mask_grid;
    CudaDevicePtr<float2> spectrum, mask_spectrum;
    CudaDevicePtr<char> fft_work;
    std::mutex lock;
    int engines = 0;

    RefinementGPULane(const RefinementGPUSetup &setup, size_t fft_work_bytes) {
        const auto sync = CudaAlloc::Synchronous;
        grid = CudaDevicePtr<float>(GridPoints(setup.grid), sync);
        spectrum = CudaDevicePtr<float2>(ComplexPoints(setup.grid), sync);
        fft_work = CudaDevicePtr<char>(std::max<size_t>(fft_work_bytes, 1), sync);
        if (!SameGrid(setup.mask_grid, setup.grid)) {
            mask_grid = CudaDevicePtr<float>(GridPoints(setup.mask_grid), sync);
            mask_spectrum = CudaDevicePtr<float2>(ComplexPoints(setup.mask_grid), sync);
        }
    }
};

// What the engines on one card share: the reflections' tables, which nothing writes once they are up, the
// bulk-solvent mask's buffers, which one engine at a time masks with (under `mask_lock`, on the card's own stream),
// and the lanes of working buffers, each engine on one of them.
struct RefinementGPUCard {
    int device = 0;
    CudaStream stream;
    CudaDevicePtr<int> ops, row_hkl;
    CudaDevicePtr<float> fo;
    CudaDevicePtr<char> work, centric;
    CudaDevicePtr<int> mate;   // a twin law's (setup.mate non-empty)
    CudaDevicePtr<int> scatter_point, scatter_start, scatter_term, scatter_phase;
    CudaDevicePtr<double2> phase;   // exp(2 pi i h.t) over the integer h.t from phase_min
    int phase_min = 0;
    std::unique_ptr<ModelMaskGPU> mask;   // on `stream`
    std::mutex mask_lock;
    std::mutex lanes_lock;
    std::vector<std::unique_ptr<RefinementGPULane>> lanes;

    // A lane for an engine: a new one where `own` (or there is none yet), else the one with the fewest engines.
    // On the card's device.
    RefinementGPULane *Lane(const RefinementGPUSetup &setup, size_t fft_work_bytes, bool own) {
        const std::lock_guard l(lanes_lock);
        RefinementGPULane *lane = nullptr;
        if (own || lanes.empty()) {
            lanes.push_back(std::make_unique<RefinementGPULane>(setup, fft_work_bytes));
            lane = lanes.back().get();
        } else
            for (const auto &x : lanes)
                if (lane == nullptr || x->engines < lane->engines)
                    lane = x.get();
        lane->engines++;
        return lane;
    }

    RefinementGPUCard(int dev, const RefinementGPUSetup &setup) : device(dev) {
        const auto sync = CudaAlloc::Synchronous;
        const size_t nr = std::max<size_t>(setup.rows.size(), 1);
        ops = CudaDevicePtr<int>(std::max<size_t>(12 * setup.sym_ops.size(), 1), sync);
        row_hkl = CudaDevicePtr<int>(3 * nr, sync);
        fo = CudaDevicePtr<float>(nr, sync);
        work = CudaDevicePtr<char>(nr, sync);
        centric = CudaDevicePtr<char>(nr, sync);
        if (!setup.mate.empty())
            mate = CudaDevicePtr<int>(nr, 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);
        scatter_phase = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_term.size(), 1), sync);
        int phase_max = 0;
        PhaseRange(setup, phase_min, phase_max);
        const int phase_count = phase_max - phase_min + 1;
        phase = CudaDevicePtr<double2>(phase_count, sync);
        mask = std::make_unique<ModelMaskGPU>(stream, GridPoints(setup.mask_grid));
        ModelMaskGrid mg;
        mg.nu = setup.mask_grid.nu;
        mg.nv = setup.mask_grid.nv;
        mg.nw = setup.mask_grid.nw;
        std::copy(setup.mask_grid.orth, setup.mask_grid.orth + 9, mg.orth);
        mg.volume = setup.volume;
        mask->SetGrid(mg, setup.mask_ops);

        // the rotations in whole units (gemmi's are den times an integer matrix), the translations as given
        std::vector<int> flat_ops;
        for (const auto &o : setup.sym_ops) {
            for (int i = 0; i < 9; i++)
                flat_ops.push_back(o[i] / setup.den);
            flat_ops.insert(flat_ops.end(), o.begin() + 9, o.end());
        }
        std::vector<int> hkl;
        for (const auto &h : setup.rows)
            hkl.insert(hkl.end(), h.begin(), h.end());
        // each scatter term as row * 2 + conj, and its phase's place in the table
        const int n_ops = static_cast<int>(setup.sym_ops.size());
        std::vector<int> term, term_phase;
        for (const int t : setup.scatter_term) {
            const int m = (t >> 1) / n_ops, o = (t >> 1) % n_ops;
            term.push_back(m * 2 + (t & 1));
            term_phase.push_back(HT(setup.rows[m], setup.sym_ops[o]) - phase_min);
        }
        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);
        if (!setup.mate.empty())
            Upload(mate, setup.mate, stream);
        Upload(scatter_point, setup.scatter_point, stream);
        Upload(scatter_start, setup.scatter_start, stream);
        Upload(scatter_term, term, stream);
        Upload(scatter_phase, term_phase, stream);
        phase_table_kernel<<<blocks(phase_count), THREADS, 0, stream>>>(phase_min, phase_count, setup.den, phase);
        cuda_err(cudaGetLastError());
        cuda_err(cudaStreamSynchronize(stream));
    }

    ~RefinementGPUCard() { cudaStreamSynchronize(stream); }
};

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;
        int version = 0;   // counts SetRestraints(): a recording of the set is of one version
        RecordedLaunches launches;
    };

    const RefinementGPUSetup &setup;
    std::shared_ptr<RefinementGPUCard> card;
    size_t work_bytes;
    int device = 0;
    CudaStream stream;
    StreamPoller waiter;   // the engine's waits on the stream, polled
    int n_atoms = 0;

    CudaDevicePtr<float4> pos;
    // What comes back, pinned, so that a copy back queues without waiting: the X-ray target and the restraints'
    // sums, their gradients, and the restraints' positions going up.
    CudaHostPtr<double> host_sums;
    CudaHostPtr<float4> host_grad, host_xyzb, host_restraint_grad;
    CudaHostPtr<float> host_grad_u;   // six per atom
    CudaDevicePtr<ModelMaskAtom> mask_atoms;
    // the lane's working buffers (used under its lock)
    RefinementGPULane *lane;
    float *grid;
    float2 *spectrum;
    char *fft_work;
    // The mask's grid, spectrum and transform where its grid is not the structure factors' (grid, spectrum,
    // forward otherwise).
    float *mask_grid;
    float2 *mask_spectrum;
    // the card's tables
    const int *ops, *row_hkl, *mate, *scatter_point, *scatter_start, *scatter_term, *scatter_phase;
    const float *fo;
    const char *work, *centric;
    const double2 *phase;
    int phase_min;
    CudaDevicePtr<double> row_scale;
    CudaDevicePtr<float> k_total, ml_d, ml_var;
    CudaDevicePtr<float2> fmask_scaled, f, fm, coef;
    // a twin law's (setup.mate non-empty): the fraction, the two parts of each row's derivative
    CudaDevicePtr<float> twin_fraction;
    CudaDevicePtr<float2> coef_own, coef_other;
    CudaDevicePtr<double> nll, nll_sum;
    CudaDevicePtr<char> reduce_temp;
    size_t reduce_bytes = 0;
    CudaDevicePtr<float4> grad;
    CudaDevicePtr<float> grad_u;
    int n_aniso = 0;   // the anisotropic atoms of the last SetDensities()
    std::unique_ptr<ModelDensityGPU> density;
    CudaFFTPlan forward, inverse, mask_forward;
    ModelDensityGPU::Geometry geom{};
    Restraints restraints[2];
    RecordedLaunches evaluate;   // QueueEvaluate() recorded

    RefinementGPUEngineImpl(int dev, const RefinementGPUSetup &s, std::shared_ptr<RefinementGPUCard> c, size_t fft_work_bytes,
                            bool own_lane)
        : setup(s), card(std::move(c)), work_bytes(fft_work_bytes), device(dev), ops(card->ops), row_hkl(card->row_hkl),
          mate(card->mate), scatter_point(card->scatter_point), scatter_start(card->scatter_start),
          scatter_term(card->scatter_term), scatter_phase(card->scatter_phase), fo(card->fo), work(card->work),
          centric(card->centric), phase(card->phase), phase_min(card->phase_min),
          lane(card->Lane(s, fft_work_bytes, own_lane)), grid(lane->grid), spectrum(lane->spectrum),
          fft_work(lane->fft_work), mask_grid(lane->mask_grid), mask_spectrum(lane->mask_spectrum) {
        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);
        host_sums = CudaHostPtr<double>(2);
        host_grad = CudaHostPtr<float4>(na);
        host_grad_u = CudaHostPtr<float>(6 * na);
        host_xyzb = CudaHostPtr<float4>(na);
        host_restraint_grad = CudaHostPtr<float4>(na);
        mask_atoms = CudaDevicePtr<ModelMaskAtom>(na, sync);
        grad = CudaDevicePtr<float4>(na, sync);
        grad_u = CudaDevicePtr<float>(6 * na, 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);
        fmask_scaled = CudaDevicePtr<float2>(nr, sync);
        f = CudaDevicePtr<float2>(nr, sync);
        fm = CudaDevicePtr<float2>(nr, sync);
        coef = CudaDevicePtr<float2>(nr, sync);
        nll = CudaDevicePtr<double>(nr, sync);
        nll_sum = CudaDevicePtr<double>(1, sync);
        if (!setup.mate.empty()) {
            twin_fraction = CudaDevicePtr<float>(1, sync);
            coef_own = CudaDevicePtr<float2>(nr, sync);
            coef_other = CudaDevicePtr<float2>(nr, sync);
        }
        reduce_bytes = ReduceBytes(nr);
        reduce_temp = CudaDevicePtr<char>(std::max<size_t>(reduce_bytes, 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));

        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 (!SameGrid(setup.mask_grid, setup.grid))
            mask_forward = CudaFFTPlan(ForwardShape(setup.mask_grid), stream, fft_work);
        if (forward.WorkBytes() > work_bytes || inverse.WorkBytes() > work_bytes ||
            (mask_grid != nullptr && mask_forward.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))); }
    float MaskNorm() const {
        return static_cast<float>(setup.volume / static_cast<double>(GridPoints(setup.mask_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)));
    }

    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()),
            phase, phase_min, row_hkl, row_scale, rows, evaluate, k_total, fmask_scaled, ml_d, ml_var, fo, work, centric,
            f, nll, coef);
        cuda_err(cudaGetLastError());
    }

    // The X-ray target into host_sums[0] and, with `grad`, its gradient into host_grad: queued, not waited for.
    void QueueEvaluate(bool grad) {
        const int rows = static_cast<int>(setup.rows.size());
        Transform();
        if (setup.mate.empty())
            Compose(true);
        else if (rows > 0) {
            Compose(false);
            twin_target_kernel<<<blocks(rows), THREADS, 0, stream>>>(f, fmask_scaled, mate, twin_fraction, rows, k_total,
                                                                     ml_d, ml_var, fo, work, centric, nll, coef_own,
                                                                     coef_other);
            cuda_err(cudaGetLastError());
            twin_gather_kernel<<<blocks(rows), THREADS, 0, stream>>>(coef_own, coef_other, mate, row_scale, rows, coef);
            cuda_err(cudaGetLastError());
        }
        if (rows > 0) {
            size_t temp = reduce_bytes;
            cuda_err(cub::DeviceReduce::Sum(reduce_temp.get(), temp, nll.get(), nll_sum.get(), rows, stream));
            cuda_err(cudaMemcpyAsync(host_sums.get(), nll_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));
        }
        if (grad) {
            const ModelDensityGrid &g = setup.grid;
            cuda_err(cudaMemsetAsync(spectrum, 0, ComplexPoints(g) * sizeof(float2), stream));
            const int points = static_cast<int>(setup.scatter_point.size());
            if (points > 0) {
                scatter_kernel<<<blocks(points), THREADS, 0, stream>>>(coef, phase, scatter_point, scatter_start,
                                                                       scatter_term, scatter_phase, points, spectrum);
                cuda_err(cudaGetLastError());
            }
            cufft_err(inverse.C2R(reinterpret_cast<cufftComplex *>(spectrum), this->grid));
            if (n_atoms > 0) {
                const size_t threads = static_cast<size_t>(n_atoms) * WARP;
                atom_gradient_kernel<<<blocks(threads), THREADS, 0, stream>>>(density->DeviceAtoms(), pos,
                                                                              density->DeviceBoxes(), n_atoms, geom,
                                                                              this->grid, Norm(), this->grad,
                                                                              n_aniso > 0 ? this->grad_u.get() : nullptr);
                cuda_err(cudaGetLastError());
                cuda_err(cudaMemcpyAsync(host_grad.get(), this->grad, n_atoms * sizeof(float4), cudaMemcpyDeviceToHost,
                                         stream));
                if (n_aniso > 0)
                    cuda_err(cudaMemcpyAsync(host_grad_u.get(), this->grad_u, 6 * n_atoms * sizeof(float),
                                             cudaMemcpyDeviceToHost, stream));
            }
        }
    }

    // Restraint set `which` at host_xyzb into host_sums[1] and, with `grad`, its gradient into
    // host_restraint_grad: queued, not waited for.
    void QueueRestraints(int which, bool grad) {
        Restraints &r = restraints[which];
        if (r.atoms > 0)
            cuda_err(cudaMemcpyAsync(r.xyzb, host_xyzb.get(), r.atoms * sizeof(float4), cudaMemcpyHostToDevice, stream));
        if (r.records > 0) {
            restraint_kernel<<<blocks(r.records), THREADS, 0, 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, stream));
            cuda_err(cudaMemcpyAsync(host_sums.get() + 1, r.sum, sizeof(double), cudaMemcpyDeviceToHost, stream));
        }
        if (grad && r.atoms > 0) {
            if (r.records > 0) {
                slot_gather_kernel<<<blocks(r.atoms), THREADS, 0, 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), stream));
            cuda_err(cudaMemcpyAsync(host_restraint_grad.get(), r.grad, r.atoms * sizeof(float4), cudaMemcpyDeviceToHost,
                                     stream));
        }
    }

    // QueueEvaluate() and QueueRestraints() as their recordings (the sums that no launch writes zeroed here, on
    // the host, so that a replay finds them so).
    void LaunchEvaluate(bool grad) {
        host_sums[0] = 0;
        evaluate.Launch(stream, 4L * n_atoms + (n_aniso > 0 ? 2 : 0) + (grad ? 1 : 0), [&]() { QueueEvaluate(grad); });
    }

    void LaunchRestraints(int which, bool grad) {
        Restraints &r = restraints[which];
        host_sums[1] = 0;
        r.launches.Launch(stream, 2L * r.version + (grad ? 1 : 0), [&]() { QueueRestraints(which, grad); });
    }
};

namespace {

void CopyOut(const float4 *from, int n, std::vector<std::array<float, 4>> &to) {
    to.resize(n);
    std::copy(reinterpret_cast<const std::array<float, 4> *>(from), reinterpret_cast<const std::array<float, 4> *>(from) + n,
              to.begin());
}

}  // namespace

size_t RefinementGPUEngine::FFTWorkBytes(int device, const RefinementGPUSetup &setup) {
    DeviceGuard guard(device);
    size_t bytes = std::max(CudaFFTPlan::WorkBytes(ForwardShape(setup.grid)),
                            CudaFFTPlan::WorkBytes(InverseShape(setup.grid)));
    if (!SameGrid(setup.mask_grid, setup.grid))
        bytes = std::max(bytes, CudaFFTPlan::WorkBytes(ForwardShape(setup.mask_grid)));
    return bytes;
}

size_t RefinementGPUEngine::LaneBytes(const RefinementGPUSetup &s, size_t fft_work_bytes) {
    size_t bytes = GridPoints(s.grid) * sizeof(float) + ComplexPoints(s.grid) * sizeof(float2) + fft_work_bytes;
    if (!SameGrid(s.mask_grid, s.grid))
        bytes += GridPoints(s.mask_grid) * sizeof(float) + ComplexPoints(s.mask_grid) * sizeof(float2);
    return bytes;
}

size_t RefinementGPUEngine::CardBytes(const RefinementGPUSetup &s) {
    const size_t nr = s.rows.size() + 1;
    size_t bytes = 12 * s.sym_ops.size() * sizeof(int) + nr * (3 * sizeof(int) + sizeof(float) + 2);
    if (!s.mate.empty())
        bytes += nr * sizeof(int);
    bytes += (s.scatter_point.size() + s.scatter_start.size() + 2 * s.scatter_term.size()) * sizeof(int);
    int phase_min = 0, phase_max = 0;
    PhaseRange(s, phase_min, phase_max);
    bytes += static_cast<size_t>(phase_max - phase_min + 1) * sizeof(double2);
    bytes += ModelMaskGPU::DeviceBytes(GridPoints(s.mask_grid));
    return bytes;
}

size_t RefinementGPUEngine::DeviceBytes(const RefinementGPUSetup &s, size_t records, size_t slots) {
    const size_t na = s.max_atoms + 1, nr = s.rows.size() + 1;
    size_t bytes = na * (2 * sizeof(float4) + 6 * sizeof(float) + 3 * sizeof(int) + sizeof(ModelMaskAtom) +
                         sizeof(ModelDensityAtom));
    bytes += nr * (2 * sizeof(double) + 3 * sizeof(float) + 4 * sizeof(float2));
    if (!s.mate.empty())
        bytes += nr * 2 * sizeof(float2);
    bytes += ModelDensityGPU::DeviceBytes(s.max_atoms, s.max_pairs, ModelDensityGPU::Bricks(s.grid.nu, s.grid.nv, s.grid.nw));
    // the two restraint sets of SetRestraints(), together up to these records (plane atoms counted as records) and
    // gradient slots
    bytes += records * (sizeof(RestraintRecord) + sizeof(double)) + slots * (sizeof(float4) + sizeof(int)) +
             2 * na * (2 * sizeof(float4) + sizeof(int));
    return bytes;
}

std::shared_ptr<RefinementGPUCard> RefinementGPUEngine::MakeCard(int device, const RefinementGPUSetup &setup) {
    DeviceGuard guard(device);
    return std::make_shared<RefinementGPUCard>(device, setup);
}

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

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(e.waiter.Wait(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, float twin_fraction) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    Upload(e.k_total, k_total, e.stream);
    const std::vector<float> fraction{twin_fraction};   // kept until the wait below
    if (e.twin_fraction)
        Upload(e.twin_fraction, fraction, 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(e.waiter.Wait(e.stream));
}

void RefinementGPUEngine::SetDensities(const ModelDensityAtom *atoms, size_t n) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    if (n > e.setup.max_atoms)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: atoms over the reserve");
    e.density->SetAtoms(e.setup.grid, atoms, n);
    e.n_atoms = static_cast<int>(n);
    e.n_aniso = 0;
    for (size_t i = 0; i < n; i++)
        e.n_aniso += atoms[i].aniso != 0;
}

void RefinementGPUEngine::SetPositions(const std::array<float, 4> *pos, size_t n) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    if (static_cast<int>(n) != e.n_atoms)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: positions do not match the atoms");
    if (n > 0)
        cuda_err(cudaMemcpyAsync(e.pos, pos, n * sizeof(float4), cudaMemcpyHostToDevice, e.stream));
}

void RefinementGPUEngine::Fcalc(std::vector<std::array<float, 2>> &f) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    const std::lock_guard lock(e.lane->lock);
    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(e.waiter.Wait(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));
    const bool own = e.mask_grid != nullptr;
    float *grid = own ? e.mask_grid : e.grid;
    float2 *spectrum = own ? e.mask_spectrum : e.spectrum;
    const std::lock_guard lock(e.lane->lock);
    // the mask on the card's stream, once the atoms are up
    cuda_err(e.waiter.Wait(e.stream));
    {
        const std::lock_guard mask_lock(e.card->mask_lock);
        e.card->mask->Compute(e.mask_atoms, static_cast<int>(mask_atoms.size()), grid);
        cuda_err(cudaStreamSynchronize(e.card->stream));
    }
    cufft_err((own ? e.mask_forward : e.forward).R2C(grid, reinterpret_cast<cufftComplex *>(spectrum)));
    const int rows = static_cast<int>(e.setup.rows.size());
    f.resize(rows);
    if (rows > 0) {
        const ModelDensityGrid &g = e.setup.mask_grid;
        fmask_kernel<<<blocks(rows), THREADS, 0, e.stream>>>(spectrum, e.MaskNorm(), e.row_hkl, rows, g.nu, g.nv, g.nw,
                                                             e.fm);
        cuda_err(cudaGetLastError());
        cuda_err(cudaMemcpyAsync(f.data(), e.fm, rows * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
    }
    cuda_err(e.waiter.Wait(e.stream));
}

double RefinementGPUEngine::Evaluate(bool grad) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    const std::lock_guard lock(e.lane->lock);
    e.LaunchEvaluate(grad);
    cuda_err(e.waiter.Wait(e.stream));
    return e.host_sums[0];
}

const std::array<float, 4> *RefinementGPUEngine::HostGradient() const {
    return reinterpret_cast<const std::array<float, 4> *>(impl_->host_grad.get());
}

const float *RefinementGPUEngine::HostUGradient() const {
    return impl_->host_grad_u.get();
}

const std::array<float, 4> *RefinementGPUEngine::HostRestraintGradient() const {
    return reinterpret_cast<const std::array<float, 4> *>(impl_->host_restraint_grad.get());
}

double RefinementGPUEngine::EvaluateWithRestraints(bool grad, int which,
                                                   const std::function<void(std::array<float, 4> *)> &xyzb,
                                                   bool restraint_grad, double &restraints) {
    RefinementGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    RefinementGPUEngineImpl::Restraints &r = e.restraints[which];
    if (r.atoms != e.n_atoms)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: restraint atoms do not match");
    const std::lock_guard lock(e.lane->lock);
    e.LaunchEvaluate(grad);
    xyzb(reinterpret_cast<std::array<float, 4> *>(e.host_xyzb.get()));
    e.LaunchRestraints(which, restraint_grad);
    cuda_err(e.waiter.Wait(e.stream));
    restraints = e.host_sums[1];
    return e.host_sums[0];
}

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);
    r.version++;
    cuda_err(e.waiter.Wait(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");
    // the staging buffer free again
    cuda_err(e.waiter.Wait(e.stream));
    std::copy(xyzb.begin(), xyzb.end(), reinterpret_cast<std::array<float, 4> *>(e.host_xyzb.get()));
    e.LaunchRestraints(which, grad != nullptr);
    cuda_err(e.waiter.Wait(e.stream));
    if (grad != nullptr)
        CopyOut(e.host_restraint_grad, r.atoms, *grad);
    return e.host_sums[1];
}