Jungfraujoch/image_analysis/structure_refinement/ModelDensityGPU.cu

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

#include "ModelDensityGPU.h"

#include <algorithm>
#include <cmath>

#include <cub/cub.cuh>

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

namespace {

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

constexpr int BRICK = 8;                  // the gather's bricks are BRICK^3 grid points, one block each
constexpr int GATHER_TILE = 128;           // atoms staged in shared memory at a time
constexpr int MAX_BRICKS_PER_AXIS = 16;   // an atom's box touches at most this many bricks along an axis
constexpr int THREADS = 256;

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

// gemmi's unsafe_expapprox() (formfact.hpp), the exponential the density is computed with.
__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))));
}

// The distinct bricks the points c - d ... c + d of one axis fall in, wrapped into the cell. A grid size
// that is not a multiple of BRICK leaves the last brick partial, which is why this is done on wrapped
// points and not in brick coordinates. -1 where there are more than MAX_BRICKS_PER_AXIS of them.
__device__ int axis_bricks(int c, int d, int n, int *out) {
    int k = 0;
    for (int p = c - d; p <= c + d; p++) {
        const int b = imod(p, n) / BRICK;
        bool seen = false;
        for (int j = 0; j < k; j++)
            seen = seen || out[j] == b;
        if (seen)
            continue;
        if (k == MAX_BRICKS_PER_AXIS)
            return -1;
        out[k++] = b;
    }
    return k;
}

// The bricks atom i's box touches, per axis. The box is gemmi's: `box` points either side of the
// point nearest the atom.
__device__ void atom_bricks(const float4 &pos, const int *box, int nu, int nv, int nw, int *bu, int &nbu, int *bv,
                            int &nbv, int *bw, int &nbw) {
    nbu = axis_bricks(__float2int_rn(pos.x), box[0], nu, bu);
    nbv = axis_bricks(__float2int_rn(pos.y), box[1], nv, bv);
    nbw = axis_bricks(__float2int_rn(pos.z), box[2], nw, bw);
}

__global__ void count_pairs_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw, int *count,
                                   int *overflow) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i > n)
        return;
    if (i == n) {
        count[n] = 0;   // so the exclusive scan's last entry is the total
        return;
    }
    int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c;
    atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c);
    if (a < 0 || b < 0 || c < 0) {
        *overflow = 1;
        a = b = c = 0;
    }
    count[i] = a * b * c;
}

// The (brick, atom) pairs, atom-major: sorted stably by brick, each brick's atoms stay in model order.
__global__ void fill_pairs_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw,
                                  const int *offset, int *key, int *value) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n)
        return;
    int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c;
    atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c);
    const int nbu = (nu + BRICK - 1) / BRICK, nbv = (nv + BRICK - 1) / BRICK;
    int o = offset[i];
    for (int k = 0; k < c; k++)
        for (int j = 0; j < b; j++)
            for (int l = 0; l < a; l++) {
                key[o] = bu[l] + nbu * (bv[j] + nbv * bw[k]);
                value[o] = i;
                o++;
            }
}

__global__ void brick_ranges_kernel(const int *key, int npairs, int *start, int *end) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= npairs)
        return;
    const int k = key[i];
    if (i == 0 || key[i - 1] != k)
        start[k] = i;
    if (i == npairs - 1 || key[i + 1] != k)
        end[k] = i + 1;
}

// What the gather needs of one staged atom: its isotropic Gaussians (an anisotropic atom's matrices are read
// from the atoms in global memory), and its box as the brick sees it - a bit per point of the brick along each
// axis for whether the point is in the box, and for the points that are, what gemmi's box offset makes of them:
// the point's unwrapped box index is base + t, less n where t reaches wrap.
struct GatherAtom {
    float a[5], b[5];
    float occ, max_r2;
    int index, aniso;
    int masks;         // bits 0-7: u, 8-15: v, 16-23: w
    int base[3], wrap[3];
    float nctr[3];
};

// The point p's offset in the box along an axis of n points, 0 to span where it is in it.
__device__ __forceinline__ int box_offset(int p, int first, int n) {
    const int m = p - first;
    return m < 0 ? m + n : m;
}

// The density of the model on the grid, as UseDensityBoxesByPlane() and AtomDensity (ModelGrid.cpp) put
// it there, bit for bit: every point adds, in model order, each atom whose box holds it and whose radius
// reaches it, with the CPU's single-precision arithmetic in the CPU's order (no a*b+c is contracted:
// CMakeLists.txt compiles device code with --fmad=false). One block per brick; the brick's atoms are
// staged in shared memory. Supports() leaves only boxes narrower than the cell, so a point is in an atom's
// box at most once, at one index of the box's own unwrapped range.
__global__ void gather_kernel(const ModelDensityAtom *atoms, const float4 *pos, const int *box, const int *pair_atom,
                              const int *brick_start, const int *brick_end, ModelDensityGPU::Geometry g,
                              float *grid) {
    const int nbu = (g.nu + BRICK - 1) / BRICK, nbv = (g.nv + BRICK - 1) / BRICK;
    const int brick = blockIdx.x;
    const int s = brick_start[brick], e = brick_end[brick];   // s == e: an empty brick, written as zeros
    const int origin[3] = {(brick % nbu) * BRICK, (brick / nbu % nbv) * BRICK, (brick / (nbu * nbv)) * BRICK};
    const int tu = threadIdx.x, tv = threadIdx.y, tw = threadIdx.z;
    const int u = origin[0] + tu, v = origin[1] + tv, w = origin[2] + tw;
    const bool inside = u < g.nu && v < g.nv && w < g.nw;
    const int tid = tu + BRICK * (tv + BRICK * tw);
    const int n3[3] = {g.nu, g.nv, g.nw};

    __shared__ GatherAtom sh[GATHER_TILE];
    float acc = 0.0f;
    for (int c = s; c < e; c += GATHER_TILE) {
        const int m = min(GATHER_TILE, e - c);
        __syncthreads();
        if (tid < m) {
            const int a = pair_atom[c + tid];
            const ModelDensityAtom &at = atoms[a];
            const float4 p = pos[a];
            const float nctr[3] = {p.x, p.y, p.z};
            GatherAtom &r = sh[tid];
            for (int q = 0; q < 5; q++) {
                r.a[q] = at.a[q];
                r.b[q] = at.b[q][0];
            }
            r.occ = at.occ;
            r.max_r2 = at.max_r2;
            r.index = a;
            r.aniso = at.aniso;
            r.masks = 0;
            for (int k = 0; k < 3; k++) {
                const int d = box[3 * a + k];
                const int lo = __float2int_rn(nctr[k]) - d, first = imod(lo, n3[k]), span = 2 * d;
                const int m0 = box_offset(origin[k], first, n3[k]);
                for (int t = 0; t < BRICK; t++) {
                    const int mt = m0 + t;
                    if ((mt < n3[k] ? mt : mt - n3[k]) <= span)
                        r.masks |= 1 << (8 * k + t);
                }
                r.base[k] = lo + m0;
                r.wrap[k] = n3[k] - m0;
                r.nctr[k] = nctr[k];
            }
        }
        __syncthreads();
        if (!inside)
            continue;
        for (int k = 0; k < m; k++) {
            const GatherAtom &at = sh[k];
            if (!((at.masks >> tu) & (at.masks >> (8 + tv)) & (at.masks >> (16 + tw)) & 1))
                continue;
            const int pu = at.base[0] + tu - (tu >= at.wrap[0] ? g.nu : 0);
            const int pv = at.base[1] + tv - (tv >= at.wrap[1] ? g.nv : 0);
            const int pw = at.base[2] + tw - (tw >= at.wrap[2] ? g.nw : 0);
            // f = centre - (u, v, w) in the box's own (unwrapped) indices, then in A, as the CPU walks it
            const float fx = at.nctr[0] - static_cast<float>(pu);
            const float fy = at.nctr[1] - static_cast<float>(pv);
            const float fz = at.nctr[2] - static_cast<float>(pw);
            const float z = g.m22 * fz;
            const float y = g.m11 * fy + g.m12 * fz;
            const float yz = y * y + z * z;
            const float x = g.m00 * fx + (g.m01 * fy + g.m02 * fz);
            const float r2 = yz + x * x;
            if (r2 > at.max_r2)
                continue;
            float density = 0.0f;
            if (!at.aniso) {
                for (int q = 0; q < 5; q++)
                    density += at.a[q] * expapprox(fmaxf(at.b[q] * r2, -88.f));
            } else {
                for (int q = 0; q < 5; q++) {
                    const float *bq = atoms[at.index].b[q];
                    // gemmi's SMat33::r_u_r(), in its order
                    const float rur = x * x * bq[0] + y * y * bq[1] + z * z * bq[2] +
                                      2 * (x * y * bq[3] + x * z * bq[4] + y * z * bq[5]);
                    density += at.a[q] * expapprox(fmaxf(rur, -88.f));
                }
            }
            acc += at.occ * density;
        }
    }
    if (inside)
        grid[u + static_cast<size_t>(g.nu) * (v + static_cast<size_t>(g.nv) * w)] = acc;
}

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

// cub's temporary storage for the pair scan and sort at the capacity.
size_t CubBytes(size_t atoms, size_t pairs) {
    size_t scan = 0, sort = 0;
    cuda_err(cub::DeviceScan::ExclusiveSum(nullptr, scan, static_cast<int *>(nullptr), static_cast<int *>(nullptr),
                                           static_cast<int>(atoms + 1)));
    cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, sort, static_cast<int *>(nullptr), static_cast<int *>(nullptr),
                                             static_cast<int *>(nullptr), static_cast<int *>(nullptr),
                                             static_cast<int>(pairs)));
    return std::max(scan, sort);
}

}  // namespace

// The most bricks the 2 d + 1 wrapped points of a box can fall in along an axis of n points: one more
// than the points span for where they start in a brick, and one more again where they wrap past a last
// brick that is partial (n = 17: the points 15, 16, 0 fall in bricks 1, 2 and 0).
size_t ModelDensityGPU::AxisBrickBound(int d, int n) {
    const int nb = (n + BRICK - 1) / BRICK;
    return std::min(nb, (2 * d + 1 + BRICK - 1) / BRICK + 2);
}

size_t ModelDensityGPU::Bricks(int nu, int nv, int nw) {
    return static_cast<size_t>((nu + BRICK - 1) / BRICK) * ((nv + BRICK - 1) / BRICK) * ((nw + BRICK - 1) / BRICK);
}

void ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const ModelDensityAtom *atoms, size_t n, int *box) {
    const int n3[3] = {grid.nu, grid.nv, grid.nw};
    double spacing[3];
    for (int k = 0; k < 3; k++)
        spacing[k] = 1.0 / (n3[k] * std::sqrt(grid.frac[3 * k] * grid.frac[3 * k] + grid.frac[3 * k + 1] * grid.frac[3 * k + 1] +
                                              grid.frac[3 * k + 2] * grid.frac[3 * k + 2]));
    for (size_t i = 0; i < n; i++)
        for (int k = 0; k < 3; k++)
            box[3 * i + k] = static_cast<int>(std::ceil(atoms[i].radius / spacing[k]));
}

std::vector<int> ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atoms) {
    std::vector<int> box(3 * atoms.size());
    Boxes(grid, atoms.data(), atoms.size(), box.data());
    return box;
}

size_t ModelDensityGPU::PairBound(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atoms) {
    const std::vector<int> box = Boxes(grid, atoms);
    const int n3[3] = {grid.nu, grid.nv, grid.nw};
    size_t pairs = 0;
    for (size_t i = 0; i < atoms.size(); i++) {
        size_t product = 1;
        for (int k = 0; k < 3; k++)
            product *= AxisBrickBound(box[3 * i + k], n3[k]);
        pairs += product;
    }
    return pairs;
}

bool ModelDensityGPU::Supports(const ModelDensityGrid &grid, const int *box, size_t n) {
    const int n3[3] = {grid.nu, grid.nv, grid.nw};
    for (size_t i = 0; i < 3 * n; i++)
        if (2 * box[i] + 1 > n3[i % 3] || AxisBrickBound(box[i], n3[i % 3]) > MAX_BRICKS_PER_AXIS)
            return false;
    return true;
}

bool ModelDensityGPU::Supports(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atoms) {
    const std::vector<int> box = Boxes(grid, atoms);
    return Supports(grid, box.data(), atoms.size());
}

size_t ModelDensityGPU::DeviceBytes(size_t max_atoms, size_t max_pairs, size_t max_bricks) {
    const size_t na = max_atoms + 1;
    return na * (sizeof(ModelDensityAtom) + 3 * sizeof(int) + 2 * sizeof(int)) + 4 * max_pairs * sizeof(int) +
           2 * max_bricks * sizeof(int) + CubBytes(na, max_pairs + 1);
}

ModelDensityGPU::ModelDensityGPU(cudaStream_t stream, size_t max_atoms_, size_t max_pairs_, size_t max_bricks_)
    : stream(stream), max_atoms(max_atoms_), max_pairs(max_pairs_), max_bricks(max_bricks_) {
    const auto sync = CudaAlloc::Synchronous;
    const size_t na = std::max<size_t>(max_atoms, 1), np = std::max<size_t>(max_pairs, 1),
                 nb = std::max<size_t>(max_bricks, 1);
    atoms = CudaDevicePtr<ModelDensityAtom>(na, sync);
    box = CudaDevicePtr<int>(3 * na, sync);
    overflow = CudaDevicePtr<int>(1, sync);
    count = CudaDevicePtr<int>(na + 1, sync);
    offset = CudaDevicePtr<int>(na + 1, sync);
    key = CudaDevicePtr<int>(np, sync);
    value = CudaDevicePtr<int>(np, sync);
    key_sorted = CudaDevicePtr<int>(np, sync);
    value_sorted = CudaDevicePtr<int>(np, sync);
    brick_start = CudaDevicePtr<int>(nb, sync);
    brick_end = CudaDevicePtr<int>(nb, sync);
    host_count = CudaHostPtr<int>(2);
    cub_bytes = CubBytes(na, np);
    cub_temp = CudaDevicePtr<char>(std::max<size_t>(cub_bytes, 1), sync);
    host_box = CudaHostPtr<int>(3 * na);
}

void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atom_list) {
    SetAtoms(grid, atom_list.data(), atom_list.size());
}

void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const ModelDensityAtom *atom_list, size_t n) {
    if (n > max_atoms || Bricks(grid.nu, grid.nv, grid.nw) > max_bricks)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: grid or atoms over the reserve");
    // The boxes into the pinned staging buffer, so that their copy and the atoms' (pinned where the caller keeps
    // them so) go up by DMA instead of through the driver's staging of pageable memory.
    Boxes(grid, atom_list, n, host_box);
    // An atom whose box is wider than the cell would reach a point through more than one image, which
    // gemmi's box walk adds and the gather's box test does not; such a cell is left to the CPU.
    if (!Supports(grid, host_box, n))
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: an atom wider than the cell");
    geom.nu = grid.nu;
    geom.nv = grid.nv;
    geom.nw = grid.nw;
    // As gemmi's Grid sets orth_n (multiply_by_diagonal with 1/n) and ModelDensityGeometry() rounds it.
    const double inv_n[3] = {1.0 / grid.nu, 1.0 / grid.nv, 1.0 / grid.nw};
    geom.m00 = static_cast<float>(grid.orth[0] * inv_n[0]);
    geom.m01 = static_cast<float>(grid.orth[1] * inv_n[1]);
    geom.m02 = static_cast<float>(grid.orth[2] * inv_n[2]);
    geom.m11 = static_cast<float>(grid.orth[4] * inv_n[1]);
    geom.m12 = static_cast<float>(grid.orth[5] * inv_n[2]);
    geom.m22 = static_cast<float>(grid.orth[8] * inv_n[2]);
    n_atoms = static_cast<int>(n);
    if (n_atoms > 0) {
        cuda_err(cudaMemcpyAsync(atoms, atom_list, n * sizeof(ModelDensityAtom), cudaMemcpyHostToDevice, stream));
        cuda_err(cudaMemcpyAsync(box, host_box, 3 * n * sizeof(int), cudaMemcpyHostToDevice, stream));
    }
    // The uploads read the caller's atoms, which it may change once this returns.
    cuda_err(cudaStreamSynchronize(stream));
}

void ModelDensityGPU::Compute(const float4 *d_pos, float *d_grid) {
    cuda_err(cudaMemsetAsync(overflow, 0, sizeof(int), stream));
    count_pairs_kernel<<<blocks(n_atoms + 1), THREADS, 0, stream>>>(d_pos, box, n_atoms, geom.nu, geom.nv, geom.nw,
                                                                    count, overflow);
    cuda_err(cudaGetLastError());
    size_t temp = cub_bytes;
    cuda_err(cub::DeviceScan::ExclusiveSum(cub_temp.get(), temp, count.get(), offset.get(), n_atoms + 1, stream));
    // the count and the flag into pinned memory, the stream polled for them: this wait is on every evaluation
    cuda_err(cudaMemcpyAsync(host_count.get(), offset.get() + n_atoms, sizeof(int), cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaMemcpyAsync(host_count.get() + 1, overflow.get(), sizeof(int), cudaMemcpyDeviceToHost, stream));
    cudaError_t err;
    while ((err = cudaStreamQuery(stream)) == cudaErrorNotReady) {}
    cuda_err(err);
    const int npairs = host_count[0];
    const int over = host_count[1];
    if (over)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: an atom's box over the bricks per axis");
    if (static_cast<size_t>(npairs) > max_pairs)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: gather pairs over the reserve");
    fill_pairs_kernel<<<blocks(n_atoms), THREADS, 0, stream>>>(d_pos, box, n_atoms, geom.nu, geom.nv, geom.nw,
                                                               offset, key, value);
    cuda_err(cudaGetLastError());
    const int nb = static_cast<int>(Bricks(geom.nu, geom.nv, geom.nw));
    int bits = 1;
    while ((1 << bits) < nb)
        bits++;
    temp = cub_bytes;
    cuda_err(cub::DeviceRadixSort::SortPairs(cub_temp.get(), temp, key.get(), key_sorted.get(), value.get(),
                                             value_sorted.get(), npairs, 0, bits, stream));
    cuda_err(cudaMemsetAsync(brick_start, 0, nb * sizeof(int), stream));
    cuda_err(cudaMemsetAsync(brick_end, 0, nb * sizeof(int), stream));
    if (npairs > 0) {
        brick_ranges_kernel<<<blocks(npairs), THREADS, 0, stream>>>(key_sorted, npairs, brick_start, brick_end);
        cuda_err(cudaGetLastError());
    }
    gather_kernel<<<nb, dim3(BRICK, BRICK, BRICK), 0, stream>>>(atoms, d_pos, box, value_sorted, brick_start, brick_end,
                                                                geom, d_grid);
    cuda_err(cudaGetLastError());
}