rotation merge: the fulls' frame and group CSRs built on the card

Every merge downloaded the fulls' frame and group keys and built both CSRs by serial counting
sorts on the host - four passes over the fulls, two of them scattering into an array as long
as the distinct hkl (15 M on 4nwv) - and uploaded them again. They are now stable radix sorts
of the resident fulls on the card, with the counts summed there; the host gets the counts per
frame and the group CSR it walks itself. The same permutations (stable sorts list each frame's
and each group's fulls in index order, as the counting sorts did): p.hkl identical on cytc
and 4nwv.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3

image_analysis/scale_merge/RotationScaleMerge.cpp

@@ -6619,26 +6619,11 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
         g_full.assign(n_frames, 1.0);
 
         if (scale_fulls && nf > 0) {
-            // Frame + group CSRs over the emit-ordered fulls, built by counting sort on the host (stable,
-            // deterministic). frame is always in [0, n_frames); group is <0 for absent/out-of-range fulls.
-            std::vector<int32_t> ff(nf), fg(nf);
-            gpu_->GetFullsKeys(ff.data(), fg.data());
-            std::vector<int32_t> f_start(n_frames, 0), f_count(n_frames, 0), f_perm(nf);
-            for (int i = 0; i < nf; ++i) ++f_count[ff[i]];
-            for (int f = 1; f < n_frames; ++f) f_start[f] = f_start[f - 1] + f_count[f - 1];
-            { std::vector<int32_t> fill = f_start; for (int i = 0; i < nf; ++i) f_perm[fill[ff[i]]++] = i; }
-            gpu_->SetFullsFrameCSR(f_perm.data(), nf, f_start.data(), f_count.data());
-
-            std::vector<int32_t> g_count(n_groups, 0), g_start(n_groups, 0);
-            for (int i = 0; i < nf; ++i) if (fg[i] >= 0) ++g_count[fg[i]];
-            int acc = 0;
-            for (int g = 0; g < n_groups; ++g) { g_start[g] = acc; acc += g_count[g]; }
-            std::vector<int32_t> g_perm(acc);
-            { std::vector<int32_t> fill = g_start; for (int i = 0; i < nf; ++i) if (fg[i] >= 0) g_perm[fill[fg[i]]++] = i; }
-            gpu_->SetFullsGroups(g_perm.data(), acc, g_start.data(), g_count.data());
-            fulls_group_start = g_start;
-            fulls_group_start.push_back(acc);
-            fulls_group_perm = std::move(g_perm);
+            // Frame + group CSRs over the emit-ordered fulls, built on the device by stable sorts: each
+            // frame's and each group's fulls in index order. frame is always in [0, n_frames); group is <0
+            // for absent/out-of-range fulls, which no group lists.
+            std::vector<int32_t> f_count(n_frames);
+            gpu_->BuildFullsCSRs(f_count.data(), fulls_group_start, fulls_group_perm);
             fulls_csr_built = true;
 
             gpu_->ResetFullsScale();

image_analysis/scale_merge/RotationScaleMergeGPU.cu

@@ -2034,33 +2034,6 @@ void RotationScaleMergeGPU::GetFulls(int32_t *h, int32_t *k, int32_t *l, float *
     dl(clipped, dd.f_clipped.get(), n * sizeof(uint8_t));
 }
 
-void RotationScaleMergeGPU::GetFullsKeys(int32_t *frame, int32_t *group) 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(int32_t);
-    CopyAndWait(frame, d.f_frame.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_frame");
-    CopyAndWait(group, d.f_group.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_group");
-}
-
-void RotationScaleMergeGPU::SetFullsFrameCSR(const int32_t *frame_perm, int n_perm,
-                                             const int32_t *frame_start, const int32_t *frame_count) {
-    DeviceGuard guard(impl_->device, impl_->available);
-    auto &d = *impl_;
-    d.Upload(d.f_frame_perm, frame_perm, n_perm);
-    d.Upload(d.f_frame_start, frame_start, d.n_frames);
-    d.Upload(d.f_frame_count, frame_count, d.n_frames);
-}
-
-void RotationScaleMergeGPU::SetFullsGroups(const int32_t *gperm, int n_gperm,
-                                           const int32_t *gstart, const int32_t *gcount) {
-    DeviceGuard guard(impl_->device, impl_->available);
-    auto &d = *impl_;
-    d.Upload(d.f_gperm, gperm, n_gperm);
-    d.Upload(d.f_gstart, gstart, d.n_groups);
-    d.Upload(d.f_gcount, gcount, d.n_groups);
-}
-
 void RotationScaleMergeGPU::ResetFullsScale() {
     DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
@@ -2611,6 +2584,19 @@ namespace {
 
     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.
@@ -3118,3 +3104,65 @@ void RotationScaleMergeGPU::IngestFinish(float *corr) {
     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);
+    const auto stable_sort = [&](const int32_t *key_in, int32_t *key_out, const int32_t *val_in, int32_t *val_out) {
+        size_t bytes = 0;
+        CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key_in, key_out, val_in, val_out, 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, key_in, key_out, val_in, val_out, nf, 0, 32, d.s()),
+                  "sort");
+        CudaCheck(cudaStreamSynchronize(d.s()), "sort");
+    };
+    // 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);
+    if (nf > 0) IotaKernel<<<(nf + BLK - 1) / BLK, BLK, 0, d.s()>>>(nf, idx.get());
+    CudaCheck(cudaGetLastError(), "IotaKernel launch");
+
+    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) stable_sort(d.f_frame.get(), key.get(), idx.get(), d.f_frame_perm.get());
+
+    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.get(), sorted.get(), idx.get(), by_group.get());
+    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");
+    }
+}

image_analysis/scale_merge/RotationScaleMergeGPU.h

@@ -243,18 +243,11 @@ public:
 
     // --- scale the resident fulls on the device (Unity model), no round-trip ---
 
-    // Download the fulls' frame and ASU-group keys (emit order) so the host can build the frame/group CSRs
-    // with a counting sort (deterministic, no GPU stable-sort) and hand them back below.
-    void GetFullsKeys(int32_t *frame, int32_t *group) const;
-
-    // The fulls' per-frame CSR: frame_perm groups the emit-ordered fulls by frame (frame_start/count length
-    // n_frames index it), so FitPerFrameG can scale the fulls without physically reordering them.
-    void SetFullsFrameCSR(const int32_t *frame_perm, int n_perm,
-                          const int32_t *frame_start, const int32_t *frame_count);
-
-    // The fulls' per-ASU-group CSR (group-ordered permutation of the fulls with group>=0, + its CSR).
-    void SetFullsGroups(const int32_t *gperm, int n_gperm,
-                        const int32_t *gstart, const int32_t *gcount);
+    // The fulls' per-frame CSR (each frame's fulls, for the per-frame fit) and per-ASU-group CSR, built on
+    // the device by stable sorts of the emit-ordered fulls, so that each frame and each group lists its
+    // fulls in index order as a counting sort does. Hands back the fulls per frame (length n_frames) and
+    // the group CSR for the host's own walks: gstart (n_groups + 1 entries, the last the total) and gperm.
+    void BuildFullsCSRs(int32_t *frame_count, std::vector<int32_t> &gstart, std::vector<int32_t> &gperm);
 
     // Start the Unity scaling of the resident fulls: corr, partiality, prescaling and zeta all 1, no
     // frame fitted yet. Once before the FitFullsScale rounds.
@@ -262,7 +255,7 @@ public:
 
     // One Unity fit of the resident fulls (reduce group means -> per-frame LS G and its information,
     // GetG / GetInfo), leaving the fulls' working corr as it is: the host smooths the fitted scales
-    // and applies them with SmoothFullsCorr. Requires SetFullsFrameCSR + SetFullsGroups and
+    // and applies them with SmoothFullsCorr. Requires BuildFullsCSRs and
     // ResetFullsScale.
     void FitFullsScale(double min_partiality);