Jungfraujoch/image_analysis/structure_refinement/RigidBodyGPU.cu

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

#include "RigidBodyGPUEngine.h"

#include <algorithm>
#include <cmath>
#include <map>
#include <string>

#include <cufft.h>

#include "ModelDensityGPU.h"
#include "ModelMaskGPU.h"
#include "ModelScaleGPU.h"
#include "../../common/JFJochException.h"
#include "../indexing/CUDAMemHelpers.h"

// Every kernel below is deterministic: nothing is summed with a floating-point atomic, each grid point
// adds its atoms in model order, and every reduction runs in a fixed order. The same placement gives the
// same residuals, bit for bit, on every run and whichever engine runs it.

namespace {

void cuda_err(cudaError_t val) {
    if (val != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}

// The calling thread on the engine's card for one call, and back on its own card after: a caller's later work
// - a function of the search, or memory it allocates - runs on the card the caller chose, not on whichever card
// the last engine it borrowed is on.
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;
};

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 MAX_SCALE_PARAMS = 7;       // k_overall + up to six B* constraints

// The grid of a zone.
struct GridGeom {
    int nu, nv, nw;
};

// Where the body is: x -> R x_rel + t, and the fractionalization.
struct Placement {
    double r[9];
    double t[3];
    double frac[9];
};

// The scale a residual or a Jacobian row is taken at.
struct Scale {
    double k_overall;
    double b_star[6];       // u11 u22 u33 u12 u13 u23
    double k_sol, b_sol;
    int n_constraints;
    double constraints[6][6];
};

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

// b_star.r_u_r(h), as gemmi's SMat33.
__device__ __forceinline__ double r_u_r(const double *u, const int *h) {
    const double x = h[0], y = h[1], z = h[2];
    return x * x * u[0] + y * y * u[1] + z * z * u[2] + 2 * (x * y * u[3] + x * z * u[4] + y * z * u[5]);
}

// Each atom placed: its grid coordinates (fractional times the grid size, not wrapped, as the CPU's density
// walk centres an atom) for the gather and, for a mask atom, its wrapped fractional coordinates and mask
// radius for the mask.
__global__ void place_kernel(const double *rel, int n, Placement p, int nu, int nv, int nw, const int *mask_slot,
                             const float *mask_radius, float4 *pos, double *mask_atoms) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n)
        return;
    const double x = rel[3 * i], y = rel[3 * i + 1], z = rel[3 * i + 2];
    const double c[3] = {p.r[0] * x + p.r[1] * y + p.r[2] * z + p.t[0],
                         p.r[3] * x + p.r[4] * y + p.r[5] * z + p.t[1],
                         p.r[6] * x + p.r[7] * y + p.r[8] * z + p.t[2]};
    const int n3[3] = {nu, nv, nw};
    double f[3];
    float g[3];
    for (int k = 0; k < 3; k++) {
        f[k] = p.frac[3 * k] * c[0] + p.frac[3 * k + 1] * c[1] + p.frac[3 * k + 2] * c[2];
        g[k] = static_cast<float>(f[k] * n3[k]);
        f[k] -= floor(f[k]);
        if (f[k] >= 1.0)   // a tiny negative coordinate wraps to 1 - epsilon, which rounds to 1
            f[k] = 0.0;
    }
    pos[i] = make_float4(g[0], g[1], g[2], 0.0f);
    const int slot = mask_slot[i];
    if (slot >= 0) {
        for (int k = 0; k < 3; k++)
            mask_atoms[4 * slot + k] = f[k];
        mask_atoms[4 * slot + 3] = mask_radius[i];
    }
}

// F at the zone's rows composed from the transform of one copy (SymmetryComposition::Compose), with
// dF/dt where `df_dt` is not null. `f1` is cuFFT's forward transform, which carries exp(-2 pi i h.x):
// conjugated and scaled by V/N it is gemmi's F1.
__global__ void compose_kernel(const float2 *f1, float norm, const RigidBodyGPUTerm *terms, const double *row_scale,
                               int rows, int ops, const double2 *fixed, double2 *f, double2 *df_dt) {
    const int m = blockIdx.x * blockDim.x + threadIdx.x;
    if (m >= rows)
        return;
    double sr = 0, si = 0, d[3][2] = {};
    for (int o = 0; o < ops; o++) {
        const RigidBodyGPUTerm t = terms[static_cast<size_t>(m) * ops + o];
        const float2 out = f1[t.index];
        double vr = out.x * norm, vi = -out.y * norm;   // F1 at the index read
        if (t.conj)
            vi = -vi;
        const double tr = t.phase[0] * vr - t.phase[1] * vi, ti = t.phase[0] * vi + t.phase[1] * vr;
        sr += tr;
        si += ti;
        for (int k = 0; k < 3; k++) {
            d[k][0] += tr * t.s[k];
            d[k][1] += ti * t.s[k];
        }
    }
    const double scale = row_scale[m];
    f[m] = fixed != nullptr ? make_double2(scale * sr + fixed[m].x, scale * si + fixed[m].y)
                            : make_double2(scale * sr, scale * si);
    if (df_dt != nullptr) {
        // times 2 pi i
        const double two_pi = 2 * 3.14159265358979323846;
        for (int k = 0; k < 3; k++)
            df_dt[3 * static_cast<size_t>(m) + k] = make_double2(-scale * two_pi * d[k][1], scale * two_pi * d[k][0]);
    }
}

__global__ void add_kernel(const double2 *a, int n, double2 *sum) {
    const int m = blockIdx.x * blockDim.x + threadIdx.x;
    if (m < n)
        sum[m] = make_double2(sum[m].x + a[m].x, sum[m].y + a[m].y);
}

// Fmask at each row, read at h directly (the mask is symmetric): gemmi's get_value_by_hkl() on the
// transform of the mask, taking the Friedel mate where h is outside the stored half.
__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);
}

// Scaling<float>::Point's fcmol and fmask of each point.
__global__ void pack_points_kernel(const int *point_obs, const int *obs_row, int n, const double2 *f,
                                   const float2 *fmask, float2 *fcmol_p, float2 *fmask_p) {
    const int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p >= n)
        return;
    const int m = obs_row[point_obs[p]];
    fcmol_p[p] = make_float2(static_cast<float>(f[m].x), static_cast<float>(f[m].y));
    fmask_p[p] = fmask[m];
}

// gemmi's Scaling::get_fcalc() of a row: fcmol + (float) solvent scale * fmask, in float.
__device__ float2 solvent_fcalc(const double2 &f, const float2 &fm, double stol2, const Scale &sc) {
    const float solvent = static_cast<float>(sc.k_sol * exp(-sc.b_sol * stol2));
    return make_float2(static_cast<float>(f.x) + solvent * fm.x, static_cast<float>(f.y) + solvent * fm.y);
}

// The residual of each observation: scale_data() applied to the row, then (Fobs - |F|) / <Fobs>.
__global__ void residual_kernel(const int *obs_row, const float *obs_fobs, int n, const double2 *f,
                                const float2 *fmask, const int *row_hkl, const double *row_stol2, Scale sc,
                                double f_mean, double *residual) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n)
        return;
    const int m = obs_row[i];
    if (m < 0) {
        residual[i] = 0.0;
        return;
    }
    float2 value = solvent_fcalc(f[m], fmask[m], row_stol2[m], sc);
    const float k = static_cast<float>(sc.k_overall * exp(-0.25 * r_u_r(sc.b_star, row_hkl + 3 * m)));
    value.x *= k;
    value.y *= k;
    residual[i] = (obs_fobs[i] - hypotf(value.x, value.y)) / f_mean;
}

// One row of the fixed-scale Jacobian and of the scale parameters' own columns (RigidBodyTarget::Jacobian).
__global__ void jacobian_kernel(const int *obs_row, int n, const double2 *f, const double2 *df_dt,
                                const double2 *shifted, size_t rows, const float2 *fmask, const int *row_hkl,
                                const double *row_stol2, Scale sc, double f_mean, double step, double *jacobian,
                                double *jk) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n)
        return;
    const int p = 1 + sc.n_constraints;
    const int m = obs_row[i];
    if (m < 0) {
        for (int j = 0; j < 6; j++)
            jacobian[6 * static_cast<size_t>(i) + j] = 0.0;
        for (int a = 0; a < p; a++)
            jk[static_cast<size_t>(i) * p + a] = 0.0;
        return;
    }
    const int *h = row_hkl + 3 * m;
    // compute_value_and_derivatives() with k_sol and b_sol fixed
    const float2 fcalc = solvent_fcalc(f[m], fmask[m], row_stol2[m], sc);
    const double fcalc_abs = hypotf(fcalc.x, fcalc.y);
    const double kaniso = exp(-0.25 * r_u_r(sc.b_star, h));
    const double fe = fcalc_abs * kaniso;
    const double y = sc.k_overall * fe;
    const double x = h[0], yy = h[1], z = h[2];
    const double du[6] = {-0.25 * y * (x * x), -0.25 * y * (yy * yy), -0.25 * y * (z * z),
                          -0.5 * y * (x * yy), -0.5 * y * (x * z), -0.5 * y * (yy * z)};
    jk[static_cast<size_t>(i) * p] = -fe / f_mean;
    for (int c = 0; c < sc.n_constraints; c++) {
        double dot = 0;
        for (int q = 0; q < 6; q++)
            dot += sc.constraints[c][q] * du[q];
        jk[static_cast<size_t>(i) * p + 1 + c] = -dot / f_mean;
    }
    // d|F_scaled|/dq = K Re(conj(F_total) dFcalc/dq) / |F_total|
    const double k = sc.k_overall * kaniso / (fcalc_abs * f_mean);
    for (int j = 0; j < 6; j++) {
        double2 d;
        if (j < 3) {
            const double2 s = shifted[j * rows + m];
            d = make_double2((s.x - f[m].x) / step, (s.y - f[m].y) / step);
        } else {
            d = df_dt[3 * static_cast<size_t>(m) + j - 3];
        }
        jacobian[6 * static_cast<size_t>(i) + j] = -k * (fcalc.x * d.x + fcalc.y * d.y);
    }
}

// out[b] = sum over i of a[i * lda + ia] * c[i * ldc + ic] for the pair b = (ia, ic): one block per pair,
// a fixed stride per thread and a fixed shared-memory tree, so the sum is the same on every run.
__global__ void products_kernel(const double *a, int lda, const double *c, int ldc, int n, int nc, double *out) {
    const int ia = blockIdx.x / nc, ic = blockIdx.x % nc;
    double sum = 0;
    for (int i = threadIdx.x; i < n; i += THREADS)
        sum += a[static_cast<size_t>(i) * lda + ia] * c[static_cast<size_t>(i) * ldc + ic];
    __shared__ double sh[THREADS];
    sh[threadIdx.x] = sum;
    __syncthreads();
    for (int s = THREADS / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s)
            sh[threadIdx.x] += sh[threadIdx.x + s];
        __syncthreads();
    }
    if (threadIdx.x == 0)
        out[blockIdx.x] = sh[0];
}

// J - J_k x, the scale re-fit projected out.
__global__ void project_kernel(double *jacobian, const double *jk, int p, const double *x, int n) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n)
        return;
    for (int j = 0; j < 6; j++) {
        double v = jacobian[6 * static_cast<size_t>(i) + j];
        for (int a = 0; a < p; a++)
            v -= jk[static_cast<size_t>(i) * p + a] * x[a * 6 + j];
        jacobian[6 * static_cast<size_t>(i) + j] = v;
    }
}

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

// An (nu, nv, nw) R2C transform, u fastest, `batch` grids back to back.
CudaFFTShape GridShape(int nu, int nv, int nw, int batch) {
    return {.n = {nw, nv, nu}, .idist = nu * nv * nw, .odist = (nu / 2 + 1) * nv * nw, .type = CUFFT_R2C,
            .batch = batch};
}

int GreatestStreamPriority() {
    int least = 0, greatest = 0;
    cuda_err(cudaDeviceGetStreamPriorityRange(&least, &greatest));
    return greatest;
}

}  // namespace

// The zone's grid as the density gather takes it.
static ModelDensityGrid DensityGrid(const RigidBodyGPUZone &zone) {
    ModelDensityGrid g;
    g.nu = zone.nu;
    g.nv = zone.nv;
    g.nw = zone.nw;
    std::copy(zone.orth, zone.orth + 9, g.orth);
    std::copy(zone.frac, zone.frac + 9, g.frac);
    return g;
}

size_t RigidBodyGPUEngine::AxisBrickBound(int d, int n) {
    return ModelDensityGPU::AxisBrickBound(d, n);
}

struct RigidBodyGPUEngineImpl {
    int device = 0;
    RigidBodyGPUCapacity cap;
    // At the device's highest priority: a validation can run beside a merge, whose kernels would
    // otherwise fill the card ahead of a fit's short ones, one evaluation at a time.
    CudaStream stream{cudaStreamNonBlocking, GreatestStreamPriority()};

    CudaDevicePtr<double> rel;
    CudaDevicePtr<float4> pos;
    CudaDevicePtr<int> mask_slot;
    CudaDevicePtr<float> mask_radius;
    CudaDevicePtr<double> mask_atoms;   // 4 per mask atom: fractional x, y, z in [0,1), radius
    CudaDevicePtr<double> copies_mask_atoms;   // the same for the body and fixed copies of it, grown as needed
    size_t copies_mask_capacity = 0;

    CudaDevicePtr<float> grid;          // 3 grids back to back
    CudaDevicePtr<float2> spectrum;     // their 3 transforms
    CudaDevicePtr<char> fft_work;

    CudaDevicePtr<RigidBodyGPUTerm> terms;
    CudaDevicePtr<int> row_hkl;
    CudaDevicePtr<double> row_scale, row_stol2;
    CudaDevicePtr<double2> f, df_dt, shifted, fixed;
    bool has_fixed = false;   // the zone's fixed Fcalc is in `fixed`
    CudaDevicePtr<float2> fmask;

    CudaDevicePtr<int> obs_row, point_obs;
    CudaDevicePtr<float> obs_fobs;
    CudaDevicePtr<double> residual, jacobian, jk;
    CudaDevicePtr<float2> point_fcmol, point_fmask;
    CudaDevicePtr<double> products, x;

    std::map<std::array<int, 4>, CudaFFTPlan> plans;
    std::unique_ptr<ModelDensityGPU> density;
    std::unique_ptr<ModelMaskGPU> mask;
    std::unique_ptr<ModelScaleGPU> scale;

    // The zone.
    int n_atoms = 0;
    GridGeom geom{};
    double frac[9] = {};
    double volume = 0;
    size_t rows = 0, ops = 1, n_obs = 0, n_points = 0, n_constraints = 0, n_mask_atoms = 0;
    std::vector<std::array<double, 6>> constraints;
    double f_mean = 1;

    explicit RigidBodyGPUEngineImpl(const RigidBodyGPUCapacity &c, int dev)
        : device(dev), cap(c) {
        // A card busy with a probe pass beside the run is waited for, not given up on (RigidBodyGPUPool).
        const auto sync = CudaAlloc::SynchronousAfterWorkBeside;
        const size_t na = std::max<size_t>(cap.atoms, 1);
        rel = CudaDevicePtr<double>(3 * na, sync);
        pos = CudaDevicePtr<float4>(na, sync);
        mask_slot = CudaDevicePtr<int>(na, sync);
        mask_radius = CudaDevicePtr<float>(na, sync);
        mask_atoms = CudaDevicePtr<double>(4 * na, sync);
        grid = CudaDevicePtr<float>(3 * cap.grid_points, sync);
        spectrum = CudaDevicePtr<float2>(3 * cap.complex_points, sync);
        fft_work = CudaDevicePtr<char>(std::max<size_t>(cap.fft_work_bytes, 1), sync);
        const size_t nr = std::max<size_t>(cap.rows, 1), no = std::max<size_t>(cap.observations, 1);
        terms = CudaDevicePtr<RigidBodyGPUTerm>(std::max<size_t>(cap.terms, 1), sync);
        row_hkl = CudaDevicePtr<int>(3 * nr, sync);
        row_scale = CudaDevicePtr<double>(nr, sync);
        row_stol2 = CudaDevicePtr<double>(nr, sync);
        f = CudaDevicePtr<double2>(nr, sync);
        df_dt = CudaDevicePtr<double2>(3 * nr, sync);
        shifted = CudaDevicePtr<double2>(3 * nr, sync);
        fixed = CudaDevicePtr<double2>(nr, sync);
        fmask = CudaDevicePtr<float2>(nr, sync);
        obs_row = CudaDevicePtr<int>(no, sync);
        point_obs = CudaDevicePtr<int>(no, sync);
        obs_fobs = CudaDevicePtr<float>(no, sync);
        residual = CudaDevicePtr<double>(no, sync);
        jacobian = CudaDevicePtr<double>(6 * no, sync);
        jk = CudaDevicePtr<double>(MAX_SCALE_PARAMS * no, sync);
        point_fcmol = CudaDevicePtr<float2>(no, sync);
        point_fmask = CudaDevicePtr<float2>(no, sync);
        products = CudaDevicePtr<double>(MAX_SCALE_PARAMS * (MAX_SCALE_PARAMS + 6), sync);
        x = CudaDevicePtr<double>(MAX_SCALE_PARAMS * 6, sync);
        density = std::make_unique<ModelDensityGPU>(stream, cap.atoms, cap.pairs, cap.bricks);
        mask = std::make_unique<ModelMaskGPU>(stream, std::max<size_t>(cap.grid_points, 1));
        scale = std::make_unique<ModelScaleGPU>(stream, no);
    }

    // The points' fcmol and fmask at this evaluation.
    void PackPoints() {
        if (n_points == 0)
            return;
        pack_points_kernel<<<blocks(n_points), THREADS, 0, stream>>>(point_obs, obs_row, static_cast<int>(n_points),
                                                                    f, fmask, point_fcmol, point_fmask);
        cuda_err(cudaGetLastError());
    }

    ~RigidBodyGPUEngineImpl() {
        cudaSetDevice(device);
        cudaStreamSynchronize(stream);
    }

    cufftHandle Plan(int batch) {
        const std::array<int, 4> k{geom.nu, geom.nv, geom.nw, batch};
        auto it = plans.find(k);
        if (it != plans.end())
            return it->second;
        CudaFFTPlan plan(GridShape(geom.nu, geom.nv, geom.nw, batch), stream, fft_work);
        if (plan.WorkBytes() > cap.fft_work_bytes)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: cuFFT work area over the reserve");
        return plans.emplace(k, std::move(plan)).first->second;
    }

    size_t GridPoints() const { return static_cast<size_t>(geom.nu) * geom.nv * geom.nw; }
    size_t ComplexPoints() const { return static_cast<size_t>(geom.nu / 2 + 1) * geom.nv * geom.nw; }
    float Norm() const { return static_cast<float>(volume / static_cast<double>(GridPoints())); }

    // The body placed at (R, t) and its density gathered into grid `slot`.
    void Density(const double r[9], const double t[3], int slot) {
        Placement p{};
        std::copy(r, r + 9, p.r);
        std::copy(t, t + 3, p.t);
        std::copy(frac, frac + 9, p.frac);
        place_kernel<<<blocks(n_atoms), THREADS, 0, stream>>>(rel, n_atoms, p, geom.nu, geom.nv, geom.nw, mask_slot,
                                                            mask_radius, pos, mask_atoms);
        cuda_err(cudaGetLastError());
        density->Compute(pos, grid.get() + slot * GridPoints());
    }

    void Compose(int slot, double2 *out, double2 *derivative) {
        compose_kernel<<<blocks(rows), THREADS, 0, stream>>>(spectrum.get() + slot * ComplexPoints(), Norm(), terms,
                                                            row_scale, static_cast<int>(rows), static_cast<int>(ops),
                                                            has_fixed ? fixed.get() : nullptr, out, derivative);
        cuda_err(cudaGetLastError());
    }

    Scale MakeScale(double k_overall, const double b_star[6], double k_sol, double b_sol) const {
        Scale s{};
        s.k_overall = k_overall;
        std::copy(b_star, b_star + 6, s.b_star);
        s.k_sol = k_sol;
        s.b_sol = b_sol;
        s.n_constraints = static_cast<int>(n_constraints);
        for (size_t c = 0; c < n_constraints; c++)
            for (int q = 0; q < 6; q++)
                s.constraints[c][q] = constraints[c][q];
        return s;
    }
};

size_t RigidBodyGPUEngine::Bricks(int nu, int nv, int nw) {
    return ModelDensityGPU::Bricks(nu, nv, nw);
}

size_t RigidBodyGPUEngine::PairBound(const RigidBodyGPUZone &zone) {
    return ModelDensityGPU::PairBound(DensityGrid(zone), zone.atoms);
}

bool RigidBodyGPUEngine::Supports(const RigidBodyGPUZone &zone) {
    return ModelDensityGPU::Supports(DensityGrid(zone), zone.atoms);
}

int RigidBodyGPUEngine::CurrentDevice() {
    int device = 0;
    cuda_err(cudaGetDevice(&device));
    return device;
}

size_t RigidBodyGPUEngine::FFTWorkBytes(int nu, int nv, int nw) {
    size_t largest = 0;
    for (int batch : {1, 3})
        largest = std::max(largest, CudaFFTPlan::WorkBytes(GridShape(nu, nv, nw, batch)));
    return largest;
}

size_t RigidBodyGPUEngine::DeviceBytes(const RigidBodyGPUCapacity &c) {
    const size_t na = c.atoms + 1, no = c.observations + 1, nr = c.rows + 1;
    size_t bytes = na * (7 * sizeof(double) + sizeof(float4) + sizeof(int) + sizeof(float));
    bytes += 3 * c.grid_points * sizeof(float) + 3 * c.complex_points * sizeof(float2) + c.fft_work_bytes;
    bytes += ModelDensityGPU::DeviceBytes(c.atoms, c.pairs, c.bricks);
    bytes += c.terms * sizeof(RigidBodyGPUTerm);
    bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 8 * sizeof(double2) + sizeof(float2));
    bytes += no * (2 * sizeof(int) + sizeof(float) + (1 + 6 + MAX_SCALE_PARAMS) * sizeof(double) + 2 * sizeof(float2));
    bytes += ModelMaskGPU::DeviceBytes(c.grid_points) + ModelScaleGPU::DeviceBytes(no);
    return bytes;
}

RigidBodyGPUEngine::RigidBodyGPUEngine(const RigidBodyGPUCapacity &capacity, int device) {
    const DeviceGuard guard(device);
    impl_ = std::make_unique<RigidBodyGPUEngineImpl>(capacity, device);
}

// Released on its own card, which need not be the one the destroying thread is on.
RigidBodyGPUEngine::~RigidBodyGPUEngine() {
    int current = 0;
    cudaGetDevice(&current);
    cudaSetDevice(impl_->device);
    impl_.reset();
    cudaSetDevice(current);
}

int RigidBodyGPUEngine::Device() const {
    return impl_->device;
}

void RigidBodyGPUEngine::SetBody(const std::vector<std::array<double, 3>> &relative) {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    if (relative.size() > e.cap.atoms)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: more atoms than the reserve");
    e.n_atoms = static_cast<int>(relative.size());
    cuda_err(cudaMemcpyAsync(e.rel, relative.data(), relative.size() * 3 * sizeof(double), cudaMemcpyHostToDevice,
                             e.stream));
    cuda_err(cudaStreamSynchronize(e.stream));
}

void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    const size_t np = static_cast<size_t>(zone.nu) * zone.nv * zone.nw;
    const size_t nc = static_cast<size_t>(zone.nu / 2 + 1) * zone.nv * zone.nw;
    if (static_cast<int>(zone.atoms.size()) != e.n_atoms || np > e.cap.grid_points || nc > e.cap.complex_points ||
        Bricks(zone.nu, zone.nv, zone.nw) > e.cap.bricks || zone.row_hkl.size() > e.cap.rows ||
        zone.terms.size() > e.cap.terms || zone.obs_row.size() > e.cap.observations ||
        zone.constraints.size() + 1 > MAX_SCALE_PARAMS)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: zone over the engine's reserve");

    e.geom.nu = zone.nu;
    e.geom.nv = zone.nv;
    e.geom.nw = zone.nw;
    std::copy(zone.frac, zone.frac + 9, e.frac);
    e.volume = zone.volume;
    e.rows = zone.row_hkl.size();
    e.ops = zone.ops;
    e.n_obs = zone.obs_row.size();
    e.n_points = zone.point_obs.size();
    e.n_mask_atoms = zone.mask_atom.size();
    e.n_constraints = zone.constraints.size();
    e.constraints = zone.constraints;
    e.f_mean = zone.f_mean;
    e.has_fixed = false;

    std::vector<int> slot(zone.atoms.size(), -1);
    std::vector<float> radius(zone.atoms.size(), 0.0f);
    for (size_t j = 0; j < zone.mask_atom.size(); j++) {
        slot[zone.mask_atom[j]] = static_cast<int>(j);
        radius[zone.mask_atom[j]] = zone.mask_radius[j];
    }

    auto up = [&](void *dst, const void *src, size_t bytes) {
        if (bytes > 0)
            cuda_err(cudaMemcpyAsync(dst, src, bytes, cudaMemcpyHostToDevice, e.stream));
    };
    up(e.mask_slot, slot.data(), slot.size() * sizeof(int));
    up(e.mask_radius, radius.data(), radius.size() * sizeof(float));
    up(e.terms, zone.terms.data(), zone.terms.size() * sizeof(RigidBodyGPUTerm));
    up(e.row_hkl, zone.row_hkl.data(), zone.row_hkl.size() * 3 * sizeof(int));
    up(e.row_scale, zone.row_scale.data(), zone.row_scale.size() * sizeof(double));
    up(e.row_stol2, zone.row_stol2.data(), zone.row_stol2.size() * sizeof(double));
    up(e.obs_row, zone.obs_row.data(), zone.obs_row.size() * sizeof(int));
    up(e.obs_fobs, zone.obs_fobs.data(), zone.obs_fobs.size() * sizeof(float));
    up(e.point_obs, zone.point_obs.data(), zone.point_obs.size() * sizeof(int));
    ModelMaskGrid mask_grid;
    mask_grid.nu = zone.nu;
    mask_grid.nv = zone.nv;
    mask_grid.nw = zone.nw;
    std::copy(zone.orth, zone.orth + 9, mask_grid.orth);
    mask_grid.volume = zone.volume;
    std::vector<ModelMaskOp> images(zone.images.size());
    for (size_t i = 0; i < images.size(); i++) {
        std::copy(zone.images[i].begin(), zone.images[i].begin() + 9, images[i].rot);
        std::copy(zone.images[i].begin() + 9, zone.images[i].end(), images[i].tran);
    }
    e.mask->SetGrid(mask_grid, images);
    e.density->SetAtoms(DensityGrid(zone), zone.atoms);
    e.scale->SetPoints(zone.point_hkl, zone.point_stol2, zone.point_fobs, zone.point_sigma, zone.constraints, zone.frac);
    // The previous zone's plans go back to be lent to another engine; this one takes the zone's own.
    e.plans.clear();
    e.Plan(1);
    e.Plan(3);
    cuda_err(cudaStreamSynchronize(e.stream));
}

void RigidBodyGPUEngine::Fcalc(const double rotation[9], const double translation[3]) {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    e.Density(rotation, translation, 0);
    cufft_err(cufftExecR2C(e.Plan(1), e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
    e.Compose(0, e.f, e.df_dt);
}

void RigidBodyGPUEngine::SetFixedCopies(const std::vector<std::array<double, 12>> &fixed,
                                        const std::array<double, 12> &moving) {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    const size_t copies = fixed.size() + 1;
    if (e.copies_mask_capacity < copies * e.n_mask_atoms) {
        e.copies_mask_capacity = copies * e.n_mask_atoms;
        e.copies_mask_atoms = CudaDevicePtr<double>(4 * std::max<size_t>(e.copies_mask_capacity, 1), CudaAlloc::Synchronous);
    }
    cuda_err(cudaMemsetAsync(e.fixed, 0, e.rows * sizeof(double2), e.stream));
    for (size_t c = 0; c < copies; c++) {
        const std::array<double, 12> &x = c < fixed.size() ? fixed[c] : moving;
        Placement p{};
        std::copy(x.begin(), x.begin() + 9, p.r);
        std::copy(x.begin() + 9, x.end(), p.t);
        std::copy(e.frac, e.frac + 9, p.frac);
        place_kernel<<<blocks(e.n_atoms), THREADS, 0, e.stream>>>(e.rel, e.n_atoms, p, e.geom.nu, e.geom.nv, e.geom.nw,
                                                                e.mask_slot, e.mask_radius, e.pos,
                                                                e.copies_mask_atoms.get() + 4 * c * e.n_mask_atoms);
        cuda_err(cudaGetLastError());
        if (c < fixed.size()) {
            e.density->Compute(e.pos, e.grid.get());
            cufft_err(cufftExecR2C(e.Plan(1), e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
            e.has_fixed = false;
            e.Compose(0, e.shifted, nullptr);
            add_kernel<<<blocks(e.rows), THREADS, 0, e.stream>>>(e.shifted, static_cast<int>(e.rows), e.fixed);
            cuda_err(cudaGetLastError());
        }
    }
    float *mask_grid = e.grid.get() + e.GridPoints();
    cufftComplex *mask_spectrum = reinterpret_cast<cufftComplex *>(e.spectrum.get() + e.ComplexPoints());
    e.mask->Compute(reinterpret_cast<const ModelMaskAtom *>(e.copies_mask_atoms.get()),
                    static_cast<int>(copies * e.n_mask_atoms), mask_grid);
    cufft_err(cufftExecR2C(e.Plan(1), mask_grid, mask_spectrum));
    fmask_kernel<<<blocks(e.rows), THREADS, 0, e.stream>>>(reinterpret_cast<float2 *>(mask_spectrum), e.Norm(),
                                                          e.row_hkl, static_cast<int>(e.rows), e.geom.nu, e.geom.nv,
                                                          e.geom.nw, e.fmask);
    cuda_err(cudaGetLastError());
    cuda_err(cudaStreamSynchronize(e.stream));
    e.has_fixed = true;
}

void RigidBodyGPUEngine::Fmask() {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    float *mask_grid = e.grid.get() + e.GridPoints();
    cufftComplex *mask_spectrum = reinterpret_cast<cufftComplex *>(e.spectrum.get() + e.ComplexPoints());
    // The mask atoms Fcalc() placed, four doubles each, as ModelMaskAtom lays them out.
    e.mask->Compute(reinterpret_cast<const ModelMaskAtom *>(e.mask_atoms.get()), static_cast<int>(e.n_mask_atoms),
                    mask_grid);
    cufft_err(cufftExecR2C(e.Plan(1), mask_grid, mask_spectrum));
    fmask_kernel<<<blocks(e.rows), THREADS, 0, e.stream>>>(reinterpret_cast<float2 *>(mask_spectrum), e.Norm(),
                                                          e.row_hkl, static_cast<int>(e.rows), e.geom.nu, e.geom.nv,
                                                          e.geom.nw, e.fmask);
    cuda_err(cudaGetLastError());
}

void RigidBodyGPUEngine::DownloadPoints(std::vector<std::array<float, 2>> &fcmol,
                                        std::vector<std::array<float, 2>> &fmask) {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    fcmol.resize(e.n_points);
    fmask.resize(e.n_points);
    if (e.n_points == 0)
        return;
    e.PackPoints();
    cuda_err(cudaMemcpyAsync(fcmol.data(), e.point_fcmol, e.n_points * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
    cuda_err(cudaMemcpyAsync(fmask.data(), e.point_fmask, e.n_points * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
    cuda_err(cudaStreamSynchronize(e.stream));
}

void RigidBodyGPUEngine::FitSolvent(double b_hi, double &k_sol, double &b_sol) {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    e.PackPoints();
    const ModelSolventFit fit = e.scale->FitSolvent(e.point_fcmol, e.point_fmask, b_hi);
    k_sol = fit.k_sol;
    b_sol = fit.b_sol;
}

void RigidBodyGPUEngine::FitScale(double k_sol, double b_sol, double &k_overall, double b_star[6]) {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    e.PackPoints();
    const ModelScaleParams p = e.scale->Fit(e.point_fcmol, e.point_fmask, k_sol, b_sol);
    k_overall = p.k_overall;
    std::copy(p.b_star, p.b_star + 6, b_star);
}

void RigidBodyGPUEngine::Residuals(double k_overall, const double b_star[6], double k_sol, double b_sol,
                                   double *residuals) {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    const Scale sc = e.MakeScale(k_overall, b_star, k_sol, b_sol);
    residual_kernel<<<blocks(e.n_obs), THREADS, 0, e.stream>>>(e.obs_row, e.obs_fobs, static_cast<int>(e.n_obs), e.f,
                                                              e.fmask, e.row_hkl, e.row_stol2, sc, e.f_mean,
                                                              e.residual);
    cuda_err(cudaGetLastError());
    cuda_err(cudaMemcpyAsync(residuals, e.residual, e.n_obs * sizeof(double), cudaMemcpyDeviceToHost, e.stream));
    cuda_err(cudaStreamSynchronize(e.stream));
}

void RigidBodyGPUEngine::Jacobian(const double rotation[3][9], const double translation[3][3], double step,
                                  double k_overall, const double b_star[6], double k_sol, double b_sol,
                                  std::vector<double> &jtj, std::vector<double> &jtq) {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    for (int j = 0; j < 3; j++)
        e.Density(rotation[j], translation[j], j);
    cufft_err(cufftExecR2C(e.Plan(3), e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
    for (int j = 0; j < 3; j++)
        e.Compose(j, e.shifted.get() + j * e.rows, nullptr);

    const Scale sc = e.MakeScale(k_overall, b_star, k_sol, b_sol);
    const int p = 1 + static_cast<int>(e.n_constraints);
    const int n = static_cast<int>(e.n_obs);
    jacobian_kernel<<<blocks(e.n_obs), THREADS, 0, e.stream>>>(e.obs_row, n, e.f, e.df_dt, e.shifted, e.rows, e.fmask,
                                                              e.row_hkl, e.row_stol2, sc, e.f_mean, step, e.jacobian,
                                                              e.jk);
    cuda_err(cudaGetLastError());
    products_kernel<<<p * p, THREADS, 0, e.stream>>>(e.jk, p, e.jk, p, n, p, e.products);
    cuda_err(cudaGetLastError());
    products_kernel<<<p * 6, THREADS, 0, e.stream>>>(e.jk, p, e.jacobian, 6, n, 6, e.products.get() + p * p);
    cuda_err(cudaGetLastError());
    std::vector<double> out(p * p + p * 6);
    cuda_err(cudaMemcpyAsync(out.data(), e.products, out.size() * sizeof(double), cudaMemcpyDeviceToHost, e.stream));
    cuda_err(cudaStreamSynchronize(e.stream));
    jtj.assign(out.begin(), out.begin() + p * p);
    jtq.assign(out.begin() + p * p, out.end());
}

void RigidBodyGPUEngine::ProjectJacobian(const std::vector<double> &x, double *jacobian) {
    RigidBodyGPUEngineImpl &e = *impl_;
    const DeviceGuard guard(e.device);
    const int p = 1 + static_cast<int>(e.n_constraints);
    cuda_err(cudaMemcpyAsync(e.x, x.data(), x.size() * sizeof(double), cudaMemcpyHostToDevice, e.stream));
    project_kernel<<<blocks(e.n_obs), THREADS, 0, e.stream>>>(e.jacobian, e.jk, p, e.x, static_cast<int>(e.n_obs));
    cuda_err(cudaGetLastError());
    cuda_err(cudaMemcpyAsync(jacobian, e.jacobian, e.n_obs * 6 * sizeof(double), cudaMemcpyDeviceToHost, e.stream));
    cuda_err(cudaStreamSynchronize(e.stream));
}