Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu

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

#include "BraggIntegrationEngineGPU.h"

using namespace bragg_engine;

namespace {

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

// Fixed scalars passed by value to every kernel (mirrors BraggIntegrationEngine's members).
struct BraggGpuParams {
    int   W, H;
    float r1_sq, r2, r2_sq, r3, r3_sq;
    int   R, G, GG;
    float bkg_clip_nsigma;   // high-side background sigma-clip multiplier (0 = no clip)
    int   empirical;      // ProfileEmpirical vs ProfileGaussian
    int   use_ellipse;
    float bw_sigma;
    float c_radial;
    float F_px;
    float beam_x, beam_y;
    float bkg_trim;       // idea 1: symmetric trimmed-mean background fraction (0 = plain ring mean)
    BraggStencilParams stencil;   // per-reflection signal/background geometry (BraggStencil.h)
    int   rad_w;          // radial-background window held in shared memory, in bins of one pixel
    int   n_kern;         // radial-background kernels in the table
    int   overlap;        // 0 = off, 1 = reject, 2 = exclude (OverlapMode)
    int   boxsum_reject;  // BoxSum mode under Reject: drop on the disk-AREA fraction (see the CPU engine)
    int   exclude_px;     // drop a neighbour's pixels from the disk: Exclude, and not a box sum
    float claim_sq, inv_claim;   // how far a reflection claims pixels in the owner map
    float minpk;          // least readable/clean profile fraction that is kept (XDS MINPK)
    int   partial_ok;     // keep a signal disk with unreadable pixels: anything but a box sum
    float peak_frac;      // most of the profile's peak value an unreadable pixel may carry
};

__device__ inline bool valid(int32_t v) { return v != INT32_MIN && v != INT32_MAX; }

// Learned second moments, (sum v*rad^2, sum v*tan^2, sum v) per shell plus one global slot at N_SHELL.
constexpr int MOM_STRIDE = 3;

// --- Mark the r2 signal region of every predicted reflection (race-free: all writes are 1). The
//     region is the INNER stencil ellipse in the neighbour's own frame; see the CPU engine.
//     The same sweep fills the owner map when an overlap treatment is on: a pixel two signal regions
//     share belongs to the nearer centre, which an atomicMin over (distance, index) settles without
//     any ordering or sorting (BraggStencil.h). The claim disk sits inside the ellipse this loop
//     already walks, so ownership costs one atomic on the pixels already visited. ---
// Mark the per-pixel reflection mask and the overlap owner, or - with clear set - put the very same
// pixels back the way they were. The two directions share one kernel because they have to visit
// exactly the same set: what is written here is what has to be unwritten, and a second kernel that
// recomputed the boxes for itself would be free to drift from this one.
__global__ void mark_mask(const float *px_x, const float *px_y, uint8_t *mask, uint32_t *owner,
                          BraggGpuParams p, int n, int clear) {
    const int i = blockIdx.x;
    if (i >= n) return;
    const float cx = px_x[i], cy = px_y[i];
    const BraggStencil st = MakeBraggStencil(cx, cy, p.stencil);
    const int x0 = max(0, (int) floorf(cx - st.ex_in - 1.0f));
    const int x1 = min(p.W - 1, (int) ceilf(cx + st.ex_in + 1.0f));
    const int y0 = max(0, (int) floorf(cy - st.ey_in - 1.0f));
    const int y1 = min(p.H - 1, (int) ceilf(cy + st.ey_in + 1.0f));
    const int bw = x1 - x0 + 1, bh = y1 - y0 + 1;
    if (bw <= 0 || bh <= 0) return;
    for (int t = threadIdx.x; t < bw * bh; t += blockDim.x) {
        const int x = x0 + t % bw, y = y0 + t / bw;
        const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
        if (d.inner < p.r2_sq) mask[y * p.W + x] = clear ? (uint8_t) 0 : (uint8_t) 1;
        if (p.overlap && d.signal < p.claim_sq) {
            // Clearing writes one fixed value, so the blocks that share a pixel cannot disagree and
            // no atomic is needed; marking keeps the nearest centre and does.
            if (clear)
                owner[y * p.W + x] = BRAGG_OWNER_NONE;
            else
                atomicMin(&owner[y * p.W + x], BraggOwnerKey(sqrtf(d.signal), p.inv_claim, i));
        }
    }
}

// Sum one value across the warp so a single lane does the shared-memory atomic.
//
// Every block-wide accumulation here has all 128 lanes targeting one address, and a shared float or
// 64-bit atomicAdd has no native instruction on Turing OR Ada - it compiles to a compare-and-swap
// retry loop, so those 128 lanes serialise into 128 retries. Reducing within the warp first leaves
// four atomics per block instead of 128. Every call site below sits after a loop that all threads
// reach, so the full-warp mask is the right one.
//
// The integer sums are unchanged by this: addition is associative. The float and double ones change
// in their last bits, and become MORE reproducible - a fixed shuffle tree replaces whatever order
// the atomics happened to arrive in.
template <typename T>
__device__ __forceinline__ T warp_sum(T v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_down_sync(0xffffffffu, v, off);
    return v;
}

#define WARP_ATOMIC_ADD(dst, val) do {                     \
        const auto _w = warp_sum(val);                     \
        if ((threadIdx.x & 31u) == 0u) atomicAdd(&(dst), _w); \
    } while (0)

// Sum one float across the whole block, in a fixed order. warp_sum gives each warp its own total
// with a fixed shuffle tree; the warps then leave those totals in slots of their own and every
// thread adds the slots up by warp index. WARP_ATOMIC_ADD is the right tool for an INTEGER
// accumulator, where addition is associative, but on a float it adds the same handful of numbers in
// whatever order the warps happened to arrive - and s_num / s_den below are the numerator and
// denominator of the fitted intensity, so that order reached the answer. Every caller is reached by
// the whole block, as the WARP_ATOMIC_ADDs it replaces already were, and each concurrent quantity
// needs slots of its own.
constexpr int BRAGG_WARPS_MAX = 8;   // the kernels launch 128 threads; room to 256
__device__ __forceinline__ float block_sum(float v, float *slots) {
    const float w = warp_sum(v);
    if ((threadIdx.x & 31u) == 0u) slots[threadIdx.x >> 5] = w;
    __syncthreads();
    float s = 0.0f;
    const int nwarps = (blockDim.x + 31) >> 5;
    for (int j = 0; j < nwarps; ++j) s += slots[j];
    return s;
}

// Fixed-point scale for the learned profile and its moments. Those accumulators are sums of
// v = (px - bkg) / I over every strong reflection of the frame, and a float atomicAdd adds them in
// whatever order the blocks happen to arrive - so the profile, and every intensity fitted through
// it, moved between two runs of the same command on the same image. In fixed point the sum is
// integer addition, which is associative, exactly as the ring statistics and the box sums above
// already are. 2^20 puts the quantum at 1e-6 of one I-normalised pixel, far below the Poisson noise
// of the pixel it came from, and leaves the widest sum this kernel can build - every pixel of every
// reflection, times the r1^2 moment arm - some five orders inside a signed 64-bit accumulator.
constexpr double PROFILE_FIXED = 1048576.0;   // 2^20

// Signed fixed-point value carried as two's complement, so the accumulators can use the unsigned
// 64-bit atomicAdd (the only 64-bit one CUDA offers) - the same convention as s_Isum above.
__device__ __forceinline__ unsigned long long profile_fixed(float v) {
    return (unsigned long long) (long long) llrintf(v * (float) PROFILE_FIXED);
}
__device__ __forceinline__ float profile_float(unsigned long long v) {
    return (float) ((double) (long long) v / PROFILE_FIXED);
}

// --- Pass A box-sum: rough I / background / centroid / strong flag, one block per reflection. ---
__global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
                       const int32_t *img, const uint8_t *mask, const uint32_t *owner,
                       BraggGpuParams p, int n,
                       int *cx_o, int *cy_o, float *I_o, float *sigma_o, float *bkg_o,
                       float *bkgvar_o, float *varbkg_o, float *obsx_o, float *obsy_o, uint8_t *ok_o, uint8_t *strong_o,
                       uint8_t *hasobs_o, unsigned long long *invd2mm,
                       float *isum_o, int *ninner_o, int *rbin_o, int *kbin_o,
                       unsigned long long *rad_sum, int *rad_cnt, int n_rad) {
    const int i = blockIdx.x;
    if (i >= n) return;

    __shared__ unsigned long long s_Isum, s_Ix, s_Iy;
    __shared__ int s_ninner, s_ninner_valid, s_nbkg, s_ndisk, s_nown;
    __shared__ unsigned long long s_bkgsum;
    __shared__ int s_accept, s_full;
    __shared__ double s_bkg, s_thr;
    __shared__ unsigned long long s_clipsum;
    __shared__ long long s_thr_i;
    __shared__ int s_clipn;
    __shared__ float s_r0;
    // The stencil spans r0 +- the radial semi-axis of the outer ellipse, so the window has to be
    // sized from the widest aperture on the detector - which the host does, into p.rad_w. A fixed
    // 32-bin window silently dropped the outer bins as soon as anything was elongated.
    // The radial sums are INTEGERS, for the reason reduce_rings_shared gives in the spot finder: a
    // preprocessed pixel is an exact int32, so the sum is exact in 64 bits and integer addition is
    // associative - the radial background curve is then the same whatever order the blocks arrive in.
    // With float accumulators it moved in its last bits between runs of the same command, and it is
    // subtracted from every reflection's background.
    extern __shared__ unsigned long long s_rad[];
    unsigned long long *s_radv = s_rad;                       // signed value carried as two's complement
    int   *s_radn = (int *) (s_rad + p.rad_w);
    __shared__ int   s_radbase;
    if (threadIdx.x == 0) {
        const float rx = px_x[i] - p.beam_x, ry = px_y[i] - p.beam_y;
        s_r0 = sqrtf(rx * rx + ry * ry);
        s_Isum = 0; s_Ix = 0; s_Iy = 0;
        s_ninner = 0; s_ninner_valid = 0; s_nbkg = 0; s_bkgsum = 0;
        s_ndisk = 0; s_nown = 0;
    }
    __syncthreads();

    const float cx = px_x[i], cy = px_y[i];
    // Every thread builds the same stencil from the same inputs; it is a few flops against a whole
    // bounding box of pixel reads, so it costs less than staging it through shared memory.
    const BraggStencil st = MakeBraggStencil(cx, cy, p.stencil);
    const int x0 = max(0, (int) floorf(cx - st.ex_out - 1.0f));
    const int x1 = min(p.W - 1, (int) ceilf(cx + st.ex_out + 1.0f));
    const int y0 = max(0, (int) floorf(cy - st.ey_out - 1.0f));
    const int y1 = min(p.H - 1, (int) ceilf(cy + st.ey_out + 1.0f));
    const int bw = x1 - x0 + 1, bh = y1 - y0 + 1;
    const int area = (bw > 0 && bh > 0) ? bw * bh : 0;

    long long l_Isum = 0, l_Ix = 0, l_Iy = 0;
    int l_ni = 0, l_niv = 0, l_nb = 0, l_nd = 0, l_no = 0;
    long long l_bkg = 0;
    for (int t = threadIdx.x; t < area; t += blockDim.x) {
        const int x = x0 + t % bw, y = y0 + t / bw;
        const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
        // The window is the bounding box of the outer ellipse, so nearly half of it is neither the
        // signal disk nor the background ring. Reading the pixel only once it is known to be wanted
        // keeps those slots from fetching a cache line for nothing.
        if (d.signal < p.r1_sq) {
            // A pixel a nearer neighbour owns carries that neighbour's flux; see the CPU engine.
            ++l_nd;
            if (p.overlap) {
                if (BraggOwnedBy(owner[y * p.W + x], i)) ++l_no;
                else if (p.exclude_px) continue;
            }
            ++l_ni;
            const int32_t px = img[y * p.W + x];
            if (valid(px)) { l_Isum += px; l_Ix += (long long) x * px; l_Iy += (long long) y * px; ++l_niv; }
        } else if (d.inner >= p.r2_sq && d.outer < p.r3_sq) {
            if (mask[y * p.W + x]) continue;
            const int32_t px = img[y * p.W + x];
            if (!valid(px)) continue;
            l_bkg += px; ++l_nb;
        }
    }
    WARP_ATOMIC_ADD(s_Isum, (unsigned long long) l_Isum);
    WARP_ATOMIC_ADD(s_Ix, (unsigned long long) l_Ix);
    WARP_ATOMIC_ADD(s_Iy, (unsigned long long) l_Iy);
    WARP_ATOMIC_ADD(s_ninner, l_ni);
    WARP_ATOMIC_ADD(s_ninner_valid, l_niv);
    WARP_ATOMIC_ADD(s_nbkg, l_nb);
    WARP_ATOMIC_ADD(s_bkgsum, (unsigned long long) l_bkg);
    WARP_ATOMIC_ADD(s_ndisk, l_nd);
    WARP_ATOMIC_ADD(s_nown, l_no);
    __syncthreads();

    for (int t = threadIdx.x; t < p.rad_w; t += blockDim.x) { s_radv[t] = 0; s_radn[t] = 0; }
    if (threadIdx.x == 0) s_radbase = (int) lroundf(s_r0) - p.rad_w / 2;
    if (threadIdx.x == 0) {
        // A hole in the signal disk no longer discards the reflection in the profile modes - the fit
        // renormalises to the pixels it can read and Pass B cuts on how much of the profile survived
        // (XDS's MINPK). A box sum has no profile to renormalise with. See the CPU engine.
        s_full = (s_ninner_valid == s_ninner) ? 1 : 0;
        s_accept = ((s_full || p.partial_ok) && s_nbkg > 5) ? 1 : 0;
        s_bkg = s_accept ? ((double) (long long) s_bkgsum / (double) s_nbkg) : 0.0;
        s_thr = s_bkg + (double) p.bkg_clip_nsigma * sqrt(fmax(s_bkg, 1.0));
        // The pixel is an integer, so comparing it against the floor of the threshold accepts
        // exactly the same set - and does it with an integer compare instead of a widening
        // conversion and a double comparison, per ring pixel.
        s_thr_i = (long long) floor(s_thr);
        s_clipsum = 0; s_clipn = 0;
    }
    __syncthreads();

    // Trimmed-mean background (idea 1): collect the annulus into shared memory, bitonic-sort the block, and
    // average the middle (1 - 2*bkg_trim) fraction - robust to the high-side contamination that biases the
    // plain ring mean. Falls back to the flat mean when the ring exceeds the shared buffer.
    __shared__ int s_bvals[BKG_TRIM_MAX];
    __shared__ int s_bn;
    if (threadIdx.x == 0) s_bn = 0;
    __syncthreads();
    const bool do_trim = s_accept && p.bkg_trim > 0.0f && s_nbkg > 5 && s_nbkg <= BKG_TRIM_MAX;
    if (do_trim) {
        for (int t = threadIdx.x; t < area; t += blockDim.x) {
            const int x = x0 + t % bw, y = y0 + t / bw;
            const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
            if (!(d.inner >= p.r2_sq && d.outer < p.r3_sq)) continue;
            if (mask[y * p.W + x]) continue;
            const int32_t px = img[y * p.W + x];
            if (!valid(px)) continue;
            const int slot = atomicAdd(&s_bn, 1);
            if (slot < BKG_TRIM_MAX) s_bvals[slot] = px;
        }
        __syncthreads();
        const int nb = min(s_bn, BKG_TRIM_MAX);
        int n2 = 1;
        while (n2 < nb) n2 <<= 1;
        for (int t = threadIdx.x + nb; t < n2; t += blockDim.x) s_bvals[t] = INT32_MAX;   // pad to power of 2
        __syncthreads();
        for (int k = 2; k <= n2; k <<= 1)                       // ascending bitonic sort of s_bvals[0..n2)
            for (int j = k >> 1; j > 0; j >>= 1) {
                for (int idx = threadIdx.x; idx < n2; idx += blockDim.x) {
                    const int ixj = idx ^ j;
                    if (ixj > idx) {
                        const bool up = ((idx & k) == 0);
                        const int a = s_bvals[idx], b = s_bvals[ixj];
                        if ((up && a > b) || (!up && a < b)) { s_bvals[idx] = b; s_bvals[ixj] = a; }
                    }
                }
                __syncthreads();
            }
        if (threadIdx.x == 0) {
            const int lo = (int) (nb * p.bkg_trim), hi = nb - lo;
            if (hi > lo) {
                double s = 0.0;
                for (int t = lo; t < hi; ++t) s += (double) s_bvals[t];
                s_bkg = s / (double) (hi - lo);
            }
        }
        __syncthreads();
    }

    // Second ring pass for the high-side sigma-clip (re-reads the annulus; avoids storing bkg values).
    if (s_accept && p.bkg_clip_nsigma > 0.0f && !do_trim) {
        long long c_l = 0; int cn_l = 0;
        for (int t = threadIdx.x; t < area; t += blockDim.x) {
            const int x = x0 + t % bw, y = y0 + t / bw;
            const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
            if (!(d.inner >= p.r2_sq && d.outer < p.r3_sq)) continue;
            if (mask[y * p.W + x]) continue;
            const int32_t px = img[y * p.W + x];
            if (!valid(px)) continue;
            if ((long long) px <= s_thr_i) {
                c_l += px; ++cn_l;
                if (n_rad > 0) {
                    // The radial curve is binned on the TRUE detector radius, so this is the
                    // unshrunk radial projection - no per-pixel sqrt.
                    const int idx = (int) lroundf(s_r0 + d.rad) - s_radbase;
                    if (idx >= 0 && idx < p.rad_w) {
                        atomicAdd(&s_radv[idx], (unsigned long long) (long long) px);   // shared, not global
                        atomicAdd(&s_radn[idx], 1);
                    }
                }
            }
        }
        WARP_ATOMIC_ADD(s_clipsum, (unsigned long long) c_l); WARP_ATOMIC_ADD(s_clipn, cn_l);
    }
    __syncthreads();

    if (n_rad > 0) {
        for (int t = threadIdx.x; t < p.rad_w; t += blockDim.x) {
            if (s_radn[t] == 0) continue;
            const int b = min(max(s_radbase + t, 0), n_rad - 1);
            atomicAdd(&rad_sum[b], s_radv[t]);   // two's complement, summed as unsigned
            atomicAdd(&rad_cnt[b], s_radn[t]);
        }
        __syncthreads();
    }

    if (threadIdx.x != 0) return;
    if (!s_accept) { ok_o[i] = 0; strong_o[i] = 0; hasobs_o[i] = 0; return; }
    // A box sum has no profile to renormalise a disk it has taken pixels out of, so dropping the
    // reflection is the only overlap treatment it has (see the CPU engine).
    if (p.boxsum_reject && (float) s_nown < p.minpk * (float) s_ndisk) {
        ok_o[i] = 0; strong_o[i] = 0; hasobs_o[i] = 0; return;
    }

    double bkg = s_bkg;
    int n_bkg_used = s_nbkg;   // pixels behind the FINAL background value (trim/clip shrink it)
    if (do_trim) {
        const int nb = min(s_bn, BKG_TRIM_MAX);
        const int lo = (int) (nb * p.bkg_trim), hi = nb - lo;
        if (hi > lo) n_bkg_used = hi - lo;
    }
    if (p.bkg_clip_nsigma > 0.0f && s_clipn > 5) { bkg = (double) (long long) s_clipsum / (double) s_clipn; n_bkg_used = s_clipn; }
    const long long Isum = (long long) s_Isum;
    // The sum is over the pixels actually READ, so that is the count the background is subtracted
    // with; with nothing missing it is the whole disk, exactly as before.
    const double I = (double) Isum - (double) s_ninner_valid * bkg;
    // See the CPU engine: bkg is estimated from n_bkg_used pixels and subtracted n_inner times, so
    // var(I) = Isum + n_inner^2 * bkg/n_bkg_used. Both engines must agree.
    const double bkg_var = bkg / (double) n_bkg_used;
    const double var_bkg_term = (double) s_ninner_valid * (double) s_ninner_valid * bkg_var;
    double sigma = 1.0;
    uint8_t hasobs = 0; double ox = 0.0, oy = 0.0;
    if (Isum > 0) {
        sigma = fmax(sigma, sqrt((double) Isum + var_bkg_term));
        ox = (double) (long long) s_Ix / (double) Isum;
        oy = (double) (long long) s_Iy / (double) Isum;
        // A disk with a hole gives a centroid pulled away from it; see the CPU engine.
        hasobs = s_full ? 1 : 0;
    }
    cx_o[i] = (int) lroundf(cx);
    cy_o[i] = (int) lroundf(cy);
    I_o[i] = (float) I; sigma_o[i] = (float) sigma; bkg_o[i] = (float) bkg;
    bkgvar_o[i] = (float) bkg_var;
    varbkg_o[i] = (float) ((double) s_ninner_valid * bkg + var_bkg_term);
    isum_o[i] = (float) Isum;
    ninner_o[i] = s_ninner_valid;
    rbin_o[i] = min(max((int) lroundf(s_r0), 0), n_rad > 0 ? n_rad - 1 : 0);
    kbin_o[i] = BraggStencilKernelIndex(st, p.n_kern);
    obsx_o[i] = (float) ox; obsy_o[i] = (float) oy; hasobs_o[i] = hasobs;
    ok_o[i] = 1;
    // The profile and its resolution shells are learned from COMPLETE reflections only; see the CPU
    // engine.
    strong_o[i] = (s_full && sigma > 0.0 && I / sigma >= STRONG_I_OVER_SIGMA) ? 1 : 0;

    const float d = dd[i];
    if (s_full && d > 0.0f) {
        // Positive doubles keep IEEE bit-pattern ordering, so atomicMin/Max on the ull view works.
        const unsigned long long b = (unsigned long long) __double_as_longlong(1.0 / ((double) d * d));
        atomicMin(&invd2mm[0], b);
        atomicMax(&invd2mm[1], b);
    }
}

// Resolution shell of one reflection from the global inv-d^2 range (mirrors CPU shell_of). Computed
// inline in learn_profile and fit so no separate shell array/kernel is needed.
__device__ inline int compute_shell(float d, const unsigned long long *invd2mm) {
    const unsigned long long mn = invd2mm[0], mx = invd2mm[1];
    if (!(d > 0.0f) || mx <= mn) return 0;
    const double invd2 = 1.0 / ((double) d * d);
    const double dmn = __longlong_as_double((long long) mn), dmx = __longlong_as_double((long long) mx);
    const int s = (int) ((invd2 - dmn) / (dmx - dmn) * N_SHELL);
    return s < 0 ? 0 : (s >= N_SHELL ? N_SHELL - 1 : s);
}

// --- Zero the profile accumulators and seed the inv-d^2 range, in one launch (replaces a handful
//     of small cudaMemsetAsync calls, which matter when kernel-launch latency is high). ---
__global__ void reset(unsigned long long *shell_grid, unsigned long long *global_grid,
                      unsigned long long *mom, int *shell_n, int *global_n,
                      unsigned long long *invd2mm, int GG) {
    for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < N_SHELL * GG; k += blockDim.x * gridDim.x)
        shell_grid[k] = 0;
    for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < GG; k += blockDim.x * gridDim.x)
        global_grid[k] = 0;
    if (blockIdx.x == 0 && threadIdx.x < MOM_STRIDE * (N_SHELL + 1)) mom[threadIdx.x] = 0;
    if (blockIdx.x == 0 && threadIdx.x < N_SHELL) shell_n[threadIdx.x] = 0;
    if (blockIdx.x == 0 && threadIdx.x == 0) {
        *global_n = 0;
        invd2mm[0] = ~0ull;   // min seed
        invd2mm[1] = 0ull;    // max seed
    }
}

// --- Learn the profile: each strong spot adds its bkg-subtracted, I-normalised grid to its shell
//     (and the global grid), and its second moments to the same shell's moment slot. The grid stays
//     in the DETECTOR frame (that is where the empirical profile is applied); the moments are taken in
//     the spot's own radial/tangential frame, because a detector-frame stack is azimuthally averaged
//     and cannot tell a radially smeared spot from a tangentially wide one (see the CPU engine).
//     One block per reflection. ---
__global__ void learn_profile(const int32_t *img, const uint32_t *owner,
                              const float *px_x, const float *px_y,
                              const int *cx_a, const int *cy_a, const float *dd,
                              const unsigned long long *invd2mm,
                              const float *I_a, const float *bkg_a, const uint8_t *ok_a, const uint8_t *strong_a,
                              unsigned long long *shell_grid, unsigned long long *global_grid,
                              unsigned long long *mom, int *shell_n, int *global_n,
                              BraggGpuParams p, int n) {
    const int i = blockIdx.x;
    if (i >= n || !ok_a[i] || !strong_a[i]) return;
    const float I = I_a[i];
    if (!(I > 0.0f)) return;
    const int cx = cx_a[i], cy = cy_a[i];
    // The shell belongs to the reflection, not the pixel, so one thread works it out for the whole
    // block: it is two software double-precision divisions, and all 128 threads were doing both.
    __shared__ int s_sh;
    if (threadIdx.x == 0) s_sh = compute_shell(dd[i], invd2mm);
    __syncthreads();
    const int sh = s_sh;
    const float bkg = bkg_a[i];
    const float rx = px_x[i] - p.beam_x, ry = px_y[i] - p.beam_y;
    const float Rpx = sqrtf(rx * rx + ry * ry);
    const float ux = Rpx > 1e-6f ? rx / Rpx : 1.0f, uy = Rpx > 1e-6f ? ry / Rpx : 0.0f;
    unsigned long long *sg = shell_grid + (size_t) sh * p.GG;
    __shared__ unsigned long long s_rad, s_tan, s_w;
    if (threadIdx.x == 0) { s_rad = 0; s_tan = 0; s_w = 0; }
    __syncthreads();
    float l_rad = 0.0f, l_tan = 0.0f, l_w = 0.0f;
    for (int k = threadIdx.x; k < p.GG; k += blockDim.x) {
        const int dx = k % p.G - p.R, dy = k / p.G - p.R;
        const int x = cx + dx, y = cy + dy;
        if (x < 0 || y < 0 || x >= p.W || y >= p.H) continue;
        const int32_t px = img[y * p.W + x];
        if (!valid(px)) continue;
        if (p.overlap == 2 && !BraggOwnedBy(owner[y * p.W + x], i)) continue;
        const float v = ((float) px - bkg) / I;
        const unsigned long long vq = profile_fixed(v);
        atomicAdd(&sg[k], vq);
        atomicAdd(&global_grid[k], vq);
        if ((float) (dx * dx + dy * dy) >= p.r1_sq) continue;
        const float rad = dx * ux + dy * uy, tn = -dx * uy + dy * ux;
        l_rad += v * rad * rad; l_tan += v * tn * tn; l_w += v;
    }
    // Reduce in shared memory first: every block adds to the same handful of global words, so one
    // atomic per warp instead of one per thread. A thread's own l_* sum is over a fixed set of k, so
    // it is the same every run; quantising here makes everything above it associative as well.
    WARP_ATOMIC_ADD(s_rad, profile_fixed(l_rad));
    WARP_ATOMIC_ADD(s_tan, profile_fixed(l_tan));
    WARP_ATOMIC_ADD(s_w,   profile_fixed(l_w));
    __syncthreads();
    if (threadIdx.x == 0) {
        atomicAdd(&mom[MOM_STRIDE * sh + 0], s_rad);
        atomicAdd(&mom[MOM_STRIDE * sh + 1], s_tan);
        atomicAdd(&mom[MOM_STRIDE * sh + 2], s_w);
        atomicAdd(&mom[MOM_STRIDE * N_SHELL + 0], s_rad);
        atomicAdd(&mom[MOM_STRIDE * N_SHELL + 1], s_tan);
        atomicAdd(&mom[MOM_STRIDE * N_SHELL + 2], s_w);
        atomicAdd(&shell_n[sh], 1);
        atomicAdd(global_n, 1);
    }
}

// --- Turn the learned moments into radial/tangential variances and, for empirical, normalise the
//     learned grid into a profile. One block per grid: blocks [0,N_SHELL) are the shells, block
//     N_SHELL is the global one. ---
__global__ void build_profiles(const unsigned long long *shell_grid, const unsigned long long *global_grid,
                               const unsigned long long *mom,
                               const int *shell_n, const int *global_n,
                               float *shell_P, float *global_P, float *sigma2_r, float *sigma2_t,
                               BraggGpuParams p) {
    const int b = blockIdx.x;
    const unsigned long long *grid; int nstrong; float *P;
    if (b < N_SHELL) { grid = shell_grid + (size_t) b * p.GG; nstrong = shell_n[b]; P = shell_P + (size_t) b * p.GG; }
    else            { grid = global_grid; nstrong = *global_n; P = global_P; }

    if (threadIdx.x == 0) {
        const float w = profile_float(mom[MOM_STRIDE * b + 2]);
        sigma2_r[b] = w > 0.0f ? fmaxf(0.25f, profile_float(mom[MOM_STRIDE * b + 0]) / w) : 1.0f;
        sigma2_t[b] = w > 0.0f ? fmaxf(0.25f, profile_float(mom[MOM_STRIDE * b + 1]) / w) : 1.0f;
    }
    if (!p.empirical) return;

    // The normalisation total is summed in fixed point as well. It divides every cell of the profile,
    // so a float atomicAdd here - 128 lanes on one address, in arrival order - reached every intensity
    // the profile fits, which is most of what was left moving between two runs of the same command.
    __shared__ unsigned long long s_sum_q;
    if (threadIdx.x == 0) s_sum_q = 0;
    __syncthreads();
    unsigned long long l_sum = 0;
    for (int k = threadIdx.x; k < p.GG; k += blockDim.x) {
        const long long gq = (long long) grid[k];
        const unsigned long long g = gq > 0 ? (unsigned long long) gq : 0ull;   // a profile has to be non-negative
        P[k] = profile_float(g);
        l_sum += g;
    }
    WARP_ATOMIC_ADD(s_sum_q, l_sum);
    __syncthreads();
    const float s_sum = profile_float(s_sum_q);
    const bool normalise = nstrong > 0 && s_sum > 0.0f;
    for (int k = threadIdx.x; k < p.GG; k += blockDim.x) P[k] = normalise ? P[k] / s_sum : 0.0f;
}

// --- Radial background curvature correction: one thread per reflection, no pixel reads.
//     The disk and the annulus are concentric, so a background linear in position cancels between
//     them; what is left is the curvature of the radial background. k_diff (annulus minus disk
//     histogram over radial offset) turns that into one short dot product. An empty radial bin
//     contributes the reflection's own background, so an empty neighbourhood gives exactly zero
//     correction because the kernel weights sum to zero. The kernel depends on how far this
//     reflection's ring was elongated, so k_diff is a table and kbin_a picks the row.
//     Mirrors BraggIntegrationEngineCPU. ---
__global__ void radial_correct(const unsigned long long *rad_sum, const int *rad_cnt, int n_rad,
                               const float *k_diff, int k_len, int k_off,
                               const float *isum_a, const int *ninner_a, const int *rbin_a,
                               const int *kbin_a,
                               const uint8_t *ok_a, float *bkg_o, float *I_o, int n) {
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n || !ok_a[i]) return;
    const float bkg = bkg_o[i];
    const float *kern = k_diff + (size_t) kbin_a[i] * k_len;
    float corr = 0.0f;
    for (int k = 0; k < k_len; ++k) {
        int b = rbin_a[i] + k - k_off;
        b = min(max(b, 0), n_rad - 1);
        const float v = rad_cnt[b] > 0 ? (float) (long long) rad_sum[b] / (float) rad_cnt[b] : bkg;
        corr += kern[k] * v;
    }
    const float bkg_new = bkg - corr;
    bkg_o[i] = bkg_new;
    I_o[i] = isum_a[i] - (float) ninner_a[i] * bkg_new;
}

// --- Pass B Kabsch profile fit: I = sum P(c-B)/v over sum P^2/v, v = B + max(I,0)P (iterate). The
//     reweighting is the Kabsch/Otwinowski iteration: Kabsch, Acta Cryst D66, 133-144 (2010);
//     Otwinowski & Minor, Methods Enzymol 276, 307-326 (1997).
//     One block per reflection; the (possibly elongated) profile is built in shared memory. ---
__global__ void fit(const int32_t *img, const uint32_t *owner, const float *px_x, const float *px_y,
                    const int *cx_a, const int *cy_a, const float *dd, const unsigned long long *invd2mm,
                    const float *I_seed, const float *sigma_seed, const float *bkg_a,
                    const float *bkgvar_a, const float *varbkg_seed, const uint8_t *ok_a,
                    const float *shell_P, const float *global_P,
                    const float *sigma2_r, const float *sigma2_t, const int *shell_n,
                    float *I_o, float *sigma_o, float *varbkg_o, uint8_t *ok_o, BraggGpuParams p, int n) {
    const int i = blockIdx.x;
    if (i >= n) return;
    extern __shared__ float Pbuf[];
    __shared__ float s_I;
    __shared__ float s_slot[5][BRAGG_WARPS_MAX];   // per-warp partials, see block_sum
    // Max reductions. The profile is non-negative, and for non-negative floats the IEEE bit pattern
    // orders exactly as the value does, so an integer atomicMax on that pattern is an exact float max.
    __shared__ int s_ppeak_i, s_plost_i;
    __shared__ int s_Rf, s_Gf;

    if (!ok_a[i]) { if (threadIdx.x == 0) ok_o[i] = 0; return; }

    const int cx = cx_a[i], cy = cy_a[i];
    // As in learn_profile: one thread per block, not one per thread. Two double-precision divisions
    // on a card whose double throughput is a sixty-fourth of its single is worth doing once.
    __shared__ int s_sh;
    if (threadIdx.x == 0) s_sh = compute_shell(dd[i], invd2mm);
    __syncthreads();
    const int sh = s_sh;
    const bool use_shell = shell_n[sh] >= MIN_STRONG_PER_SHELL;   // else fall back to the global profile
    const float bkg = bkg_a[i];

    if (p.empirical) {
        const float *Psrc = use_shell ? (shell_P + (size_t) sh * p.GG) : global_P;
        if (threadIdx.x == 0) { s_Rf = p.R; s_Gf = p.G; }
        __syncthreads();
        for (int k = threadIdx.x; k < p.GG; k += blockDim.x) Pbuf[k] = Psrc[k];
        __syncthreads();
    } else {
        const int si = use_shell ? sh : N_SHELL;   // N_SHELL is the global slot
        const float s2t = sigma2_t[si];
        const float rx = px_x[i] - p.beam_x, ry = px_y[i] - p.beam_y;
        const float Rpx = sqrtf(rx * rx + ry * ry);
        const float tan2t = Rpx / p.F_px;
        float s2r = s2t, ux = 1.0f, uy = 0.0f;
        bool elong = false;
        if (p.use_ellipse) {
            // Measured radial excess over the tangential width, floored by the analytic bandwidth +
            // parallax/capture term (see the CPU engine).
            const float sbw = p.bw_sigma * Rpx;
            const float radial_extra = fmaxf(sigma2_r[si] - s2t, sbw * sbw + p.c_radial * tan2t * tan2t);
            if (Rpx > 1e-6f && radial_extra > 0.25f) { ux = rx / Rpx; uy = ry / Rpx; s2r = s2t + radial_extra; elong = true; }
        }
        const int Rf = elong ? min(3 * p.R, (int) ceilf(p.r2 + 2.0f * sqrtf(s2r))) : p.R;
        const int Gf = 2 * Rf + 1;
        if (threadIdx.x == 0) { s_Rf = Rf; s_Gf = Gf; }
        __syncthreads();
        const float fx = px_x[i] - cx, fy = px_y[i] - cy;
        // The two widths are the same for every cell of this reflection, so their reciprocals are
        // taken once rather than per cell.
        const float inv_2s2r = 0.5f / s2r, inv_2s2t = 0.5f / s2t;
        float l_gs = 0.0f;
        for (int k = threadIdx.x; k < Gf * Gf; k += blockDim.x) {
            const float ex = (k % Gf - Rf) - fx, ey = (k / Gf - Rf) - fy;
            const float rad = ex * ux + ey * uy, tn = -ex * uy + ey * ux;
            const float g = expf(-rad * rad * inv_2s2r - tn * tn * inv_2s2t);
            Pbuf[k] = g; l_gs += g;
        }
        const float gsum = block_sum(l_gs, s_slot[0]);
        // Normalising by a reciprocal rather than dividing per cell: the divisor is the same for
        // every cell of the reflection, and a division here is thirteen instructions of a
        // twenty-instruction loop.
        const float inv_gs = __frcp_rn(gsum);
        for (int k = threadIdx.x; k < Gf * Gf; k += blockDim.x) Pbuf[k] *= inv_gs;
        __syncthreads();
    }

    const int Rf = s_Rf, Gf = s_Gf, GfGf = Gf * Gf;
    const float B = fmaxf(bkg, (float) PIXEL_VARIANCE_FLOOR);

    // How much of the expected profile the fit can actually see: s_pvalid over the whole grid against
    // s_pgrid, the mass that falls on the detector at all (XDS's MINPK), s_pown the same over
    // neighbour-owned pixels (what Reject cuts on), and s_mread / s_mall over the r1 disk alone,
    // which is what the summation seed the runaway guard compares against actually summed. See the
    // CPU engine.
    if (threadIdx.x == 0) { s_ppeak_i = 0; s_plost_i = 0; }
    __syncthreads();
    float l_pgrid = 0.0f, l_pvalid = 0.0f, l_pown = 0.0f, l_mall = 0.0f, l_mread = 0.0f;
    float l_ppeak = 0.0f, l_plost = 0.0f;
    for (int k = threadIdx.x; k < GfGf; k += blockDim.x) {
        const float Pp = Pbuf[k];
        if (Pp <= 0.0f) continue;
        const int dx = k % Gf - Rf, dy = k / Gf - Rf;
        const int x = cx + dx, y = cy + dy;
        if (x < 0 || y < 0 || x >= p.W || y >= p.H) continue;
        const bool in_disk = (float) (dx * dx + dy * dy) < p.r1_sq;
        l_pgrid += Pp;
        l_ppeak = fmaxf(l_ppeak, Pp);
        if (in_disk) l_mall += Pp;
        if (!valid(img[y * p.W + x])) {
            l_plost = fmaxf(l_plost, Pp);
            continue;
        }
        l_pvalid += Pp;
        const bool own = !p.overlap || BraggOwnedBy(owner[y * p.W + x], i);
        // Zeroing the profile here is how Exclude drops the pixel: the fit skips any cell with
        // P <= 0 already, and P is not renormalised, so the fitted amplitude comes out on the
        // scale of the WHOLE profile - the renormalisation is the estimator's own doing.
        if (!own && p.overlap == 2) Pbuf[k] = 0.0f;
        if (own) l_pown += Pp;
        if (in_disk && (own || p.overlap != 2)) l_mread += Pp;
    }
    atomicMax(&s_ppeak_i, __float_as_int(l_ppeak)); atomicMax(&s_plost_i, __float_as_int(l_plost));
    const float s_pgrid = block_sum(l_pgrid, s_slot[0]);
    const float s_pvalid = block_sum(l_pvalid, s_slot[1]);
    const float s_pown = block_sum(l_pown, s_slot[2]);
    const float s_mall = block_sum(l_mall, s_slot[3]);
    const float s_mread = block_sum(l_mread, s_slot[4]);
    if (s_pvalid < p.minpk * s_pgrid) {
        if (threadIdx.x == 0) ok_o[i] = 0;
        return;
    }
    // A hole in the profile's PEAK is a different defect from a hole in its wings. See the CPU engine.
    if (__int_as_float(s_plost_i) > p.peak_frac * __int_as_float(s_ppeak_i)) {
        if (threadIdx.x == 0) ok_o[i] = 0;
        return;
    }
    if (p.overlap == 1 && s_pown < p.minpk) {
        if (threadIdx.x == 0) ok_o[i] = 0;
        return;
    }

    if (threadIdx.x == 0) s_I = I_seed[i];
    __syncthreads();

    float s_num = 0.0f, s_den = 0.0f, s_wsum = 0.0f;
    for (int iter = 0; iter < 4; ++iter) {
        const float Ihere = s_I;
        float l_num = 0.0f, l_den = 0.0f, l_wsum = 0.0f;
        for (int k = threadIdx.x; k < GfGf; k += blockDim.x) {
            const float Pp = Pbuf[k];
            if (Pp <= 0.0f) continue;
            const int x = cx + (k % Gf - Rf), y = cy + (k / Gf - Rf);
            if (x < 0 || y < 0 || x >= p.W || y >= p.H) continue;
            const int32_t px = img[y * p.W + x];
            if (!valid(px)) continue;
            const float v = fmaxf(B + Ihere * Pp, (float) WEIGHT_VARIANCE_MIN_FRACTION * B);
            // One reciprocal for the three weights. Written as three divisions the compiler emits
            // the whole correctly-rounded sequence three times over - the same estimate, the same
            // refinement, the same check - for a weight that only has to be a weight.
            const float iv = __frcp_rn(v);
            l_num += Pp * ((float) px - bkg) * iv;
            l_den += Pp * Pp * iv;
            l_wsum += Pp * iv;
        }
        s_num = block_sum(l_num, s_slot[0]);
        s_den = block_sum(l_den, s_slot[1]);
        s_wsum = block_sum(l_wsum, s_slot[2]);
        if (threadIdx.x == 0 && s_den > 0.0f) s_I = s_num / s_den;
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        if (s_den > 0.0f) {
            // 1/s_den takes the background as exact; dI/dbkg = -s_wsum/s_den adds the background
            // estimate's own error (see the CPU engine).
            const float wr = s_wsum / s_den;
            float I = s_I, sigma = sqrtf(1.0f / s_den + wr * wr * bkgvar_a[i]);
            // The signal part to remove is I itself, not max(0, I) - see the CPU engine.
            float var_bkg = fmaxf(0.0f, 1.0f / s_den - I + wr * wr * bkgvar_a[i]);
            // Guard against profile-fit runaways (see the CPU engine): fall back to the summation seed
            // when the profile result diverges from it. Pixels missing from both - excluded to a
            // neighbour, or unreadable - scale the full-profile intensity back to the disk the seed
            // read; nothing dropped gives exactly 1.
            const float gs = s_mall > 0.0f ? s_mread / s_mall : 1.0f;
            if (fabsf(I * gs - I_seed[i]) > (float) PROFILE_SUMMATION_MAX_NSIGMA * sigma_seed[i]) {
                I = I_seed[i];
                sigma = sigma_seed[i];
                var_bkg = varbkg_seed[i];
            }
            I_o[i] = I; sigma_o[i] = sigma; varbkg_o[i] = var_bkg; ok_o[i] = 1;
        } else ok_o[i] = 0;
    }
}

}   // namespace

BraggIntegrationEngineGPU::BraggIntegrationEngineGPU(const DiffractionExperiment &experiment,
                                                     std::shared_ptr<CudaStream> stream)
    : BraggIntegrationEngine(experiment),
      stream(std::move(stream)),
      d_mask(npixel),
      d_shell_grid(static_cast<size_t>(bragg_engine::N_SHELL) * GG),
      d_global_grid(GG),
      d_shell_P(static_cast<size_t>(bragg_engine::N_SHELL) * GG),
      d_global_P(GG),
      d_mom(MOM_STRIDE * (bragg_engine::N_SHELL + 1)),
      d_sigma2_r(bragg_engine::N_SHELL + 1),
      d_sigma2_t(bragg_engine::N_SHELL + 1),
      d_shell_n(bragg_engine::N_SHELL),
      d_global_n(1),
      d_invd2(2) {
    threads = 128;

    // Fit profile grid: R for empirical / box, up to 3R (radially elongated) for the Gaussian.
    const int max_Rf = empirical ? R : 3 * R;
    const int max_Gf = 2 * max_Rf + 1;
    fit_shared_bytes = static_cast<size_t>(max_Gf) * max_Gf * sizeof(float);

    // Shared radial window of the background curve in boxsum. The stencil spans r0 +- the radial
    // semi-axis of the outer ellipse, widest at the far corner of the detector, and the window
    // covers twice that so a reflection anywhere on the detector fits.
    rad_w = 2 * (static_cast<int>(std::ceil(r3 + BraggStencilGrow_px(static_cast<float>(r_max), stencil))) + 1) + 1;

    // The pixels one reflection can mark: the box mark_mask walks, at the widest aperture on the
    // detector. Run() weighs that against the frame to decide how to clear the mask afterwards.
    const int mark_half = static_cast<int>(std::ceil(r2 + BraggStencilGrow_px(static_cast<float>(r_max), stencil))) + 1;
    mask_box_px = static_cast<size_t>(2 * mark_half + 1) * static_cast<size_t>(2 * mark_half + 1);
    boxsum_shared_bytes = static_cast<size_t>(rad_w) * (sizeof(unsigned long long) + sizeof(int));

    // The current device, not device 0: workers are pinned round-robin across GPUs, so device 0's
    // shared-memory size can belong to a different card than the one these kernels launch on.
    int device = 0;
    cuda_err(cudaGetDevice(&device));
    cudaDeviceProp prop{};
    cuda_err(cudaGetDeviceProperties(&prop, device));
    if (fit_shared_bytes > prop.sharedMemPerBlock)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                              "BraggIntegrationEngineGPU: profile grid exceeds shared memory (r2 too large)");
    // boxsum's dynamic window sits alongside its static shared arrays (the trimmed-mean buffer
    // dominates them), so leave room for those rather than budgeting the whole block.
    if (boxsum_shared_bytes + sizeof(int) * BKG_TRIM_MAX + 256 > prop.sharedMemPerBlock)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                              "BraggIntegrationEngineGPU: background ring exceeds shared memory");

    // Radial background curve: one bin per pixel of distance from the beam, out to the far corner.
    // Allocated whenever the correction COULD run, so the auto mode can turn it on for an individual
    // image; whether it runs for a given image is decided in Run() from the per-image bkg_radial.
    if (overlap != OverlapMode::Off)
        d_owner = CudaDevicePtr<uint32_t>(npixel);

    if (bkg_radial || bkg_radial_auto) {
        n_rad = static_cast<int>(std::ceil(r_max)) + 2;
        d_rad_sum = CudaDevicePtr<unsigned long long>(n_rad);
        d_rad_cnt = CudaDevicePtr<int>(n_rad);
        d_k_diff = CudaDevicePtr<float>(k_diff.size());
        cuda_err(cudaMemcpy(d_k_diff, k_diff.data(), sizeof(float) * k_diff.size(),
                            cudaMemcpyHostToDevice));
    }
}

void BraggIntegrationEngineGPU::EnsureCapacity(size_t n) {
    if (n <= capacity)
        return;
    // Grow with slack. cudaMalloc/cudaFree take a device-wide lock in the driver, so growing to
    // exactly n makes every image that sets a new reflection-count record stall all other workers.
    const size_t new_capacity = std::max(n, capacity + capacity / 2);

    d_px_x = CudaDevicePtr<float>(new_capacity);
    d_px_y = CudaDevicePtr<float>(new_capacity);
    d_d = CudaDevicePtr<float>(new_capacity);
    d_cx = CudaDevicePtr<int>(new_capacity);
    d_cy = CudaDevicePtr<int>(new_capacity);
    d_I = CudaDevicePtr<float>(new_capacity);
    d_sigma = CudaDevicePtr<float>(new_capacity);
    d_bkg = CudaDevicePtr<float>(new_capacity);
    d_bkg_var = CudaDevicePtr<float>(new_capacity);
    d_var_bkg = CudaDevicePtr<float>(new_capacity);
    d_isum = CudaDevicePtr<float>(new_capacity);
    d_ninner = CudaDevicePtr<int>(new_capacity);
    d_rbin = CudaDevicePtr<int>(new_capacity);
    d_kbin = CudaDevicePtr<int>(new_capacity);
    d_obs_x = CudaDevicePtr<float>(new_capacity);
    d_obs_y = CudaDevicePtr<float>(new_capacity);
    d_ok = CudaDevicePtr<uint8_t>(new_capacity);
    d_strong = CudaDevicePtr<uint8_t>(new_capacity);
    d_has_obs = CudaDevicePtr<uint8_t>(new_capacity);

    h_px_x = CudaHostPtr<float>(new_capacity);
    h_px_y = CudaHostPtr<float>(new_capacity);
    h_d = CudaHostPtr<float>(new_capacity);
    h_I = CudaHostPtr<float>(new_capacity);
    h_sigma = CudaHostPtr<float>(new_capacity);
    h_bkg = CudaHostPtr<float>(new_capacity);
    h_var_bkg = CudaHostPtr<float>(new_capacity);
    h_obs_x = CudaHostPtr<float>(new_capacity);
    h_obs_y = CudaHostPtr<float>(new_capacity);
    h_ok = CudaHostPtr<uint8_t>(new_capacity);
    h_has_obs = CudaHostPtr<uint8_t>(new_capacity);
    capacity = new_capacity;
}

std::vector<Reflection> BraggIntegrationEngineGPU::Run(const ImagePreprocessorBuffer &image,
                                                       const std::vector<Reflection> &predicted,
                                                       size_t npredicted, int64_t image_number) {
    std::vector<BraggFitResult> results(npredicted);
    if (image.size() != npixel || npredicted == 0)
        return Finalize(predicted, npredicted, results, image_number);

    const int32_t *img = image.getGPUBuffer();
    if (img == nullptr)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "BraggIntegrationEngineGPU: image buffer is not on the GPU");

    EnsureCapacity(npredicted);
    const int n = static_cast<int>(npredicted);
    for (size_t i = 0; i < npredicted; ++i) {
        h_px_x[i] = predicted[i].predicted_x;
        h_px_y[i] = predicted[i].predicted_y;
        h_d[i] = predicted[i].d;
    }
    cuda_err(cudaMemcpyAsync(d_px_x, h_px_x.get(), sizeof(float) * npredicted, cudaMemcpyHostToDevice, *stream));
    cuda_err(cudaMemcpyAsync(d_px_y, h_px_y.get(), sizeof(float) * npredicted, cudaMemcpyHostToDevice, *stream));
    cuda_err(cudaMemcpyAsync(d_d, h_d.get(), sizeof(float) * npredicted, cudaMemcpyHostToDevice, *stream));

    BraggGpuParams p{
        .W = static_cast<int>(xpixel), .H = static_cast<int>(ypixel),
        .r1_sq = r1_sq, .r2 = r2, .r2_sq = r2_sq, .r3 = r3, .r3_sq = r3_sq,
        .R = R, .G = G, .GG = GG,
        .bkg_clip_nsigma = mode != IntegratorMode::BoxSum ? bkg_clip_nsigma : 0.0f,
        .empirical = empirical ? 1 : 0,
        .use_ellipse = use_ellipse ? 1 : 0,
        .bw_sigma = static_cast<float>(bw_sigma), .c_radial = static_cast<float>(c_radial),
        .F_px = static_cast<float>(F_px),
        .beam_x = beam_x, .beam_y = beam_y,
        .bkg_trim = bkg_trim,   // effective trim fraction (0 for stills), set by the base ctor from settings
        .stencil = stencil, .rad_w = rad_w, .n_kern = n_kern,
        .overlap = overlap == OverlapMode::Reject ? 1 : (overlap == OverlapMode::Exclude ? 2 : 0),
        .boxsum_reject = (overlap == OverlapMode::Reject && mode == IntegratorMode::BoxSum) ? 1 : 0,
        .exclude_px = (overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum) ? 1 : 0,
        .claim_sq = claim * claim, .inv_claim = inv_claim, .minpk = overlap_min_peak,
        .partial_ok = mode != IntegratorMode::BoxSum ? 1 : 0,
        .peak_frac = static_cast<float>(MINPK_MAX_MISSING_PEAK),
    };

    // Whether the radial correction runs for THIS image. n_rad only says the buffers exist - under
    // the auto mode they are allocated for every image and used for the ones the ice score selects.
    const int rad_n = bkg_radial ? n_rad : 0;

    // The mask and the owner map can be taken back either by clearing the whole frame or by revisiting
    // the boxes that were marked, and which is cheaper depends on the detector: on 18 Mpx the marks are
    // a twentieth of the frame and clearing all of it costs 109 us of card time per image - three
    // quarters of what the two integration kernels themselves cost - while on 2.5 Mpx with twelve
    // thousand predictions the boxes cover more than the frame does. So compare the two. The factor of
    // two is the measured penalty for scattering the writes over boxes instead of streaming them:
    // 0.73 TB/s against 1.2 TB/s for the memset.
    const bool clear_by_box = static_cast<size_t>(n) * mask_box_px * 2 < npixel;

    // Whoever wrote the buffers last left them clean, so nothing has to be cleared here; `dirty` says
    // the previous call did not get that far - or chose the frame-wide clear, or is the first call -
    // and the frame has to be taken back wholesale before this image marks it.
    if (dirty) {
        cuda_err(cudaMemsetAsync(d_mask, 0, npixel, *stream));
        if (p.overlap)
            cuda_err(cudaMemsetAsync(d_owner, 0xff, sizeof(uint32_t) * npixel, *stream));   // BRAGG_OWNER_NONE
    }
    if (rad_n > 0) {
        cuda_err(cudaMemsetAsync(d_rad_sum, 0, sizeof(unsigned long long) * n_rad, *stream));
        cuda_err(cudaMemsetAsync(d_rad_cnt, 0, sizeof(int) * n_rad, *stream));
    }
    reset<<<32, 256, 0, *stream>>>(d_shell_grid, d_global_grid, d_mom, d_shell_n, d_global_n, d_invd2, GG);
    dirty = true;
    mark_mask<<<n, threads, 0, *stream>>>(d_px_x, d_px_y, d_mask, d_owner, p, n, 0);
    boxsum<<<n, threads, boxsum_shared_bytes, *stream>>>(d_px_x, d_px_y, d_d, img, d_mask, d_owner, p, n,
                                       d_cx, d_cy, d_I, d_sigma, d_bkg, d_bkg_var, d_var_bkg, d_obs_x, d_obs_y,
                                       d_ok, d_strong, d_has_obs, d_invd2,
                                       d_isum, d_ninner, d_rbin, d_kbin,
                                       d_rad_sum, d_rad_cnt, rad_n);

    // Correct the flat annulus background for the curvature of the radial background before anything
    // downstream (profile fit, variance) reads it.
    if (rad_n > 0)
        radial_correct<<<(n + threads - 1) / threads, threads, 0, *stream>>>(
            d_rad_sum, d_rad_cnt, rad_n, d_k_diff, k_len, k_off,
            d_isum, d_ninner, d_rbin, d_kbin, d_ok, d_bkg, d_I, n);

    if (mode != IntegratorMode::BoxSum) {
        // Pass B: learn (shell computed inline) -> build -> fit.
        learn_profile<<<n, threads, 0, *stream>>>(img, d_owner, d_px_x, d_px_y, d_cx, d_cy, d_d, d_invd2,
                                                  d_I, d_bkg, d_ok, d_strong,
                                                  d_shell_grid, d_global_grid, d_mom, d_shell_n, d_global_n, p, n);
        build_profiles<<<bragg_engine::N_SHELL + 1, threads, 0, *stream>>>(
            d_shell_grid, d_global_grid, d_mom, d_shell_n, d_global_n,
            d_shell_P, d_global_P, d_sigma2_r, d_sigma2_t, p);
        fit<<<n, threads, fit_shared_bytes, *stream>>>(img, d_owner, d_px_x, d_px_y, d_cx, d_cy, d_d, d_invd2,
                                                       d_I, d_sigma, d_bkg, d_bkg_var, d_var_bkg, d_ok, d_shell_P, d_global_P,
                                                       d_sigma2_r, d_sigma2_t, d_shell_n,
                                                       d_I, d_sigma, d_var_bkg, d_ok, p, n);
    }

    // Nothing below reads the mask or the owner map, so put them back now, over the boxes that were
    // marked rather than over the frame - where that is the cheaper of the two.
    if (clear_by_box)
        mark_mask<<<n, threads, 0, *stream>>>(d_px_x, d_px_y, d_mask, d_owner, p, n, 1);

    cuda_err(cudaMemcpyAsync(h_I.get(), d_I, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(h_sigma.get(), d_sigma, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(h_bkg.get(), d_bkg, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(h_var_bkg.get(), d_var_bkg, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(h_ok.get(), d_ok, sizeof(uint8_t) * npredicted, cudaMemcpyDeviceToHost, *stream));
    // Pass A always fills the box-sum centroid (observed spot position), so copy it back in all modes -
    // post-refinement uses it as the observed position (beam-centre / distance).
    cuda_err(cudaMemcpyAsync(h_obs_x.get(), d_obs_x, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(h_obs_y.get(), d_obs_y, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(h_has_obs.get(), d_has_obs, sizeof(uint8_t) * npredicted, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaStreamSynchronize(*stream));
    // The clearing pass has run, so the buffers are clean again for the next image.
    if (clear_by_box)
        dirty = false;

    for (size_t i = 0; i < npredicted; ++i) {
        if (!h_ok[i]) continue;
        results[i].I = h_I[i];
        results[i].sigma = h_sigma[i];
        results[i].bkg = h_bkg[i];
        results[i].var_bkg = h_var_bkg[i];
        results[i].ok = true;
        if (h_has_obs[i]) {
            results[i].observed_x = h_obs_x[i];
            results[i].observed_y = h_obs_y[i];
            results[i].has_observed = true;
        }
    }
    return Finalize(predicted, npredicted, results, image_number);
}