Merge branch 'sg-arms' into rc175: every merge starts from G = 1; the all-observation arm merged and searched beside the first

Exact: P4 p.hkl identical to the base on 39 noise-subset sets; concurrent arms p.hkl md5 and MTZ data identical
to P4 on 12 sets at default -N and -N 8 (incl. 5 sets where the arms disagree). Side arm on card 2 with 3+ cards,
else on the last card when 2 x observations x 256 B fits its total memory (decided up front).

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

common/CUDAWrapper.cpp

@@ -61,6 +61,10 @@ std::vector<std::string> get_gpu_names() {
     return {};
 }
 
+size_t get_gpu_total_memory(int32_t) {
+    return 0;
+}
+
 void set_gpu(int32_t dev_id) {}
 
 void pin_gpu() {}

common/CUDAWrapper.cu

@@ -49,6 +49,13 @@ std::vector<std::string> get_gpu_names() {
     return names;
 }
 
+size_t get_gpu_total_memory(int32_t dev_id) {
+    cudaDeviceProp prop{};
+    if (cudaGetDeviceProperties(&prop, dev_id) != cudaSuccess)
+        return 0;
+    return prop.totalGlobalMem;
+}
+
 void set_gpu(int32_t dev_id) {
     auto dev_count = get_gpu_count();
 

common/CUDAWrapper.h

@@ -4,6 +4,7 @@
 #pragma once
 
 #include <chrono>
+#include <cstddef>
 #include <cstdint>
 #include <string>
 #include <vector>
@@ -17,6 +18,11 @@ std::vector<std::string> get_gpu_names();
 // The same list collapsed for a person: "4x NVIDIA A100-SXM4-80GB", or several such groups separated
 // by ", " on a mixed machine. Empty when no GPU is visible.
 std::string get_gpu_description();
+
+// The device's total memory in bytes - not what is free at the moment. 0 without CUDA, or when the
+// device cannot be queried.
+size_t get_gpu_total_memory(int32_t dev_id);
+
 void set_gpu(int32_t dev_id);
 
 // Pin the calling thread to the next GPU in round-robin order, using a process-wide counter

image_analysis/scale_merge/RotationScaleMerge.cpp

@@ -706,6 +706,25 @@ void RotationScaleMerge::StageObservations() {
 }
 #endif
 
+std::unique_ptr<RotationScaleMerge> RotationScaleMerge::CloneIngest(int device, Logger &log) const {
+    TimingMark mark(*logger, "rsm ingest copy");
+    std::unique_ptr<RotationScaleMerge> c(new RotationScaleMerge(*this));
+    c->logger = &log;
+    c->gpu_device = device;
+    c->write_back_per_frame_scale = false;
+    // Of the ingest arrays, IngestDevice reads only this one (StageObservations); the rest served
+    // IngestHost alone.
+    if (ingest_partiality) {
+        c->ingest_partiality.reset(new float[n_partials_obs]);
+        std::copy_n(ingest_partiality.get(), n_partials_obs, c->ingest_partiality.get());
+    }
+#ifdef JFJOCH_USE_CUDA
+    if (gpu_)
+        c->gpu_ = std::make_unique<RotationScaleMergeGPU>(device);
+#endif
+    return c;
+}
+
 void RotationScaleMerge::IngestDevice() {
     const ScopedGpuPin pin(gpu_device);   // see IngestHost
     TimingMark mark(*logger, "rsm ingest (device)");
@@ -6211,9 +6230,8 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st
     // returned is the one run last, so that every per-pass state the merge leaves behind is that arm's.
     //
     // Where the prior is to scale the partials and that merge resolves its error model, the other arm
-    // cannot change the choice, so it is not made. The partials' arm then starts from G = 1, which is
-    // what the fulls-alone arm would have left it (it is the only state that arm leaves for the next),
-    // so the merge is the one the pair would have returned.
+    // cannot change the choice, so it is not made. Every arm starts from G = 1 (RunArm), so the merge
+    // is the one the pair would have returned.
     if (rocking_event_frames_at_start < 0)
         rocking_event_frames_at_start = RockingEventFrames(&rocking_events_at_start);
     const double events_per_frame = n_frames > 0 ? static_cast<double>(rocking_events_at_start) / n_frames : 0.0;
@@ -6221,7 +6239,6 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st
     if (for_search && !scale_partials_choice)
         return RunArm(prior, for_search, full_stats, measure_cc_before_corrections);
     if (!scale_partials_choice && prior) {
-        std::fill(g_partial.begin(), g_partial.end(), 1.0);
         Result with = RunArm(true, for_search, full_stats, measure_cc_before_corrections);
         if (ErrorModelResolved(with)) {
             scale_partials_choice = true;
@@ -6270,6 +6287,12 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
             });
     }
 
+    // ... and from G = 1, so that no merge starts from another's scales. A frame this merge does not
+    // fit would otherwise keep the G an earlier merge fitted it with, and the fulls' CC1/2 weight
+    // reads it (frame_scale_fulls): the merge would depend on which merges ran before it on this
+    // engine, and two merges could not run side by side on two engines with the same result.
+    std::fill(g_partial.begin(), g_partial.end(), 1.0);
+
     TimingMark section_mark(*logger, "rsm scale partials and filter frames");
     // --- 1. Per-frame partial scaling (Rotation model, per-image G only). ---
     const int n_groups = ComputeAsuGroups(keygen);   // one ASU grouping, shared by partials and fulls
@@ -6325,7 +6348,6 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
     }
 #endif
     if (!scale_partials) {
-        std::fill(g_partial.begin(), g_partial.end(), 1.0);
         frame_scaled_scratch.assign(n_frames, 1);
         partial_loop.converged = true;
     } else if (!scaled_on_gpu) {
@@ -6462,6 +6484,9 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
         ReducePartialGroupMeans(n_groups, partial_mean);
         ComputePerFrameCC(partial_mean, cc, cc_n);
     }
+    last_cc = cc;
+    last_cc_n = cc_n;
+    last_frame_scaled = partial_scaled;
     if (write_back_per_frame_scale)
         FinalizePerFrameScale(cc, cc_n, partial_scaled);
 

image_analysis/scale_merge/RotationScaleMerge.h

@@ -151,6 +151,27 @@ public:
     // was written.
     void SetWriteBackPerFrameScale(bool on) { write_back_per_frame_scale = on; }
 
+    // Write back what the last Run() would have written with the write-back on: for a merge that ran
+    // beside others with it off, at the point where it would have run in series.
+    void WriteBackPerFrameScale() { FinalizePerFrameScale(last_cc, last_cc_n, last_frame_scaled); }
+
+    // A second engine on this one's ingest, on `device` and logging to `log`, for a merge to run beside
+    // this engine's: the host half of the ingest is copied, and the copy's IngestDevice() is still to
+    // be called. Call between IngestHost() and IngestDevice(). No merge depends on the merges run
+    // before it (see RunArm), so a merge on the copy gives what the same merge gives on this engine.
+    // The copy does not write its per-frame scales back (WriteBackPerFrameScale does).
+    [[nodiscard]] std::unique_ptr<RotationScaleMerge> CloneIngest(int device, Logger &log) const;
+
+    // The partial observations the ingest kept.
+    [[nodiscard]] int64_t ObservationCount() const { return n_partials_obs; }
+
+    // The card memory an engine may hold during a search merge, per observation the ingest kept: what
+    // a caller goes by when it decides whether a second engine fits beside one already on the card.
+    // Between what two engines merging side by side were measured at, ~150 bytes per observation each
+    // (20.6M observations, 4.2M fulls), and the layout's worst case of ~280 (71 bytes per observation,
+    // 83 per full, ~130 per ASU group, with a full and a group for every observation).
+    static constexpr size_t DEVICE_BYTES_PER_OBSERVATION = 256;
+
     // Merge only the first n frames of the outcomes (by default all of them). Set before Ingest().
     void SetFrameLimit(int n) { frame_limit = n; }
 
@@ -180,6 +201,17 @@ public:
     [[nodiscard]] double GetSearchMinZeta() const { return search_min_zeta; }
 
 private:
+    // A unique_ptr that a copy of the engine starts without (see CloneIngest).
+    template <class T> struct NotCopied : std::unique_ptr<T> {
+        using std::unique_ptr<T>::unique_ptr;
+        using std::unique_ptr<T>::operator=;
+        NotCopied() = default;
+        NotCopied(NotCopied &&) noexcept = default;
+        NotCopied &operator=(NotCopied &&) noexcept = default;
+        NotCopied(const NotCopied &) {}
+    };
+    RotationScaleMerge(const RotationScaleMerge &) = default;   // CloneIngest only
+
     static constexpr uint8_t OBS_CLIPPED = 1;
     static constexpr uint8_t OBS_OVERLOADED = 2;
 
@@ -296,11 +328,11 @@ private:
     // BuildInRangeObservations and handed back at the end of Ingest. Plain arrays, not sized vectors, for
     // the reason the staging buffers are (see Ingest): the build writes every element, and a sized vector
     // would first zero all of them on one thread.
-    std::unique_ptr<float[]> ingest_delta_phi, ingest_partiality, ingest_image_number;
+    NotCopied<float[]> ingest_delta_phi, ingest_partiality, ingest_image_number;
     // The rocking-event usability test (RockingEventFrames) on the ingest-time values - which is what
     // that walk reads on the resident path, where the host corr is never refreshed (see
     // partials_released). Resident path only.
-    std::unique_ptr<uint8_t[]> ingest_rock_ok;
+    NotCopied<uint8_t[]> ingest_rock_ok;
     // Which source reflections the resolution range kept (indexed like the sort keys: frame by frame,
     // from ingest_src_start), so that the i-th kept one of a frame can be found again.
     std::vector<uint8_t> ingest_keep;
@@ -329,7 +361,7 @@ private:
     std::vector<int32_t> frame_start, frame_count; // CSR ranges of `partials` per frame
     std::vector<uint8_t> frame_cell_ok;           // per-frame cell-consistency mask (1 = kept)
     std::vector<uint8_t> finite_ok;               // per-obs AcceptReflection finiteness (immutable; 1 = kept)
-    std::vector<double> g_partial;                // per-frame partial scale G (the RESIDUAL after the flux)
+    std::vector<double> g_partial;                // per-frame partial scale G (the RESIDUAL after the flux); 1 at the start of every merge
     std::vector<double> frame_flux;               // per-frame incident flux, run median = 1 (see the .cpp)
     // corr as Ingest built it, kept for the whole life of the object and put back at the start of every
     // Run. A pass that starts anywhere else is not the pass that would have run on its own, and the
@@ -509,7 +541,7 @@ private:
     // GPU engine: the whole hot path (scaling, combine, scale-fulls, per-frame CC, smooth-G, merge +
     // error model) runs on the device, resident, when a GPU is present. Null / inactive otherwise, with
     // the CPU loops as the bit-parity fallback. Built in Ingest.
-    std::unique_ptr<RotationScaleMergeGPU> gpu_;
+    NotCopied<RotationScaleMergeGPU> gpu_;
     bool gpu_active_ = false;
     // The fulls of this pass live on the device (the GPU combine ran), so the correction surfaces can
     // build their terms there. Set by RunArm.
@@ -519,6 +551,10 @@ private:
     std::vector<float> surface_corr;
 #endif
     bool write_back_per_frame_scale = true;   // see SetWriteBackPerFrameScale
+    // What the last Run wrote back, or would have (WriteBackPerFrameScale).
+    std::vector<double> last_cc;
+    std::vector<int64_t> last_cc_n;
+    std::vector<uint8_t> last_frame_scaled;
     int frame_limit = -1;                     // see SetFrameLimit
     int gpu_device = 0;                       // see SetGpuDevice
     bool export_scaled_fulls = false;         // see SetExportScaledFulls
@@ -557,7 +593,7 @@ private:
     // which ComputeAsuGroups builds for the GPU reduction anyway. ReducePartialGroupMeans sums each
     // group over it, on all threads and in the order the serial loop would.
     std::vector<int32_t> partial_group_start, partial_group_count;
-    std::unique_ptr<int32_t[]> partial_group_perm;
+    NotCopied<int32_t[]> partial_group_perm;
 
     // The CPU partial scaling loop's working copy of the partials. It passes over them several times
     // per iteration and reads a handful of fields each time, so it works on these rather than on the

rugnux/RugnuxScaleMerge.cpp

@@ -958,6 +958,15 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
                     100.0 * *bkg_starved_fraction_);
         return true;
     }
+    // The all-observation arm of the space-group search (below) is the same ingest merged with the
+    // Lorentz filter off. No merge depends on the merges made before it (RotationScaleMerge::RunArm),
+    // so it is merged on a copy of the engine, started here and read where the search needs it: beside
+    // the P1 merge, the re-scaled search merge and the first arm. Started whenever the search will run;
+    // a search that turns out not to need it (no point group from the first arm) does not read it.
+    // Declared before the future, which waits for the merge when it goes, so the engine outlives it.
+    std::unique_ptr<RotationScaleMerge> rsm_all;
+    Logger all_arm_log = Logger::Buffered();
+    std::future<RotationScaleMerge::Result> all_arm_ahead;
     if (is_rotation) {
         if (experiment_.GetRefineRotationWedgeInScaling()
             || rot_ss.GetRotationWedgeForScaling().has_value())
@@ -985,6 +994,32 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
                         config_.nthreads, logger, config_.observation_dump_path);
             rsm->IngestHost();
         }
+        // On a card of its own from three cards up (card 0 has this engine, card 1 the quality guard's
+        // and the speculative passes). Otherwise on the last card, beside the engine already there -
+        // this one, or the guard's, which merges at most as many observations - and only where two
+        // engines of this size fit the card's memory, judged on its total, not on what happens to be
+        // free: the answer must not depend on the moment. On one card that is the small and middling
+        // data sets; the large ones merge their arms one after the other.
+        const int32_t cards = get_gpu_count();
+        if (!geometry_prepass && !experiment_.GetGemmiSpaceGroup().has_value() && rsm->GetSearchMinZeta() > 0.0
+            && cards > 0) {
+            const int32_t device = std::min(cards - 1, 2);
+            const size_t engine_bytes = static_cast<size_t>(rsm->ObservationCount())
+                                        * RotationScaleMerge::DEVICE_BYTES_PER_OBSERVATION;
+            if (cards >= 3 || 2 * engine_bytes <= get_gpu_total_memory(device)) {
+                rsm_all = rsm->CloneIngest(device, all_arm_log);
+                rsm_all->SetSearchMinZeta(0.0);
+                rsm_all->SetExportScaledFulls(false);   // as scale_and_merge sets them for a search merge
+                rsm_all->SetFrenchWilson(false);
+                all_arm_ahead = std::async(std::launch::async, [&rsm_all, &all_arm_log, device] {
+                    pin_gpu(device);   // the card and its NUMA node, as the guard's merge has them
+                    TimingMark mark(all_arm_log, "merge P1, all observations");
+                    rsm_all->IngestDevice();
+                    return rsm_all->Run(/*for_search=*/true, /*full_stats=*/true,
+                                        /*measure_cc_before_corrections=*/false);
+                });
+            }
+        }
         rsm->IngestDevice();
     }
 
@@ -1001,6 +1036,28 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
                 indexer->GetIntegrationOutcome(), config_.nthreads,
                 /*release_reflections=*/true, logger);
 
+    // What a rotation merge leaves in the result, taken in the order the merges are made - which, for
+    // the all-observation arm merged beside the others, is where the search reads it.
+    const auto take_merge = [&](const std::string &label, bool for_search, RotationScaleMerge::Result r)
+            -> ScaleMergeResult {
+        result.error_model_isa = r.isa;
+        result.error_model_isa_asymptotic = r.isa_asymptotic;
+        result.error_model_isa_resolved = r.isa_resolved;
+        if (r.resolution_fit_A) result.resolution_fit_A = r.resolution_fit_A;
+        result.error_model_a = r.error_model_a;
+        result.error_model_b = r.error_model_b;
+        result.scaling_iterations_partials = r.scaling_iterations_partials;
+        result.scaling_iterations_fulls = r.scaling_iterations_fulls;
+        result.scaling_converged = r.scaling_converged;
+        if (for_search)
+            result.search_scaling_converged = r.scaling_converged;
+        Note(fmt::format("merge in {}: {} unique", label, r.merged.size()));
+        return ScaleMergeResult{std::move(r.merged), std::move(r.statistics),
+                                r.cc_half_before_corrections,
+                                std::move(r.reference_statistics), r.reference_isa,
+                                std::move(r.scaled_fulls)};
+    };
+
     // measure_cc: whether anything reads the result's cc_half_before_corrections, which costs a
     // merge of its own - only the first merge, which the two-pass quality guard reads.
     auto scale_and_merge = [&](const std::string &label, bool for_search, bool measure_cc = false)
@@ -1017,24 +1074,9 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
             // Amplitudes are made once, on the merge that is written, where the anisotropy tensor
             // their prior uses is known (see below) - not on every merge the tail makes on its way.
             rsm->SetFrenchWilson(false);
-            auto r = rsm->Run(for_search, /*full_stats=*/!geometry_prepass,
-                              /*measure_cc_before_corrections=*/measure_cc);
-            result.error_model_isa = r.isa;
-            result.error_model_isa_asymptotic = r.isa_asymptotic;
-            result.error_model_isa_resolved = r.isa_resolved;
-            if (r.resolution_fit_A) result.resolution_fit_A = r.resolution_fit_A;
-            result.error_model_a = r.error_model_a;
-            result.error_model_b = r.error_model_b;
-            result.scaling_iterations_partials = r.scaling_iterations_partials;
-            result.scaling_iterations_fulls = r.scaling_iterations_fulls;
-            result.scaling_converged = r.scaling_converged;
-            if (for_search)
-                result.search_scaling_converged = r.scaling_converged;
-            Note(fmt::format("merge in {}: {} unique", label, r.merged.size()));
-            return ScaleMergeResult{std::move(r.merged), std::move(r.statistics),
-                                    r.cc_half_before_corrections,
-                                    std::move(r.reference_statistics), r.reference_isa,
-                                    std::move(r.scaled_fulls)};
+            return take_merge(label, for_search,
+                              rsm->Run(for_search, /*full_stats=*/!geometry_prepass,
+                                       /*measure_cc_before_corrections=*/measure_cc));
         }
         // Stills (rotation goes through RotationScaleMerge above): self-scale each image against the
         // running merge with ScaleOnTheFly (fixed partiality), then merge directly. This runs even
@@ -1347,6 +1389,30 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
         // Set per merge, not once: the two arms below are scaled separately and their ISa differ by up
         // to a factor of six. 0 (b not measurable) leaves the cut as the caller set it.
         sg_opts.merge_isa = result.error_model_isa;
+        // The all-observation merge made beside everything since the ingest (see rsm_all), and its
+        // search started beside this one's: a search reads its merge and its options and nothing
+        // else. Whether the all-observation arm is wanted at all is known only from the first arm
+        // (below), so nothing of it reaches the run before then. A merge that failed beside the
+        // others is made again where it is wanted, in series, as it would have been without it.
+        std::optional<RotationScaleMerge::Result> all_merge;
+        std::future<SearchSpaceGroupResult> all_search_ahead;
+        if (all_arm_ahead.valid()) {
+            TimingMark wait_mark(logger, "wait for the all-observation merge");
+            try {
+                all_merge = all_arm_ahead.get();
+            } catch (const std::exception &e) {
+                logger.Warning("Space-group search: the all-observation merge beside the others failed ({}); "
+                               "it is made again in series", e.what());
+            }
+            wait_mark.Stop();
+        }
+        if (all_merge) {
+            SearchSpaceGroupOptions all_opts = sg_opts;
+            all_opts.merge_isa = all_merge->isa;   // as take_merge leaves it below
+            all_search_ahead = std::async(std::launch::async, [&merged = all_merge->merged, all_opts] {
+                return SearchSpaceGroup(merged, all_opts);
+            });
+        }
         TimingMark first_arm_mark(logger, "SearchSpaceGroup (first arm)");
         auto sg_search = SearchSpaceGroup(sm.merged, sg_opts);
         first_arm_mark.Stop();
@@ -1380,15 +1446,28 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
         // all-observation arm on both twin gates and promoted to R32 by the filtered one. So the
         // losing arm's refusal is logged and carried to the report below; the rule is unchanged.
         if (rsm && rsm->GetSearchMinZeta() > 0.0 && !sg_search.point_group_hm.empty()) {
-            const double zeta = rsm->GetSearchMinZeta();
-            rsm->SetSearchMinZeta(0.0);
-            auto sm_all = scale_and_merge("P1, all observations", true);
-            rsm->SetSearchMinZeta(zeta);
+            ScaleMergeResult sm_all;
+            SearchSpaceGroupResult alt;
             merged_filtered_isa = sg_opts.merge_isa;     // the filtered arm's, before it is replaced
-            sg_opts.merge_isa = result.error_model_isa;   // this arm's own error model
-            TimingMark all_arm_mark(logger, "SearchSpaceGroup (all-observation arm)");
-            auto alt = SearchSpaceGroup(sm_all.merged, sg_opts);
-            all_arm_mark.Stop();
+            if (all_merge) {
+                TimingMark all_arm_mark(logger, "SearchSpaceGroup (all-observation arm, wait)");
+                alt = all_search_ahead.get();
+                all_arm_mark.Stop();
+                phase("Scale/combine/merge (P1, all observations)");
+                all_arm_log.ReplayInto(logger);
+                rsm_all->WriteBackPerFrameScale();
+                sm_all = take_merge("P1, all observations", true, std::move(*all_merge));
+                sg_opts.merge_isa = result.error_model_isa;   // this arm's own error model
+            } else {
+                const double zeta = rsm->GetSearchMinZeta();
+                rsm->SetSearchMinZeta(0.0);
+                sm_all = scale_and_merge("P1, all observations", true);
+                rsm->SetSearchMinZeta(zeta);
+                sg_opts.merge_isa = result.error_model_isa;   // this arm's own error model
+                TimingMark all_arm_mark(logger, "SearchSpaceGroup (all-observation arm)");
+                alt = SearchSpaceGroup(sm_all.merged, sg_opts);
+                all_arm_mark.Stop();
+            }
             // The order of the point group each arm confirmed. Read from the search, not from the
             // space group it went on to pick: Stage B leaves best_space_group unset when no
             // candidate has enough absences to be eligible, which is common on the filtered arm
@@ -1535,6 +1614,13 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
             sm = std::move(sm_all);
         }
 
+        // The arm merged beside the search is done with - or was not needed - before anything below
+        // changes the experiment it reads, and its card memory goes back before the merges below.
+        if (all_search_ahead.valid())
+            all_search_ahead.wait();
+        all_merge.reset();
+        rsm_all.reset();
+
         // Everything from here to the adopted group's merge: the lattice re-asks and the choices
         // between the candidates the arms left.
         TimingMark decisions_mark(logger, "space-group decisions");