Merge stillskhz exact commit into rc175: stills image loop 3.6-3.9x faster, output unchanged

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

image_analysis/IndexAndRefine.cpp

@@ -171,7 +171,20 @@ IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetryRotat
     return outcome;
 }
 
-IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetry(DataMessage &msg) {
+IndexerResult IndexAndRefine::RunIndexer(const std::vector<Coord> &recip, FrameIndexerCalls *frame_calls) {
+    if (frame_calls)
+        for (const auto &[asked, answer] : frame_calls->calls)
+            if (asked == recip)
+                return answer;
+    auto res = indexer_->Run(experiment, recip);
+    // A failure says nothing about the list, so it is not an answer to keep.
+    if (frame_calls && !res.error)
+        frame_calls->calls.emplace_back(recip, res);
+    return res;
+}
+
+IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetry(DataMessage &msg,
+                                                                            FrameIndexerCalls *frame_calls) {
     auto indexing_start_time = std::chrono::steady_clock::now();
 
     IndexingOutcome outcome(experiment);
@@ -209,7 +222,7 @@ IndexAndRefine::IndexingOutcome IndexAndRefine::DetermineLatticeAndSymmetry(Data
                     break;
             }
         }
-        auto res = indexer_->Run(experiment, recip);
+        auto res = RunIndexer(recip, frame_calls);
         if (res.error)
             RecordIndexerError(res.error);
         any_executed |= res.executed;
@@ -811,7 +824,7 @@ SupercellTrend IndexAndRefine::GetSupercellTrend() {
 
 std::optional<IndexAndRefine::IndexingOutcome>
 IndexAndRefine::DetermineRefineAnalyze(DataMessage &msg, const SpotFindingSettings &spot_finding_settings,
-                                      int64_t *spots_on_lattice) {
+                                      int64_t *spots_on_lattice, FrameIndexerCalls *frame_calls) {
     if (!indexer_ || !spot_finding_settings.indexing)
         return std::nullopt;
 
@@ -820,7 +833,7 @@ IndexAndRefine::DetermineRefineAnalyze(DataMessage &msg, const SpotFindingSettin
     if (rotation_indexer)
         outcome = DetermineLatticeAndSymmetryRotation(msg);
     else
-        outcome = DetermineLatticeAndSymmetry(msg);
+        outcome = DetermineLatticeAndSymmetry(msg, frame_calls);
 
     if (!outcome.lattice_candidate)
         return std::nullopt;
@@ -849,17 +862,19 @@ void IndexAndRefine::ProcessImage(DataMessage &msg,
                                   const SpotFindingSettings &spot_finding_settings,
                                   BraggPrediction &prediction,
                                   const BraggIntegrateFn &integrate,
-                                  const BraggIntegrateFn &probe_integrate) {
-    auto outcome = DetermineRefineAnalyze(msg, spot_finding_settings);
+                                  const BraggIntegrateFn &probe_integrate,
+                                  FrameIndexerCalls *frame_calls) {
+    auto outcome = DetermineRefineAnalyze(msg, spot_finding_settings, nullptr, frame_calls);
     if (outcome && spot_finding_settings.quick_integration)
         QuickPredictAndIntegrate(msg, spot_finding_settings, prediction, integrate, *outcome, probe_integrate);
 }
 
-bool IndexAndRefine::IndexFrameOnly(DataMessage &msg, const SpotFindingSettings &spot_finding_settings) {
+bool IndexAndRefine::IndexFrameOnly(DataMessage &msg, const SpotFindingSettings &spot_finding_settings,
+                                    FrameIndexerCalls *frame_calls) {
     // The per-frame verdict, not "was there anything to integrate": this is what the rotation first pass
     // scores candidate lattices on, and on rotation a frame can be integrable without indexing on its
     // own (AnalyzeIndexing), which would score a sparse frame for every candidate alike.
-    return DetermineRefineAnalyze(msg, spot_finding_settings).has_value()
+    return DetermineRefineAnalyze(msg, spot_finding_settings, nullptr, frame_calls).has_value()
            && msg.indexing_result.value_or(false);
 }
 

image_analysis/IndexAndRefine.h

@@ -75,6 +75,15 @@ struct SupercellProbeFit {
 };
 SupercellProbeFit FitSupercellProbe(const SupercellProbeClass &c);
 
+// The indexer calls one frame has made so far, kept by the caller across that frame's own calls. The stills
+// min-pix choice indexes one frame up to four times (three min-pix, then the winner again), and the seed of
+// strongest spots the indexer is handed is often the same list each time - the strongest spots are the
+// large ones every min-pix keeps. The indexer answers the same list the same way, so it is asked once.
+// Only within one frame: the geometry and settings the answer depends on do not change inside it.
+struct FrameIndexerCalls {
+    std::vector<std::pair<std::vector<Coord>, IndexerResult>> calls;
+};
+
 class IndexAndRefine {
     // When false, the current image's result is still returned via the outgoing message, but the
     // whole-run integration_outcome vector is not retained (viewer live/interactive use, which never
@@ -158,14 +167,17 @@ class IndexAndRefine {
                              const BraggPredictionSettings &settings);
 
     IndexingOutcome DetermineLatticeAndSymmetryRotation(DataMessage &msg);
-    IndexingOutcome DetermineLatticeAndSymmetry(DataMessage &msg);
+    IndexingOutcome DetermineLatticeAndSymmetry(DataMessage &msg, FrameIndexerCalls *frame_calls);
+    // indexer_->Run, answered from frame_calls where this list was already asked about on the frame.
+    IndexerResult RunIndexer(const std::vector<Coord> &recip, FrameIndexerCalls *frame_calls);
     // Shared indexing path: determine the lattice/symmetry, refine geometry, and run AnalyzeIndexing.
     // Returns the outcome (ready for integration) when the frame indexes, nullopt otherwise. Both the
     // real per-image ProcessImage and the first-pass scheme validation go through this, so they cannot
     // diverge.
     std::optional<IndexingOutcome> DetermineRefineAnalyze(DataMessage &msg,
                                                           const SpotFindingSettings &spot_finding_settings,
-                                                          int64_t *spots_on_lattice = nullptr);
+                                                          int64_t *spots_on_lattice = nullptr,
+                                                          FrameIndexerCalls *frame_calls = nullptr);
     void RefineGeometryIfNeeded(DataMessage &msg, IndexingOutcome &outcome);
     void QuickPredictAndIntegrate(DataMessage &msg,
                                  const SpotFindingSettings &spot_finding_settings,
@@ -218,9 +230,11 @@ public:
     // is a further step gated on quick_integration.
     // probe_integrate integrates the supercell probe's reflections (ProbeSupercell); a caller that
     // gives none takes no part in the probe.
+    // frame_calls, where given, carries the indexer calls this frame already made (FrameIndexerCalls).
     void ProcessImage(DataMessage &msg, const SpotFindingSettings &settings,
                       BraggPrediction &prediction, const BraggIntegrateFn &integrate,
-                      const BraggIntegrateFn &probe_integrate = {});
+                      const BraggIntegrateFn &probe_integrate = {},
+                      FrameIndexerCalls *frame_calls = nullptr);
 
     // Does the lattice each frame is refined on have an index-2 superstructure? On the given frames,
     // after the frame's own integration, predict the 2a x 2b x 2c supercell of its PRIMITIVE lattice
@@ -236,7 +250,8 @@ public:
     [[nodiscard]] std::optional<CrystalLattice> GetSupercellProbePrimitive();
     // Index a single frame (no integration) with the current forced rotation lattice; used to score
     // first-pass sampling schemes on the real per-image path. Returns whether the frame indexed.
-    bool IndexFrameOnly(DataMessage &msg, const SpotFindingSettings &settings);
+    bool IndexFrameOnly(DataMessage &msg, const SpotFindingSettings &settings,
+                        FrameIndexerCalls *frame_calls = nullptr);
     // How many of this frame's spots lie on the forced rotation lattice, whatever the per-frame gate
     // then makes of them. The first pass's acceptance test counts spots over the validation set rather
     // than frames, so it needs the count from the sparse frames IndexFrameOnly refuses too.

image_analysis/MXAnalysisWithoutFPGA.cpp

@@ -260,6 +260,7 @@ void MXAnalysisWithoutFPGA::Analyze(DataMessage &output,
             float indexing_time_s = 0.0f;
 
             SpotFindingSettings s = spot_finding_settings;
+            FrameIndexerCalls frame_calls;
             std::vector<DiffractionSpot> best_spots;
             int best_mp = 0;
             double best_score = -1.0;
@@ -270,7 +271,7 @@ void MXAnalysisWithoutFPGA::Analyze(DataMessage &output,
                 spot_finding_time_s +=
                         std::chrono::duration<float>(std::chrono::steady_clock::now() - extract_start_time).count();
                 SpotAnalyze(experiment, s, spots, output);
-                const bool indexed = indexer.IndexFrameOnly(output, s);
+                const bool indexed = indexer.IndexFrameOnly(output, s, &frame_calls);
                 indexing_time_s += output.indexing_time_s.value_or(0.0f);
                 if (indexed) {
                     const double n_idx = static_cast<double>(output.spot_count_indexed.value_or(0));
@@ -288,7 +289,7 @@ void MXAnalysisWithoutFPGA::Analyze(DataMessage &output,
                 s.min_pix_per_spot = best_mp;
                 SpotAnalyze(experiment, s, best_spots, output);
                 decide_radial_background();
-                indexer.ProcessImage(output, s, *prediction, integrate_fn, probe_integrate_fn);
+                indexer.ProcessImage(output, s, *prediction, integrate_fn, probe_integrate_fn, &frame_calls);
                 indexing_time_s += output.indexing_time_s.value_or(0.0f);
             }
             // Each indexer call reports only its own time, so the escalation's total is summed here.

image_analysis/indexing/FFTIndexer.cpp

@@ -148,7 +148,13 @@ std::vector<CrystalLattice> FFTIndexer::ReduceResults(const std::vector<Coord> &
     // most triples of one lattice Niggli-reduce to the same cell (keeps the refine set small).
     const size_t n = std::min<size_t>(results.size(), 64);
     const size_t n_short = std::min<size_t>(n, 12);
-    std::vector<UnitCell> candidate_cells;   // parallel to `candidates`, see the dedup scan below
+    // The cells accepted so far, ordered by a, for the dedup test below. A triple duplicates an accepted
+    // cell only if is_close holds, which needs the two a to agree within 2% of the new one; so only the
+    // accepted cells in that window of a are asked, with the window a little wider than the tolerance so
+    // rounding cannot leave a match outside it. The answer - is there any accepted cell close to this
+    // one - is the one a scan over all of them gives. A cell that is not finite is close to nothing, so
+    // it is never a duplicate and never needs to be found.
+    std::vector<UnitCell> candidate_cells;
     for (size_t i = 0; i < n_short; i++) {
         for (size_t j = i + 1; j < n_short; j++) {
             for (size_t k = j + 1; k < n; k++) {
@@ -166,14 +172,22 @@ std::vector<CrystalLattice> FFTIndexer::ReduceResults(const std::vector<Coord> &
                 if (reduced.VolumeFraction() < MIN_BASIS_VOLUME_FRACTION)
                     continue;
 
+                const auto by_a = [](const UnitCell &x, float a) { return x.a < a; };
                 bool duplicate = false;
-                for (const auto &c : candidate_cells)
-                    if (c.is_close(uc, 0.02f, 1.0f)) { duplicate = true; break; }
+                if (uc.is_finite()) {
+                    const float window = 0.021f * std::max(std::fabs(uc.a), 1e-6f);
+                    for (auto c = std::lower_bound(candidate_cells.begin(), candidate_cells.end(),
+                                                   uc.a - window, by_a);
+                         c != candidate_cells.end() && c->a <= uc.a + window; ++c)
+                        if (c->is_close(uc, 0.02f, 1.0f)) { duplicate = true; break; }
+                }
                 if (!duplicate) {
                     // Keep each accepted candidate's cell rather than re-deriving it on the next
-                    // triple: GetUnitCell costs three acos, and this scan runs over every candidate
-                    // accepted so far, for every triple.
-                    candidate_cells.push_back(uc);
+                    // triple: GetUnitCell costs three acos.
+                    if (uc.is_finite())
+                        candidate_cells.insert(std::upper_bound(candidate_cells.begin(), candidate_cells.end(), uc.a,
+                                                                [](float a, const UnitCell &x) { return a < x.a; }),
+                                               uc);
                     candidates.emplace_back(std::move(reduced));
                 }
             }
@@ -185,25 +199,35 @@ std::vector<CrystalLattice> FFTIndexer::ReduceResults(const std::vector<Coord> &
 
 std::vector<Coord> FFTIndexer::FilterFFTResults(size_t max_vectors,
                                                std::vector<float> *magnitudes) const {
-    std::multimap<float, FFTResult> fft_result_map;
-
     // No legitimate row is non-finite (lengths are bin*coeff or -1, magnitudes guarded
-    // prominences), but a corrupted result buffer would put NaN into the length sort below,
+    // prominences), but a corrupted result buffer would put NaN into the magnitude sort below,
     // and std::sort with a NaN key walks out of bounds. Drop such rows instead.
+    std::vector<int> by_magnitude;
+    by_magnitude.reserve(direction_vectors.size());
     for (int i = 0; i < direction_vectors.size(); i++) {
         if (!std::isfinite(result_fft[i].magnitude) || !std::isfinite(result_fft[i].length))
             continue;
-        fft_result_map.insert(std::make_pair(result_fft[i].magnitude, result_fft[i]));
+        by_magnitude.push_back(i);
     }
 
+    // Strongest first, and among equal magnitudes the later direction first - the order a multimap
+    // keyed on the magnitude, walked from its end, gives. Only the head of that order is read in the
+    // normal case, so only the head is sorted; the tail is sorted on the rare walk that reaches it.
+    const auto stronger = [this](int a, int b) {
+        const float ma = result_fft[a].magnitude, mb = result_fft[b].magnitude;
+        if (ma != mb)
+            return ma > mb;
+        return a > b;
+    };
+    const size_t head = std::min(by_magnitude.size(), std::max<size_t>(4 * max_vectors, 256));
+    std::partial_sort(by_magnitude.begin(), by_magnitude.begin() + head, by_magnitude.end(), stronger);
+
     std::vector<FFTResult> fft_result_filtered;
-    int count = 0;
     // `it` outlives the loop: the extra scan at the end of the function carries on from the
     // first direction the budget did not reach.
-    auto it = fft_result_map.rbegin();
-    for (; it != fft_result_map.rend() && count < max_vectors; ++it, ++count) {
-        fft_result_filtered.emplace_back(it->second);
-    }
+    size_t it = 0;
+    for (; it < by_magnitude.size() && it < max_vectors; ++it)
+        fft_result_filtered.emplace_back(result_fft[by_magnitude[it]]);
 
     std::vector<Coord> ret;
 
@@ -302,8 +326,11 @@ std::vector<Coord> FFTIndexer::FilterFFTResults(size_t max_vectors,
     // results[0..8] and a five-vector shortlist leaves most of that window unusable.
     constexpr int EXTRA_DISTINCT_DIRECTIONS = 4;
 
-    for (int extra = 0; it != fft_result_map.rend() && extra < EXTRA_DISTINCT_DIRECTIONS; ++it) {
-        Coord dir = direction_vectors.at(it->second.direction);
+    for (int extra = 0; it < by_magnitude.size() && extra < EXTRA_DISTINCT_DIRECTIONS; ++it) {
+        if (it == head)
+            std::sort(by_magnitude.begin() + head, by_magnitude.end(), stronger);
+        const FFTResult &r = result_fft[by_magnitude[it]];
+        Coord dir = direction_vectors.at(r.direction);
 
         bool distinct = true;
         for (const auto &v: ret) {
@@ -314,9 +341,9 @@ std::vector<Coord> FFTIndexer::FilterFFTResults(size_t max_vectors,
         }
 
         if (distinct) {
-            ret.push_back(dir * it->second.length);
+            ret.push_back(dir * r.length);
             if (magnitudes)
-                magnitudes->push_back(it->second.magnitude);
+                magnitudes->push_back(r.magnitude);
             extra++;
         }
     }

image_analysis/indexing/FFTIndexerGPU.cu

@@ -10,8 +10,38 @@ __device__ __host__ inline float complex_abs(const cufftComplex &z) {
     return sqrtf(z.x * z.x + z.y * z.y);
 }
 
+// The magnitudes of the transform, transposed to bin-major order: mag_t[bin * directions + direction].
+// The R2C output is direction-major, and the peak search gives each direction to one thread, so read
+// straight from it neighbouring lanes are a whole spectrum apart - every load a separate sector, three
+// streams per thread. Transposed, a warp reads 32 consecutive floats per bin. On an idle card that buys
+// nothing (the search is latency-bound either way, and the transpose costs its own pass); with the
+// stills indexer pool's streams sharing the card it is the difference between 67 and 98 frames/s
+// (measured). A 32x32 tile in shared memory keeps both the read and the write coalesced. The search took
+// complex_abs of the same values, so the magnitudes are bit for bit the ones it computed itself.
+__global__ void magnitude_transpose(const cufftComplex *__restrict__ d_output,
+                                    const int out_len,
+                                    const int directions_size,
+                                    float *__restrict__ mag_t) {
+    __shared__ float tile[32][33];
+    const int bin0 = blockIdx.x * 32;
+    const int dir0 = blockIdx.y * 32;
+    for (int r = threadIdx.y; r < 32; r += blockDim.y) {
+        const int dir = dir0 + r;
+        const int bin = bin0 + threadIdx.x;
+        if (dir < directions_size && bin < out_len)
+            tile[r][threadIdx.x] = complex_abs(d_output[static_cast<size_t>(dir) * out_len + bin]);
+    }
+    __syncthreads();
+    for (int r = threadIdx.y; r < 32; r += blockDim.y) {
+        const int bin = bin0 + r;
+        const int dir = dir0 + threadIdx.x;
+        if (dir < directions_size && bin < out_len)
+            mag_t[static_cast<size_t>(bin) * directions_size + dir] = tile[threadIdx.x][r];
+    }
+}
+
 __global__ void calculate_fft_result(
-    const cufftComplex *__restrict__ d_output,
+    const float *__restrict__ mag_t,
     const float max_length_A,
     const float min_length_A,
     const int histogram_size,
@@ -22,7 +52,8 @@ __global__ void calculate_fft_result(
 
     if (i < directions_size) {
         const int out_len = (histogram_size / 2) + 1;
-        size_t offset = static_cast<size_t>(out_len) * i;
+        // |transform| of bin k of this direction, from the bin-major array (magnitude_transpose)
+        const auto mag_at = [&](int k) { return mag_t[static_cast<size_t>(k) * directions_size + i]; };
         float len_coeff = 2.0f * max_length_A / static_cast<float>(histogram_size);
 
         // Pick the peak by PROMINENCE above a running-mean background of half-width bg_half:
@@ -34,7 +65,7 @@ __global__ void calculate_fft_result(
         int wlo = 0;
         int whi = min(bg_half, out_len - 1);
         for (int k = wlo; k <= whi; ++k)
-            winsum += complex_abs(d_output[offset + k]);
+            winsum += mag_at(k);
 
         float best_prom = 0.0f;
         FFTResult result{.magnitude = 0.0f, .direction = i, .length = -1};
@@ -51,13 +82,13 @@ __global__ void calculate_fft_result(
                 want_lo = max(0, want_lo - (want_hi - (out_len - 1)));
                 want_hi = out_len - 1;
             }
-            while (whi < want_hi) { ++whi; winsum += complex_abs(d_output[offset + whi]); }
-            while (wlo < want_lo) { winsum -= complex_abs(d_output[offset + wlo]); ++wlo; }
+            while (whi < want_hi) { ++whi; winsum += mag_at(whi); }
+            while (wlo < want_lo) { winsum -= mag_at(wlo); ++wlo; }
 
             const float len = len_coeff * static_cast<float>(j);
             if (len <= min_length_A) continue;
 
-            const float mag = complex_abs(d_output[offset + j]);
+            const float mag = mag_at(j);
             const float bg = static_cast<float>(winsum / static_cast<double>(whi - wlo + 1));
             const float prom = mag - bg;
             if (prom > best_prom) {
@@ -239,7 +270,15 @@ void FFTIndexerGPU::ExecuteFFT(const std::vector<Coord> &coord, size_t nspots) {
     const double len_coeff = 2.0 * static_cast<double>(max_length_A) / static_cast<double>(histogram_size);
     const int bg_half = std::max(1, static_cast<int>(std::lround(15.0 / len_coeff)));
 
-    calculate_fft_result<<<l_gridDim, l_blockDim, 0, stream>>>(d_output_fft,
+    const int out_len = static_cast<int>(histogram_size / 2 + 1);
+    const dim3 tr_grid((out_len + 31) / 32, (direction_vectors.size() + 31) / 32);
+    // The histogram has been transformed and is not read again, so its buffer - histogram_size floats
+    // per direction, more than the out_len magnitudes need - holds them.
+    float *d_mag_t = d_input_fft;
+    magnitude_transpose<<<tr_grid, dim3(32, 8), 0, stream>>>(d_output_fft, out_len,
+                                                             direction_vectors.size(), d_mag_t);
+    cuda_err(cudaGetLastError());
+    calculate_fft_result<<<l_gridDim, l_blockDim, 0, stream>>>(d_mag_t,
                                                                max_length_A, min_length_A, histogram_size,
                                                                bg_half,
                                                                direction_vectors.size(), d_result_fft);

rugnux/Rugnux.cpp

@@ -926,6 +926,14 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b
             // the setting's limit and throw before a single image is read.
             indexing_settings.RefineThreads(std::clamp(config_.nthreads / 2, 1, MAX_REFINE_THREADS));
             RaiseFFTBoundForKnownCell(indexing_settings, experiment_);
+        } else if (!experiment_.IsRotationIndexing()) {
+            // Stills index every frame, and each image worker waits on the pool until its frame is
+            // indexed, so the default four threads were what the whole loop ran at: the post-processing
+            // of the FFT candidates is CPU work, and four threads of it held sixteen workers back.
+            // Measured on a 16 Mpx serial set, de-novo FFT: 83 frames/s with 4, 125 with 8, 133 with 16 -
+            // by eight the card is what limits it. Each GPU FFT indexer holds ~0.5 GB of the card
+            // (sixteen of them took the run to 12.7 GB), so eight rather than one per worker.
+            indexing_settings.IndexingThreads(std::clamp(config_.nthreads, 1, 8));
         }
         indexing_bound_A = indexing_settings.GetFFT_MaxUnitCell_A();
         indexer_pool = std::make_unique<IndexerThreadPool>(indexing_settings, IndexerConstruction::OnFirstUse,