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

#include "RotationScaleMergeGPU.h"
#include "OutlierBand.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <string>
#include <type_traits>
#include <vector>
#include <cuda_runtime.h>
#include <math_constants.h>
#include <cub/device/device_radix_sort.cuh>
#include <cub/device/device_run_length_encode.cuh>
#include <cub/device/device_scan.cuh>

#include "../indexing/CUDAMemHelpers.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
#include "../../common/ParallelFor.h"
#include "../../common/Reflection.h"

namespace {
    constexpr int BLK = 256;
    constexpr int MIN_REFLECTIONS = 20;

    // Every kernel and copy here is queued on the instance's own stream (Impl::stream), not on the
    // legacy NULL stream: the merge and the image analysis of a probe pass beside it would otherwise
    // wait for each other's work at every launch and every synchronisation. As in BeamCenterFFTGPU no buffer comes from the pool. A
    // pooled buffer is freed with cudaFreeAsync on the thread's allocation stream, which is not
    // ordered after this one. Each entry point below waits for its own work before it returns, so no
    // free here has yet overtaken a read, but that holds only by that convention, and not at all for a
    // free while unwinding from a failed call; nor can compute-sanitizer --track-stream-ordered-races
    // see it, and it reported every reassigned merge buffer as a use-after-free. cudaFree
    // synchronises the device first.
    constexpr CudaAlloc ALLOC = CudaAlloc::SynchronousAfterWorkBeside;

    __device__ __forceinline__ double SafeInvD(double x, double fallback) {
        return (isfinite(x) && x != 0.0) ? 1.0 / x : fallback;
    }

    // Block reduction of a double, deterministic for a fixed thread->element mapping (fixed order,
    // no atomics). Returns the sum on thread 0; `s` is BLK doubles of shared scratch.
    __device__ double BlockReduceSum(double v, double *s) {
        const int t = threadIdx.x;
        s[t] = v;
        __syncthreads();
        for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
            if (t < stride) s[t] += s[t + stride];
            __syncthreads();
        }
        return s[0];
    }

    // The four immutable fields of an observation that the scaling, the combine and the grouping read
    // together, in one 16-byte record: the kernels that reach an observation through a permutation
    // (the group reduction above all) then fetch one sector for it instead of one per field.
    struct alignas(16) ObsRecord { float I, sigma, partiality, prescaling_corr; };   // floats 0, 1, 2, 3

    // A partial's flags, one byte: on an ice ring, then RotationScaleMerge::OBS_CLIPPED and
    // OBS_OVERLOADED one bit up. A full's on_ice byte is 0 or OBS_ON_ICE_FLAG, so the kernels shared by
    // the two read either the same way.
    constexpr uint8_t OBS_ON_ICE_FLAG = 1;
    constexpr uint8_t OBS_CLIPPED_FLAG = 2;
    constexpr uint8_t OBS_OVERLOADED_FLAG = 4;

    // An observation's frame. The fulls carry theirs (frame); the partials lie in frame order, so theirs
    // is the last frame starting at or before them in the frame CSR (start, n_frames entries) - an empty
    // frame starts where the next one does, so it is never the last.
    struct FrameOf {
        const int32_t *frame;
        const int32_t *start;
        int n_frames;
        __device__ __forceinline__ int operator()(int i) const {
            if (frame) return frame[i];
            int lo = 0, hi = n_frames - 1;
            while (lo < hi) {
                const int mid = (lo + hi + 1) / 2;
                if (start[mid] <= i) lo = mid;
                else hi = mid - 1;
            }
            return lo;
        }
    };

    // The observation's inverse-variance weight in the scaling reference (0 = not in it), and its
    // scaled intensity. The same filter as the host's ReferenceWeight.
    __device__ __forceinline__ double ReferenceWeight(float I, float sigma, float partiality, float corr,
                                                      double min_partiality, float &I_corr) {
        if (!(corr > 0.0f) || !isfinite(corr)) return 0.0;
        if (partiality < min_partiality) return 0.0;
        const float Ic = I * corr;
        const float sigma_corr = sigma * corr;
        if (!isfinite(Ic) || !isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0;
        I_corr = Ic;
        return 1.0 / (double(sigma_corr) * sigma_corr);
    }

    // One thread per ASU group (grid-stride): inverse-variance mean of I*corr over the group's contiguous,
    // fixed-order segment of group_perm. Groups are small (avg tens of obs), so a whole block per group
    // wastes threads; summing each group in one thread avoids launch/sync overhead and stays deterministic
    // (fixed group_perm order). Matches the CPU ReduceGroupMeans filter (no ice/cell mask - scaling ref).
    __global__ void ReduceGroupMeansKernel(int n_groups, double min_partiality,
                                           const int32_t *__restrict__ group_perm,
                                           const int32_t *__restrict__ group_start,
                                           const int32_t *__restrict__ group_count,
                                           const ObsRecord *__restrict__ obs,
                                           const float *__restrict__ corr,
                                           double *__restrict__ group_mean) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < n_groups; g += gridDim.x * blockDim.x) {
            const int lo = group_start[g], hi = group_start[g] + group_count[g];
            double sw = 0.0, swI = 0.0;
            for (int p = lo; p < hi; ++p) {
                const int i = group_perm[p];
                const ObsRecord o = obs[i];
                float I_corr;
                const double w = ReferenceWeight(o.I, o.sigma, o.partiality, corr[i], min_partiality, I_corr);
                if (w <= 0.0) continue;
                sw += w;
                swI += w * I_corr;
            }
            group_mean[g] = sw > 0.0 ? swI / sw : NAN;
        }
    }

    // The scale-fit coefficient of an observation (rotation model), recomputed each scaling iteration
    // once the group means are known: coeff = partiality * (1/prescaling_corr) * mean[group], or false
    // where the observation is not one the reference is built from (the host RunArm says why, where it fills ScalingFit).
    __device__ __forceinline__ bool ScaleCoeff(int i, double min_partiality, const int32_t *__restrict__ group,
                                               const ObsRecord &o, const float *__restrict__ zeta,
                                               const uint8_t *__restrict__ on_ice,
                                               const double *__restrict__ group_mean, float &coeff) {
        const int g = group[i];
        if (!((g >= 0) && !(on_ice[i] & OBS_ON_ICE_FLAG) && isfinite(zeta[i]) && zeta[i] > 0.0f && o.partiality >= min_partiality))
            return false;
        const double mean = group_mean[g];
        if (!isfinite(mean)) return false;
        coeff = float(double(o.partiality) * SafeInvD(o.prescaling_corr, 1.0) * mean);
        return true;
    }

    // One block per frame: the per-frame scale G as the weighted least-squares slope over the frame's
    // contiguous obs - the same plain fit as the host ScaleFromSums (which says why it is not robust).
    // Leaves g/scaled untouched for frames with fewer than min_obs usable observations. `perm` (null
    // for the partials, whose arrays are already frame-contiguous) maps a position in the frame's
    // [lo,hi) range to the obs index, so the same kernel scales the fulls (emit-ordered) through a
    // frame-grouping permutation without physically reordering the fulls arrays. info[f] receives the
    // fit's information on G, sum w^2 c^2 (0 for a frame left alone). The weight is 1/sigma, as the
    // CPU's ScalingFit::w2 is its square.
    __global__ void FitPerFrameGKernel(int n_frames,
                                       const int32_t *__restrict__ frame_start,
                                       const int32_t *__restrict__ frame_count,
                                       const int32_t *__restrict__ perm, double min_partiality,
                                       const int32_t *__restrict__ group, const ObsRecord *__restrict__ obs,
                                       const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice,
                                       const double *__restrict__ group_mean, long min_obs,
                                       double *__restrict__ g, uint8_t *__restrict__ scaled,
                                       double *__restrict__ info) {
        const int f = blockIdx.x;
        if (f >= n_frames) return;
        const int lo = frame_start[f], hi = frame_start[f] + frame_count[f];
        __shared__ double sh[BLK];
        if (threadIdx.x == 0) info[f] = 0.0;

        // The sums are formed in the same pass as the count; a frame with too few observations
        // discards them.
        long cnt_local = 0;
        double num = 0.0, den = 0.0;
        for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
            const int a = perm ? perm[i] : i;
            const ObsRecord o = obs[a];
            float coeff_f;
            if (!ScaleCoeff(a, min_partiality, group, o, zeta, on_ice, group_mean, coeff_f)) continue;
            ++cnt_local;
            const double coeff = coeff_f;
            const double w = SafeInvD(o.sigma, 1.0);
            const double w2 = w * w;
            num += w2 * coeff * double(o.I);
            den += w2 * coeff * coeff;
        }
        const double cnt = BlockReduceSum(double(cnt_local), sh);
        __shared__ double s_cnt;
        if (threadIdx.x == 0) s_cnt = cnt;
        __syncthreads();
        if (s_cnt < min_obs) return;   // leave g[f]/scaled[f] as-is

        const double tnum = BlockReduceSum(num, sh); __syncthreads();
        const double tden = BlockReduceSum(den, sh);
        if (threadIdx.x == 0) {
            const double G = tden > 0.0 ? tnum / tden : NAN;
            g[f] = isfinite(G) ? fmax(0.0, G) : 1.0;
            scaled[f] = 1;
            info[f] = tden;
        }
    }

    // One block per frame: Pearson CC of (I*corr) vs the merged group mean over the frame's partials,
    // == FinalizePerFrameScale's per-frame loop. Diagnostic only (per-image scaling table), so the tree
    // reduction's ~ulp difference from the CPU is immaterial; deterministic run-to-run.
    __global__ void PerFrameCCKernel(int n_frames, double min_partiality,
                                     const int32_t *__restrict__ frame_start,
                                     const int32_t *__restrict__ frame_count,
                                     const ObsRecord *__restrict__ obs, const float *__restrict__ corr,
                                     const uint8_t *__restrict__ flags, const int32_t *__restrict__ group,
                                     const double *__restrict__ group_mean,
                                     double *__restrict__ cc_out, int64_t *__restrict__ cc_n_out) {
        const int f = blockIdx.x;
        if (f >= n_frames) return;
        const int lo = frame_start[f], hi = frame_start[f] + frame_count[f];
        __shared__ double sh[BLK];
        double sx = 0, sy = 0, sx2 = 0, sy2 = 0, sxy = 0;
        long nl = 0;
        for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
            if (flags[i] & OBS_ON_ICE_FLAG) continue;
            const int g = group[i];
            if (g < 0) continue;
            const ObsRecord o = obs[i];
            if (o.partiality < min_partiality) continue;
            const float c = corr[i];
            if (!isfinite(o.I) || !isfinite(c) || !(c > 0.0f)) continue;
            if (!isfinite(o.sigma) || !(o.sigma > 0.0f)) continue;
            const double mean = group_mean[g];
            if (!isfinite(mean)) continue;
            const double img = double(o.I) * c;
            sx += img; sy += mean; sx2 += img * img; sy2 += mean * mean; sxy += img * mean; ++nl;
        }
        const double tsx = BlockReduceSum(sx, sh); __syncthreads();
        const double tsy = BlockReduceSum(sy, sh); __syncthreads();
        const double tsx2 = BlockReduceSum(sx2, sh); __syncthreads();
        const double tsy2 = BlockReduceSum(sy2, sh); __syncthreads();
        const double tsxy = BlockReduceSum(sxy, sh); __syncthreads();
        const double tn = BlockReduceSum(double(nl), sh);
        if (threadIdx.x == 0) {
            cc_out[f] = NAN; cc_n_out[f] = 0;
            if (tn >= MIN_REFLECTIONS) {
                const double cov = tsxy - tsx * tsy / tn;
                const double vx = tsx2 - tsx * tsx / tn;
                const double vy = tsy2 - tsy * tsy / tn;
                if (vx > 0.0 && vy > 0.0) { cc_out[f] = cov / sqrt(vx * vy); cc_n_out[f] = int64_t(tn); }
            }
        }
    }

    // SmoothG corr adjust: corr[i] *= ratio[frame[i]] for frames flagged apply, in double then stored
    // as float - matching CPU SmoothG's `corr = float(corr * (g/g_smooth))`. Grid-stride, resident corr.
    __global__ void SmoothCorrKernel(int n_obs, FrameOf frame_of,
                                     const uint8_t *__restrict__ apply, const double *__restrict__ ratio,
                                     float *__restrict__ corr) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
            const int f = frame_of(i);
            if (!apply[f]) continue;
            const float c = corr[i];
            if (isfinite(c)) corr[i] = float(double(c) * ratio[f]);
        }
    }

    // Zero corr where the rocking geometry is too tangential for the de-novo search, counting what that
    // removed from the merge. zeta is compared in double, as the host does with a double threshold. The
    // count is an integer sum, so the order the atomics land in cannot change it.
    __global__ void FilterZetaKernel(int n_obs, double min_zeta, const float *__restrict__ zeta,
                                     float *__restrict__ corr, unsigned long long *dropped) {
        unsigned long long local = 0;
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
            const float z = zeta[i];
            if (isfinite(z) && double(z) >= min_zeta) continue;
            const float c = corr[i];
            if (isfinite(c) && c > 0.0f) ++local;
            corr[i] = 0.0f;
        }
        if (local > 0) atomicAdd(dropped, local);
    }

    // Zero corr on the rejected frames (--min-image-cc), grid-stride over the resident corr.
    __global__ void FilterFrameKernel(int n_obs, FrameOf frame_of,
                                      const uint8_t *__restrict__ reject, float *__restrict__ corr) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x)
            if (reject[frame_of(i)]) corr[i] = 0.0f;
    }

    // corr = prescaling_corr / (partiality * G[frame]) for fitted frames; unchanged otherwise (grid-stride).
    __global__ void UpdateCorrKernel(int n_obs, FrameOf frame_of,
                                     const ObsRecord *__restrict__ obs,
                                     const double *__restrict__ g, const uint8_t *__restrict__ scaled,
                                     float *__restrict__ corr) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
            const int f = frame_of(i);
            if (!scaled[f]) continue;
            const ObsRecord o = obs[i];
            const double denom = double(o.partiality) * g[f];
            corr[i] = (isfinite(double(o.prescaling_corr)) && isfinite(denom) && denom > 0.0)
                          ? float(o.prescaling_corr / denom) : NAN;
        }
    }

    // std::max / std::min return (a<b)?b:a and (b<a)?b:a - reproduce that exactly (NOT fmax/fmin, which
    // differ on NaN) so the combine matches the CPU path bit-for-bit on the same inputs.
    __device__ __forceinline__ double Dmax(double a, double b) { return (a < b) ? b : a; }
    __device__ __forceinline__ double Dmin(double a, double b) { return (b < a) ? b : a; }
    __device__ __forceinline__ float  Fmax(float a, float b)   { return (a < b) ? b : a; }

    // A partial is usable for the combine iff its corr and (I, sigma) are finite with corr>0, sigma>0.
    __device__ __forceinline__ bool CombineUsable(int i, const ObsRecord *obs, const float *corr) {
        const float c = corr[i];
        if (!(c > 0.0f) || !isfinite(c)) return false;
        return isfinite(obs[i].I) && isfinite(obs[i].sigma) && obs[i].sigma > 0.0f;
    }


    // All device pointers + scalars the combine kernels need, passed by value.
    struct CombineParams {
        int n_runs;
        double min_partiality, capture_uncertainty_coeff, min_captured_fraction;
        // RotationScaleMerge::max_frame_gap, passed in rather than recomputed here so that the host
        // and the device compare the same float and the combine stays bit-identical on both paths.
        float max_frame_gap;
        const ObsRecord *__restrict__ obs;
        const float *__restrict__ corr, *__restrict__ bkg, *__restrict__ var_bkg,
                    *__restrict__ image_number, *__restrict__ d;
        FrameOf frame_of;
        const uint8_t *__restrict__ flags;
        const int32_t *__restrict__ perm, *__restrict__ rr_start, *__restrict__ rr_count,
                      *__restrict__ rr_h, *__restrict__ rr_k, *__restrict__ rr_l,
                      *__restrict__ rr_group;
        int32_t *rr_nevents;                // count pass output
        int32_t *rr_noverloaded;            // count pass output: events dropped for a saturated pixel
        const int32_t *rr_offset;           // emit pass: per-run base offset into the fulls arrays
        int32_t *f_run, *f_peak, *f_frame, *f_group;
        float *f_I, *f_sigma, *f_d, *f_img, *f_var_bkg, *f_var_per_I, *f_capture;
        uint8_t *f_on_ice, *f_clipped;
    };

    // One thread's raw-hkl run: split its usable partials (already in image_number order within the run)
    // into rocking events, and for each event pool background, seed F and run the 3-iter de-biased Poisson
    // reweight - the exact objective of RotationScaleMerge::Combine::process_rawrun. When Emit, write the
    // resulting full at rr_offset[r] + (event index); otherwise just count the emitted events. Both modes
    // run the identical accept test, so the count pass predicts the emit pass exactly.
    template <bool Emit>
    __device__ void CombineRawRun(int r, const CombineParams &p) {
        const int lo = p.rr_start[r];
        const int hi = lo + p.rr_count[r];
        const int group = p.rr_group[r];

        int n_emit = 0, n_overloaded = 0;
        int cursor = lo;
        while (cursor < hi) {
            while (cursor < hi && !CombineUsable(p.perm[cursor], p.obs, p.corr)) ++cursor;
            if (cursor >= hi) break;
            const int ev_start = cursor;             // first usable position of the event
            int ev_end = cursor;                     // last usable position (inclusive), extended below
            float last_img = p.image_number[p.perm[cursor]];
            int probe = cursor + 1;
            while (probe < hi) {
                while (probe < hi && !CombineUsable(p.perm[probe], p.obs, p.corr)) ++probe;
                if (probe >= hi) break;
                const float img = p.image_number[p.perm[probe]];
                if (img - last_img > p.max_frame_gap) break;
                last_img = img;
                ev_end = probe;
                ++probe;
            }
            cursor = ev_end + 1;

            // Pass A: pooled background = mean of the event members' finite backgrounds.

            double pooled_bkg = 0.0;
            int n_pool = 0;
            for (int m = ev_start; m <= ev_end; ++m) {
                const int i = p.perm[m];
                if (!CombineUsable(i, p.obs, p.corr)) continue;
                const float b = p.bkg[i];
                if (isfinite(b)) { pooled_bkg += b; ++n_pool; }
            }
            pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0;

            auto pooled_I = [&](int i) -> double {
                const double n_bkg = Dmax(0.0, double(p.obs[i].sigma) * p.obs[i].sigma - p.obs[i].I)
                                     / Fmax(p.bkg[i], 1.0f);
                return double(p.obs[i].I) + n_bkg * (double(p.bkg[i]) - pooled_bkg);
            };

            // Pass B: seed F (inverse-variance mean of pooled_I*corr), plus peak / d / on_ice / partiality.
            double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0;
            float d = NAN;
            const int first = p.perm[ev_start];
            int peak_outcome = p.frame_of(first);
            float peak_frame = p.image_number[first];
            int peak_obs = first;
            float peak_partiality = -1.0f;
            const bool on_ice = p.flags[first] & OBS_ON_ICE_FLAG;
            uint8_t flags = 0;
            for (int m = ev_start; m <= ev_end; ++m) flags |= p.flags[p.perm[m]];
            // An overloaded event is not a measurement of the reflection; see the host combine.
            if (flags & OBS_OVERLOADED_FLAG) { ++n_overloaded; continue; }
            const bool clipped = flags & OBS_CLIPPED_FLAG;
            for (int m = ev_start; m <= ev_end; ++m) {
                const int i = p.perm[m];
                if (!CombineUsable(i, p.obs, p.corr)) continue;
                const double sigma_corr = double(p.obs[i].sigma) * p.corr[i];
                const double w = 1.0 / (sigma_corr * sigma_corr);
                sum_w += w;
                sum_wI += w * pooled_I(i) * p.corr[i];
                sum_partiality += p.obs[i].partiality;
                if (p.obs[i].partiality > peak_partiality) {
                    peak_partiality = p.obs[i].partiality;
                    peak_outcome = p.frame_of(i);
                    peak_frame = p.image_number[i];
                    peak_obs = i;
                }
                if (!isfinite(d) && isfinite(p.d[i]) && p.d[i] > 0.0f) d = p.d[i];
            }
            double F = sum_wI / sum_w;

            // Pass C: 3 de-biased Poisson reweights (variance = bkg part + corr*max(0,F)), plus the
            // full's var(I) = var_bkg + var_per_I*I for the merge (see the host combine).
            double sum_wb = 0.0, sum_cwb = 0.0;
            for (int iter = 0; iter < 3; ++iter) {
                sum_w = 0.0; sum_wI = 0.0; sum_wb = 0.0; sum_cwb = 0.0;
                for (int m = ev_start; m <= ev_end; ++m) {
                    const int i = p.perm[m];
                    if (!CombineUsable(i, p.obs, p.corr)) continue;
                    const double corr = p.corr[i];
                    const double I_corr = pooled_I(i) * corr;
                    const double sigma_corr = double(p.obs[i].sigma) * corr;
                    // The integrator's own non-signal variance (see the host combine).
                    const double bkg_var = corr * corr * (double) p.var_bkg[i];
                    const double a_var = bkg_var > 0.0 ? bkg_var : sigma_corr * sigma_corr;
                    double var = a_var + corr * Dmax(0.0, F);
                    const double w = 1.0 / var;
                    sum_w += w;
                    sum_wI += w * I_corr;
                    // a_var has no F in it, so these two are the same in every reweight and only
                    // the last round's values are ever read. Two thirds of them were two divisions
                    // each, thrown away.
                    if (iter == 2) {
                        sum_wb += 1.0 / a_var;
                        sum_cwb += corr / (a_var * a_var);
                    }
                }
                F = sum_wI / sum_w;
            }
            const double var_bkg_full = 1.0 / sum_wb;

            if (sum_w <= 0.0 || sum_partiality < p.min_partiality
                             || sum_partiality < p.min_captured_fraction)
                continue;

            double sigma_full = 1.0 / sqrt(sum_w);
            const double capture = p.capture_uncertainty_coeff * (1.0 - Dmin(1.0, sum_partiality));
            if (p.capture_uncertainty_coeff > 0.0) {
                const double extra = capture * Dmax(0.0, F);
                sigma_full = sqrt(sigma_full * sigma_full + extra * extra);
            }

            if (Emit) {
                const int o = p.rr_offset[r] + n_emit;
                p.f_run[o] = r;
                p.f_I[o] = float(F);
                p.f_sigma[o] = float(sigma_full);
                p.f_var_bkg[o] = float(var_bkg_full);
                p.f_var_per_I[o] = float(var_bkg_full * var_bkg_full * sum_cwb);
                p.f_capture[o] = float(capture);
                p.f_d[o] = d;
                p.f_img[o] = peak_frame;
                p.f_peak[o] = peak_obs;
                p.f_frame[o] = peak_outcome;
                p.f_on_ice[o] = on_ice ? 1 : 0;
                p.f_clipped[o] = clipped ? 1 : 0;
                p.f_group[o] = group;
            }
            ++n_emit;
        }

        if (!Emit) {
            p.rr_nevents[r] = n_emit;
            p.rr_noverloaded[r] = n_overloaded;
        }
    }

    template <bool Emit>
    __global__ void CombineKernel(CombineParams p) {
        for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < p.n_runs; r += gridDim.x * blockDim.x)
            CombineRawRun<Emit>(r, p);
    }

    // ===== the per-space-group grouping, built on the device from the raw-hkl runs =====

    // The host's finiteness test (finite_ok in RotationScaleMerge::Ingest), on the same uploaded values.
    __device__ __forceinline__ bool ObsFinite(int i, const ObsRecord *obs) {
        const ObsRecord o = obs[i];
        return isfinite(o.I) && isfinite(o.prescaling_corr) && o.prescaling_corr != 0.0f && isfinite(o.sigma)
               && o.sigma > 0.0f;
    }

    struct GroupParams {
        int n_runs, n_groups;
        const int32_t *perm, *rr_start, *rr_count, *rr_group;
        const int32_t *sorted_run, *group_first;   // group g = runs sorted_run[group_first[g] .. [g+1])
        const ObsRecord *obs;
        int32_t *group, *group_start, *group_count, *group_perm;
    };

    // Every observation's group: its raw-hkl run's, or -1 when the run has none or the observation
    // failed the finiteness test. Runs own disjoint observations.
    __global__ void StampGroupsKernel(GroupParams p) {
        for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < p.n_runs; r += gridDim.x * blockDim.x) {
            const int g = p.rr_group[r];
            for (int q = p.rr_start[r]; q < p.rr_start[r] + p.rr_count[r]; ++q) {
                const int i = p.perm[q];
                p.group[i] = (g >= 0 && ObsFinite(i, p.obs)) ? g : -1;
            }
        }
    }

    __global__ void CountGroupsKernel(GroupParams p) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            int c = 0;
            for (int j = p.group_first[g]; j < p.group_first[g + 1]; ++j) {
                const int r = p.sorted_run[j];
                for (int q = p.rr_start[r]; q < p.rr_start[r] + p.rr_count[r]; ++q)
                    c += ObsFinite(p.perm[q], p.obs) ? 1 : 0;
            }
            p.group_count[g] = c;
        }
    }

    // Each group gathers its observations and puts them in ascending index order (an in-place heap
    // sort: a group is tens of observations, rarely thousands) - the order the host's stable counting
    // sort over the observations produced, which is unique because the indices are distinct.
    __global__ void FillGroupsKernel(GroupParams p) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            int32_t *a = p.group_perm + p.group_start[g];
            int n = 0;
            for (int j = p.group_first[g]; j < p.group_first[g + 1]; ++j) {
                const int r = p.sorted_run[j];
                for (int q = p.rr_start[r]; q < p.rr_start[r] + p.rr_count[r]; ++q) {
                    const int i = p.perm[q];
                    if (ObsFinite(i, p.obs)) a[n++] = i;
                }
            }
            auto sift = [a](int root, int end) {
                while (2 * root + 1 < end) {
                    int child = 2 * root + 1;
                    if (child + 1 < end && a[child] < a[child + 1]) ++child;
                    if (a[root] >= a[child]) return;
                    const int32_t t = a[root]; a[root] = a[child]; a[child] = t;
                    root = child;
                }
            };
            for (int k = n / 2 - 1; k >= 0; --k) sift(k, n);
            for (int end = n - 1; end > 0; --end) {
                const int32_t t = a[0]; a[0] = a[end]; a[end] = t;
                sift(0, end);
            }
        }
    }

    // One field of a slice of the records, from a contiguous array: rec[i].field = v[i], the field being
    // the record's float number `field`.
    __global__ void SetRecordFieldKernel(int n, const float *__restrict__ v, ObsRecord *__restrict__ rec, int field) {
        float *f = reinterpret_cast<float *>(rec);
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
            f[size_t(i) * 4 + field] = v[i];
    }

    // The fulls' records for their Unity scaling: partiality and prescaling_corr 1.
    __global__ void FullsRecordKernel(int n, const float *__restrict__ I, const float *__restrict__ sigma,
                                      ObsRecord *__restrict__ obs) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
            obs[i] = {I[i], sigma[i], 1.0f, 1.0f};
    }

    template <typename T>
    __global__ void FillKernel(T *p, int n, T v) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) p[i] = v;
    }


    // ===== correction-surface terms built on the device from the resident fulls =====
    // ApplyCellSurface's host build, step for step: which fulls give a term, the terms in fulls order,
    // the group CSR over them, each subset's per-(block, cell) segments and the shells' weights. Every
    // order the host makes is a stable counting sort, which a stable radix sort of the indices
    // reproduces; the only arithmetic (1/d^2, the shell weights) is rounded as the host rounds it.

    // Per full: 1 if it gives a term (all, even frame, odd frame) - the host's `usable(o) && cell >= 0`.
    // A full's partiality is 1 (the combine sets it so on both paths).
    __global__ void SurfaceKeepKernel(int n, const int32_t *group, const float *corr, const uint8_t *on_ice,
                                      const int32_t *frame, const int32_t *cell, double min_partiality,
                                      int32_t *keep, int32_t *keep_even, int32_t *keep_odd) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n) return;
        const float c = corr[i];
        const bool k = group[i] >= 0 && c > 0.0f && isfinite(c) && 1.0f >= min_partiality && !on_ice[i]
                       && cell[i] >= 0;
        const bool odd = frame[i] & 1;
        keep[i] = k; keep_even[i] = k && !odd; keep_odd[i] = k && odd;
    }

    // The term of each kept full at its place (the scans of the keep flags), its frame parity, and its
    // place in the even / odd lists.
    __global__ void SurfaceWriteTermKernel(int n, const int32_t *keep_at, const int32_t *even_at,
                                           const int32_t *odd_at, const int32_t *last_flags,
                                           const float *I, const float *sigma, const float *corr, const float *d,
                                           const int32_t *cell, const int32_t *group, const int32_t *frame,
                                           RotationScaleMergeGPU::SurfaceTerm *term, uint8_t *parity,
                                           int32_t *idx_even, int32_t *idx_odd, uint32_t *group_key) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n) return;
        // A full is kept when its exclusive-scan position steps up to the next one (the last full by
        // its own flag).
        const bool kept = i + 1 < n ? keep_at[i + 1] > keep_at[i] : last_flags[0] != 0;
        if (!kept) return;
        const int w = keep_at[i];
        term[w] = {I[i], sigma[i], corr[i], d[i], cell[i], group[i]};
        group_key[w] = static_cast<uint32_t>(group[i]);
        const uint8_t par = static_cast<uint8_t>(frame[i] & 1);
        parity[w] = par;
        if (par) idx_odd[odd_at[i]] = w;
        else idx_even[even_at[i]] = w;
    }

    __global__ void IotaKernel(int n, int32_t *v) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i < n) v[i] = i;
    }

    // Counting into bins: integer adds, so the counts do not depend on their order.
    __global__ void CountKeysKernel(int n, const uint32_t *key, int32_t *count) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i < n) atomicAdd(&count[key[i]], 1);
    }

    // The block-and-cell key of position k of a subset of n terms: block b holds positions
    // [n b / nb, n (b + 1) / nb), as the host's ParallelFor over blocks cuts them.
    __global__ void SurfaceSubsetKeyKernel(int n, int nb, int ncell, const int32_t *sel,
                                           const RotationScaleMergeGPU::SurfaceTerm *term, uint32_t *key) {
        const int k = blockIdx.x * blockDim.x + threadIdx.x;
        if (k >= n) return;
        auto lo = [n, nb](int b) { return static_cast<int64_t>(n) * b / nb; };
        int b = static_cast<int>(static_cast<int64_t>(k) * nb / n);
        while (b + 1 < nb && lo(b + 1) <= k) ++b;
        while (b > 0 && lo(b) > k) --b;
        key[k] = static_cast<uint32_t>(b * ncell + term[sel[k]].cell);
    }

    // 1/d^2 of each term, as the host's shells take it.
    __global__ void SurfaceS2Kernel(int n, const RotationScaleMergeGPU::SurfaceTerm *term, float *s2) {
        const int k = blockIdx.x * blockDim.x + threadIdx.x;
        if (k >= n) return;
        const float d = term[k].d;
        s2[k] = d > 0.0f ? __fdiv_rn(1.0f, __fmul_rn(d, d)) : 0.0f;
    }

    // The shell of each term (how many edges are at or below its s^2 - the host's upper_bound) and
    // the (cell, shell) bin its weight goes to.
    __global__ void SurfaceShellKernel(int n, const float *s2, const float *edge, int n_edge, int nshell,
                                       const RotationScaleMergeGPU::SurfaceTerm *term, uint8_t *shell,
                                       uint32_t *bin) {
        const int k = blockIdx.x * blockDim.x + threadIdx.x;
        if (k >= n) return;
        int sh = 0;
        while (sh < n_edge && edge[sh] <= s2[k]) ++sh;
        shell[k] = static_cast<uint8_t>(sh);
        bin[k] = static_cast<uint32_t>(term[k].cell * nshell + sh);
    }

    // The shell of each group: that of its last term in fulls order (the host overwrites it term by term).
    __global__ void SurfaceGroupShellKernel(int n_groups, const int32_t *gstart, const int32_t *gperm,
                                            const uint8_t *shell, int32_t *group_shell) {
        const int g = blockIdx.x * blockDim.x + threadIdx.x;
        if (g >= n_groups) return;
        group_shell[g] = gstart[g + 1] > gstart[g] ? shell[gperm[gstart[g + 1] - 1]] : 0;
    }

    // One thread per (cell, shell) bin: the statistical weight of its terms, in fulls order.
    __global__ void SurfaceShellWeightKernel(int n_bins, const int32_t *bin_start, const int32_t *order,
                                             const RotationScaleMergeGPU::SurfaceTerm *term, double *shw) {
        const int b = blockIdx.x * blockDim.x + threadIdx.x;
        if (b >= n_bins) return;
        double w = 0.0;
        for (int q = bin_start[b]; q < bin_start[b + 1]; ++q) {
            const auto &t = term[order[q]];
            const double sc = __dmul_rn(double(t.sigma), double(t.corr));
            if (!(sc > 0.0)) continue;
            w = __dadd_rn(w, 1.0 / __dmul_rn(sc, sc));
        }
        shw[b] = w;
    }

    // ===== error-model + merge reductions over the resident, scaled fulls (mirror MergeAndStats) =====

    struct MergeParams {
        int n_groups;
        double min_partiality, error_model_a, error_model_b, reject_nsigma;
        double reject_up, reject_down;   // the multiplicative part of the outlier band (OutlierBand.h)
        int for_search, error_model_active, reject_outliers;
        const float *I, *sigma, *corr, *d, *reject_median, *var_bkg, *var_per_I, *capture;
        const float *reject_var_add;   // per group: the shell's measured share of dI^2/4
        const int32_t *group, *frame;
        const uint8_t *on_ice, *frame_cell_ok, *half;
        const uint8_t *hand, *has_hands;   // per-full Bijvoet hand; per-group "two hands to split"
        const int32_t *gperm, *gstart, *gcount;
        const double *em_mean, *merged_I, *cc_factor;
        // outputs
        double *sw, *swI, *em_mean_out;
        double *swh, *swIh;      // per (group, hand), indexed 2*g + hand
        int32_t *cnt, *cnth;
        int32_t *sample_at;            // per full: 1 where it gives a sample, then (scanned) where it goes
        ErrorModelSample *sample;      // the samples, compacted in fulls order
        double *a_swI, *a_sw, *a_swIh0, *a_swIh1, *a_swh0, *a_swh1, *a_swht0, *a_swht1, *a_d;
        int32_t *a_nh0, *a_nh1, *a_rejected;
        uint8_t *a_on_ice;
        double *r_absdev, *r_sumI, *r_wabsdev, *r_wsumI, *r_sumv, *r_sumv2;
        int32_t *r_n, *r_nusable;
        uint8_t *rejected_obs;   // per-full: set by MergeAccum (outlier-rejected), read by MergeRmeas
    };

    // A full's partiality: the device merges the fulls it scaled, in the Unity model.
    constexpr float FULL_PARTIALITY = 1.0f;

    // A full passes the merge / error-model filter (mirrors MergeAndStats::usable_merge).
    __device__ __forceinline__ bool MergeUsable(int i, const MergeParams &p) {
        const int g = p.group[i];
        if (g < 0) return false;
        if (!p.frame_cell_ok[p.frame[i]]) return false;
        const float c = p.corr[i];
        if (!(c > 0.0f) || !isfinite(c)) return false;
        if (p.for_search && p.on_ice[i]) return false;
        if (FULL_PARTIALITY < p.min_partiality) return false;
        const float I_corr = p.I[i] * c, sigma_corr = p.sigma[i] * c;
        return isfinite(I_corr) && isfinite(sigma_corr) && sigma_corr > 0.0f;
    }

    // The looser R_meas filter (no ice / for_search - Mask = cell only).
    __device__ __forceinline__ bool RmeasUsable(int i, const MergeParams &p) {
        const int g = p.group[i];
        if (g < 0) return false;
        if (!p.frame_cell_ok[p.frame[i]]) return false;
        const float c = p.corr[i];
        if (!(c > 0.0f) || !isfinite(c)) return false;
        if (FULL_PARTIALITY < p.min_partiality) return false;
        const float I_corr = p.I[i] * c, sigma_corr = p.sigma[i] * c;
        return isfinite(I_corr) && isfinite(sigma_corr) && sigma_corr > 0.0f;
    }

    // One thread per group: inverse-variance sums over the group's usable fulls + the group mean (>=2).
    __global__ void MergeEmStatsKernel(MergeParams p) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            const int lo = p.gstart[g], hi = lo + p.gcount[g];
            double sw = 0.0, swI = 0.0, swh[2] = {0.0, 0.0}, swIh[2] = {0.0, 0.0};
            int cnt = 0, cnth[2] = {0, 0};
            const bool split = p.has_hands && p.has_hands[g];
            for (int q = lo; q < hi; ++q) {
                const int i = p.gperm[q];
                if (!MergeUsable(i, p)) continue;
                const double sigma_corr = double(p.sigma[i]) * p.corr[i];
                const double w = 1.0 / (sigma_corr * sigma_corr);
                const double wI = w * (double(p.I[i]) * p.corr[i]);
                sw += w; swI += wI; ++cnt;
                if (split) {
                    const int hh = p.hand[i];
                    swh[hh] += w; swIh[hh] += wI; ++cnth[hh];
                }
            }
            if (p.swh) {
                p.swh[2 * g] = swh[0]; p.swh[2 * g + 1] = swh[1];
                p.swIh[2 * g] = swIh[0]; p.swIh[2 * g + 1] = swIh[1];
                p.cnth[2 * g] = cnth[0]; p.cnth[2 * g + 1] = cnth[1];
            }
            p.sw[g] = sw; p.swI[g] = swI; p.cnt[g] = cnt;
            p.em_mean_out[g] = (cnt >= 2 && sw > 0.0) ? swI / sw : NAN;
        }
    }

    // The host's counting_variance (RotationScaleMerge::MergeAndStats), to the bit: each product
    // rounded on its own, as the host build rounds it, and summed in the host's order.
    __device__ __forceinline__ double CountingVarianceExact(float corr, float var_bkg, float var_per_I, float capture,
                                                            double I_for_b, double own_s2) {
        const double I_exp = 0.0 < I_for_b ? I_for_b : 0.0;
        const double c = corr, cap = capture;
        double base = __dmul_rn(__dmul_rn(c, c), double(var_bkg));
        base = __dadd_rn(base, __dmul_rn(__dmul_rn(c, double(var_per_I)), I_exp));
        base = __dadd_rn(base, __dmul_rn(__dmul_rn(__dmul_rn(cap, cap), I_exp), I_exp));
        return base > 0.0 ? base : own_s2;
    }

    // The host's model_sigma, to the bit (see CountingVarianceExact).
    __device__ __forceinline__ float ModelSigmaExact(int i, const MergeParams &p, double I_for_b, float sigma_corr) {
        if (!p.error_model_active) return sigma_corr;
        const double a_var = CountingVarianceExact(p.corr[i], p.var_bkg[i], p.var_per_I[i], p.capture[i], I_for_b,
                                                   __dmul_rn(double(sigma_corr), double(sigma_corr)));
        const double bI = __dmul_rn(p.error_model_b, I_for_b);
        const double v = __dadd_rn(__dmul_rn(p.error_model_a, a_var), __dmul_rn(bI, bI));
        return v > 0.0 ? float(sqrt(v)) : sigma_corr;
    }

    // One thread per group: what the pooled outlier test's widening is measured from (the host's
    // reject_var_add block) - the group's d (its first usable full's), the excess dI^2 - var(dI) of its
    // two hands under the error model's sigmas (NaN where a hand is missing), and the d range of the
    // fulls that entered. Summed over the group's fulls in index order, as the host sums them.
    __global__ void BijvoetExcessKernel(MergeParams p, float *group_d, double *excess, float *d_lo, float *d_hi) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            const int lo = p.gstart[g], hi = lo + p.gcount[g];
            double sw[2] = {0.0, 0.0}, swI[2] = {0.0, 0.0};
            float gd = NAN, dmin = INFINITY, dmax = 0.0f;
            bool first = true;
            for (int q = lo; q < hi; ++q) {
                const int i = p.gperm[q];
                if (!MergeUsable(i, p)) continue;
                if (first) { gd = p.d[i]; first = false; }
                if (!p.has_hands[g]) continue;
                const float I_corr = p.I[i] * p.corr[i];
                const double I_for_b = isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr);
                const float sc = ModelSigmaExact(i, p, I_for_b, p.sigma[i] * p.corr[i]);
                if (!(sc > 0.0f) || !isfinite(sc)) continue;
                const double w = 1.0 / __dmul_rn(double(sc), double(sc));
                const int h = p.hand[i];
                sw[h] = __dadd_rn(sw[h], w);
                swI[h] = __dadd_rn(swI[h], __dmul_rn(w, double(I_corr)));
                const float d = p.d[i];
                if (isfinite(d) && d > 0.0f) { dmin = fminf(dmin, d); dmax = fmaxf(dmax, d); }
            }
            group_d[g] = gd; d_lo[g] = dmin; d_hi[g] = dmax;
            excess[g] = NAN;
            if (sw[0] > 0.0 && sw[1] > 0.0 && isfinite(gd)) {
                const double di = swI[0] / sw[0] - swI[1] / sw[1];
                excess[g] = __dmul_rn(di, di) - (1.0 / sw[0] + 1.0 / sw[1]);
            }
        }
    }

    // One thread per group: the sums the strong-reflection ISa asymptote is read from (MergeAndStats),
    // over the group's usable fulls in index order, with the host's rounding.
    __global__ void GroupScatterKernel(MergeParams p, double *sum, double *sum_sq, double *sum_var, int32_t *n) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            const int lo = p.gstart[g], hi = lo + p.gcount[g];
            double s = 0.0, s2 = 0.0, sv = 0.0;
            int cnt = 0;
            for (int q = lo; q < hi; ++q) {
                const int i = p.gperm[q];
                if (!MergeUsable(i, p)) continue;
                const double I_corr = __dmul_rn(double(p.I[i]), double(p.corr[i]));
                const double sigma_corr = __dmul_rn(double(p.sigma[i]), double(p.corr[i]));
                s = __dadd_rn(s, I_corr);
                s2 = __dadd_rn(s2, __dmul_rn(I_corr, I_corr));
                sv = __dadd_rn(sv, CountingVarianceExact(p.corr[i], p.var_bkg[i], p.var_per_I[i], p.capture[i],
                                                         p.em_mean[g], __dmul_rn(sigma_corr, sigma_corr)));
                ++cnt;
            }
            sum[g] = s; sum_sq[g] = s2; sum_var[g] = sv; n[g] = cnt;
        }
    }

    // The host's ObservationKey (RotationScaleMerge.cpp): integer arithmetic, so the same on both.
    __device__ __forceinline__ uint64_t ObservationKeyDevice(int32_t h, int32_t k, int32_t l, int32_t frame) {
        uint64_t z = uint64_t(uint32_t(h)) * 0x9E3779B97F4A7C15ULL
                   ^ uint64_t(uint32_t(k)) * 0xC2B2AE3D27D4EB4FULL
                   ^ uint64_t(uint32_t(l)) * 0x165667B19E3779F9ULL
                   ^ uint64_t(uint32_t(frame)) * 0x27D4EB2F165667C5ULL;
        z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
        z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
        return z ^ (z >> 31);
    }

    // One thread per full: its observation key, and whether the merge can use it.
    // A full's raw hkl is its run's.
    __global__ void ObservationKeyKernel(int n_obs, MergeParams p, const int32_t *run, const int32_t *h,
                                         const int32_t *k, const int32_t *l, uint64_t *key, uint8_t *usable) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
            const int r = run[i];
            key[i] = ObservationKeyDevice(h[r], k[r], l[r], p.frame[i]);
            usable[i] = MergeUsable(i, p) ? 1 : 0;
        }
    }

    // One thread per group: the CC1/2 half of each usable full, the parity of its key's rank among the
    // group's usable fulls (the host's AssignHalvesOfMembers). Unusable fulls keep half 0.
    __global__ void HalvesKernel(MergeParams p, const uint64_t *key, const uint8_t *usable, uint8_t *half) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            const int lo = p.gstart[g], hi = lo + p.gcount[g];
            for (int a = lo; a < hi; ++a) {
                const int i = p.gperm[a];
                if (!usable[i]) continue;
                int rank = 0;
                for (int b = lo; b < hi; ++b) {
                    const int j = p.gperm[b];
                    if (usable[j] && key[j] < key[i]) ++rank;
                }
                half[i] = uint8_t(rank & 1);
            }
        }
    }

    // The leverage-corrected error-model sample of full i; false where the full gives none.
    __device__ __forceinline__ bool EmSample(int i, const MergeParams &p, ErrorModelSample &out) {
        if (!MergeUsable(i, p)) return false;
        const int g = p.group[i];
        if (p.cnt[g] < 2) return false;
        // The hand's own mean where the hand has two of its own, the pooled pair where it has
        // not - see the host obs_hand block. em_mean (the merge weights) is untouched.
        const int sub = 2 * g + (p.hand ? p.hand[i] : 0);
        const bool on_hand = p.swh && p.has_hands[g] && p.cnth[sub] >= 2 && p.swh[sub] > 0.0;
        const double mean = on_hand ? p.swIh[sub] / p.swh[sub] : p.em_mean[g];
        if (!isfinite(mean)) return false;
        const double sigma_corr = double(p.sigma[i]) * p.corr[i];
        const double s2 = sigma_corr * sigma_corr;
        const double factor = 1.0 - (1.0 / s2) / (on_hand ? p.swh[sub] : p.sw[g]);
        if (factor < 0.05) return false;
        const double resid = double(p.I[i]) * p.corr[i] - mean;
        out.s2 = CountingVarianceExact(p.corr[i], p.var_bkg[i], p.var_per_I[i], p.capture[i], p.em_mean[g], s2);
        out.I2 = mean * mean;
        out.dev2 = resid * resid / factor;
        out.d = p.d[i];
        return true;
    }

    // One thread per full: whether it gives an error-model sample (the scan turns it into where).
    __global__ void MergeSampleFlagKernel(int n_obs, MergeParams p) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
            ErrorModelSample s;
            p.sample_at[i] = EmSample(i, p, s) ? 1 : 0;
        }
    }

    // One thread per full: its sample, at its place among the samples in fulls order.
    __global__ void MergeSampleWriteKernel(int n_obs, MergeParams p) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
            ErrorModelSample s;
            if (EmSample(i, p, s)) p.sample[p.sample_at[i]] = s;
        }
    }

    // The full's sigma under the error model with its variance at intensity I_for_b - the host model_sigma.
    // At the reflection's expected intensity the merge weight cannot know this full's own fluctuation.
    __device__ __forceinline__ float ModelSigma(int i, const MergeParams &p, double I_for_b, float sigma_raw) {
        const double bi = p.error_model_b * I_for_b;
        const double c = p.corr[i];
        double a_var = double(sigma_raw) * sigma_raw;
        const double I_exp = Dmax(0.0, I_for_b);
        const double cap = p.capture[i];
        const double base = c * c * p.var_bkg[i] + c * p.var_per_I[i] * I_exp + cap * cap * I_exp * I_exp;
        if (base > 0.0) a_var = base;
        const double v = p.error_model_a * a_var + bi * bi;
        return v > 0.0 ? float(sqrt(v)) : sigma_raw;
    }

    // One thread per group: the merge accumulators (inv-var sums + deterministic half-sets), with the
    // error-model-corrected sigma. Mirrors MergeAndStats' merge loop (reject path stays on the host).
    __global__ void MergeAccumKernel(MergeParams p) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            const int lo = p.gstart[g], hi = lo + p.gcount[g];
            double swI = 0, sw = 0, swIh0 = 0, swIh1 = 0, swh0 = 0, swh1 = 0, swht0 = 0, swht1 = 0, dd = NAN;
            int nh0 = 0, nh1 = 0, rejected = 0, onice = 0;
            const float rmed = p.reject_median[g];
            // The pooled cut, widened by the shell's measured Bijvoet variance - see the host loop.
            const double v_add = p.reject_var_add ? p.reject_var_add[g] : 0.0;
            for (int q = lo; q < hi; ++q) {
                const int i = p.gperm[q];
                if (!MergeUsable(i, p)) continue;
                if (p.rejected_obs[i]) { ++rejected; continue; }   // by the host's Wilson test
                const float I_corr = p.I[i] * p.corr[i];
                const float sigma_raw = p.sigma[i] * p.corr[i];
                const float sigma_corr = p.error_model_active
                        ? ModelSigma(i, p, isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr), sigma_raw)
                        : sigma_raw;
                if (p.reject_outliers && p.error_model_active && isfinite(rmed)) {
                    const double I_b = isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr);
                    if (OutsideOutlierBand(I_corr, rmed, p.reject_nsigma, double(sigma_corr) * sigma_corr + v_add,
                                           I_b, p.error_model_b * I_b, p.reject_up, p.reject_down)) {
                        ++rejected; p.rejected_obs[i] = 1; continue;
                    }
                }
                const double w = 1.0 / (double(sigma_corr) * sigma_corr);
                const double wI = w * I_corr;
                // The half-sets are assigned once on the host (AssignHalvesByRank), not decided here:
                // a rank within the reflection cannot be read off a running count, and computing it
                // in one place is what makes this kernel and the host merge loop agree exactly.
                const int half = p.half[i];
                const double wt = w * p.cc_factor[p.frame[i]];   // see the host merge loop (cc_weight)
                swI += wI; sw += w;
                if (half) { swIh1 += wI; swh1 += w; swht1 += wt; ++nh1; }
                else { swIh0 += wI; swh0 += w; swht0 += wt; ++nh0; }
                if (p.on_ice[i]) onice = 1;   // see the host merge loop
                if (!isfinite(dd) && isfinite(p.d[i]) && p.d[i] > 0.0f) dd = p.d[i];
            }
            p.a_swI[g] = swI; p.a_sw[g] = sw; p.a_swIh0[g] = swIh0; p.a_swIh1[g] = swIh1;
            p.a_swh0[g] = swh0; p.a_swh1[g] = swh1; p.a_swht0[g] = swht0; p.a_swht1[g] = swht1;
            p.a_nh0[g] = nh0; p.a_nh1[g] = nh1; p.a_d[g] = dd;
            p.a_rejected[g] = rejected; p.a_on_ice[g] = uint8_t(onice);
        }
    }

    // One thread per group: R_meas accumulators (sum|I_corr - merged_I|, sum I_corr, the same weighted by
    // v = the observation's merge weight 1/sigma^2 under the error model (as MergeAccumKernel weights it),
    // sum v, sum v^2, n) + the count of observations this looser walk accepted.
    // Mirrors MergeAndStats' R_meas re-walk (cell-only filter).
    // That count is NOT the per-shell total_observations - it is wider than the merge, so it would
    // over-report multiplicity; the host uses it only to skip empty groups.
    __global__ void MergeRmeasKernel(MergeParams p) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
            const int lo = p.gstart[g], hi = lo + p.gcount[g];
            const double mI = p.merged_I[g];
            const bool have = isfinite(mI);
            double absdev = 0, sumI = 0, wabsdev = 0, wsumI = 0, sumv = 0, sumv2 = 0;
            int n = 0, nusable = 0;
            for (int q = lo; q < hi; ++q) {
                const int i = p.gperm[q];
                if (!RmeasUsable(i, p)) continue;
                if (p.rejected_obs[i]) continue;   // outlier-rejected in the merge -> also out of R_meas (XDS convention)
                if (!isfinite(p.d[i]) || !(p.d[i] > 0.0f)) continue;   // host counts only fulls with a shell
                ++nusable;
                if (have) {
                    const double I_corr = double(p.I[i]) * p.corr[i];
                    const float sigma_raw = p.sigma[i] * p.corr[i];
                    const double sc = p.error_model_active
                            ? ModelSigma(i, p, isfinite(p.em_mean[g]) ? p.em_mean[g] : I_corr, sigma_raw)
                            : sigma_raw;
                    const double v = 1.0 / (sc * sc);
                    const double dev = fabs(I_corr - mI);
                    absdev += dev; sumI += I_corr;
                    wabsdev += v * dev; wsumI += v * I_corr; sumv += v; sumv2 += v * v; ++n;
                }
            }
            p.r_absdev[g] = absdev; p.r_sumI[g] = sumI; p.r_wabsdev[g] = wabsdev; p.r_wsumI[g] = wsumI;
            p.r_sumv[g] = sumv; p.r_sumv2[g] = sumv2;
            p.r_n[g] = n; p.r_nusable[g] = nusable;
        }
    }

    // --- correction-surface fit (ApplyCellSurface) ---
    // The host fit is the reference here: these kernels form the same sums from the same terms in the
    // same order, and every rounding is spelled out (__dmul_rn / __dadd_rn round each step on its own,
    // fma rounds once) to be the one the host build makes - nvcc would otherwise contract a multiply
    // and an add wherever it sees them, and GCC at -march=x86-64-v3 does so only in some of them. The
    // products are the host's (I * corr) * a and (sigma * corr) * a.
    using SurfaceTerm = RotationScaleMergeGPU::SurfaceTerm;

    __device__ __forceinline__ double SurfaceIs(const SurfaceTerm &t, double a) {
        return __dmul_rn(__dmul_rn(double(t.I), double(t.corr)), a);
    }

    __device__ __forceinline__ double SurfaceSigma(const SurfaceTerm &t, double a) {
        return __dmul_rn(__dmul_rn(double(t.sigma), double(t.corr)), a);
    }

    // One thread per ASU group: the group's inverse-variance sums over its terms of frame parity `parity`
    // (< 0 = all), in fulls order - the host's `reference`. The host builds that loop twice, and only
    // the parity-filtered copy fuses the multiply-add of swI; the unfiltered one rounds the product.
    __global__ void SurfaceReferenceKernel(int n_groups, int parity, const int32_t *__restrict__ gperm,
                                           const int32_t *__restrict__ gstart,
                                           const SurfaceTerm *__restrict__ term,
                                           const uint8_t *__restrict__ term_parity,
                                           const double *__restrict__ A,
                                           double *__restrict__ sw, double *__restrict__ swI) {
        for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < n_groups; g += gridDim.x * blockDim.x) {
            double s_w = 0.0, s_wI = 0.0;
            for (int k = gstart[g]; k < gstart[g + 1]; ++k) {
                const int i = gperm[k];
                if (parity >= 0 && term_parity[i] != parity) continue;
                const SurfaceTerm t = term[i];
                const double a = A[t.cell];
                const double Is = SurfaceIs(t, a), sc = SurfaceSigma(t, a);
                const double w = 1.0 / __dmul_rn(sc, sc);
                s_w = __dadd_rn(s_w, w);
                s_wI = parity >= 0 ? fma(Is, w, s_wI) : __dadd_rn(s_wI, __dmul_rn(Is, w));
            }
            sw[g] = s_w; swI[g] = s_wI;
        }
    }

    // The fit's sums of one round. Each term's contribution is formed on its own thread (one per term,
    // in segment order), and only the two multiply-adds that sum them are left to the thread of each
    // (reduction block, cell) - the per-block accumulators of the host fit, one slot at a time, walking
    // its segment in term order. A term the host skips is marked by a NaN Iref (a kept one is finite).
    __global__ void SurfaceFitTermKernel(int n, const int32_t *__restrict__ perm,
                                         const SurfaceTerm *__restrict__ term,
                                         const double *__restrict__ A,
                                         const double *__restrict__ sw, const double *__restrict__ swI,
                                         double *__restrict__ w_Is, double *__restrict__ w_Iref,
                                         double *__restrict__ Iref_out) {
        for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < n; k += gridDim.x * blockDim.x) {
            const SurfaceTerm t = term[perm[k]];
            Iref_out[k] = NAN;
            if (sw[t.group] <= 0.0) continue;
            const double Iref = swI[t.group] / sw[t.group], a = A[t.cell];
            const double Is = SurfaceIs(t, a), sc = SurfaceSigma(t, a);
            if (!isfinite(Iref) || Iref <= 0.0 || !(sc > 0.0)) continue;
            const double w = 1.0 / __dmul_rn(sc, sc);
            w_Is[k] = __dmul_rn(w, Is);
            w_Iref[k] = __dmul_rn(w, Iref);
            Iref_out[k] = Iref;
        }
    }

    __global__ void SurfaceFitSegmentKernel(int n_seg, const int32_t *__restrict__ seg_start,
                                            const double *__restrict__ w_Is, const double *__restrict__ w_Iref,
                                            const double *__restrict__ Iref,
                                            double *__restrict__ tcross, double *__restrict__ tref2) {
        for (int s = blockIdx.x * blockDim.x + threadIdx.x; s < n_seg; s += gridDim.x * blockDim.x) {
            double xcross = 0.0, xref2 = 0.0;
            for (int k = seg_start[s]; k < seg_start[s + 1]; ++k) {
                if (isnan(Iref[k])) continue;
                xcross = fma(w_Is[k], Iref[k], xcross);
                xref2 = fma(w_Iref[k], Iref[k], xref2);
            }
            tcross[s] = xcross; tref2[s] = xref2;
        }
    }

    // One thread per cell: the blocks' sums added up in block order, as the host adds its slots.
    __global__ void SurfaceFitCellKernel(int n_blocks, int ncell, const double *__restrict__ tcross,
                                         const double *__restrict__ tref2,
                                         double *__restrict__ cross, double *__restrict__ ref2) {
        for (int c = blockIdx.x * blockDim.x + threadIdx.x; c < ncell; c += gridDim.x * blockDim.x) {
            double sc = 0.0, sr = 0.0;
            for (int b = 0; b < n_blocks; ++b) {
                sc = __dadd_rn(sc, tcross[b * ncell + c]);
                sr = __dadd_rn(sr, tref2[b * ncell + c]);
            }
            cross[c] = sc; ref2[c] = sr;
        }
    }

    void CudaCheck(cudaError_t e, const char *what);

    // A copy on the instance's stream, waited for - what cudaMemcpy on the NULL stream was, without also
    // waiting for every other stream on the card.
    void CopyAndWait(void *dst, const void *src, size_t bytes, cudaMemcpyKind kind, cudaStream_t s,
                     const char *what) {
        CudaCheck(cudaMemcpyAsync(dst, src, bytes, kind, s), what);
        CudaCheck(cudaStreamSynchronize(s), what);
    }

    void CudaCheck(cudaError_t e, const char *what) {
        if (e != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                                  std::string("RotationScaleMergeGPU: ") + what + ": " + cudaGetErrorString(e));
    }

    // Device bytes per observation in the arrays SetPartialsLayout allocates: ten floats (the record's
    // I, sigma, partiality and prescaling_corr; zeta, corr, bkg, var_bkg, image_number, d) and the flags
    // byte. The merge's other buffers come on top of these.
    constexpr size_t OBS_BYTES = 10 * 4 + 1;

    // Every buffer here is sized by the data set, and the scaling does not fall back to the CPU
    // mid-run: the two paths differ in the last bits, so the result would depend on the card. A
    // buffer that does not fit ends the run, and says why and what to do instead.
    [[noreturn]] void ThrowTooLargeForGPU(int64_t n_obs) {
        size_t free_bytes = 0, total_bytes = 0;
        cudaMemGetInfo(&free_bytes, &total_bytes);
        cudaGetLastError();
        char msg[768];
        snprintf(msg, sizeof(msg),
                 "Scaling %lld partial observations needs more GPU memory than this card has: %.1f GB for "
                 "their per-observation arrays alone (%zu bytes each), before the merge's other buffers, "
                 "on a %.1f GB card. This data set is too large for GPU scaling on this card - run it on "
                 "the CPU (rugnux --no-gpu ..., or a CPU-only build) on a machine with plenty of "
                 "host memory: the CPU path keeps all of it in RAM, which for a set this large can be "
                 "tens of GB",
                 static_cast<long long>(n_obs), double(n_obs) * OBS_BYTES / 1e9, OBS_BYTES,
                 double(total_bytes) / 1e9);
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, msg);
    }

    // A buffer that did not fit although the per-observation arrays fit the card, with nothing else in
    // this process running on the GPU any more: what holds the rest is this run or another program.
    [[noreturn]] void ThrowOutOfGPUMemory(int64_t n_obs, size_t bytes) {
        size_t free_bytes = 0, total_bytes = 0;
        get_gpu_free_memory(free_bytes, total_bytes);
        cudaGetLastError();
        char msg[768];
        snprintf(msg, sizeof(msg),
                 "Scaling %lld partial observations ran out of GPU memory: a buffer of %.1f MB did not fit, "
                 "with %.1f GB of the %.1f GB card free (the per-observation arrays take %.1f GB). Either "
                 "another program is using this GPU, or this data set is too large for GPU scaling on this "
                 "card - run it on the CPU (rugnux --no-gpu ..., or a CPU-only build) on a "
                 "machine with plenty of host memory",
                 static_cast<long long>(n_obs), double(bytes) / 1e6, double(free_bytes) / 1e9,
                 double(total_bytes) / 1e9, double(n_obs) * OBS_BYTES / 1e9);
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, msg);
    }

    // How long a merge waits for GPU work running beside it (GPUWorkBeside) to give its memory back.
    // That work is a probe pass of seconds, so this long means it is stuck, not slow.
    constexpr auto GPU_BUSY_TIMEOUT = std::chrono::minutes(10);

    [[noreturn]] void ThrowGPUBusy(int64_t n_obs) {
        char msg[512];
        snprintf(msg, sizeof(msg),
                 "GPU busy: scaling %lld partial observations waited %lld minutes for GPU work running "
                 "beside it in this process to give back its memory, and it did not. The data set fits "
                 "this card; run it again, or on the CPU (rugnux --no-gpu ...)",
                 static_cast<long long>(n_obs),
                 static_cast<long long>(std::chrono::duration_cast<std::chrono::minutes>(GPU_BUSY_TIMEOUT).count()));
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, msg);
    }
}

struct RotationScaleMergeGPU::Impl {
    // First, so it goes last: the buffers below are freed before the stream their work ran on.
    std::unique_ptr<CudaStream> stream;
    cudaStream_t s() const { return stream->get(); }
    int device = 0;   // the GPU this instance's buffers live on
    bool available = false;
    int n_obs = 0, n_frames = 0, n_groups = 0;

    // The arrays fit the card (SetPartialsLayout checks that first), but not necessarily beside other
    // GPU work of this process: a probe pass started on a copy of the run holds several GB of engines
    // for a few seconds while this merge runs. So a buffer that does not fit waits for that work to
    // end and is asked for again - the same buffer, so what the merge computes does not depend on the
    // wait. It is asked for again also when nothing is running beside it any more: work that ended
    // between the failure and the check has just given its memory back.
    //
    // The failed allocation is handled here, so its error must not stay behind as this thread's last
    // error: the next kernel launch's cudaGetLastError() would report it over work that went fine.
    //
    // Plain Synchronous, not ALLOC: this function waits for the work beside itself, with its own limit.
    template <typename T>
    CudaDevicePtr<T> Alloc(size_t n) const {
        try {
            return CudaDevicePtr<T>(n, CudaAlloc::Synchronous);
        } catch (const JFJochException &) {
            cuda_throw_if_context_lost();   // not memory: waiting would only hide the real error
            cuda_clear_error();
        }
        if (!wait_for_gpu_work_beside(GPU_BUSY_TIMEOUT))
            ThrowGPUBusy(n_obs);
        try {
            return CudaDevicePtr<T>(n, CudaAlloc::Synchronous);
        } catch (const JFJochException &) {
            ThrowOutOfGPUMemory(n_obs, n * sizeof(T));
        }
    }

    template <typename T>
    void Upload(CudaDevicePtr<T> &dst, const T *src, int n) const {
        dst = Alloc<T>(std::max(1, n));
        if (n > 0)
            CopyAndWait(dst.get(), src, size_t(n) * sizeof(T), cudaMemcpyHostToDevice, s(), "upload");
    }

    // immutable per-obs
    CudaDevicePtr<ObsRecord> obs;
    void UploadRecordChunk(int field, int offset, int count, const float *v);
    // The per-observation arrays for n observations, the per-frame ones for n_frames and the frame CSR
    // (SetPartialsLayout, the resident ingest and CopyObservationsFrom all lay them out the same way).
    void AllocObservations(int n);
    // (frame_start null: the CSR is left to the caller to fill.)
    void AllocFrames(const int32_t *frame_start, const int32_t *frame_count);
    // The previous grouping's fulls, their CSRs and what the merge left of them, before a new grouping
    // or a new combine: nothing reads them again, and on a large cell they are gigabytes beside the
    // arrays that replace them.
    void ReleaseFulls();
    CudaDevicePtr<float> zeta;
    CudaDevicePtr<uint8_t> flags;   // OBS_ON_ICE_FLAG | OBS_CLIPPED_FLAG | OBS_OVERLOADED_FLAG
    // A partial's frame, from the frame CSR (the partials lie in frame order).
    FrameOf PartialFrame() const { return FrameOf{nullptr, frame_start.get(), n_frames}; }
    CudaDevicePtr<float> corr;               // mutable, resident across iterations
    CudaDevicePtr<int32_t> frame_start, frame_count;
    // per space group
    CudaDevicePtr<int32_t> group, group_perm, group_start, group_count;
    // scratch
    CudaDevicePtr<double> group_mean, g, info;
    CudaDevicePtr<uint8_t> scaled;
    CudaDevicePtr<double> cc;        // per-frame CC (diagnostic), length n_frames
    CudaDevicePtr<int64_t> cc_n;
    CudaDevicePtr<uint8_t> smooth_apply;   // per-frame smooth-G apply flag + ratio, length n_frames
    CudaDevicePtr<double> smooth_ratio;
    CudaDevicePtr<uint8_t> filter_reject;  // per-frame --min-image-cc rejection flag, length n_frames
    // merge / error-model reductions over the resident fulls (reuse the fulls group CSR f_gperm/...)
    // What one merge keeps between its calls; every other merge buffer lives only in the call that
    // needs it.
    CudaDevicePtr<uint8_t> frame_cell_ok;
    CudaDevicePtr<double> m_em_mean;                // per group (n_groups)
    CudaDevicePtr<uint8_t> m_hand;                  // per full: Bijvoet hand (0 = I(+))
    CudaDevicePtr<uint8_t> m_has_hands;             // per group: acentric, and the merge pools the hands
    CudaDevicePtr<uint8_t> m_rejected;   // per-full outlier-rejected flag (MergeAccum -> MergeRmeas)
    // merge accum per group, from MergeAccum to MergeAccumRelease
    CudaDevicePtr<double> a_swI, a_sw, a_swIh0, a_swIh1, a_swh0, a_swh1, a_swht0, a_swht1, a_d;
    CudaDevicePtr<int32_t> a_nh0, a_nh1, a_rejected;
    CudaDevicePtr<uint8_t> a_on_ice;     // per group: any contributing full on an ice ring
    int merge_for_search = 0;   // filter context for one MergeAndStats call
    double merge_min_part = 0.0;
    double merge_em_a = 1.0, merge_em_b = 0.0;   // the error model of the last MergeAccum, for MergeRmeas
    int merge_em_active = 0;

    // combine: extra per-obs inputs + the one-time raw-hkl run layout
    CudaDevicePtr<float> bkg, var_bkg, image_number, d_obs;
    int n_runs = 0, n_perm = 0;
    CudaDevicePtr<int32_t> perm, rr_start, rr_count, rr_h, rr_k, rr_l;
    // combine: resident fulls SoA (rebuilt each Combine)
    int n_fulls = 0;
    // run: the raw-hkl run a full is of; peak: the partial its position is taken from
    CudaDevicePtr<int32_t> f_run, f_peak, f_frame, f_group;
    CudaDevicePtr<float> f_I, f_sigma, f_d, f_img, f_var_bkg, f_var_per_I, f_capture;
    CudaDevicePtr<uint8_t> f_on_ice, f_clipped;
    // scale-fulls (Unity model, kept resident): all-ones partiality/zeta and records with partiality and
    // prescaling_corr 1, so the shared scaling kernels yield coeff=mean, plus the working corr and the
    // fulls frame/group CSRs (built on the host from the small f_frame/f_group key arrays, over the
    // emit-ordered fulls).
    CudaDevicePtr<float> f_corr, f_zeta;
    CudaDevicePtr<ObsRecord> f_obs;
    CudaDevicePtr<int32_t> f_frame_perm, f_frame_start, f_frame_count;
    CudaDevicePtr<int32_t> f_gperm, f_gstart, f_gcount;
    // correction-surface fit (one ApplyCellSurface call at a time): its terms and their group CSR, the
    // three subsets' per-(block, cell) segments, the surface and the sums of the round
    int s_ncell = 0, s_n_groups = 0;
    int s_n_blocks[3] = {0, 0, 0};
    int s_n_sel[3] = {0, 0, 0};
    size_t s_slots = 0;   // capacity of s_tcross / s_tref2
    CudaDevicePtr<RotationScaleMergeGPU::SurfaceTerm> s_term;
    CudaDevicePtr<uint8_t> s_parity;
    CudaDevicePtr<int32_t> s_gperm, s_gstart;
    CudaDevicePtr<int32_t> s_perm[3], s_seg_start[3];
    CudaDevicePtr<double> s_A, s_sw, s_swI, s_tcross, s_tref2, s_cross, s_ref2;
    CudaDevicePtr<double> s_w_Is, s_w_Iref, s_Iref;   // per term of a subset, in segment order
    // the resident ingest (IngestUpload .. IngestFinish): each frame's first reflection (host, length
    // n_frames + 1); per reflection what the sort and the cut read - its raw hkl as the sort key, its
    // frame position, resolution and background - until they have, and the place the cut gives it (as
    // an exclusive prefix over the kept ones, n_keep in all) until the observations are taken from the
    // reflections; per observation its delta_phi as integrated until the smoothing has used it
    std::vector<int32_t> in_src_start;
    CudaDevicePtr<uint64_t> in_key;
    CudaDevicePtr<float> in_img, in_d, in_bkg;
    CudaDevicePtr<int32_t> in_new_idx;
    int in_n_keep = 0;
    CudaDevicePtr<float> in_dphi;
    // The reflections of the ingest's frames, through two pinned buffers of whole frames, each filled on
    // all threads while the other is on its way to the card; take(cnt, base, src) takes cnt of them,
    // the first being reflection base, apart on the card.
    template <typename Take>
    void StageReflections(const Reflection *const *refl, const int32_t *count, size_t nthreads, Take take);
    RotationScaleMergeGPU::IngestLayout in_layout;
};

// Set the device this instance's memory lives on for the duration of a call, and put the caller's
// back afterwards. CUDA's current device is per-thread, so without this a single set_gpu() in the
// constructor silently re-pins the calling thread for the rest of its life - and, worse, the
// destructor would free several gigabytes against whatever device happened to be current then.
// CudaDevicePtr records no device of its own, so every entry point needs this.
namespace {
struct DeviceGuard {
    int prev = 0;
    bool active = false;
    explicit DeviceGuard(int device, bool enable) : active(enable) {
        if (!active)
            return;
        cudaGetDevice(&prev);
        if (prev != device)
            set_gpu(device);
    }
    ~DeviceGuard() {
        if (active)
            cudaSetDevice(prev);   // not set_gpu: a destructor must not throw
    }
};
}  // namespace

RotationScaleMergeGPU::RotationScaleMergeGPU(int device) : impl_(std::make_unique<Impl>()) {
    if (get_gpu_count() > 0) {
        // One instance, one device, recorded so the guard below can put the caller's device back
        // instead of leaving the thread moved.
        impl_->device = device;
        DeviceGuard guard(impl_->device, true);
        impl_->stream = std::make_unique<CudaStream>();
        impl_->available = true;
    }
}

RotationScaleMergeGPU::~RotationScaleMergeGPU() {
    DeviceGuard guard(impl_->device, impl_->available);
    impl_.reset();
}

bool RotationScaleMergeGPU::Available() const { return impl_->available; }

void RotationScaleMergeGPU::SetPartialsLayout(int n_obs, int n_frames,
                                              const int32_t *frame_start, const int32_t *frame_count) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    d.n_obs = n_obs;
    d.n_frames = n_frames;
    // Known now, before anything is allocated: arrays that could not fit even an empty card fail here,
    // not after gigabytes of them have been.
    size_t free_bytes = 0, total_bytes = 0;
    CudaCheck(cudaMemGetInfo(&free_bytes, &total_bytes), "memory info");
    if (size_t(std::max(0, n_obs)) * OBS_BYTES > total_bytes)
        ThrowTooLargeForGPU(n_obs);
    d.AllocFrames(frame_start, frame_count);
    d.AllocObservations(n_obs);
}

void RotationScaleMergeGPU::Impl::AllocObservations(int n_in) {
    // What the image passes' engines handed back to the allocation pool is still reserved by it when
    // the scaling starts - 5 to 7 GB of a 16 GB card on a large data set, measured - and the scaling's
    // buffers are not pooled, so it goes back to the driver here rather than only once an allocation
    // has failed for want of it.
    cuda_trim_allocation_pool();
    const int n = std::max(1, n_in);
    obs = Alloc<ObsRecord>(n);
    zeta = Alloc<float>(n); corr = Alloc<float>(n);
    bkg = Alloc<float>(n); var_bkg = Alloc<float>(n);
    image_number = Alloc<float>(n); d_obs = Alloc<float>(n);
    flags = Alloc<uint8_t>(n);
}

void RotationScaleMergeGPU::Impl::AllocFrames(const int32_t *fstart, const int32_t *fcount) {
    frame_start = Alloc<int32_t>(std::max(1, n_frames));
    frame_count = Alloc<int32_t>(std::max(1, n_frames));
    if (n_frames > 0 && fstart) {
        CopyAndWait(frame_start.get(), fstart, size_t(n_frames) * sizeof(int32_t), cudaMemcpyHostToDevice, s(),
                    "frame start");
        CopyAndWait(frame_count.get(), fcount, size_t(n_frames) * sizeof(int32_t), cudaMemcpyHostToDevice, s(),
                    "frame count");
    }
    g = Alloc<double>(std::max(1, n_frames));
    scaled = Alloc<uint8_t>(std::max(1, n_frames));
    info = Alloc<double>(std::max(1, n_frames));
    cc = Alloc<double>(std::max(1, n_frames));
    cc_n = Alloc<int64_t>(std::max(1, n_frames));
}

void RotationScaleMergeGPU::Impl::ReleaseFulls() {
    n_fulls = 0;
    f_run = {}; f_peak = {}; f_frame = {}; f_group = {};
    f_I = {}; f_sigma = {}; f_d = {}; f_img = {};
    f_var_bkg = {}; f_var_per_I = {}; f_capture = {}; f_on_ice = {}; f_clipped = {};
    f_corr = {}; f_zeta = {}; f_obs = {};
    f_frame_perm = {}; f_frame_start = {}; f_frame_count = {};
    f_gperm = {}; f_gstart = {}; f_gcount = {};
    m_em_mean = {}; m_hand = {}; m_has_hands = {}; m_rejected = {};
    a_swI = {}; a_sw = {}; a_swIh0 = {}; a_swIh1 = {}; a_swh0 = {}; a_swh1 = {}; a_swht0 = {}; a_swht1 = {};
    a_d = {}; a_nh0 = {}; a_nh1 = {}; a_rejected = {}; a_on_ice = {};
}

namespace {
    template <typename T>
    void UploadChunk(CudaDevicePtr<T> &dst, int offset, int count, const T *v, cudaStream_t s) {
        if (count > 0)
            CopyAndWait(dst.get() + offset, v, size_t(count) * sizeof(T), cudaMemcpyHostToDevice, s,
                        "upload chunk");
    }

}

// A slice of one field of the records (`field` its float number), staged contiguously and then put
// in place on the device: a strided copy from pageable memory goes element by element.
void RotationScaleMergeGPU::Impl::UploadRecordChunk(int field, int offset, int count, const float *v) {
    if (count == 0)
        return;
    auto staged = Alloc<float>(count);
    CopyAndWait(staged.get(), v, size_t(count) * sizeof(float), cudaMemcpyHostToDevice, s(), "upload record chunk");
    SetRecordFieldKernel<<<std::min(65535, (count + BLK - 1) / BLK), BLK, 0, s()>>>(count, staged.get(),
        obs.get() + offset, field);
    CudaCheck(cudaGetLastError(), "record field launch");
    CudaCheck(cudaStreamSynchronize(s()), "upload record chunk");
}

void RotationScaleMergeGPU::SetObsField(ObsField f, int offset, int count, const float *v) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    switch (f) {
        case ObsField::I:               d.UploadRecordChunk(0, offset, count, v); break;
        case ObsField::Sigma:           d.UploadRecordChunk(1, offset, count, v); break;
        case ObsField::Partiality:      d.UploadRecordChunk(2, offset, count, v); break;
        case ObsField::PrescalingCorr:  d.UploadRecordChunk(3, offset, count, v); break;
        case ObsField::Zeta:            UploadChunk(d.zeta, offset, count, v, impl_->s()); break;
        case ObsField::Corr0:           UploadChunk(d.corr, offset, count, v, impl_->s()); break;
        case ObsField::Bkg:             UploadChunk(d.bkg, offset, count, v, impl_->s()); break;
        case ObsField::VarBkg:          UploadChunk(d.var_bkg, offset, count, v, impl_->s()); break;
        case ObsField::ImageNumber:     UploadChunk(d.image_number, offset, count, v, impl_->s()); break;
        case ObsField::D:               UploadChunk(d.d_obs, offset, count, v, impl_->s()); break;
    }
}

void RotationScaleMergeGPU::CopyObservationsFrom(const RotationScaleMergeGPU &from) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const auto &f = *from.impl_;
    d.n_obs = f.n_obs;
    d.n_frames = f.n_frames;
    d.n_runs = f.n_runs;
    d.n_perm = f.n_perm;
    d.AllocFrames(nullptr, nullptr);
    d.AllocObservations(d.n_obs);
    const auto copy = [&](auto &dst, const auto &src, size_t n) {
        CudaCheck(cudaMemcpyPeerAsync(dst.get(), d.device, src.get(), f.device,
                                      n * sizeof(*src.get()), d.s()), "copy observations");
    };
    copy(d.frame_start, f.frame_start, d.n_frames); copy(d.frame_count, f.frame_count, d.n_frames);
    const size_t n = d.n_obs;
    copy(d.obs, f.obs, n); copy(d.zeta, f.zeta, n); copy(d.corr, f.corr, n);
    copy(d.bkg, f.bkg, n); copy(d.var_bkg, f.var_bkg, n); copy(d.image_number, f.image_number, n);
    copy(d.d_obs, f.d_obs, n);
    copy(d.flags, f.flags, n);
    const size_t nr = std::max(1, d.n_runs);
    d.perm = d.Alloc<int32_t>(std::max(1, d.n_perm));
    d.rr_start = d.Alloc<int32_t>(nr); d.rr_count = d.Alloc<int32_t>(nr);
    d.rr_h = d.Alloc<int32_t>(nr); d.rr_k = d.Alloc<int32_t>(nr); d.rr_l = d.Alloc<int32_t>(nr);
    copy(d.perm, f.perm, d.n_perm);
    copy(d.rr_start, f.rr_start, d.n_runs); copy(d.rr_count, f.rr_count, d.n_runs);
    copy(d.rr_h, f.rr_h, d.n_runs); copy(d.rr_k, f.rr_k, d.n_runs); copy(d.rr_l, f.rr_l, d.n_runs);
    CudaCheck(cudaStreamSynchronize(d.s()), "copy observations");
}

void RotationScaleMergeGPU::SetObsFlags(int offset, int count, const uint8_t *flags) {
    DeviceGuard guard(impl_->device, impl_->available);
    UploadChunk(impl_->flags, offset, count, flags, impl_->s());
}

void RotationScaleMergeGPU::BuildGroups(int n_groups, const int32_t *rawrun_group, int n_sorted,
                                        const int32_t *sorted_run, const int32_t *group_first) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    d.n_groups = n_groups;
    // The previous grouping goes before this one is allocated: the old CSRs (Combine frees the partials'
    // one) and the fulls grouped by it.
    d.ReleaseFulls();
    d.group_perm = {}; d.group_start = {}; d.group_count = {}; d.group_mean = {};
    CudaDevicePtr<int32_t> rr_group, sorted, first;
    d.Upload(rr_group, rawrun_group, d.n_runs);
    d.Upload(sorted, sorted_run, n_sorted);
    d.Upload(first, group_first, n_groups + 1);
    if (d.group.get() == nullptr) d.group = d.Alloc<int32_t>(std::max(1, d.n_obs));
    d.group_count = d.Alloc<int32_t>(std::max(1, n_groups));
    d.group_start = d.Alloc<int32_t>(std::max(1, n_groups));

    GroupParams p{};
    p.n_runs = d.n_runs; p.n_groups = n_groups;
    p.perm = d.perm.get(); p.rr_start = d.rr_start.get(); p.rr_count = d.rr_count.get(); p.rr_group = rr_group.get();
    p.sorted_run = sorted.get(); p.group_first = first.get();
    p.obs = d.obs.get();
    p.group = d.group.get(); p.group_start = d.group_start.get(); p.group_count = d.group_count.get();
    const int run_blocks = std::min(65535, (d.n_runs + BLK - 1) / BLK);
    const int grp_blocks = std::min(65535, (n_groups + BLK - 1) / BLK);
    StampGroupsKernel<<<std::max(1, run_blocks), BLK, 0, impl_->s()>>>(p);
    CudaCheck(cudaGetLastError(), "StampGroupsKernel launch");
    CountGroupsKernel<<<std::max(1, grp_blocks), BLK, 0, impl_->s()>>>(p);
    CudaCheck(cudaGetLastError(), "CountGroupsKernel launch");

    // The starts are a prefix over the counts: a few megabytes there and back, summed in order.
    std::vector<int32_t> count(n_groups), start(n_groups);
    if (n_groups > 0)
        CopyAndWait(count.data(), d.group_count.get(), size_t(n_groups) * sizeof(int32_t),
                    cudaMemcpyDeviceToHost, impl_->s(), "download group counts");
    int acc = 0;
    for (int g = 0; g < n_groups; ++g) { start[g] = acc; acc += count[g]; }
    if (n_groups > 0)
        CopyAndWait(d.group_start.get(), start.data(), size_t(n_groups) * sizeof(int32_t),
                    cudaMemcpyHostToDevice, impl_->s(), "upload group starts");
    d.group_perm = d.Alloc<int32_t>(std::max(1, acc));
    p.group_perm = d.group_perm.get();
    FillGroupsKernel<<<std::max(1, grp_blocks), BLK, 0, impl_->s()>>>(p);
    CudaCheck(cudaGetLastError(), "FillGroupsKernel launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "groups sync");
    d.group_mean = d.Alloc<double>(std::max(1, n_groups));
}

void RotationScaleMergeGPU::SetCorr(const float *corr) {
    DeviceGuard guard(impl_->device, impl_->available);
    CopyAndWait(impl_->corr.get(), corr, size_t(impl_->n_obs) * sizeof(float),
                         cudaMemcpyHostToDevice, impl_->s(), "upload corr");
}

void RotationScaleMergeGPU::ScalePartials(int iters, double min_partiality, bool /*has_d_min*/) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    // Reset per call: the host keeps the G of a frame across calls (RunScalingLoop), so a frame this
    // call did not fit must read as unfitted, not as fitted with the value of the call before.
    CudaCheck(cudaMemsetAsync(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t), impl_->s()), "memset scaled");
    CudaCheck(cudaMemsetAsync(d.g.get(), 0, size_t(d.n_frames) * sizeof(double), impl_->s()), "memset g"); // unscaled g unused
    const int upd_blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
    const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
    for (int it = 0; it < iters; ++it) {
        ReduceGroupMeansKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(d.n_groups, min_partiality,
            d.group_perm.get(), d.group_start.get(), d.group_count.get(),
            d.obs.get(), d.corr.get(), d.group_mean.get());
        CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
        FitPerFrameGKernel<<<d.n_frames, BLK, 0, impl_->s()>>>(d.n_frames, d.frame_start.get(), d.frame_count.get(),
            nullptr, min_partiality, d.group.get(), d.obs.get(), d.zeta.get(), d.flags.get(), d.group_mean.get(),
            long(MIN_REFLECTIONS), d.g.get(), d.scaled.get(), d.info.get());
        CudaCheck(cudaGetLastError(), "FitPerFrameGKernel launch");
        UpdateCorrKernel<<<upd_blocks, BLK, 0, impl_->s()>>>(d.n_obs, d.PartialFrame(), d.obs.get(),
            d.g.get(), d.scaled.get(), d.corr.get());
    }
    CudaCheck(cudaGetLastError(), "kernel launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "scale sync");
}

void RotationScaleMergeGPU::GetCorr(float *corr_out) const {
    DeviceGuard guard(impl_->device, impl_->available);
    CopyAndWait(corr_out, impl_->corr.get(), size_t(impl_->n_obs) * sizeof(float),
                         cudaMemcpyDeviceToHost, impl_->s(), "download corr");
}

void RotationScaleMergeGPU::GetInfo(double *info_out) const {
    DeviceGuard guard(impl_->device, impl_->available);
    CopyAndWait(info_out, impl_->info.get(), size_t(impl_->n_frames) * sizeof(double),
                         cudaMemcpyDeviceToHost, impl_->s(), "download info");
}

void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const {
    DeviceGuard guard(impl_->device, impl_->available);
    CopyAndWait(g_out, impl_->g.get(), size_t(impl_->n_frames) * sizeof(double),
                         cudaMemcpyDeviceToHost, impl_->s(), "download g");
    CopyAndWait(scaled_out, impl_->scaled.get(), size_t(impl_->n_frames) * sizeof(uint8_t),
                         cudaMemcpyDeviceToHost, impl_->s(), "download scaled");
}

void RotationScaleMergeGPU::SetFrameCellOk(const uint8_t *frame_cell_ok) {
    DeviceGuard guard(impl_->device, impl_->available);
    impl_->Upload(impl_->frame_cell_ok, frame_cell_ok, impl_->n_frames);
}

// The per-group inv-var mean (em_mean) + the per-full leverage-corrected error-model samples over the
// resident+scaled fulls. Stashes the filter context for the later MergeAccum/MergeRmeas calls.
void RotationScaleMergeGPU::MergeEmSamples(bool for_search, double min_partiality,
                                           const uint8_t *hand, const uint8_t *has_hands,
                                           double *em_mean_out, int32_t *cnt_out,
                                           std::vector<ErrorModelSample> &samples) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int ng = d.n_groups, nf = d.n_fulls;
    d.merge_for_search = for_search ? 1 : 0; d.merge_min_part = min_partiality;
    // Only em_mean and the hands outlive this call.
    d.m_em_mean = d.Alloc<double>(std::max(1, ng));
    CudaDevicePtr<double> sw = d.Alloc<double>(std::max(1, ng)), swI = d.Alloc<double>(std::max(1, ng));
    CudaDevicePtr<int32_t> cnt = d.Alloc<int32_t>(std::max(1, ng));
    // The error model is fitted on the Bijvoet hands (see the host obs_hand block). Absent - a merge
    // that already separates them - the hand sums are not allocated and the kernels take the pooled
    // group, which is what they did before.
    CudaDevicePtr<double> swh, swIh;
    CudaDevicePtr<int32_t> cnth;
    d.m_hand = {}; d.m_has_hands = {};
    if (hand && has_hands) {
        swh = d.Alloc<double>(std::max(1, 2 * ng)); swIh = d.Alloc<double>(std::max(1, 2 * ng));
        cnth = d.Alloc<int32_t>(std::max(1, 2 * ng));
        if (nf > 0) d.Upload(d.m_hand, hand, nf);
        if (ng > 0) d.Upload(d.m_has_hands, has_hands, ng);
    }
    CudaDevicePtr<int32_t> sample_at = d.Alloc<int32_t>(std::max(1, nf));

    MergeParams p{};
    p.n_groups = ng; p.min_partiality = min_partiality;
    p.for_search = d.merge_for_search;
    p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get();
    p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
    p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get();
    p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
    p.em_mean = d.m_em_mean.get(); p.sw = sw.get(); p.swI = swI.get();
    p.em_mean_out = d.m_em_mean.get(); p.cnt = cnt.get();
    p.swh = swh.get(); p.swIh = swIh.get(); p.cnth = cnth.get();
    p.hand = d.m_hand.get(); p.has_hands = d.m_has_hands.get();
    p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get(); p.capture = d.f_capture.get();
    p.sample_at = sample_at.get();

    const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
    const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
    MergeEmStatsKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p);
    MergeSampleFlagKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(nf, p);
    CudaCheck(cudaGetLastError(), "merge em/samples launch");
    // Where each sample goes: the fulls that give one, in order, as the host's compaction lists them.
    int n_samples = 0;
    if (nf > 0) {
        int32_t last_flag = 0;
        CopyAndWait(&last_flag, sample_at.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost,
                    impl_->s(), "dl last sample flag");
        size_t bytes = 0;
        CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, sample_at.get(), sample_at.get(), nf,
                                                impl_->s()), "sample scan size");
        auto scratch = d.Alloc<uint8_t>(std::max<size_t>(bytes, 1));
        CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, sample_at.get(), sample_at.get(), nf,
                                                impl_->s()), "sample scan");
        CopyAndWait(&n_samples, sample_at.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost,
                    impl_->s(), "dl sample count");
        n_samples += last_flag;
    }
    CudaDevicePtr<ErrorModelSample> sample = d.Alloc<ErrorModelSample>(std::max(1, n_samples));
    p.sample = sample.get();
    MergeSampleWriteKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(nf, p);
    CudaCheck(cudaGetLastError(), "merge samples write launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "merge em/samples sync");
    CopyAndWait(em_mean_out, d.m_em_mean.get(), size_t(ng) * sizeof(double),
                         cudaMemcpyDeviceToHost, impl_->s(), "dl em_mean");
    CopyAndWait(cnt_out, cnt.get(), size_t(ng) * sizeof(int32_t),
                         cudaMemcpyDeviceToHost, impl_->s(), "dl cnt");
    samples.resize(n_samples);
    if (n_samples > 0)
        CopyAndWait(samples.data(), sample.get(), size_t(n_samples) * sizeof(ErrorModelSample),
                    cudaMemcpyDeviceToHost, impl_->s(), "dl samples");
}

void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_b, bool error_model_active,
                                       bool reject_outliers, double reject_nsigma, double reject_mult_sd,
                                       const float *reject_median,
                                       const float *reject_var_add,
                                       const uint8_t *half, const double *frame_cc_factor,
                                       uint8_t *rejected_obs) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int ng = d.n_groups;
    d.a_swI = d.Alloc<double>(std::max(1, ng)); d.a_sw = d.Alloc<double>(std::max(1, ng));
    d.a_swIh0 = d.Alloc<double>(std::max(1, ng)); d.a_swIh1 = d.Alloc<double>(std::max(1, ng));
    d.a_swh0 = d.Alloc<double>(std::max(1, ng)); d.a_swh1 = d.Alloc<double>(std::max(1, ng));
    d.a_swht0 = d.Alloc<double>(std::max(1, ng)); d.a_swht1 = d.Alloc<double>(std::max(1, ng));
    d.a_nh0 = d.Alloc<int32_t>(std::max(1, ng)); d.a_nh1 = d.Alloc<int32_t>(std::max(1, ng));
    d.a_d = d.Alloc<double>(std::max(1, ng)); d.a_rejected = d.Alloc<int32_t>(std::max(1, ng));
    d.a_on_ice = d.Alloc<uint8_t>(std::max(1, ng));
    d.Upload(d.m_rejected, rejected_obs, d.n_fulls);
    CudaDevicePtr<float> median, var_add;
    d.Upload(median, reject_median, ng);
    if (reject_var_add && ng > 0) d.Upload(var_add, reject_var_add, ng);
    CudaDevicePtr<uint8_t> d_half;
    d.Upload(d_half, half, d.n_fulls);
    CudaDevicePtr<double> cc_factor;
    d.Upload(cc_factor, frame_cc_factor, d.n_frames);

    d.merge_em_a = error_model_a; d.merge_em_b = error_model_b; d.merge_em_active = error_model_active ? 1 : 0;
    MergeParams p{};
    p.n_groups = ng; p.min_partiality = d.merge_min_part;
    p.for_search = d.merge_for_search;
    p.error_model_a = error_model_a; p.error_model_b = error_model_b;
    p.error_model_active = error_model_active ? 1 : 0;
    p.reject_outliers = reject_outliers ? 1 : 0; p.reject_nsigma = reject_nsigma;
    p.reject_up = std::expm1(reject_nsigma * reject_mult_sd);
    p.reject_down = -std::expm1(-reject_nsigma * reject_mult_sd);
    p.reject_median = median.get();
    p.reject_var_add = var_add.get();
    p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get();
    p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
    p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get();
    p.half = d_half.get();
    p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get(); p.capture = d.f_capture.get();
    p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
    p.em_mean = d.m_em_mean.get(); p.cc_factor = cc_factor.get();
    p.a_swI = d.a_swI.get(); p.a_sw = d.a_sw.get(); p.a_swIh0 = d.a_swIh0.get(); p.a_swIh1 = d.a_swIh1.get();
    p.a_swh0 = d.a_swh0.get(); p.a_swh1 = d.a_swh1.get();
    p.a_swht0 = d.a_swht0.get(); p.a_swht1 = d.a_swht1.get(); p.a_nh0 = d.a_nh0.get(); p.a_nh1 = d.a_nh1.get();
    p.a_d = d.a_d.get(); p.a_rejected = d.a_rejected.get(); p.a_on_ice = d.a_on_ice.get();
    p.rejected_obs = d.m_rejected.get();

    const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
    MergeAccumKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p);
    CudaCheck(cudaGetLastError(), "merge accum launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "merge accum sync");
    CopyAndWait(rejected_obs, d.m_rejected.get(), size_t(d.n_fulls) * sizeof(uint8_t), cudaMemcpyDeviceToHost, impl_->s(),
              "dl rejected_obs");
}

void RotationScaleMergeGPU::MergeAccumRange(int g0, int n, double *swI, double *sw, double *swIh0,
                                            double *swIh1, double *swh0, double *swh1, double *swh_typ0,
                                            double *swh_typ1, int32_t *nh0, int32_t *nh1, double *d_out,
                                            int32_t *rejected, uint8_t *on_ice_out) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    auto dl = [&](void *h, const auto &s) {
        CopyAndWait(h, s.get() + g0, size_t(n) * sizeof(*s.get()), cudaMemcpyDeviceToHost, impl_->s(),
                  "dl accum"); };
    dl(swI, d.a_swI); dl(sw, d.a_sw); dl(swIh0, d.a_swIh0); dl(swIh1, d.a_swIh1);
    dl(swh0, d.a_swh0); dl(swh1, d.a_swh1); dl(swh_typ0, d.a_swht0); dl(swh_typ1, d.a_swht1);
    dl(d_out, d.a_d);
    dl(nh0, d.a_nh0); dl(nh1, d.a_nh1); dl(rejected, d.a_rejected); dl(on_ice_out, d.a_on_ice);
}

void RotationScaleMergeGPU::MergeAccumRelease() {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    d.a_swI = {}; d.a_sw = {}; d.a_swIh0 = {}; d.a_swIh1 = {}; d.a_swh0 = {}; d.a_swh1 = {};
    d.a_swht0 = {}; d.a_swht1 = {}; d.a_d = {}; d.a_nh0 = {}; d.a_nh1 = {}; d.a_rejected = {}; d.a_on_ice = {};
}

void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, double *sumI, double *wabsdev,
                                       double *wsumI, double *sumv, double *sumv2, int32_t *n,
                                       int32_t *nusable) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int ng = d.n_groups;
    CudaDevicePtr<double> d_merged_I;
    d.Upload(d_merged_I, merged_I, ng);
    CudaDevicePtr<double> r_absdev = d.Alloc<double>(std::max(1, ng)), r_sumI = d.Alloc<double>(std::max(1, ng)),
                          r_wabsdev = d.Alloc<double>(std::max(1, ng)), r_wsumI = d.Alloc<double>(std::max(1, ng)),
                          r_sumv = d.Alloc<double>(std::max(1, ng)), r_sumv2 = d.Alloc<double>(std::max(1, ng));
    CudaDevicePtr<int32_t> r_n = d.Alloc<int32_t>(std::max(1, ng)), r_nusable = d.Alloc<int32_t>(std::max(1, ng));

    MergeParams p{};
    p.n_groups = ng; p.min_partiality = d.merge_min_part;
    p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get();
    p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get(); p.frame_cell_ok = d.frame_cell_ok.get();
    p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
    p.merged_I = d_merged_I.get();
    p.r_absdev = r_absdev.get(); p.r_sumI = r_sumI.get(); p.r_n = r_n.get(); p.r_nusable = r_nusable.get();
    p.r_wabsdev = r_wabsdev.get(); p.r_wsumI = r_wsumI.get();
    p.r_sumv = r_sumv.get(); p.r_sumv2 = r_sumv2.get();
    p.error_model_a = d.merge_em_a; p.error_model_b = d.merge_em_b; p.error_model_active = d.merge_em_active;
    p.em_mean = d.m_em_mean.get(); p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get();
    p.capture = d.f_capture.get();
    p.rejected_obs = d.m_rejected.get();

    const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
    MergeRmeasKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p);
    CudaCheck(cudaGetLastError(), "merge rmeas launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "merge rmeas sync");
    CopyAndWait(absdev, r_absdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl absdev");
    CopyAndWait(sumI, r_sumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumI");
    CopyAndWait(wabsdev, r_wabsdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl wabsdev");
    CopyAndWait(wsumI, r_wsumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl wsumI");
    CopyAndWait(sumv, r_sumv.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumv");
    CopyAndWait(sumv2, r_sumv2.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumv2");
    CopyAndWait(n, r_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl rn");
    CopyAndWait(nusable, r_nusable.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl rnusable");
}

void RotationScaleMergeGPU::MergeBijvoetExcess(double error_model_a, double error_model_b, float *group_d,
                                               double *excess, float &d_lo, float &d_hi) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int ng = d.n_groups;
    auto g_d = d.Alloc<float>(std::max(1, ng)), g_lo = d.Alloc<float>(std::max(1, ng)),
         g_hi = d.Alloc<float>(std::max(1, ng));
    auto g_excess = d.Alloc<double>(std::max(1, ng));

    MergeParams p{};
    p.n_groups = ng; p.min_partiality = d.merge_min_part; p.for_search = d.merge_for_search;
    p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get();
    p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
    p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get();
    p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
    p.hand = d.m_hand.get(); p.has_hands = d.m_has_hands.get();
    p.error_model_a = error_model_a; p.error_model_b = error_model_b; p.error_model_active = 1;
    p.em_mean = d.m_em_mean.get(); p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get();
    p.capture = d.f_capture.get();

    const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
    BijvoetExcessKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p, g_d.get(), g_excess.get(), g_lo.get(), g_hi.get());
    CudaCheck(cudaGetLastError(), "Bijvoet excess launch");
    std::vector<float> lo(ng), hi(ng);
    CopyAndWait(group_d, g_d.get(), size_t(ng) * sizeof(float), cudaMemcpyDeviceToHost, impl_->s(), "dl group d");
    CopyAndWait(excess, g_excess.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(),
                "dl excess");
    CopyAndWait(lo.data(), g_lo.get(), size_t(ng) * sizeof(float), cudaMemcpyDeviceToHost, impl_->s(), "dl d lo");
    CopyAndWait(hi.data(), g_hi.get(), size_t(ng) * sizeof(float), cudaMemcpyDeviceToHost, impl_->s(), "dl d hi");
    d_lo = INFINITY; d_hi = 0.0f;
    for (int g = 0; g < ng; ++g) { d_lo = std::min(d_lo, lo[g]); d_hi = std::max(d_hi, hi[g]); }
}

void RotationScaleMergeGPU::MergeGroupScatter(double *sum, double *sum_sq, double *sum_var, int32_t *n) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int ng = d.n_groups;
    auto g_sum = d.Alloc<double>(std::max(1, ng)), g_sum_sq = d.Alloc<double>(std::max(1, ng)),
         g_sum_var = d.Alloc<double>(std::max(1, ng));
    auto g_n = d.Alloc<int32_t>(std::max(1, ng));

    MergeParams p{};
    p.n_groups = ng; p.min_partiality = d.merge_min_part; p.for_search = d.merge_for_search;
    p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get();
    p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
    p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get();
    p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
    p.em_mean = d.m_em_mean.get(); p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get();
    p.capture = d.f_capture.get();

    const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
    GroupScatterKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p, g_sum.get(), g_sum_sq.get(), g_sum_var.get(), g_n.get());
    CudaCheck(cudaGetLastError(), "group scatter launch");
    CopyAndWait(sum, g_sum.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sum");
    CopyAndWait(sum_sq, g_sum_sq.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sum_sq");
    CopyAndWait(sum_var, g_sum_var.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(),
                "dl sum_var");
    CopyAndWait(n, g_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl n");
}

void RotationScaleMergeGPU::MergeHalves(uint64_t *obs_key, uint8_t *half) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int ng = d.n_groups, nf = d.n_fulls;
    if (nf == 0) return;
    auto key = d.Alloc<uint64_t>(nf);
    auto usable = d.Alloc<uint8_t>(nf), h = d.Alloc<uint8_t>(nf);

    MergeParams p{};
    p.n_groups = ng; p.min_partiality = d.merge_min_part; p.for_search = d.merge_for_search;
    p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get();
    p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
    p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get();
    p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();

    const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
    const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
    ObservationKeyKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(nf, p, d.f_run.get(), d.rr_h.get(), d.rr_k.get(),
                                                             d.rr_l.get(), key.get(), usable.get());
    CudaCheck(cudaGetLastError(), "observation key launch");
    CudaCheck(cudaMemsetAsync(h.get(), 0, size_t(nf), impl_->s()), "memset halves");
    HalvesKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p, key.get(), usable.get(), h.get());
    CudaCheck(cudaGetLastError(), "halves launch");
    CopyAndWait(obs_key, key.get(), size_t(nf) * sizeof(uint64_t), cudaMemcpyDeviceToHost, impl_->s(), "dl keys");
    CopyAndWait(half, h.get(), size_t(nf), cudaMemcpyDeviceToHost, impl_->s(), "dl halves");
}

void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    d.Upload(d.smooth_apply, apply, d.n_frames);
    d.Upload(d.smooth_ratio, ratio, d.n_frames);
    const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
    SmoothCorrKernel<<<blocks, BLK, 0, impl_->s()>>>(d.n_obs, d.PartialFrame(), d.smooth_apply.get(),
                                      d.smooth_ratio.get(), d.corr.get());
    CudaCheck(cudaGetLastError(), "smooth corr launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "smooth corr sync");
}

void RotationScaleMergeGPU::SmoothFullsCorr(const uint8_t *apply, const double *ratio) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    if (d.n_fulls == 0) return;
    d.Upload(d.smooth_apply, apply, d.n_frames);
    d.Upload(d.smooth_ratio, ratio, d.n_frames);
    const int blocks = std::min(65535, (d.n_fulls + BLK - 1) / BLK);
    SmoothCorrKernel<<<blocks, BLK, 0, impl_->s()>>>(d.n_fulls, FrameOf{d.f_frame.get(), nullptr, 0}, d.smooth_apply.get(),
                                      d.smooth_ratio.get(), d.f_corr.get());
    CudaCheck(cudaGetLastError(), "smooth fulls corr launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "smooth fulls corr sync");
}

int64_t RotationScaleMergeGPU::FilterCorrByZeta(double min_zeta) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    CudaDevicePtr<unsigned long long> dropped = d.Alloc<unsigned long long>(1);
    CudaCheck(cudaMemsetAsync(dropped.get(), 0, sizeof(unsigned long long), impl_->s()), "zero zeta drop count");
    const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
    FilterZetaKernel<<<blocks, BLK, 0, impl_->s()>>>(d.n_obs, min_zeta, d.zeta.get(), d.corr.get(), dropped.get());
    CudaCheck(cudaGetLastError(), "zeta filter launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "zeta filter sync");
    unsigned long long n = 0;
    CopyAndWait(&n, dropped.get(), sizeof(unsigned long long), cudaMemcpyDeviceToHost, impl_->s(),
              "dl zeta drop count");
    return static_cast<int64_t>(n);
}

void RotationScaleMergeGPU::FilterCorrByFrame(const uint8_t *reject) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    d.Upload(d.filter_reject, reject, d.n_frames);
    const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
    FilterFrameKernel<<<blocks, BLK, 0, impl_->s()>>>(d.n_obs, d.PartialFrame(), d.filter_reject.get(), d.corr.get());
    CudaCheck(cudaGetLastError(), "frame filter launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "frame filter sync");
}

void RotationScaleMergeGPU::ComputePartialCC(double min_partiality, double *cc_out, int64_t *cc_n_out) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
    // Post-smooth group means (reuse the scaling reduce; reads the resident, smoothed corr), then the
    // per-frame CC over the resident partials. Only the tiny per-frame cc/cc_n come back to the host.
    ReduceGroupMeansKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(d.n_groups, min_partiality,
        d.group_perm.get(), d.group_start.get(), d.group_count.get(),
        d.obs.get(), d.corr.get(), d.group_mean.get());
    CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
    PerFrameCCKernel<<<d.n_frames, BLK, 0, impl_->s()>>>(d.n_frames, min_partiality,
        d.frame_start.get(), d.frame_count.get(), d.obs.get(), d.corr.get(), d.flags.get(), d.group.get(),
        d.group_mean.get(), d.cc.get(), d.cc_n.get());
    CudaCheck(cudaGetLastError(), "partial CC launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "partial CC sync");
    CopyAndWait(cc_out, d.cc.get(), size_t(d.n_frames) * sizeof(double),
                         cudaMemcpyDeviceToHost, impl_->s(), "download cc");
    CopyAndWait(cc_n_out, d.cc_n.get(), size_t(d.n_frames) * sizeof(int64_t),
                         cudaMemcpyDeviceToHost, impl_->s(), "download cc_n");
}

void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *perm,
                                       const int32_t *rr_start, const int32_t *rr_count,
                                       const int32_t *rr_h, const int32_t *rr_k, const int32_t *rr_l) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    d.n_runs = n_runs;
    d.n_perm = n_perm;
    d.Upload(d.perm, perm, n_perm);
    d.Upload(d.rr_start, rr_start, n_runs);
    d.Upload(d.rr_count, rr_count, n_runs);
    d.Upload(d.rr_h, rr_h, n_runs);
    d.Upload(d.rr_k, rr_k, n_runs);
    d.Upload(d.rr_l, rr_l, n_runs);
}

int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_partiality,
                                   double capture_uncertainty_coeff, double min_captured_fraction,
                                   float max_frame_gap, int64_t &n_overloaded) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    // The partials' group CSR is read only by their scaling, which is over (BuildGroups makes the next
    // one); the group means stay for the fulls' scaling. The previous fulls go before the new ones are
    // allocated.
    d.group = {}; d.group_perm = {}; d.group_start = {}; d.group_count = {};
    d.ReleaseFulls();
    CudaDevicePtr<int32_t> rr_group, rr_nevents = d.Alloc<int32_t>(std::max(1, d.n_runs)),
                           rr_noverloaded = d.Alloc<int32_t>(std::max(1, d.n_runs)), rr_offset;
    d.Upload(rr_group, rawrun_group, d.n_runs);

    CombineParams p{};
    p.n_runs = d.n_runs;
    p.min_partiality = min_partiality;
    p.capture_uncertainty_coeff = capture_uncertainty_coeff;
    p.min_captured_fraction = min_captured_fraction;
    p.max_frame_gap = max_frame_gap;
    p.obs = d.obs.get(); p.corr = d.corr.get();
    p.bkg = d.bkg.get(); p.var_bkg = d.var_bkg.get(); p.image_number = d.image_number.get(); p.d = d.d_obs.get();
    p.frame_of = d.PartialFrame(); p.flags = d.flags.get();
    p.perm = d.perm.get(); p.rr_start = d.rr_start.get(); p.rr_count = d.rr_count.get();
    p.rr_h = d.rr_h.get(); p.rr_k = d.rr_k.get(); p.rr_l = d.rr_l.get(); p.rr_group = rr_group.get();
    p.rr_nevents = rr_nevents.get();
    p.rr_noverloaded = rr_noverloaded.get();

    const int blocks = std::min(65535, (d.n_runs + BLK - 1) / BLK);

    // Count pass: how many fulls each run emits.
    CombineKernel<false><<<blocks, BLK, 0, impl_->s()>>>(p);
    CudaCheck(cudaGetLastError(), "combine count launch");

    // Exclusive prefix sum on the host (deterministic) -> per-run output offset + total fulls.
    std::vector<int32_t> nevents(d.n_runs);
    CopyAndWait(nevents.data(), rr_nevents.get(), size_t(d.n_runs) * sizeof(int32_t),
                         cudaMemcpyDeviceToHost, impl_->s(), "download nevents");
    std::vector<int32_t> offset(d.n_runs);
    int64_t acc = 0;
    for (int r = 0; r < d.n_runs; ++r) { offset[r] = static_cast<int32_t>(acc); acc += nevents[r]; }
    d.n_fulls = static_cast<int>(acc);
    std::vector<int32_t> noverloaded(d.n_runs);
    CopyAndWait(noverloaded.data(), rr_noverloaded.get(), size_t(d.n_runs) * sizeof(int32_t),
                         cudaMemcpyDeviceToHost, impl_->s(), "download noverloaded");
    n_overloaded = 0;
    for (int r = 0; r < d.n_runs; ++r) n_overloaded += noverloaded[r];
    rr_nevents = {}; rr_noverloaded = {};

    // Allocate the fulls SoA and emit.
    const int nf = std::max(1, d.n_fulls);
    d.f_run = d.Alloc<int32_t>(nf); d.f_peak = d.Alloc<int32_t>(nf);
    d.f_frame = d.Alloc<int32_t>(nf); d.f_group = d.Alloc<int32_t>(nf);
    d.f_I = d.Alloc<float>(nf); d.f_sigma = d.Alloc<float>(nf);
    d.f_d = d.Alloc<float>(nf); d.f_img = d.Alloc<float>(nf);
    d.f_var_bkg = d.Alloc<float>(nf); d.f_var_per_I = d.Alloc<float>(nf); d.f_capture = d.Alloc<float>(nf);
    d.f_on_ice = d.Alloc<uint8_t>(nf); d.f_clipped = d.Alloc<uint8_t>(nf);
    d.f_corr = d.Alloc<float>(nf);
    d.f_zeta = d.Alloc<float>(nf);
    d.f_obs = d.Alloc<ObsRecord>(nf);
    d.Upload(rr_offset, offset.data(), d.n_runs);

    p.rr_offset = rr_offset.get();
    p.f_run = d.f_run.get(); p.f_peak = d.f_peak.get();
    p.f_frame = d.f_frame.get(); p.f_group = d.f_group.get();
    p.f_I = d.f_I.get(); p.f_sigma = d.f_sigma.get(); p.f_d = d.f_d.get(); p.f_img = d.f_img.get();
    p.f_var_bkg = d.f_var_bkg.get(); p.f_var_per_I = d.f_var_per_I.get(); p.f_capture = d.f_capture.get();
    p.f_on_ice = d.f_on_ice.get(); p.f_clipped = d.f_clipped.get();
    if (d.n_fulls > 0) {
        CombineKernel<true><<<blocks, BLK, 0, impl_->s()>>>(p);
        CudaCheck(cudaGetLastError(), "combine emit launch");
    }
    CudaCheck(cudaStreamSynchronize(impl_->s()), "combine sync");
    return d.n_fulls;
}

void RotationScaleMergeGPU::GetFulls(int32_t *run, float *I, float *sigma, float *d,
                                     float *image_number, int32_t *frame, uint8_t *on_ice,
                                     uint8_t *clipped, int32_t *group) const {
    DeviceGuard guard(impl_->device, impl_->available);
    const auto &dd = *impl_;
    const size_t n = static_cast<size_t>(dd.n_fulls);
    if (n == 0) return;
    auto dl = [&](void *dst, const void *src, size_t bytes) {
        CopyAndWait(dst, src, bytes, cudaMemcpyDeviceToHost, impl_->s(), "download fulls");
    };
    dl(run, dd.f_run.get(), n * sizeof(int32_t)); dl(frame, dd.f_frame.get(), n * sizeof(int32_t));
    dl(group, dd.f_group.get(), n * sizeof(int32_t));
    dl(I, dd.f_I.get(), n * sizeof(float)); dl(sigma, dd.f_sigma.get(), n * sizeof(float));
    dl(d, dd.f_d.get(), n * sizeof(float)); dl(image_number, dd.f_img.get(), n * sizeof(float));
    dl(on_ice, dd.f_on_ice.get(), n * sizeof(uint8_t));
    dl(clipped, dd.f_clipped.get(), n * sizeof(uint8_t));
}

void RotationScaleMergeGPU::ResetFullsScale() {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int nf = d.n_fulls;
    if (nf == 0) return;
    const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
    // Unity model: partiality/prescaling_corr/zeta = 1 so coeff = mean; corr starts at 1.
    FillKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(d.f_corr.get(), nf, 1.0f);
    CudaCheck(cudaGetLastError(), "FillKernel launch");
    FullsRecordKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(nf, d.f_I.get(), d.f_sigma.get(), d.f_obs.get());
    CudaCheck(cudaGetLastError(), "FullsRecordKernel launch");
    FillKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(d.f_zeta.get(), nf, 1.0f);
    CudaCheck(cudaGetLastError(), "FillKernel launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "reset fulls scale sync");
}

void RotationScaleMergeGPU::FitFullsScale(double min_partiality) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int nf = d.n_fulls;
    if (nf == 0) return;
    const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);

    // Reset per call, as ScalePartials: the host keeps the G of a frame across calls.
    CudaCheck(cudaMemsetAsync(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t), impl_->s()), "memset f scaled");
    CudaCheck(cudaMemsetAsync(d.g.get(), 0, size_t(d.n_frames) * sizeof(double), impl_->s()), "memset f g");

    ReduceGroupMeansKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(d.n_groups, min_partiality,
        d.f_gperm.get(), d.f_gstart.get(), d.f_gcount.get(),
        d.f_obs.get(), d.f_corr.get(), d.group_mean.get());
    CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
    FitPerFrameGKernel<<<d.n_frames, BLK, 0, impl_->s()>>>(d.n_frames, d.f_frame_start.get(), d.f_frame_count.get(),
        d.f_frame_perm.get(), min_partiality, d.f_group.get(), d.f_obs.get(), d.f_zeta.get(), d.f_on_ice.get(),
        d.group_mean.get(), 1L, d.g.get(), d.scaled.get(), d.info.get());
    CudaCheck(cudaGetLastError(), "scale fulls launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "scale fulls sync");
}

void RotationScaleMergeGPU::GetFullsCorr(float *corr) const {
    DeviceGuard guard(impl_->device, impl_->available);
    const auto &d = *impl_;
    if (d.n_fulls == 0) return;
    CopyAndWait(corr, d.f_corr.get(), size_t(d.n_fulls) * sizeof(float),
                         cudaMemcpyDeviceToHost, impl_->s(), "download f_corr");
}

void RotationScaleMergeGPU::GetFullsPeak(int32_t *peak) const {
    DeviceGuard guard(impl_->device, impl_->available);
    const auto &d = *impl_;
    if (d.n_fulls == 0) return;
    CopyAndWait(peak, d.f_peak.get(), size_t(d.n_fulls) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(),
                "download f_peak");
}

void RotationScaleMergeGPU::GetFullsVariance(float *var_bkg, float *var_per_I, float *capture) const {
    DeviceGuard guard(impl_->device, impl_->available);
    const auto &d = *impl_;
    if (d.n_fulls == 0) return;
    const size_t bytes = size_t(d.n_fulls) * sizeof(float);
    CopyAndWait(var_bkg, d.f_var_bkg.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_var_bkg");
    CopyAndWait(var_per_I, d.f_var_per_I.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(),
              "download f_var_per_I");
    CopyAndWait(capture, d.f_capture.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_capture");
}

void RotationScaleMergeGPU::FullsDownloaded() {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    d.f_img = {}; d.f_peak = {}; d.f_clipped = {};
    d.f_obs = {}; d.f_zeta = {}; d.f_frame_perm = {}; d.f_frame_start = {}; d.f_frame_count = {};
    d.group_mean = {};
}

void RotationScaleMergeGPU::SetFullsCorr(const float *corr) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    if (d.n_fulls == 0) return;
    CopyAndWait(d.f_corr.get(), corr, size_t(d.n_fulls) * sizeof(float),
                         cudaMemcpyHostToDevice, impl_->s(), "upload f_corr");
}

void RotationScaleMergeGPU::SurfaceSetTerms(int n_terms, const SurfaceTerm *term, const uint8_t *parity,
                                            int n_groups, const int32_t *gperm, const int32_t *gstart,
                                            int ncell) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    d.s_ncell = ncell;
    d.s_n_groups = n_groups;
    d.Upload(d.s_term, term, n_terms);
    d.Upload(d.s_parity, parity, n_terms);
    d.Upload(d.s_gperm, gperm, n_terms);
    d.Upload(d.s_gstart, gstart, n_groups + 1);
    d.s_A = d.Alloc<double>(std::max(1, ncell));
    d.s_cross = d.Alloc<double>(std::max(1, ncell));
    d.s_ref2 = d.Alloc<double>(std::max(1, ncell));
    d.s_sw = d.Alloc<double>(std::max(1, n_groups));
    d.s_swI = d.Alloc<double>(std::max(1, n_groups));
    d.s_w_Is = d.Alloc<double>(std::max(1, n_terms));
    d.s_w_Iref = d.Alloc<double>(std::max(1, n_terms));
    d.s_Iref = d.Alloc<double>(std::max(1, n_terms));
    for (int &nb : d.s_n_blocks) nb = 0;
}

void RotationScaleMergeGPU::SurfaceSetSubset(int subset, int n_blocks, const int32_t *perm,
                                             const int32_t *seg_start) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int n_seg = n_blocks * d.s_ncell;
    d.s_n_sel[subset] = n_blocks > 0 ? seg_start[n_seg] : 0;
    d.Upload(d.s_perm[subset], perm, d.s_n_sel[subset]);
    d.Upload(d.s_seg_start[subset], seg_start, n_seg + 1);
    d.s_n_blocks[subset] = n_blocks;
    if (size_t(n_seg) > d.s_slots) {
        d.s_tcross = d.Alloc<double>(n_seg);
        d.s_tref2 = d.Alloc<double>(n_seg);
        d.s_slots = n_seg;
    }
}

namespace {
    // The positions, in order, of a stable radix sort of `value` by `key` (keys below 2^bits).
    void StableSortByKey(cudaStream_t s, int n, const uint32_t *key, const int32_t *value, uint32_t *key_out,
                         int32_t *value_out, int bits, const char *what) {
        size_t bytes = 0;
        CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key, key_out, value, value_out, n, 0, bits, s), what);
        CudaDevicePtr<uint8_t> scratch(std::max<size_t>(bytes, 1), ALLOC);
        CudaCheck(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, key, key_out, value, value_out, n, 0, bits, s),
                  what);
    }

    // start[b] = how many keys are below b, for b in [0, n_bins] (a counting sort's bin starts).
    void BinStarts(cudaStream_t s, int n, const uint32_t *key, int n_bins, int32_t *start, const char *what) {
        CudaCheck(cudaMemsetAsync(start, 0, size_t(n_bins + 1) * sizeof(int32_t), s), what);
        if (n > 0) CountKeysKernel<<<(n + BLK - 1) / BLK, BLK, 0, s>>>(n, key, start);
        CudaCheck(cudaGetLastError(), what);
        size_t bytes = 0;
        CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, start, start, n_bins + 1, s), what);
        CudaDevicePtr<uint8_t> scratch(std::max<size_t>(bytes, 1), ALLOC);
        CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, start, start, n_bins + 1, s), what);
    }

    int BitsFor(int64_t n_keys) {
        int bits = 1;
        while ((int64_t(1) << bits) < n_keys) ++bits;
        return bits;
    }
}

int RotationScaleMergeGPU::SurfaceBuildTerms(const int32_t *cell, const float *corr, double min_partiality,
                                             int ncell, int n_groups, int block) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const cudaStream_t st = impl_->s();
    const int nf = d.n_fulls;
    d.s_ncell = ncell;
    d.s_n_groups = n_groups;
    for (int &nb : d.s_n_blocks) nb = 0;
    for (int &ns : d.s_n_sel) ns = 0;

    // Which fulls give a term, and where each goes in the three lists.
    int n_terms = 0, n_even = 0, n_odd = 0;
    CudaDevicePtr<int32_t> idx_even, idx_odd;
    CudaDevicePtr<uint32_t> group_key;
    {
        CudaDevicePtr<int32_t> d_cell, keep, even, odd;
        CudaDevicePtr<float> d_corr;
        d.Upload(d_cell, cell, nf);
        d.Upload(d_corr, corr, nf);
        keep = d.Alloc<int32_t>(std::max(1, nf)); even = d.Alloc<int32_t>(std::max(1, nf));
        odd = d.Alloc<int32_t>(std::max(1, nf));
        if (nf > 0) {
            SurfaceKeepKernel<<<(nf + BLK - 1) / BLK, BLK, 0, st>>>(nf, d.f_group.get(), d_corr.get(), d.f_on_ice.get(),
                d.f_frame.get(), d_cell.get(), min_partiality, keep.get(), even.get(), odd.get());
            CudaCheck(cudaGetLastError(), "surface keep launch");
        }
        // The last full's flags, before the scans turn the flags into places.
        int32_t last[3] = {0, 0, 0};
        auto d_last = d.Alloc<int32_t>(3);
        if (nf > 0) {
            CopyAndWait(&last[0], keep.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl keep");
            CopyAndWait(&last[1], even.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl keep");
            CopyAndWait(&last[2], odd.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl keep");
            size_t bytes = 0;
            CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, keep.get(), keep.get(), nf, st), "surface scan");
            auto scratch = d.Alloc<uint8_t>(std::max<size_t>(bytes, 1));
            for (int32_t *v : {keep.get(), even.get(), odd.get()})
                CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, v, v, nf, st), "surface scan");
            int32_t at[3];
            CopyAndWait(&at[0], keep.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl count");
            CopyAndWait(&at[1], even.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl count");
            CopyAndWait(&at[2], odd.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl count");
            n_terms = at[0] + last[0]; n_even = at[1] + last[1]; n_odd = at[2] + last[2];
        }
        CopyAndWait(d_last.get(), last, sizeof(last), cudaMemcpyHostToDevice, st, "ul last flags");
        d.s_term = d.Alloc<SurfaceTerm>(std::max(1, n_terms));
        d.s_parity = d.Alloc<uint8_t>(std::max(1, n_terms));
        group_key = d.Alloc<uint32_t>(std::max(1, n_terms));
        idx_even = d.Alloc<int32_t>(std::max(1, n_even));
        idx_odd = d.Alloc<int32_t>(std::max(1, n_odd));
        if (nf > 0) {
            SurfaceWriteTermKernel<<<(nf + BLK - 1) / BLK, BLK, 0, st>>>(nf, keep.get(), even.get(), odd.get(),
                d_last.get(), d.f_I.get(), d.f_sigma.get(), d_corr.get(), d.f_d.get(), d_cell.get(), d.f_group.get(),
                d.f_frame.get(), d.s_term.get(), d.s_parity.get(), idx_even.get(), idx_odd.get(), group_key.get());
            CudaCheck(cudaGetLastError(), "surface term launch");
        }
        CudaCheck(cudaStreamSynchronize(st), "surface terms sync");
    }

    // The group CSR over the terms: the terms of group g at gperm[gstart[g] .. gstart[g + 1]), in
    // fulls order.
    auto iota = d.Alloc<int32_t>(std::max(1, n_terms));
    if (n_terms > 0) IotaKernel<<<(n_terms + BLK - 1) / BLK, BLK, 0, st>>>(n_terms, iota.get());
    CudaCheck(cudaGetLastError(), "iota launch");
    d.s_gperm = d.Alloc<int32_t>(std::max(1, n_terms));
    d.s_gstart = d.Alloc<int32_t>(n_groups + 1);
    {
        auto key_sorted = d.Alloc<uint32_t>(std::max(1, n_terms));
        if (n_terms > 0)
            StableSortByKey(st, n_terms, group_key.get(), iota.get(), key_sorted.get(), d.s_gperm.get(),
                            BitsFor(n_groups), "surface group sort");
        BinStarts(st, n_terms, group_key.get(), n_groups, d.s_gstart.get(), "surface group starts");
    }
    group_key = CudaDevicePtr<uint32_t>();

    // Each subset cut into the host's reduction blocks, every block's terms ordered by cell.
    const int32_t *sel_of[3] = {idx_even.get(), idx_odd.get(), iota.get()};
    const int n_of[3] = {n_even, n_odd, n_terms};
    for (int subset = 0; subset < 3; ++subset) {
        const int n = n_of[subset];
        const int nb = ReductionBlocks(n, block);
        const int n_seg = nb * ncell;
        d.s_n_sel[subset] = n;
        d.s_n_blocks[subset] = nb;
        d.s_perm[subset] = d.Alloc<int32_t>(std::max(1, n));
        d.s_seg_start[subset] = d.Alloc<int32_t>(n_seg + 1);
        if (nb == 0) continue;
        auto key = d.Alloc<uint32_t>(n), key_sorted = d.Alloc<uint32_t>(n);
        SurfaceSubsetKeyKernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, nb, ncell, sel_of[subset], d.s_term.get(),
                                                                  key.get());
        CudaCheck(cudaGetLastError(), "surface subset key launch");
        StableSortByKey(st, n, key.get(), sel_of[subset], key_sorted.get(), d.s_perm[subset].get(), BitsFor(n_seg),
                        "surface subset sort");
        BinStarts(st, n, key.get(), n_seg, d.s_seg_start[subset].get(), "surface subset starts");
        if (size_t(n_seg) > d.s_slots) {
            d.s_tcross = d.Alloc<double>(n_seg);
            d.s_tref2 = d.Alloc<double>(n_seg);
            d.s_slots = n_seg;
        }
    }
    CudaCheck(cudaStreamSynchronize(st), "surface subsets sync");
    d.s_A = d.Alloc<double>(std::max(1, ncell));
    d.s_cross = d.Alloc<double>(std::max(1, ncell));
    d.s_ref2 = d.Alloc<double>(std::max(1, ncell));
    d.s_sw = d.Alloc<double>(std::max(1, n_groups));
    d.s_swI = d.Alloc<double>(std::max(1, n_groups));
    d.s_w_Is = d.Alloc<double>(std::max(1, n_terms));
    d.s_w_Iref = d.Alloc<double>(std::max(1, n_terms));
    d.s_Iref = d.Alloc<double>(std::max(1, n_terms));
    return n_terms;
}

void RotationScaleMergeGPU::SurfaceRelease() {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    d.s_slots = 0;
    d.s_term = {}; d.s_parity = {}; d.s_gperm = {}; d.s_gstart = {};
    for (int subset = 0; subset < 3; ++subset) {
        d.s_perm[subset] = {}; d.s_seg_start[subset] = {};
        d.s_n_blocks[subset] = d.s_n_sel[subset] = 0;
    }
    d.s_A = {}; d.s_sw = {}; d.s_swI = {}; d.s_tcross = {}; d.s_tref2 = {}; d.s_cross = {}; d.s_ref2 = {};
    d.s_w_Is = {}; d.s_w_Iref = {}; d.s_Iref = {};
}

void RotationScaleMergeGPU::SurfaceShellWeights(int n_terms, int nshell, int32_t *group_shell, double *shw) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const cudaStream_t st = impl_->s();
    const int ncell = d.s_ncell, n_groups = d.s_n_groups;
    const int n = n_terms;
    // The edges: order statistics of s^2, read off a sort.
    auto s2 = d.Alloc<float>(std::max(1, n)), s2_sorted = d.Alloc<float>(std::max(1, n));
    SurfaceS2Kernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, d.s_term.get(), s2.get());
    CudaCheck(cudaGetLastError(), "surface s2 launch");
    {
        size_t bytes = 0;
        CudaCheck(cub::DeviceRadixSort::SortKeys(nullptr, bytes, s2.get(), s2_sorted.get(), n, 0, 32, st), "s2 sort");
        auto scratch = d.Alloc<uint8_t>(std::max<size_t>(bytes, 1));
        CudaCheck(cub::DeviceRadixSort::SortKeys(scratch.get(), bytes, s2.get(), s2_sorted.get(), n, 0, 32, st),
                  "s2 sort");
    }
    std::vector<float> edge(nshell - 1);
    for (int i = 1; i < nshell; ++i)
        CopyAndWait(&edge[i - 1], s2_sorted.get() + size_t(n) * size_t(i) / size_t(nshell), sizeof(float),
                    cudaMemcpyDeviceToHost, st, "dl edge");
    s2_sorted = CudaDevicePtr<float>();
    CudaDevicePtr<float> d_edge;
    d.Upload(d_edge, edge.data(), nshell - 1);
    auto shell = d.Alloc<uint8_t>(std::max(1, n));
    auto bin = d.Alloc<uint32_t>(std::max(1, n)), bin_sorted = d.Alloc<uint32_t>(std::max(1, n));
    SurfaceShellKernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, s2.get(), d_edge.get(), nshell - 1, nshell,
                                                          d.s_term.get(), shell.get(), bin.get());
    CudaCheck(cudaGetLastError(), "surface shell launch");
    auto g_shell = d.Alloc<int32_t>(std::max(1, n_groups));
    if (n_groups > 0)
        SurfaceGroupShellKernel<<<(n_groups + BLK - 1) / BLK, BLK, 0, st>>>(n_groups, d.s_gstart.get(),
            d.s_gperm.get(), shell.get(), g_shell.get());
    CudaCheck(cudaGetLastError(), "surface group shell launch");
    // The weights: the terms of each (cell, shell) bin in fulls order, summed one bin per thread.
    const int n_bins = ncell * nshell;
    auto iota = d.Alloc<int32_t>(std::max(1, n)), order = d.Alloc<int32_t>(std::max(1, n));
    auto bin_start = d.Alloc<int32_t>(n_bins + 1);
    auto d_shw = d.Alloc<double>(std::max(1, n_bins));
    IotaKernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, iota.get());
    CudaCheck(cudaGetLastError(), "iota launch");
    StableSortByKey(st, n, bin.get(), iota.get(), bin_sorted.get(), order.get(), BitsFor(n_bins), "surface bin sort");
    BinStarts(st, n, bin.get(), n_bins, bin_start.get(), "surface bin starts");
    SurfaceShellWeightKernel<<<(n_bins + BLK - 1) / BLK, BLK, 0, st>>>(n_bins, bin_start.get(), order.get(),
                                                                      d.s_term.get(), d_shw.get());
    CudaCheck(cudaGetLastError(), "surface shell weight launch");
    CopyAndWait(shw, d_shw.get(), size_t(n_bins) * sizeof(double), cudaMemcpyDeviceToHost, st, "dl shw");
    CopyAndWait(group_shell, g_shell.get(), size_t(n_groups) * sizeof(int32_t), cudaMemcpyDeviceToHost, st,
                "dl group shell");
}

void RotationScaleMergeGPU::SurfaceReference(int parity, const double *A) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    CopyAndWait(d.s_A.get(), A, size_t(d.s_ncell) * sizeof(double), cudaMemcpyHostToDevice, impl_->s(),
                "upload surface");
    const int grp_blocks = std::min(65535, (d.s_n_groups + BLK - 1) / BLK);
    if (grp_blocks > 0)
        SurfaceReferenceKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(d.s_n_groups, parity, d.s_gperm.get(),
            d.s_gstart.get(), d.s_term.get(), d.s_parity.get(), d.s_A.get(), d.s_sw.get(), d.s_swI.get());
    CudaCheck(cudaGetLastError(), "surface reference launch");
    CudaCheck(cudaStreamSynchronize(impl_->s()), "surface reference sync");
}

void RotationScaleMergeGPU::SurfaceGetReference(double *sw, double *swI) const {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const size_t bytes = size_t(d.s_n_groups) * sizeof(double);
    CopyAndWait(sw, d.s_sw.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "dl surface sw");
    CopyAndWait(swI, d.s_swI.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "dl surface swI");
}

void RotationScaleMergeGPU::SurfaceFitSums(int subset, double *cross, double *ref2) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int ncell = d.s_ncell, nb = d.s_n_blocks[subset], n_seg = nb * ncell;
    if (nb == 0) {
        std::fill(cross, cross + ncell, 0.0);
        std::fill(ref2, ref2 + ncell, 0.0);
        return;
    }
    SurfaceFitTermKernel<<<std::min(65535, (d.s_n_sel[subset] + BLK - 1) / BLK), BLK, 0, impl_->s()>>>(
        d.s_n_sel[subset], d.s_perm[subset].get(), d.s_term.get(), d.s_A.get(), d.s_sw.get(), d.s_swI.get(),
        d.s_w_Is.get(), d.s_w_Iref.get(), d.s_Iref.get());
    CudaCheck(cudaGetLastError(), "surface fit term launch");
    SurfaceFitSegmentKernel<<<std::min(65535, (n_seg + BLK - 1) / BLK), BLK, 0, impl_->s()>>>(n_seg,
        d.s_seg_start[subset].get(), d.s_w_Is.get(), d.s_w_Iref.get(), d.s_Iref.get(),
        d.s_tcross.get(), d.s_tref2.get());
    CudaCheck(cudaGetLastError(), "surface fit segment launch");
    SurfaceFitCellKernel<<<std::min(65535, (ncell + BLK - 1) / BLK), BLK, 0, impl_->s()>>>(nb, ncell,
        d.s_tcross.get(), d.s_tref2.get(), d.s_cross.get(), d.s_ref2.get());
    CudaCheck(cudaGetLastError(), "surface cell sum launch");
    CopyAndWait(cross, d.s_cross.get(), size_t(ncell) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(),
                "dl surface cross");
    CopyAndWait(ref2, d.s_ref2.get(), size_t(ncell) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(),
                "dl surface ref2");
}

// --- the resident ingest ---
namespace {
    constexpr int REFL_WORDS = sizeof(Reflection) / sizeof(uint32_t);
    static_assert(sizeof(Reflection) % sizeof(uint32_t) == 0 && alignof(Reflection) == alignof(uint32_t),
                  "the ingest reads a reflection as whole 32-bit words");

    // The raw hkl as one sort key: each index as the sixteen bits the host's SortKey holds, offset so
    // that the unsigned order of the key is the signed order of (h, k, l).
    __device__ __forceinline__ uint64_t RawHklKey(int32_t h, int32_t k, int32_t l) {
        const auto u = [](int32_t v) { return uint64_t(uint16_t(int16_t(v)) ^ 0x8000u); };
        return (u(h) << 32) | (u(k) << 16) | u(l);
    }
    __host__ __device__ __forceinline__ int32_t RawHklIndex(uint64_t key, int shift) {
        return int16_t(uint16_t(((key >> shift) & 0xffffu) ^ 0x8000u));
    }

    // A float as an unsigned key in the order of the float comparison: -0 is +0, as the comparison
    // sees it (a NaN never reaches here - the host comparator would not be an order with one).
    __device__ __forceinline__ uint32_t FloatOrderKey(float v) {
        if (v == 0.0f) v = 0.0f;
        const uint32_t u = __float_as_uint(v);
        return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
    }

    // One block per BLK reflections: the block reads its reflections as whole words (coalesced) into
    // shared memory, and each thread then has its own at hand.
    __device__ __forceinline__ const Reflection *StageBlock(int n, const Reflection *__restrict__ refl,
                                                           uint32_t *buf) {
        const int first = blockIdx.x * BLK;
        const int cnt = min(BLK, n - first);
        const uint32_t *src = reinterpret_cast<const uint32_t *>(refl + first);
        for (int w = threadIdx.x; w < cnt * REFL_WORDS; w += BLK)
            buf[w] = src[w];
        __syncthreads();
        if (int(threadIdx.x) >= cnt)
            return nullptr;
        return reinterpret_cast<const Reflection *>(buf + threadIdx.x * REFL_WORDS);
    }

    // What the sort and the cut read of every reflection.
    struct IngestKeys {
        float *image_number, *d, *bkg;
        uint64_t *key;
    };
    __global__ void IngestKeyKernel(int n, int base, const Reflection *__restrict__ refl, IngestKeys a) {
        __shared__ uint32_t buf[BLK * REFL_WORDS];
        const Reflection *r = StageBlock(n, refl, buf);
        if (!r)
            return;
        const int i = base + blockIdx.x * BLK + threadIdx.x;
        a.image_number[i] = r->image_number;
        a.d[i] = r->d;
        a.bkg[i] = r->bkg;
        a.key[i] = RawHklKey(r->h, r->k, r->l);
    }

    struct IngestArrays {
        ObsRecord *obs;
        float *zeta, *corr, *bkg, *var_bkg, *image_number, *d, *dphi;
        uint8_t *flags;
        const int32_t *new_idx;   // the cut's place of each reflection (an exclusive prefix over the kept)
        int n, n_keep;            // reflections and observations
    };

    // The observations, each at the place the cut gave its reflection (a reflection the cut dropped has
    // the place of the next kept one). The deterministic correction is the product of its three
    // factors, in the host's order; the frame's flux is divided out once the frames are laid out.
    __global__ void IngestExtractKernel(int cnt, int base, const Reflection *__restrict__ refl, IngestArrays a) {
        __shared__ uint32_t buf[BLK * REFL_WORDS];
        const Reflection *rp = StageBlock(cnt, refl, buf);
        if (!rp)
            return;
        const Reflection &r = *rp;
        const int src = base + blockIdx.x * BLK + threadIdx.x;
        const int i = a.new_idx[src];
        if ((src + 1 < a.n ? a.new_idx[src + 1] : a.n_keep) == i)
            return;
        a.obs[i] = ObsRecord{r.I, r.sigma, r.partiality,
                             __fmul_rn(__fmul_rn(r.prescaling_corr, r.qe_corr), r.flight_corr)};
        a.zeta[i] = r.zeta;
        a.corr[i] = r.image_scale_corr;
        a.bkg[i] = r.bkg;
        a.var_bkg[i] = r.var_bkg;
        a.image_number[i] = r.image_number;
        a.d[i] = r.d;
        a.dphi[i] = r.delta_phi_deg;
        a.flags[i] = (r.on_ice_ring ? OBS_ON_ICE_FLAG : 0) | (r.clipped ? OBS_CLIPPED_FLAG : 0)
                     | (r.overloaded ? OBS_OVERLOADED_FLAG : 0);
    }


    // One thread per frame, its reflections in order: the host's sum to the bit.
    __global__ void IngestMeanBkgKernel(int n_frames, const int32_t *__restrict__ start,
                                        const float *__restrict__ bkg, double *__restrict__ mean) {
        for (int o = blockIdx.x * blockDim.x + threadIdx.x; o < n_frames; o += gridDim.x * blockDim.x) {
            double sum = 0.0;
            int n = 0;
            for (int i = start[o]; i < start[o + 1]; ++i) {
                const float b = bkg[i];
                if (isfinite(b) && b > 0.0f) { sum = __dadd_rn(sum, double(b)); ++n; }
            }
            mean[o] = n > 0 ? __ddiv_rn(sum, double(n)) : CUDART_NAN;
        }
    }

    __global__ void IngestFluxKernel(int n, FrameOf frame_of, const double *__restrict__ flux,
                                     ObsRecord *__restrict__ obs) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
            obs[i].prescaling_corr = __double2float_rn(__ddiv_rn(double(obs[i].prescaling_corr), flux[frame_of(i)]));
    }

    __global__ void IngestImageKeyKernel(int n, const float *__restrict__ image_number,
                                         uint32_t *__restrict__ key, int32_t *__restrict__ idx) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {
            key[i] = FloatOrderKey(image_number[i]);
            idx[i] = i;
        }
    }

    __global__ void IngestGatherKeyKernel(int n, const int32_t *__restrict__ idx, const uint64_t *__restrict__ key,
                                          uint64_t *__restrict__ out) {
        for (int j = blockIdx.x * blockDim.x + threadIdx.x; j < n; j += gridDim.x * blockDim.x)
            out[j] = key[idx[j]];
    }

    // Per run: its resolution (that of its first member with a usable one, as the host's run scan takes
    // it) and whether the resolution range keeps it.
    __global__ void IngestRunKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count,
                                    const int32_t *__restrict__ perm, const float *__restrict__ d_obs,
                                    bool has_d_min, double d_min, bool has_d_max, double d_max,
                                    float *__restrict__ run_d, int32_t *__restrict__ keep_run) {
        for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x) {
            float d = CUDART_NAN_F;
            for (int p = start[r]; p < start[r] + count[r]; ++p) {
                const float dp = d_obs[perm[p]];
                if (isfinite(dp) && dp > 0.0f) { d = dp; break; }
            }
            run_d[r] = d;
            keep_run[r] = (isfinite(d) && d > 0.0f && !(has_d_min && d < d_min) && !(has_d_max && d > d_max)) ? 1 : 0;
        }
    }

    __global__ void IngestKeepKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count,
                                     const int32_t *__restrict__ perm, const int32_t *__restrict__ keep_run,
                                     int32_t *__restrict__ keep) {
        for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x)
            if (keep_run[r])
                for (int p = start[r]; p < start[r] + count[r]; ++p)
                    keep[perm[p]] = 1;   // distinct runs own disjoint observations
    }

    __global__ void IngestFrameCountKernel(int n_frames, const int32_t *__restrict__ start,
                                           const int32_t *__restrict__ keep, int32_t *__restrict__ count) {
        for (int o = blockIdx.x * blockDim.x + threadIdx.x; o < n_frames; o += gridDim.x * blockDim.x) {
            int c = 0;
            for (int i = start[o]; i < start[o + 1]; ++i) c += keep[i];
            count[o] = c;
        }
    }

    // The surviving runs, in order, each owning its sorted members renumbered to the observations they
    // became; their raw hkl unpacked from the key.
    struct RunCut {
        int n_runs;
        const int32_t *start, *count, *keep_run, *run_idx, *run_at, *perm, *new_idx;
        const uint64_t *key;
        const float *d;
        int32_t *out_start, *out_count, *out_h, *out_k, *out_l, *out_perm;
        float *out_d;
    };
    __global__ void IngestCutRunsKernel(RunCut c) {
        for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < c.n_runs; r += gridDim.x * blockDim.x) {
            if (!c.keep_run[r]) continue;
            const int rr = c.run_idx[r], at = c.run_at[r];
            c.out_start[rr] = at;
            c.out_count[rr] = c.count[r];
            c.out_h[rr] = RawHklIndex(c.key[r], 32);
            c.out_k[rr] = RawHklIndex(c.key[r], 16);
            c.out_l[rr] = RawHklIndex(c.key[r], 0);
            c.out_d[rr] = c.d[r];
            for (int p = 0; p < c.count[r]; ++p)
                c.out_perm[at + p] = c.new_idx[c.perm[c.start[r] + p]];
        }
    }

    __global__ void IngestFillKernel(int n, int32_t *__restrict__ v) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
            v[i] = 1;
    }

    __global__ void IngestKeptCountKernel(int n_runs, const int32_t *__restrict__ keep_run,
                                          const int32_t *__restrict__ count, int32_t *__restrict__ kept) {
        for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x)
            kept[r] = keep_run[r] ? count[r] : 0;
    }


    int Blocks(int n) { return std::max(1, std::min(65535, (n + BLK - 1) / BLK)); }

    // --- the fulls' CSRs on the device (BuildFullsCSRs) ---
    // A full's group as a sort key: the fulls in no group sort after every group.
    __global__ void FullGroupKeyKernel(int n, const int32_t *__restrict__ group, int32_t *__restrict__ key) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
            key[i] = group[i] >= 0 ? group[i] : INT32_MAX;
    }

    // Counts per bin (integers, so in any order); a negative bin counts nowhere.
    __global__ void CountBinsKernel(int n, const int32_t *__restrict__ bin, int32_t *__restrict__ count) {
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
            if (bin[i] >= 0) atomicAdd(&count[bin[i]], 1);
    }

    // --- the ingest's geometry smoothing on the card (RotationScaleMerge::IngestOnDevice) ---
    // Explicit roundings throughout, so that no compiler contracts a product into a sum differently on
    // another card: the same bits on every GPU, and within the last bits of the host's own loops.
    struct Vec3 { float x, y, z; };
    __device__ __forceinline__ Vec3 VAdd(Vec3 a, Vec3 b) { return {__fadd_rn(a.x, b.x), __fadd_rn(a.y, b.y), __fadd_rn(a.z, b.z)}; }
    __device__ __forceinline__ Vec3 VScale(Vec3 a, float s) { return {__fmul_rn(a.x, s), __fmul_rn(a.y, s), __fmul_rn(a.z, s)}; }
    __device__ __forceinline__ float VDot(Vec3 a, Vec3 b) {
        return __fadd_rn(__fadd_rn(__fmul_rn(a.x, b.x), __fmul_rn(a.y, b.y)), __fmul_rn(a.z, b.z));
    }

    struct DeltaPhiGeometry {
        const float *recip;      // per frame: a*, b*, c* (9 floats)
        const uint8_t *ok;       // per frame
        Vec3 S0, m1, m2, m3;
        float m2_S0, m3_S0;
    };

    // RotationScaleMerge::SmoothedDeltaPhi / SolvePhi. One thread per raw-hkl run, which gives the hkl of
    // its observations; each observation is solved on its own.
    __global__ void IngestDeltaPhiKernel(int n_runs, const int32_t *__restrict__ start,
                                         const int32_t *__restrict__ count, const int32_t *__restrict__ perm,
                                         const int32_t *__restrict__ rr_h, const int32_t *__restrict__ rr_k,
                                         const int32_t *__restrict__ rr_l, FrameOf frame_of,
                                         DeltaPhiGeometry g, float *__restrict__ dphi,
                                         unsigned long long *__restrict__ n_resolved) {
        constexpr float PI_F = 3.14159265358979323846f;
        unsigned long long resolved = 0;
        for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x) {
            const float h = float(rr_h[r]), k = float(rr_k[r]), l = float(rr_l[r]);
            for (int q = start[r]; q < start[r] + count[r]; ++q) {
                const int i = perm[q];
                const int o = frame_of(i);
                if (!g.ok[o]) continue;
                const float *rc = g.recip + 9 * o;
                const Vec3 a{rc[0], rc[1], rc[2]}, b{rc[3], rc[4], rc[5]}, c{rc[6], rc[7], rc[8]};
                const Vec3 p0 = VAdd(VAdd(VScale(a, h), VScale(b, k)), VScale(c, l));
                const float phi_near = __fdiv_rn(__fmul_rn(dphi[i], PI_F), 180.0f);
                const float p0_sq = VDot(p0, p0);
                if (!(p0_sq > 0.0f)) continue;
                const float p0_m1 = VDot(p0, g.m1), p0_m2 = VDot(p0, g.m2), p0_m3 = VDot(p0, g.m3);
                const float rho_sq = __fsub_rn(p0_sq, __fmul_rn(p0_m2, p0_m2));
                const float p_m3 = __fdiv_rn(__fsub_rn(__fdiv_rn(-p0_sq, 2.0f), __fmul_rn(p0_m2, g.m2_S0)), g.m3_S0);
                if (!(rho_sq > __fmul_rn(p_m3, p_m3))) continue;
                const float root = __fsqrt_rn(__fsub_rn(rho_sq, __fmul_rn(p_m3, p_m3)));
                float best = 0.0f;
                bool have = false;
                for (const float p_m1 : {root, -root}) {
                    const float cosphi = __fdiv_rn(__fadd_rn(__fmul_rn(p_m1, p0_m1), __fmul_rn(p_m3, p0_m3)), rho_sq);
                    const float sinphi = __fdiv_rn(__fsub_rn(__fmul_rn(p_m1, p0_m3), __fmul_rn(p_m3, p0_m1)), rho_sq);
                    const float phi = -atan2f(sinphi, cosphi);
                    if (!have || fabsf(__fsub_rn(phi, phi_near)) < fabsf(__fsub_rn(best, phi_near))) { best = phi; have = true; }
                }
                if (!isfinite(best)) continue;
                dphi[i] = __fdiv_rn(__fmul_rn(best, 180.0f), PI_F);
                ++resolved;
            }
        }
        atomicAdd(n_resolved, resolved);
    }

    // One exact-Bragg angle per rocking event (RotationScaleMerge::SmoothMosaicityAndPartiality), one
    // thread per raw-hkl run.
    __global__ void IngestRelayKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count,
                                      const int32_t *__restrict__ perm, const float *__restrict__ img,
                                      float *__restrict__ dphi, double increment_deg, float max_frame_gap,
                                      unsigned long long *__restrict__ n_events,
                                      unsigned long long *__restrict__ n_partials) {
        unsigned long long ev = 0, pa = 0;
        for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x) {
            const int lo = start[r], hi = lo + count[r];
            for (int i = lo; i < hi; ) {
                int j = i + 1;
                while (j < hi && __fsub_rn(img[perm[j]], img[perm[j - 1]]) <= max_frame_gap)
                    ++j;
                if (j - i >= 2) {
                    const float f0 = img[perm[i]];
                    double sum = 0.0;
                    int m_n = 0;
                    for (int m = i; m < j; ++m) {
                        const int p = perm[m];
                        if (!isfinite(dphi[p])) continue;
                        sum = __dadd_rn(sum, __dadd_rn(double(dphi[p]),
                                                       __dmul_rn(increment_deg, double(__fsub_rn(img[p], f0)))));
                        ++m_n;
                    }
                    if (m_n >= 2) {
                        const double c = __ddiv_rn(sum, double(m_n));
                        for (int m = i; m < j; ++m) {
                            const int p = perm[m];
                            dphi[p] = __double2float_rn(__dsub_rn(c, __dmul_rn(increment_deg,
                                                                               double(__fsub_rn(img[p], f0)))));
                        }
                        ++ev;
                        pa += j - i;
                    }
                }
                i = j;
            }
        }
        atomicAdd(n_events, ev);
        atomicAdd(n_partials, pa);
    }

    // The frames each rocking event spans over the observations usable as integrated
    // (RotationScaleMerge::RockingEventFramesOver), counted into a histogram: integers, so the order the
    // runs are counted in does not matter.
    __global__ void IngestRockingKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count,
                                        const int32_t *__restrict__ perm, const float *__restrict__ img,
                                        const float *__restrict__ corr, const ObsRecord *__restrict__ obs,
                                        float max_frame_gap, int n_frames, int32_t *__restrict__ hist) {
        for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x) {
            bool open = false;
            float first = 0.0f, last = 0.0f;
            for (int q = start[r]; q < start[r] + count[r]; ++q) {
                const int i = perm[q];
                const float c = corr[i];
                const ObsRecord o = obs[i];
                if (!(isfinite(c) && c > 0.0f && isfinite(o.I) && isfinite(o.sigma) && o.sigma > 0.0f))
                    continue;
                const float im = img[i];
                if (open && __fsub_rn(im, last) <= max_frame_gap) {
                    last = im;
                    continue;
                }
                if (open)
                    atomicAdd(&hist[min(int(lroundf(last) - lroundf(first)) + 1, n_frames + 1)], 1);
                first = last = im;
                open = true;
            }
            if (open)
                atomicAdd(&hist[min(int(lroundf(last) - lroundf(first)) + 1, n_frames + 1)], 1);
        }
    }

    // The partiality recomputed from the frame's smoothed mosaicity, broadened by the energy bandwidth
    // (RotationScaleMerge::SmoothMosaicityAndPartiality), and corr brought to it.
    struct PartialityParams {
        const float *mos;          // per frame
        double wedge, bandwidth_sigma, half_wavelength;
    };
    __global__ void IngestPartialityKernel(int n, FrameOf frame_of, const float *__restrict__ zeta,
                                           const float *__restrict__ d, const float *__restrict__ dphi,
                                           PartialityParams q, ObsRecord *__restrict__ obs, float *__restrict__ corr) {
        constexpr double PI_D = 3.14159265358979323846;
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {
            const float z = zeta[i];
            float mos = q.mos[frame_of(i)];
            if (q.bandwidth_sigma > 0.0 && d[i] > 0.0f) {
                const double sin_theta = __ddiv_rn(q.half_wavelength, double(d[i]));
                const double sigma_bw = __ddiv_rn(__dmul_rn(__dmul_rn(__ddiv_rn(180.0, PI_D), q.bandwidth_sigma), sin_theta),
                                                  __dsqrt_rn(__dsub_rn(1.0, __dmul_rn(sin_theta, sin_theta))));
                mos = __double2float_rn(__dsqrt_rn(__dadd_rn(double(__fmul_rn(mos, mos)), __dmul_rn(sigma_bw, sigma_bw))));
            }
            if (!(isfinite(mos) && mos > 1e-6f && isfinite(z) && z > 0.0f && isfinite(dphi[i])))
                continue;
            const double half_wedge = __ddiv_rn(q.wedge, 2.0);
            const double c1 = __ddiv_rn(double(z), __dsqrt_rn(2.0));
            const double arg_plus = __ddiv_rn(__dmul_rn(__dadd_rn(double(dphi[i]), half_wedge), c1), double(mos));
            const double arg_minus = __ddiv_rn(__dmul_rn(__dsub_rn(double(dphi[i]), half_wedge), c1), double(mos));
            const float p_new = __double2float_rn(__ddiv_rn(__dsub_rn(erf(arg_plus), erf(arg_minus)), 2.0));
            const float p_old = obs[i].partiality;
            if (p_old > 0.0f && p_new > 0.0f)
                corr[i] = __fmul_rn(corr[i], __fdiv_rn(p_old, p_new));
            obs[i].partiality = p_new;
        }
    }
}

template <typename Take>
void RotationScaleMergeGPU::Impl::StageReflections(const Reflection *const *refl, const int32_t *count,
                                                   size_t nthreads, Take take) {
    int max_count = 0;
    for (int o = 0; o < n_frames; ++o)
        max_count = std::max(max_count, count[o]);
    // Where the card reads host memory directly (gpu_direct_host_read()), the kernel takes the pinned
    // buffer apart in place, without the upload.
    constexpr size_t STAGE_BYTES = 64 << 20;
    const bool direct = gpu_direct_host_read();
    const int cap = std::max(static_cast<int>(STAGE_BYTES / sizeof(Reflection)), max_count);
    CudaHostPtr<Reflection> stage[2] = {CudaHostPtr<Reflection>(cap), CudaHostPtr<Reflection>(cap)};
    CudaDevicePtr<Reflection> staged[2];
    if (!direct) {
        staged[0] = Alloc<Reflection>(cap);
        staged[1] = Alloc<Reflection>(cap);
    }
    CudaEvent copied[2];
    bool in_flight[2] = {false, false};
    int b = 0;
    for (int f0 = 0; f0 < n_frames; ) {
        int f1 = f0 + 1;
        while (f1 < n_frames && in_src_start[f1 + 1] - in_src_start[f0] <= cap)
            ++f1;
        const int base = in_src_start[f0];
        const int cnt = in_src_start[f1] - base;
        if (in_flight[b])
            CudaCheck(cudaEventSynchronize(copied[b]), "staging buffer");
        Reflection *dst = stage[b].get();
        ParallelFor(f1 - f0, nthreads, [&](int fo) {
            const int o = f0 + fo;
            std::copy(refl[o], refl[o] + count[o], dst + (in_src_start[o] - base));
        });
        if (cnt > 0) {
            const Reflection *src = dst;
            if (!direct) {
                CudaCheck(cudaMemcpyAsync(staged[b].get(), dst, size_t(cnt) * sizeof(Reflection),
                                          cudaMemcpyHostToDevice, s()), "upload reflections");
                CudaCheck(cudaEventRecord(copied[b], s()), "upload reflections");
                src = staged[b].get();
            }
            in_flight[b] = true;
            take(cnt, base, src);
            // Read in place, the buffer is free again only once the kernel has run.
            if (direct)
                CudaCheck(cudaEventRecord(copied[b], s()), "extract reflections");
        }
        b ^= 1;
        f0 = f1;
    }
    CudaCheck(cudaStreamSynchronize(s()), "upload reflections");
}

void RotationScaleMergeGPU::IngestUpload(int n_frames, const Reflection *const *refl, const int32_t *count,
                                         size_t nthreads) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    d.n_frames = n_frames;
    d.in_src_start.assign(n_frames + 1, 0);
    for (int o = 0; o < n_frames; ++o)
        d.in_src_start[o + 1] = d.in_src_start[o] + count[o];
    const int n = d.in_src_start[n_frames];
    d.n_obs = n;
    size_t free_bytes = 0, total_bytes = 0;
    CudaCheck(cudaMemGetInfo(&free_bytes, &total_bytes), "memory info");
    if (size_t(n) * OBS_BYTES > total_bytes)
        ThrowTooLargeForGPU(n);
    // See AllocObservations: these are the first of the scaling's buffers.
    cuda_trim_allocation_pool();
    // Only what the sort and the cut read; the observations are taken from the reflections once the cut
    // has decided which they are (IngestObservations), so the two are never on the card together.
    const int na = std::max(1, n);
    d.in_img = d.Alloc<float>(na); d.in_d = d.Alloc<float>(na); d.in_bkg = d.Alloc<float>(na);
    d.in_key = d.Alloc<uint64_t>(na);
    const IngestKeys a{d.in_img.get(), d.in_d.get(), d.in_bkg.get(), d.in_key.get()};
    d.StageReflections(refl, count, nthreads, [&](int cnt, int base, const Reflection *src) {
        IngestKeyKernel<<<(cnt + BLK - 1) / BLK, BLK, 0, d.s()>>>(cnt, base, src, a);
        CudaCheck(cudaGetLastError(), "IngestKeyKernel launch");
    });
}

void RotationScaleMergeGPU::IngestMeanBackground(double *mean_bkg) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    CudaDevicePtr<int32_t> start;
    d.Upload(start, d.in_src_start.data(), d.n_frames + 1);
    CudaDevicePtr<double> mean = d.Alloc<double>(std::max(1, d.n_frames));
    IngestMeanBkgKernel<<<Blocks(d.n_frames), BLK, 0, d.s()>>>(d.n_frames, start.get(), d.in_bkg.get(), mean.get());
    CudaCheck(cudaGetLastError(), "IngestMeanBkgKernel launch");
    if (d.n_frames > 0)
        CopyAndWait(mean_bkg, mean.get(), size_t(d.n_frames) * sizeof(double), cudaMemcpyDeviceToHost, d.s(),
                    "download mean background");
    d.in_bkg = {};
}

int RotationScaleMergeGPU::IngestSort(std::optional<double> d_min, std::optional<double> d_max) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int n = d.n_obs;
    const int na = std::max(1, n);

    // (h, k, l, image_number, index) is a total order, so two stable radix passes - by the frame position
    // with the index as the value, then by the raw hkl - give the one permutation std::sort gives on the
    // host, whatever either algorithm does on the way. Each pass goes back and forth between the two
    // buffers of its keys and of its values, as neither input is needed after it, so that cub holds no
    // copy of them of its own: the sorts are where the ingest's footprint peaks. The same stable sort, so
    // the same order; it is left in key_out and val_out.
    const auto sort_pairs = [&](auto &key_in, auto &key_out, CudaDevicePtr<int32_t> &val_in,
                                CudaDevicePtr<int32_t> &val_out, int end_bit) {
        using K = std::remove_reference_t<decltype(*key_in.get())>;
        cub::DoubleBuffer<K> keys(key_in.get(), key_out.get());
        cub::DoubleBuffer<int32_t> values(val_in.get(), val_out.get());
        size_t bytes = 0;
        CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, keys, values, n, 0, end_bit, d.s()),
                  "sort temp size");
        CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
        CudaCheck(cub::DeviceRadixSort::SortPairs(temp.get(), bytes, keys, values, n, 0, end_bit, d.s()), "sort");
        CudaCheck(cudaStreamSynchronize(d.s()), "sort");
        if (keys.Current() != key_out.get()) std::swap(key_in, key_out);
        if (values.Current() != val_out.get()) std::swap(val_in, val_out);
    };
    CudaDevicePtr<int32_t> by_image = d.Alloc<int32_t>(na);
    {
        CudaDevicePtr<uint32_t> ikey = d.Alloc<uint32_t>(na), ikey_sorted = d.Alloc<uint32_t>(na);
        CudaDevicePtr<int32_t> idx = d.Alloc<int32_t>(na);
        IngestImageKeyKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, d.in_img.get(), ikey.get(), idx.get());
        CudaCheck(cudaGetLastError(), "IngestImageKeyKernel launch");
        if (n > 0) sort_pairs(ikey, ikey_sorted, idx, by_image, 32);
    }
    d.in_img = {};
    CudaDevicePtr<uint64_t> key = d.Alloc<uint64_t>(na);
    IngestGatherKeyKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, by_image.get(), d.in_key.get(), key.get());
    CudaCheck(cudaGetLastError(), "IngestGatherKeyKernel launch");
    // The runs the sort cuts carry the raw hkl from here on.
    CudaCheck(cudaStreamSynchronize(d.s()), "gather key");
    d.in_key = {};
    CudaDevicePtr<int32_t> perm = d.Alloc<int32_t>(na);
    CudaDevicePtr<uint64_t> key_sorted = d.Alloc<uint64_t>(na);
    if (n > 0) sort_pairs(key, key_sorted, by_image, perm, 48);
    key = {}; by_image = {};

    // The raw-hkl runs of the sorted order.
    int n_runs = 0;
    CudaDevicePtr<uint64_t> run_key = d.Alloc<uint64_t>(na);
    CudaDevicePtr<int32_t> run_count = d.Alloc<int32_t>(na);
    if (n > 0) {
        CudaDevicePtr<int32_t> n_runs_dev = d.Alloc<int32_t>(1);
        size_t bytes = 0;
        CudaCheck(cub::DeviceRunLengthEncode::Encode(nullptr, bytes, key_sorted.get(), run_key.get(),
                                                     run_count.get(), n_runs_dev.get(), n, d.s()), "runs temp size");
        CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
        CudaCheck(cub::DeviceRunLengthEncode::Encode(temp.get(), bytes, key_sorted.get(), run_key.get(),
                                                     run_count.get(), n_runs_dev.get(), n, d.s()), "runs");
        CopyAndWait(&n_runs, n_runs_dev.get(), sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "run count");
    }
    key_sorted = CudaDevicePtr<uint64_t>();
    const int nra = std::max(1, n_runs);
    {
        // Cut down to the runs there are: the encode had to be given room for one per reflection.
        CudaDevicePtr<uint64_t> key_exact = d.Alloc<uint64_t>(nra);
        CudaDevicePtr<int32_t> count_exact = d.Alloc<int32_t>(nra);
        if (n_runs > 0) {
            CopyAndWait(key_exact.get(), run_key.get(), size_t(n_runs) * sizeof(uint64_t), cudaMemcpyDeviceToDevice,
                        d.s(), "runs");
            CopyAndWait(count_exact.get(), run_count.get(), size_t(n_runs) * sizeof(int32_t), cudaMemcpyDeviceToDevice,
                        d.s(), "runs");
        }
        run_key = std::move(key_exact);
        run_count = std::move(count_exact);
    }
    const auto exclusive_sum = [&](const int32_t *in, int32_t *out, int count) {
        if (count == 0) return;
        size_t bytes = 0;
        CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, in, out, count, d.s()), "scan temp size");
        CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
        CudaCheck(cub::DeviceScan::ExclusiveSum(temp.get(), bytes, in, out, count, d.s()), "scan");
        CudaCheck(cudaStreamSynchronize(d.s()), "scan");
    };
    CudaDevicePtr<int32_t> run_start = d.Alloc<int32_t>(nra);
    exclusive_sum(run_count.get(), run_start.get(), n_runs);
    CudaDevicePtr<float> run_d = d.Alloc<float>(nra);
    CudaDevicePtr<int32_t> keep_run = d.Alloc<int32_t>(nra);
    const bool limited = d_min.has_value() || d_max.has_value();
    IngestRunKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, run_start.get(), run_count.get(), perm.get(),
                                                       d.in_d.get(), d_min.has_value(), d_min.value_or(0.0),
                                                       d_max.has_value(), d_max.value_or(0.0),
                                                       run_d.get(), keep_run.get());
    CudaCheck(cudaGetLastError(), "IngestRunKernel launch");

    // Which observations survive (all of them without a range, as the host keeps them), and where each
    // run and each observation lands.
    CudaDevicePtr<int32_t> keep = d.Alloc<int32_t>(na);
    if (limited) {
        CudaCheck(cudaMemsetAsync(keep.get(), 0, size_t(na) * sizeof(int32_t), d.s()), "memset keep");
        IngestKeepKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, run_start.get(), run_count.get(), perm.get(),
                                                            keep_run.get(), keep.get());
        CudaCheck(cudaGetLastError(), "IngestKeepKernel launch");
    } else {
        IngestFillKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, keep.get());
        IngestFillKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, keep_run.get());
        CudaCheck(cudaGetLastError(), "IngestFillKernel launch");
    }
    CudaDevicePtr<int32_t> new_idx = d.Alloc<int32_t>(na);
    exclusive_sum(keep.get(), new_idx.get(), n);
    CudaDevicePtr<int32_t> run_idx = d.Alloc<int32_t>(nra), run_at = d.Alloc<int32_t>(nra);
    exclusive_sum(keep_run.get(), run_idx.get(), n_runs);
    {
        CudaDevicePtr<int32_t> kept_count = d.Alloc<int32_t>(nra);
        IngestKeptCountKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, keep_run.get(), run_count.get(),
                                                                 kept_count.get());
        CudaCheck(cudaGetLastError(), "IngestKeptCountKernel launch");
        exclusive_sum(kept_count.get(), run_at.get(), n_runs);
    }
    int32_t last_keep = 0, last_new_idx = 0, last_keep_run = 0, last_run_idx = 0;
    if (n > 0) {
        CopyAndWait(&last_keep, keep.get() + n - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "count");
        CopyAndWait(&last_new_idx, new_idx.get() + n - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "count");
    }
    if (n_runs > 0) {
        CopyAndWait(&last_keep_run, keep_run.get() + n_runs - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
                    "count");
        CopyAndWait(&last_run_idx, run_idx.get() + n_runs - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
                    "count");
    }
    const int n_keep = last_new_idx + last_keep;
    const int n_keep_run = last_run_idx + last_keep_run;

    d.n_runs = n_keep_run;
    d.n_perm = n_keep;
    const int nka = std::max(1, n_keep), nkra = std::max(1, n_keep_run);
    d.perm = d.Alloc<int32_t>(nka);
    d.rr_start = d.Alloc<int32_t>(nkra); d.rr_count = d.Alloc<int32_t>(nkra);
    d.rr_h = d.Alloc<int32_t>(nkra); d.rr_k = d.Alloc<int32_t>(nkra); d.rr_l = d.Alloc<int32_t>(nkra);
    CudaDevicePtr<float> rr_d = d.Alloc<float>(nkra);
    RunCut c{n_runs, run_start.get(), run_count.get(), keep_run.get(), run_idx.get(), run_at.get(), perm.get(),
             new_idx.get(), run_key.get(), run_d.get(),
             d.rr_start.get(), d.rr_count.get(), d.rr_h.get(), d.rr_k.get(), d.rr_l.get(), d.perm.get(), rr_d.get()};
    IngestCutRunsKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(c);
    CudaCheck(cudaGetLastError(), "IngestCutRunsKernel launch");
    CudaCheck(cudaStreamSynchronize(d.s()), "cut runs");
    // The runs are cut: what was kept of them is in rr_*, so the uncut ones and the old order go before
    // the observations are cut below, one field held twice at a time.
    perm = {}; run_key = {}; run_count = {}; run_start = {}; run_d = {}; keep_run = {}; run_idx = {}; run_at = {};
    d.in_d = {};

    // The frame CSR of what survived.
    auto &L = d.in_layout;
    L = RotationScaleMergeGPU::IngestLayout{};
    L.n_reflections = n;
    L.n_runs_all = n_runs;
    L.frame_count.assign(d.n_frames, 0);
    {
        CudaDevicePtr<int32_t> start, fcount = d.Alloc<int32_t>(std::max(1, d.n_frames));
        d.Upload(start, d.in_src_start.data(), d.n_frames + 1);
        IngestFrameCountKernel<<<Blocks(d.n_frames), BLK, 0, d.s()>>>(d.n_frames, start.get(), keep.get(),
                                                                      fcount.get());
        CudaCheck(cudaGetLastError(), "IngestFrameCountKernel launch");
        if (d.n_frames > 0)
            CopyAndWait(L.frame_count.data(), fcount.get(), size_t(d.n_frames) * sizeof(int32_t),
                        cudaMemcpyDeviceToHost, d.s(), "download frame count");
    }
    L.frame_start.assign(d.n_frames, 0);
    for (int o = 0, at = 0; o < d.n_frames; ++o) { L.frame_start[o] = at; at += L.frame_count[o]; }

    // Where each observation goes, for IngestObservations to take them from the reflections.
    d.in_new_idx = std::move(new_idx);
    d.in_n_keep = n_keep;
    d.n_obs = n_keep;
    d.AllocFrames(L.frame_start.data(), L.frame_count.data());

    L.run_start.resize(n_keep_run); L.run_count.resize(n_keep_run);
    L.run_h.resize(n_keep_run); L.run_k.resize(n_keep_run); L.run_l.resize(n_keep_run); L.run_d.resize(n_keep_run);
    const auto get = [&](auto &v, const auto &src) {
        if (!v.empty())
            CopyAndWait(v.data(), src.get(), v.size() * sizeof(v[0]), cudaMemcpyDeviceToHost, d.s(), "download runs");
    };
    get(L.run_start, d.rr_start); get(L.run_count, d.rr_count);
    get(L.run_h, d.rr_h); get(L.run_k, d.rr_k); get(L.run_l, d.rr_l); get(L.run_d, rr_d);
    return n_keep;
}

void RotationScaleMergeGPU::IngestGetLayout(IngestLayout &out) const {
    out = impl_->in_layout;
}

void RotationScaleMergeGPU::IngestObservations(const Reflection *const *refl, const int32_t *count, size_t nthreads,
                                               const double *frame_flux, float *px, float *py) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int n_keep = d.in_n_keep;
    d.AllocObservations(n_keep);
    d.in_dphi = d.Alloc<float>(std::max(1, n_keep));
    const IngestArrays a{d.obs.get(), d.zeta.get(), d.corr.get(), d.bkg.get(), d.var_bkg.get(), d.image_number.get(),
                         d.d_obs.get(), d.in_dphi.get(), d.flags.get(), d.in_new_idx.get(),
                         d.in_src_start[d.n_frames], n_keep};
    d.StageReflections(refl, count, nthreads, [&](int cnt, int base, const Reflection *src) {
        IngestExtractKernel<<<(cnt + BLK - 1) / BLK, BLK, 0, d.s()>>>(cnt, base, src, a);
        CudaCheck(cudaGetLastError(), "IngestExtractKernel launch");
    });
    // The predicted positions, which only the fulls' and only on the host read, stay there.
    const int n = d.in_src_start[d.n_frames];
    std::vector<int32_t> new_idx(n);
    if (n > 0)
        CopyAndWait(new_idx.data(), d.in_new_idx.get(), size_t(n) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
                    "download places");
    d.in_new_idx = {};
    ParallelFor(d.n_frames, nthreads, [&](int o) {
        for (int j = 0; j < count[o]; ++j) {
            const int src = d.in_src_start[o] + j, i = new_idx[src];
            if ((src + 1 < n ? new_idx[src + 1] : n_keep) == i) continue;
            px[i] = refl[o][j].predicted_x;
            py[i] = refl[o][j].predicted_y;
        }
    });

    // The kept observations are in frame order, as the reflections were (IngestSort laid out the CSR).
    CudaDevicePtr<double> flux;
    d.Upload(flux, frame_flux, d.n_frames);
    IngestFluxKernel<<<Blocks(n_keep), BLK, 0, d.s()>>>(n_keep, d.PartialFrame(), flux.get(), d.obs.get());
    CudaCheck(cudaGetLastError(), "IngestFluxKernel launch");
    CudaCheck(cudaStreamSynchronize(d.s()), "flux");
}

int64_t RotationScaleMergeGPU::IngestDeltaPhi(const DeltaPhiFrame &g) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    CudaDevicePtr<float> recip;
    CudaDevicePtr<uint8_t> ok;
    d.Upload(recip, g.recip.data(), static_cast<int>(g.recip.size()));
    d.Upload(ok, g.ok.data(), static_cast<int>(g.ok.size()));
    CudaDevicePtr<unsigned long long> n_resolved = d.Alloc<unsigned long long>(1);
    CudaCheck(cudaMemsetAsync(n_resolved.get(), 0, sizeof(unsigned long long), d.s()), "memset");
    const auto v = [](const float *a) { return Vec3{a[0], a[1], a[2]}; };
    const DeltaPhiGeometry dg{recip.get(), ok.get(), v(g.S0), v(g.m1), v(g.m2), v(g.m3), g.m2_S0, g.m3_S0};
    IngestDeltaPhiKernel<<<Blocks(d.n_runs), BLK, 0, d.s()>>>(d.n_runs, d.rr_start.get(), d.rr_count.get(),
                                                              d.perm.get(), d.rr_h.get(), d.rr_k.get(), d.rr_l.get(),
                                                              d.PartialFrame(), dg, d.in_dphi.get(), n_resolved.get());
    CudaCheck(cudaGetLastError(), "IngestDeltaPhiKernel launch");
    unsigned long long out = 0;
    CopyAndWait(&out, n_resolved.get(), sizeof(out), cudaMemcpyDeviceToHost, d.s(), "resolved count");
    return static_cast<int64_t>(out);
}

void RotationScaleMergeGPU::IngestRelayEvents(double increment_deg, float max_frame_gap, int64_t &n_events,
                                              int64_t &n_partials) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    CudaDevicePtr<unsigned long long> counts = d.Alloc<unsigned long long>(2);
    CudaCheck(cudaMemsetAsync(counts.get(), 0, 2 * sizeof(unsigned long long), d.s()), "memset");
    IngestRelayKernel<<<Blocks(d.n_runs), BLK, 0, d.s()>>>(d.n_runs, d.rr_start.get(), d.rr_count.get(), d.perm.get(),
                                                          d.image_number.get(), d.in_dphi.get(), increment_deg,
                                                          max_frame_gap, counts.get(), counts.get() + 1);
    CudaCheck(cudaGetLastError(), "IngestRelayKernel launch");
    unsigned long long out[2] = {0, 0};
    CopyAndWait(out, counts.get(), sizeof(out), cudaMemcpyDeviceToHost, d.s(), "relay counts");
    n_events = static_cast<int64_t>(out[0]);
    n_partials = static_cast<int64_t>(out[1]);
}

void RotationScaleMergeGPU::IngestRockingHistogram(float max_frame_gap, int32_t *hist) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int n_bins = d.n_frames + 2;
    CudaDevicePtr<int32_t> h = d.Alloc<int32_t>(n_bins);
    CudaCheck(cudaMemsetAsync(h.get(), 0, size_t(n_bins) * sizeof(int32_t), d.s()), "memset");
    IngestRockingKernel<<<Blocks(d.n_runs), BLK, 0, d.s()>>>(d.n_runs, d.rr_start.get(), d.rr_count.get(), d.perm.get(),
                                                            d.image_number.get(), d.corr.get(), d.obs.get(),
                                                            max_frame_gap, d.n_frames, h.get());
    CudaCheck(cudaGetLastError(), "IngestRockingKernel launch");
    CopyAndWait(hist, h.get(), size_t(n_bins) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "rocking histogram");
}

void RotationScaleMergeGPU::IngestPartiality(const float *mosaicity, double wedge_deg, double bandwidth_sigma,
                                             double half_wavelength_A) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    CudaDevicePtr<float> mos;
    d.Upload(mos, mosaicity, d.n_frames);
    const PartialityParams q{mos.get(), wedge_deg, bandwidth_sigma, half_wavelength_A};
    IngestPartialityKernel<<<Blocks(d.n_obs), BLK, 0, d.s()>>>(d.n_obs, d.PartialFrame(), d.zeta.get(), d.d_obs.get(),
                                                               d.in_dphi.get(), q, d.obs.get(), d.corr.get());
    CudaCheck(cudaGetLastError(), "IngestPartialityKernel launch");
    CudaCheck(cudaStreamSynchronize(d.s()), "partiality");
}

void RotationScaleMergeGPU::IngestFinish(float *corr) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    if (d.n_obs > 0)
        CopyAndWait(corr, d.corr.get(), size_t(d.n_obs) * sizeof(float), cudaMemcpyDeviceToHost, d.s(),
                    "download corr");
    d.in_dphi = CudaDevicePtr<float>();
    d.in_src_start.clear();
    d.in_layout = RotationScaleMergeGPU::IngestLayout{};
}

void RotationScaleMergeGPU::BuildFullsCSRs(int32_t *frame_count, std::vector<int32_t> &gstart,
                                           std::vector<int32_t> &gperm) {
    DeviceGuard guard(impl_->device, impl_->available);
    auto &d = *impl_;
    const int nf = d.n_fulls, nfr = d.n_frames, ng = d.n_groups;
    const int nfa = std::max(1, nf);
    // The index order of the keys, stably sorted, left in val_out: back and forth between the two
    // buffers of each, as no input is needed after it, so that cub holds no copy of its own.
    const auto stable_sort = [&](CudaDevicePtr<int32_t> &key_in, CudaDevicePtr<int32_t> &key_out,
                                 CudaDevicePtr<int32_t> &val_in, CudaDevicePtr<int32_t> &val_out) {
        if (nf > 0) IotaKernel<<<(nf + BLK - 1) / BLK, BLK, 0, d.s()>>>(nf, val_in.get());
        CudaCheck(cudaGetLastError(), "IotaKernel launch");
        cub::DoubleBuffer<int32_t> keys(key_in.get(), key_out.get()), values(val_in.get(), val_out.get());
        size_t bytes = 0;
        CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, keys, values, nf, 0, 32, d.s()), "sort temp size");
        CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
        CudaCheck(cub::DeviceRadixSort::SortPairs(temp.get(), bytes, keys, values, nf, 0, 32, d.s()), "sort");
        CudaCheck(cudaStreamSynchronize(d.s()), "sort");
        if (values.Current() != val_out.get()) std::swap(val_in, val_out);
    };
    // The counts come back for the starts, a prefix the host forms in order.
    const auto count = [&](const int32_t *bin, int n_bins, CudaDevicePtr<int32_t> &out, std::vector<int32_t> &host) {
        out = d.Alloc<int32_t>(std::max(1, n_bins));
        CudaCheck(cudaMemsetAsync(out.get(), 0, size_t(std::max(1, n_bins)) * sizeof(int32_t), d.s()), "memset counts");
        CountBinsKernel<<<Blocks(nf), BLK, 0, d.s()>>>(nf, bin, out.get());
        CudaCheck(cudaGetLastError(), "CountBinsKernel launch");
        host.resize(n_bins);
        if (n_bins > 0)
            CopyAndWait(host.data(), out.get(), size_t(n_bins) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
                        "download counts");
    };
    const auto prefix = [](const std::vector<int32_t> &c) {
        std::vector<int32_t> start(c.size() + 1, 0);
        for (size_t i = 0; i < c.size(); ++i) start[i + 1] = start[i] + c[i];
        return start;
    };
    CudaDevicePtr<int32_t> idx = d.Alloc<int32_t>(nfa), key = d.Alloc<int32_t>(nfa), sorted = d.Alloc<int32_t>(nfa);

    std::vector<int32_t> fcount;
    count(d.f_frame.get(), nfr, d.f_frame_count, fcount);
    std::copy(fcount.begin(), fcount.end(), frame_count);
    const std::vector<int32_t> fstart = prefix(fcount);
    d.Upload(d.f_frame_start, fstart.data(), nfr);
    d.f_frame_perm = d.Alloc<int32_t>(nfa);
    if (nf > 0) {
        CopyAndWait(key.get(), d.f_frame.get(), size_t(nf) * sizeof(int32_t), cudaMemcpyDeviceToDevice, d.s(),
                    "frame key");
        stable_sort(key, sorted, idx, d.f_frame_perm);
    }

    std::vector<int32_t> gcount;
    count(d.f_group.get(), ng, d.f_gcount, gcount);
    gstart = prefix(gcount);
    d.Upload(d.f_gstart, gstart.data(), ng);
    const int n_grouped = gstart.back();
    FullGroupKeyKernel<<<Blocks(nf), BLK, 0, d.s()>>>(nf, d.f_group.get(), key.get());
    CudaCheck(cudaGetLastError(), "FullGroupKeyKernel launch");
    CudaDevicePtr<int32_t> by_group = d.Alloc<int32_t>(nfa);
    if (nf > 0) stable_sort(key, sorted, idx, by_group);
    d.f_gperm = d.Alloc<int32_t>(std::max(1, n_grouped));
    gperm.resize(n_grouped);
    if (n_grouped > 0) {
        CopyAndWait(d.f_gperm.get(), by_group.get(), size_t(n_grouped) * sizeof(int32_t), cudaMemcpyDeviceToDevice,
                    d.s(), "group perm");
        CopyAndWait(gperm.data(), by_group.get(), size_t(n_grouped) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
                    "download group perm");
    }
}