Jungfraujoch/image_analysis/structure_refinement/ModelStructureFactorsGPU.cu

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

#include "ModelStructureFactorsGPUEngine.h"

#include <algorithm>
#include <string>

#include <cufft.h>

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

// Every kernel below is deterministic: the density is ModelDensityGPU's gather, the mask ModelMaskGPU's,
// and each reflection is composed by one thread over the operators in order. The same atoms give the same
// structure factors, bit for bit, on every run.

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;

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

// F of the crystal at each row, composed from the transform of one copy of its content (SymmetryComposition,
// RigidBodyRefine.h): F(h) = n_cen * unblur * sum over the operators x -> Rx + t of exp(+2 pi i h.t) F1(hR).
// `f1` is cuFFT's r2c transform of the one copy, u halved, which carries exp(-2 pi i h.x): conjugated and
// scaled by V/N it is gemmi's F1. An operator is gemmi's Op: rot row-major and tran, both in units of 1/den,
// and hR is its apply_to_hkl().
__global__ void compose_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, float2 *f) {
    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;
        // The half-u transform holds k_u >= 0, and F1(-k) = conj F1(k) for a real map.
        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;
        // exp(+2 pi i h.t), exact for the twelfths and quarters a translation is made of
        double s, c;
        sincospi(2.0 * (h[0] * tran[0] + h[1] * tran[1] + h[2] * tran[2]) / den, &s, &c);
        sr += c * vr - s * vi;
        si += c * vi + s * vr;
    }
    const double scale = row_scale[m];
    f[m] = make_float2(static_cast<float>(scale * sr), static_cast<float>(scale * si));
}

// F_mask at each row, read at h directly (the mask is symmetric): gemmi's prepare_asu_data() of the mask's
// transform, 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);
}

// MapFromFPhi()'s input step: conjugated, as gemmi does it (rho = 1/V sum F exp(-2 pi i h.x), and the c2r
// transform carries exp(+2 pi i h.x)), with a missing coefficient (NaN) read as zero.
__global__ void conjugate_kernel(float2 *c, size_t n) {
    const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
    if (i >= n)
        return;
    const float2 x = c[i];
    c[i] = isnan(x.y) ? make_float2(0.0f, 0.0f) : make_float2(x.x, -x.y);
}

// The c2r output, z fastest and each z row padded to 2 (nw / 2 + 1) floats, into x fastest, times 1/V.
__global__ void transpose_kernel(const float *in, int nu, int nv, int nw, float scale, float *out) {
    const size_t n = static_cast<size_t>(nu) * nv * nw;
    const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
    if (i >= n)
        return;
    const int x = static_cast<int>(i % nu);
    const int y = static_cast<int>(i / nu % nv);
    const int z = static_cast<int>(i / (static_cast<size_t>(nu) * nv));
    out[i] = in[z + static_cast<size_t>(2 * (nw / 2 + 1)) * (y + static_cast<size_t>(nv) * x)] * scale;
}

// The structure factors' transform: r2c of an (nu, nv, nw) grid stored u fastest, which halves u.
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};
}

// A map's transform: in-place c2r of an (nu, nv, nw) grid stored as gemmi's ZYX half-l grid, w fastest and
// halved, the real output padded to 2 (nw / 2 + 1) per row as cuFFT does in place.
CudaFFTShape MapShape(int nu, int nv, int nw) {
    const int idist = nu * nv * (nw / 2 + 1);
    return {.n = {nu, nv, nw}, .inembed = {nu, nv, nw / 2 + 1}, .idist = idist,
            .onembed = {nu, nv, 2 * (nw / 2 + 1)}, .odist = 2 * idist, .type = CUFFT_C2R};
}

// The calling thread on `device` for as long as this lives, and back on its own device after.
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;
}

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

}  // namespace

struct ModelStructureFactorsGPUEngineImpl {
    ModelStructureFactorsGPUSetup setup;
    size_t work_bytes;
    int device = 0;
    CudaStream stream;

    CudaDevicePtr<float4> pos;
    CudaDevicePtr<ModelMaskAtom> mask_atoms;
    CudaDevicePtr<float> grid;          // the density, then the mask (on a grid of its own size or not); a map's real grid
    CudaDevicePtr<float2> spectrum;     // their transforms; a map's coefficients, transformed in place
    CudaDevicePtr<char> fft_work;
    CudaDevicePtr<int> ops, row_hkl;
    CudaDevicePtr<double> row_scale;
    CudaDevicePtr<float2> f, fmask;
    std::unique_ptr<ModelDensityGPU> density;
    std::unique_ptr<ModelMaskGPU> mask;
    CudaFFTPlan forward, mask_forward;   // the mask's only where its grid is not the density's

    // On `dev`, which the caller has made current.
    ModelStructureFactorsGPUEngineImpl(int dev, const ModelStructureFactorsGPUSetup &s, size_t work)
        : setup(s), work_bytes(work), 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);
        mask_atoms = CudaDevicePtr<ModelMaskAtom>(na, sync);
        grid = CudaDevicePtr<float>(std::max({GridPoints(setup.grid), GridPoints(setup.mask_grid), setup.map_points}), sync);
        spectrum = CudaDevicePtr<float2>(
            std::max({ComplexPoints(setup.grid), ComplexPoints(setup.mask_grid), setup.map_complex_points}), 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);
        f = CudaDevicePtr<float2>(nr, sync);
        fmask = CudaDevicePtr<float2>(nr, 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.mask_grid));

        ModelMaskGrid mask_grid;
        mask_grid.nu = setup.mask_grid.nu;
        mask_grid.nv = setup.mask_grid.nv;
        mask_grid.nw = setup.mask_grid.nw;
        std::copy(setup.mask_grid.orth, setup.mask_grid.orth + 9, mask_grid.orth);
        mask_grid.volume = setup.volume;
        mask->SetGrid(mask_grid, setup.mask_ops);

        auto up = [&](void *dst, const void *src, size_t bytes) {
            if (bytes > 0)
                cuda_err(cudaMemcpyAsync(dst, src, bytes, cudaMemcpyHostToDevice, stream));
        };
        up(ops, setup.sym_ops.data(), setup.sym_ops.size() * 12 * sizeof(int));
        up(row_hkl, setup.rows.data(), setup.rows.size() * 3 * sizeof(int));
        up(row_scale, setup.row_scale.data(), setup.row_scale.size() * sizeof(double));

        forward = CudaFFTPlan(ForwardShape(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 ||
            (!SameGrid(setup.mask_grid, setup.grid) && mask_forward.WorkBytes() > work_bytes))
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "structure factors: cuFFT work area over the reserve");
        cuda_err(cudaStreamSynchronize(stream));
    }

    ~ModelStructureFactorsGPUEngineImpl() {
        cudaStreamSynchronize(stream);
    }

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

size_t ModelStructureFactorsGPUEngine::TotalMemory(int device) {
    DeviceGuard guard(device);
    size_t free = 0, total = 0;
    cuda_err(cudaMemGetInfo(&free, &total));
    return total;
}

size_t ModelStructureFactorsGPUEngine::FFTWorkBytes(int device, const ModelStructureFactorsGPUSetup &setup, int map_nu,
                                                    int map_nv, int map_nw) {
    DeviceGuard guard(device);
    size_t bytes = std::max(CudaFFTPlan::WorkBytes(ForwardShape(setup.grid)),
                            CudaFFTPlan::WorkBytes(MapShape(map_nu, map_nv, map_nw)));
    if (!SameGrid(setup.mask_grid, setup.grid))
        bytes = std::max(bytes, CudaFFTPlan::WorkBytes(ForwardShape(setup.mask_grid)));
    return bytes;
}

size_t ModelStructureFactorsGPUEngine::DeviceBytes(const ModelStructureFactorsGPUSetup &s, size_t fft_work_bytes) {
    const size_t na = s.max_atoms + 1, nr = s.rows.size() + 1;
    size_t bytes = na * (sizeof(float4) + sizeof(ModelMaskAtom));
    bytes += std::max({GridPoints(s.grid), GridPoints(s.mask_grid), s.map_points}) * sizeof(float);
    bytes += std::max({ComplexPoints(s.grid), ComplexPoints(s.mask_grid), s.map_complex_points}) * sizeof(float2);
    bytes += fft_work_bytes;
    bytes += 12 * s.sym_ops.size() * sizeof(int) + nr * (3 * sizeof(int) + sizeof(double) + 2 * sizeof(float2));
    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.mask_grid));
    return bytes;
}

ModelStructureFactorsGPUEngine::ModelStructureFactorsGPUEngine(int device, const ModelStructureFactorsGPUSetup &setup,
                                                               size_t fft_work_bytes) {
    DeviceGuard guard(device);
    impl_ = std::make_unique<ModelStructureFactorsGPUEngineImpl>(device, setup, fft_work_bytes);
}

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

void ModelStructureFactorsGPUEngine::Compute(const std::vector<ModelDensityAtom> &atoms,
                                             const std::vector<std::array<float, 4>> &pos,
                                             const std::vector<ModelMaskAtom> &mask_atoms,
                                             std::vector<std::array<float, 2>> &fcalc,
                                             std::vector<std::array<float, 2>> &fmask) {
    ModelStructureFactorsGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    if (pos.size() != atoms.size() || mask_atoms.size() > e.setup.max_atoms)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "structure factors: atoms over the reserve");
    const ModelDensityGrid &g = e.setup.grid;
    const int rows = static_cast<int>(e.setup.rows.size());

    e.density->SetAtoms(g, atoms);
    if (!pos.empty())
        cuda_err(cudaMemcpyAsync(e.pos, pos.data(), pos.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream));
    if (!mask_atoms.empty())
        cuda_err(cudaMemcpyAsync(e.mask_atoms, mask_atoms.data(), mask_atoms.size() * sizeof(ModelMaskAtom),
                                 cudaMemcpyHostToDevice, e.stream));

    e.density->Compute(e.pos, e.grid);
    cufft_err(e.forward.R2C(e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
    if (rows > 0) {
        compose_kernel<<<blocks(rows), THREADS, 0, e.stream>>>(e.spectrum, e.Norm(g), g.nu, g.nv, g.nw, e.ops,
                                                               static_cast<int>(e.setup.sym_ops.size()), e.setup.den,
                                                               e.row_hkl, e.row_scale, rows, e.f);
        cuda_err(cudaGetLastError());
    }

    const ModelDensityGrid &mg = e.setup.mask_grid;
    e.mask->Compute(e.mask_atoms, static_cast<int>(mask_atoms.size()), e.grid);
    CudaFFTPlan &mask_forward = SameGrid(mg, g) ? e.forward : e.mask_forward;
    cufft_err(mask_forward.R2C(e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
    if (rows > 0) {
        fmask_kernel<<<blocks(rows), THREADS, 0, e.stream>>>(e.spectrum, e.Norm(mg), e.row_hkl, rows, mg.nu, mg.nv,
                                                             mg.nw, e.fmask);
        cuda_err(cudaGetLastError());
    }

    fcalc.resize(rows);
    fmask.resize(rows);
    if (rows > 0) {
        cuda_err(cudaMemcpyAsync(fcalc.data(), e.f, rows * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
        cuda_err(cudaMemcpyAsync(fmask.data(), e.fmask, rows * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
    }
    cuda_err(cudaStreamSynchronize(e.stream));
}

void ModelStructureFactorsGPUEngine::Map(int nu, int nv, int nw, double volume, const float *coefficients, float *map) {
    ModelStructureFactorsGPUEngineImpl &e = *impl_;
    DeviceGuard guard(e.device);
    const size_t points = static_cast<size_t>(nu) * nv * nw;
    const size_t complex_points = static_cast<size_t>(nu) * nv * (nw / 2 + 1);
    if (points > e.setup.map_points || complex_points > e.setup.map_complex_points)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "map: grid over the reserve");
    // Its stream is waited for before the plan goes back, on a throw as well.
    CudaFFTPlan plan(MapShape(nu, nv, nw), e.stream, e.fft_work);
    if (plan.WorkBytes() > e.work_bytes)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "map: cuFFT work area over the reserve");
    cuda_err(cudaMemcpyAsync(e.spectrum, coefficients, complex_points * sizeof(float2), cudaMemcpyHostToDevice,
                             e.stream));
    conjugate_kernel<<<blocks(complex_points), THREADS, 0, e.stream>>>(e.spectrum, complex_points);
    cuda_err(cudaGetLastError());
    cufft_err(plan.C2R(reinterpret_cast<cufftComplex *>(e.spectrum.get()),
                           reinterpret_cast<cufftReal *>(e.spectrum.get())));
    transpose_kernel<<<blocks(points), THREADS, 0, e.stream>>>(reinterpret_cast<const float *>(e.spectrum.get()),
                                                               nu, nv, nw, static_cast<float>(1.0 / volume), e.grid);
    cuda_err(cudaGetLastError());
    cuda_err(cudaMemcpyAsync(map, e.grid, points * sizeof(float), cudaMemcpyDeviceToHost, e.stream));
    cuda_err(cudaStreamSynchronize(e.stream));
}