diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp index ef313f71f..4dab42e56 100644 --- a/image_analysis/IndexAndRefine.cpp +++ b/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 &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::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); } diff --git a/image_analysis/IndexAndRefine.h b/image_analysis/IndexAndRefine.h index e6bee48e5..686d2f7e6 100644 --- a/image_analysis/IndexAndRefine.h +++ b/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, 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 &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 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 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. diff --git a/image_analysis/MXAnalysisWithoutFPGA.cpp b/image_analysis/MXAnalysisWithoutFPGA.cpp index e7efc5ccf..2c165170b 100644 --- a/image_analysis/MXAnalysisWithoutFPGA.cpp +++ b/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 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(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(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. diff --git a/image_analysis/indexing/FFTIndexer.cpp b/image_analysis/indexing/FFTIndexer.cpp index b0aa51425..24760965f 100644 --- a/image_analysis/indexing/FFTIndexer.cpp +++ b/image_analysis/indexing/FFTIndexer.cpp @@ -148,7 +148,13 @@ std::vector FFTIndexer::ReduceResults(const std::vector & // most triples of one lattice Niggli-reduce to the same cell (keeps the refine set small). const size_t n = std::min(results.size(), 64); const size_t n_short = std::min(n, 12); - std::vector 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 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 FFTIndexer::ReduceResults(const std::vector & 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 FFTIndexer::ReduceResults(const std::vector & std::vector FFTIndexer::FilterFFTResults(size_t max_vectors, std::vector *magnitudes) const { - std::multimap 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 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(4 * max_vectors, 256)); + std::partial_sort(by_magnitude.begin(), by_magnitude.begin() + head, by_magnitude.end(), stronger); + std::vector 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 ret; @@ -302,8 +326,11 @@ std::vector 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 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++; } } diff --git a/image_analysis/indexing/FFTIndexerGPU.cu b/image_analysis/indexing/FFTIndexerGPU.cu index 7672fc1a5..36bbbe827 100644 --- a/image_analysis/indexing/FFTIndexerGPU.cu +++ b/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(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(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(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(k) * directions_size + i]; }; float len_coeff = 2.0f * max_length_A / static_cast(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(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(winsum / static_cast(whi - wlo + 1)); const float prom = mag - bg; if (prom > best_prom) { @@ -239,7 +270,15 @@ void FFTIndexerGPU::ExecuteFFT(const std::vector &coord, size_t nspots) { const double len_coeff = 2.0 * static_cast(max_length_A) / static_cast(histogram_size); const int bg_half = std::max(1, static_cast(std::lround(15.0 / len_coeff))); - calculate_fft_result<<>>(d_output_fft, + const int out_len = static_cast(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<<>>(d_output_fft, out_len, + direction_vectors.size(), d_mag_t); + cuda_err(cudaGetLastError()); + calculate_fft_result<<>>(d_mag_t, max_length_A, min_length_A, histogram_size, bg_half, direction_vectors.size(), d_result_fft); diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 7fd983129..53f284e53 100644 --- a/rugnux/Rugnux.cpp +++ b/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(indexing_settings, IndexerConstruction::OnFirstUse,