RotationScaleMerge: GPU partial-scaling loop (CUDA port, phase 1)

First stage of moving the rotation scale/merge onto the GPU. The per-frame partial-scaling loop
(inverse-variance group-mean reduction -> robust per-frame IRLS G -> corr update, x scaling_iter)
now runs in RotationScaleMergeGPU (.cu) when a GPU is present; the CPU loops remain the fallback.

The host keeps the one-time raw-hkl sort and the per-space-group gemmi ASU keying, and hands the
GPU a group-ordered permutation + CSR so the per-group reduction is a DETERMINISTIC segmented
reduction (one thread per group, fixed order, no atomics) - preserving the run-to-run determinism
just won on the CPU path (a float atomicAdd reduction would have re-introduced jitter). Reduction is
one-thread-per-group (groups average tens of obs, so a block-per-group wastes threads); the IRLS is
one block per frame with a deterministic shared-memory reduction.

Validated: bit-identical to the CPU path and deterministic run-to-run on lyso/cytC/Ins_H/pding
(P41212 ISa 7.8 CC1/2 99.7%, etc.). The scaling kernels are ~7x faster than the CPU compute
(~36 ms for 3 iters vs ~0.28 s); end-to-end scale/merge ~2.0 -> ~1.5 s. The remaining gap to the
<1 s target is the per-pass host round-trip (corr down/upload for the CPU combine + per-SG group-CSR
rebuild); phase 2 keeps the data resident by moving the 3D combine and the merge/error-model onto
the GPU too, so nothing round-trips.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

image_analysis/scale_merge/RotationScaleMerge.cpp

@@ -238,6 +238,27 @@ void RotationScaleMerge::Ingest() {
                 total, n_frames, rawrun_start.size());
 
     SmoothMosaicityAndPartiality();
+
+#ifdef JFJOCH_USE_CUDA
+    // Bring the partial-scaling loop onto the GPU when one is present. Upload the immutable per-obs
+    // fields once (corr lives on the device, refreshed each pass); the CPU keeps the sort/keying/combine.
+    gpu_ = std::make_unique<RotationScaleMergeGPU>();
+    gpu_active_ = gpu_->Available();
+    if (gpu_active_) {
+        const int n = static_cast<int>(partials.size());
+        std::vector<float> I(n), sigma(n), rlp(n), part(n), zeta(n), corr(n);
+        std::vector<uint8_t> onice(n);
+        std::vector<int32_t> frm(n);
+        for (int i = 0; i < n; ++i) {
+            const auto &o = partials[i];
+            I[i] = o.I; sigma[i] = o.sigma; rlp[i] = o.rlp; part[i] = o.partiality;
+            zeta[i] = o.zeta; onice[i] = o.on_ice; frm[i] = o.frame; corr[i] = o.corr;
+        }
+        gpu_->SetPartials(n, n_frames, I.data(), sigma.data(), rlp.data(), part.data(), zeta.data(),
+                          onice.data(), frm.data(), corr.data(), frame_start.data(), frame_count.data());
+        logger.Info("RotationScaleMerge: GPU partial-scaling active");
+    }
+#endif
 }
 
 void RotationScaleMerge::SmoothMosaicityAndPartiality() {
@@ -340,6 +361,28 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) {
             }
         }
     });
+
+#ifdef JFJOCH_USE_CUDA
+    // Group-ordered permutation (obs bucketed by ASU group, obs-index order) + its CSR, so the GPU
+    // reduction is a deterministic segmented reduction (fixed order, no atomics). Built per space group.
+    if (gpu_active_) {
+        const int n = static_cast<int>(partials.size());
+        std::vector<int32_t> group_ids(n), gcount(n_groups, 0);
+        for (int i = 0; i < n; ++i) {
+            group_ids[i] = partials[i].group;
+            if (partials[i].group >= 0) ++gcount[partials[i].group];
+        }
+        std::vector<int32_t> gstart(n_groups, 0);
+        int acc = 0;
+        for (int g = 0; g < n_groups; ++g) { gstart[g] = acc; acc += gcount[g]; }
+        std::vector<int32_t> gperm(acc), gfill = gstart;
+        for (int i = 0; i < n; ++i) {
+            const int g = partials[i].group;
+            if (g >= 0) gperm[gfill[g]++] = i;
+        }
+        gpu_->SetGroups(n_groups, group_ids.data(), gperm.data(), acc, gstart.data(), gcount.data());
+    }
+#endif
     return n_groups;
 }
 
@@ -1000,10 +1043,28 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search,
     const int n_groups = ComputeAsuGroups(keygen);   // one ASU grouping, shared by partials and fulls
     lap("group hkl");
     std::vector<double> partial_mean;
-    for (int it = 0; it < scaling_iter; ++it) {
-        ReduceGroupMeans(partials, n_groups, false, {}, partial_mean);
-        FitPerFrameG(partials, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial);
-        UpdateCorr(partials, g_partial, frame_scaled_scratch);
+    bool scaled_on_gpu = false;
+#ifdef JFJOCH_USE_CUDA
+    if (gpu_active_) {
+        // Refresh corr on the device (smooth-G mutated it on the host between passes), run the whole
+        // scaling loop on the GPU, then read corr + per-frame G/scaled back.
+        std::vector<float> corr(partials.size());
+        for (size_t i = 0; i < partials.size(); ++i) corr[i] = partials[i].corr;
+        gpu_->SetCorr(corr.data());
+        gpu_->ScalePartials(scaling_iter, SCALE_ROBUST_K, min_partiality, d_min_limit.has_value());
+        gpu_->GetCorr(corr.data());
+        frame_scaled_scratch.assign(n_frames, 0);
+        gpu_->GetG(g_partial.data(), frame_scaled_scratch.data());
+        for (size_t i = 0; i < partials.size(); ++i) partials[i].corr = corr[i];
+        scaled_on_gpu = true;
+    }
+#endif
+    if (!scaled_on_gpu) {
+        for (int it = 0; it < scaling_iter; ++it) {
+            ReduceGroupMeans(partials, n_groups, false, {}, partial_mean);
+            FitPerFrameG(partials, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial);
+            UpdateCorr(partials, g_partial, frame_scaled_scratch);
+        }
     }
     lap("scale partials");
     const std::vector<uint8_t> partial_scaled = frame_scaled_scratch;