From de911fe1266d29a0ccd5bb77b8e78aee7a3c24ff Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 13:08:44 +0200 Subject: [PATCH] rugnux: index at the measured beam centre beside the file centre's first pass, on a second GPU The beam-centre check runs a second first pass at the centre measured from the background, and reads the first pass's answer only once that pass is over. With two GPUs or more, its spot finding (on the second card), scheme feeding and FFT indexing now start beside the file centre's pass, on a copy of the experiment at the measured centre with workers, a mapping and a spot list of their own (SpotFindingWorkers). The check takes them only where the experiment (rotation angles included), SpotFindingKey at the measured centre and the spot budget are still what the run has there, and scores the schemes on the run's own indexer as before; otherwise it runs that pass itself as before. The spots at the two centres are not shared: the spot list depends on the centre (ring background, resolution, budget). With one GPU, and in the CPU build, nothing changes: on one card the two passes share the card and the reader, and side by side they saved nothing. FindFirstPassSpots is the first pass's per-frame spot finding, now shared; ExperimentKey and SpotFindingKey take an experiment other than the run's. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- rugnux/Rugnux.cpp | 27 +++-- rugnux/Rugnux.h | 13 +- rugnux/RugnuxFirstPass.cpp | 237 +++++++++++++++++++++++++++++++++---- 3 files changed, 240 insertions(+), 37 deletions(-) diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index e04cb0658..a271fe370 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -483,12 +483,12 @@ void Rugnux::CollectSpeculativeProbe() { speculative_probe_.reset(); } -std::vector Rugnux::ExperimentKey(bool rotation_angles) const { - std::vector k{experiment_.GetBeamX_pxl(), experiment_.GetBeamY_pxl(), - experiment_.GetDetectorDistance_mm(), experiment_.GetPoniRot1_rad(), - experiment_.GetPoniRot2_rad(), experiment_.GetWavelength_A(), - static_cast(experiment_.GetMaxSpotCount())}; - if (const auto g = experiment_.GetGoniometer()) { +std::vector Rugnux::ExperimentKey(const DiffractionExperiment &e, bool rotation_angles) const { + std::vector k{e.GetBeamX_pxl(), e.GetBeamY_pxl(), + e.GetDetectorDistance_mm(), e.GetPoniRot1_rad(), + e.GetPoniRot2_rad(), e.GetWavelength_A(), + static_cast(e.GetMaxSpotCount())}; + if (const auto g = e.GetGoniometer()) { const Coord a = g->GetAxis(); const auto helical = g->GetHelicalStep().value_or(Coord(-1, -1, -1)); k.insert(k.end(), {1.0, a.x, a.y, a.z, helical.x, helical.y, helical.z}); @@ -497,11 +497,11 @@ std::vector Rugnux::ExperimentKey(bool rotation_angles) const { g->GetScreeningWedge().value_or(-1.0f)}); } else k.push_back(0.0); - if (const auto uc = experiment_.GetUnitCell()) + if (const auto uc = e.GetUnitCell()) k.insert(k.end(), {1.0, uc->a, uc->b, uc->c, uc->alpha, uc->beta, uc->gamma}); else k.push_back(0.0); - const auto &x = experiment_.GetIndexingSettings(); + const auto &x = e.GetIndexingSettings(); k.insert(k.end(), {x.GetTolerance(), x.GetFFT_MinUnitCell_A(), x.GetFFT_MaxUnitCell_A(), static_cast(x.GetFFT_NumVectors()), x.GetFFT_HighResolution_A(), x.GetFFT_MinAngle_deg(), x.GetFFT_MaxAngle_deg(), x.GetUnitCellDistTolerance(), @@ -514,23 +514,24 @@ std::vector Rugnux::ExperimentKey(bool rotation_angles) const { k.insert(k.end(), {f.signal_to_noise_threshold, static_cast(f.photon_count_threshold), static_cast(f.min_pix_per_spot.value_or(-1)), static_cast(f.max_pix_per_spot), f.high_resolution_limit.value_or(-1.0f)}); - k.insert(k.end(), {static_cast(experiment_.GetSpaceGroupOrP1().number), - static_cast(experiment_.GetImageNum())}); + k.insert(k.end(), {static_cast(e.GetSpaceGroupOrP1().number), + static_cast(e.GetImageNum())}); return k; } -std::vector Rugnux::SpotFindingKey(const AzimuthalIntegrationMapping &mapping) const { +std::vector Rugnux::SpotFindingKey(const DiffractionExperiment &x, + const AzimuthalIntegrationMapping &mapping) const { // Without the rotation angles: a frame's spots are found the same under any of them, and the angle // each spot carries is stamped again where a list is taken from the store (prefetch_spots). The // probes of the rotation-scale walk differ only there, so they find the spots of a frame once. - std::vector k = ExperimentKey(/*rotation_angles=*/false); + std::vector k = ExperimentKey(x, /*rotation_angles=*/false); const auto &f = config_.spot_finding; k.insert(k.end(), {static_cast(f.enable), f.low_resolution_limit.value_or(-1.0f), f.cutoff_spot_count_low_res, f.high_res_gap_Q_recipA.value_or(-1.0f), f.ice_ring_width_Q_recipA, static_cast(f.adaptive_threshold), f.false_pixels_per_frame, static_cast(f.measured_ring_q_recipA.size())}); k.insert(k.end(), f.measured_ring_q_recipA.begin(), f.measured_ring_q_recipA.end()); - k.insert(k.end(), {static_cast(experiment_.IsDetectIceRings()), + k.insert(k.end(), {static_cast(x.IsDetectIceRings()), static_cast(pixel_mask_.GetBinaryMaskChecksum()), static_cast(mapping.GetPixelToBinChecksum()), static_cast(mapping.GetCorrectionsChecksum())}); diff --git a/rugnux/Rugnux.h b/rugnux/Rugnux.h index 97b91a4de..355e1c0b3 100644 --- a/rugnux/Rugnux.h +++ b/rugnux/Rugnux.h @@ -787,7 +787,12 @@ class Rugnux { // not evaluate it again. Shared with the copies the run makes of itself. std::shared_ptr azint_geometry_ = std::make_shared(); - [[nodiscard]] std::vector SpotFindingKey(const AzimuthalIntegrationMapping &mapping) const; + [[nodiscard]] std::vector SpotFindingKey(const AzimuthalIntegrationMapping &mapping) const { + return SpotFindingKey(experiment_, mapping); + } + // The same, for spots found at another experiment (a trial beam centre) against its mapping. + [[nodiscard]] std::vector SpotFindingKey(const DiffractionExperiment &x, + const AzimuthalIntegrationMapping &mapping) const; void AddFirstPassMemo(const FirstPassMemo &memo); // The geometry walk's next indexing probe, started on a copy of the run as soon as the canonical // pass has post-refined - it depends on nothing that pass does afterwards - so that it runs beside @@ -809,7 +814,11 @@ class Rugnux { [[nodiscard]] std::vector FirstPassInputKey() const; // The part of FirstPassInputKey that is the experiment and the spot-finding settings; without the // goniometer's angles where rotation_angles is false (SpotFindingKey). - [[nodiscard]] std::vector ExperimentKey(bool rotation_angles = true) const; + [[nodiscard]] std::vector ExperimentKey(bool rotation_angles = true) const { + return ExperimentKey(experiment_, rotation_angles); + } + // The same, of another experiment beside the run's settings. + [[nodiscard]] std::vector ExperimentKey(const DiffractionExperiment &x, bool rotation_angles = true) const; // The file's detector distance, from before the first pass of the rotation two-pass: a pass whose // distance is still this one asks the post-refinement to test "the header distance is right" // (PostRefineSettings::distance_at_header); a pass that has walked off it does not. diff --git a/rugnux/RugnuxFirstPass.cpp b/rugnux/RugnuxFirstPass.cpp index eb257aa03..599a59494 100644 --- a/rugnux/RugnuxFirstPass.cpp +++ b/rugnux/RugnuxFirstPass.cpp @@ -88,6 +88,105 @@ using namespace rugnux_internal; namespace { + // The spots of one first-pass frame. A frame's spots are a pure function of that frame and the + // settings - the engine carries nothing from one image to the next - so it does not matter which + // engine, or which thread, finds them. + std::vector FindFirstPassSpots(JFJochReader &reader, const SpotFindingSettings &settings, + const JFJochReaderDataset &dataset, + MXAnalysisWithoutFPGA &analysis, + AzimuthalIntegrationProfile &profile, + JFJochReaderRawImage &img, int ordinal, int image_idx, + Logger &logger) { + std::vector spots; + try { + if (reader.ReadRawImage(image_idx, img)) { + DataMessage m{}; + m.number = ordinal; + m.original_number = image_idx; + auto first_pass = settings; + first_pass.indexing = false; + first_pass.quick_integration = false; + m.image = img.image; + if (dataset.efficiency.size() > image_idx) + m.image_collection_efficiency = dataset.efficiency[image_idx]; + analysis.Analyze(m, profile, first_pass); + spots = std::move(m.spots); + } + } catch (const std::exception &e) { + if (IsFatalResourceError(e)) throw; + logger.Warning("First-pass spot read failed for image {}: {}", image_idx, e.what()); + } + return spots; + } + + // First-pass spot finding on workers of its own, all on one card (`device`), for work that runs + // beside the first pass. Each worker keeps its engine from one call to the next - building one + // costs far more than a frame - and gives it back on that card. + class SpotFindingWorkers { + public: + SpotFindingWorkers(JFJochReader &reader, const DiffractionExperiment &x, + const AzimuthalIntegrationMapping &mapping, const PixelMask &mask, + const SpotFindingSettings &settings, const JFJochReaderDataset &dataset, + int start_image, int stride, size_t nworkers, int device) + : reader(reader), x(x), mapping(mapping), mask(mask), settings(settings), dataset(dataset), + start_image(start_image), stride(stride), device(device), engines(std::max(nworkers, 1)) {} + + ~SpotFindingWorkers() { + OnWorkers(engines.size(), [&](size_t t) { engines[t] = Engine{}; }); + } + + // The spots of each ordinal, in the order given. + std::vector> Find(const std::vector &ordinals) { + std::vector> found(ordinals.size()); + std::atomic next{0}; + OnWorkers(std::min(engines.size(), ordinals.size()), [&](size_t t) { + Engine &e = engines[t]; + if (!e.analysis) { + e.indexer = std::make_unique(x, nullptr, /*retain_outcomes=*/false); + e.analysis = std::make_unique(x, mapping, mask, *e.indexer, + /*enable_fused_adaptive_gpu=*/true); + e.profile = std::make_unique(mapping); + } + for (size_t i = next.fetch_add(1); i < ordinals.size(); i = next.fetch_add(1)) + found[i] = FindFirstPassSpots(reader, settings, dataset, *e.analysis, *e.profile, + e.raw_image, ordinals[i], start_image + ordinals[i] * stride, + logger); + }); + return found; + } + + private: + struct Engine { + // Before the analysis engine, which page-locks these bytes for its uploads. + JFJochReaderRawImage raw_image; + std::unique_ptr indexer; + std::unique_ptr analysis; + std::unique_ptr profile; + }; + + // fn(t) on n threads of its own, on the card. + template void OnWorkers(size_t n, Fn fn) { + std::vector> futures; + for (size_t t = 0; t < n; t++) + futures.emplace_back(std::async(std::launch::async, [this, &fn, t] { + pin_gpu(device); + fn(t); + })); + for (auto &f : futures) + f.get(); + } + + JFJochReader &reader; + const DiffractionExperiment &x; + const AzimuthalIntegrationMapping &mapping; + const PixelMask &mask; + const SpotFindingSettings settings; + const JFJochReaderDataset &dataset; + const int start_image, stride, device; + std::vector engines; + Logger logger{"Rugnux"}; + }; + // Field by field, bitwise, for the spot memo's verification. bool SameSpots(const std::vector &a, const std::vector &b) { if (a.size() != b.size()) @@ -248,27 +347,8 @@ bool Rugnux::FirstPassRotationIndexing(PipelineLocals &p) { // not matter which engine, or which thread, finds them. const auto find_spots = [&](MXAnalysisWithoutFPGA &analysis, AzimuthalIntegrationProfile &profile, JFJochReaderRawImage &img, int ordinal) { - const int image_idx = start_image + ordinal * config_.stride; - std::vector spots; - try { - if (reader_.ReadRawImage(image_idx, img)) { - DataMessage m{}; - m.number = ordinal; - m.original_number = image_idx; - auto first_pass = config_.spot_finding; - first_pass.indexing = false; - first_pass.quick_integration = false; - m.image = img.image; - if (dataset->efficiency.size() > image_idx) - m.image_collection_efficiency = dataset->efficiency[image_idx]; - analysis.Analyze(m, profile, first_pass); - spots = std::move(m.spots); - } - } catch (const std::exception &e) { - if (IsFatalResourceError(e)) throw; - logger.Warning("First-pass spot read failed for image {}: {}", image_idx, e.what()); - } - return spots; + return FindFirstPassSpots(reader_, config_.spot_finding, *dataset, analysis, profile, img, + ordinal, start_image + ordinal * config_.stride, logger); }; // Find the spots of every ordinal in the list that is not cached yet, on several workers, and @@ -1200,7 +1280,97 @@ bool Rugnux::FirstPassRotationIndexing(PipelineLocals &p) { std::vector> futures; // last, so it is waited for before the rest goes }; std::optional short_spec; + // The beam-centre check below runs a second first pass at the centre measured from the + // background, and it reads the first one's answer only after that pass is over. So the spots at + // the measured centre are found, and both schemes fed and indexed, here beside the file centre's + // pass - on a copy of the experiment at that centre, with workers, a mapping and a spot list of + // their own - and taken at the check's place only if what they read is still what the run has + // there (the experiment with its rotation angles, SpotFindingKey at the measured centre, the spot + // budget); otherwise that pass is run there as before. The schemes are scored there in either + // case, on the run's own indexer. + // + // Only with a second GPU, whose card its spot finding takes: on one card the two passes share the + // card and the reader, and measured, running them side by side saved nothing. There the check + // runs its pass after the first, as it always did. + struct CentreSpeculation { + std::vector experiment_key; + std::vector key; + size_t spots_per_image = 0; + std::unique_ptr experiment; + std::unique_ptr mask; + std::unique_ptr mapping; + std::map> spots; + SchemeIndexers ris; + std::future done; // last, so it is waited for before the rest goes + }; + std::optional centre_spec; + const auto start_centre_speculation = [&] { + if (get_gpu_count() < 2) + return; + JoinBeamCenterCapture(); + if (cancelled_ || !config_.beam_center_check || !background_center_) + return; + auto &cs = centre_spec.emplace(); + cs.experiment = std::make_unique(experiment_); + cs.experiment->BeamX_pxl(background_center_->beam_x_pxl).BeamY_pxl(background_center_->beam_y_pxl); + cs.mask = std::make_unique(pixel_mask_); + cs.mapping = std::make_unique(*cs.experiment, *cs.mask, *azint_geometry_); + cs.experiment_key = ExperimentKey(*cs.experiment); + cs.key = SpotFindingKey(*cs.experiment, *cs.mapping); + cs.spots_per_image = first_pass_spots_per_image; + { + // What an earlier pass already found at this centre under these settings. + std::lock_guard lock(first_pass_spots_->m); + for (const auto &[k, spots] : first_pass_spots_->entries) + if (k == cs.key) + cs.spots = spots; + } + // At this pass's angles, as prefetch_spots takes them. + const DiffractionGeometry spot_geometry = cs.experiment->GetDiffractionGeometry(); + for (auto &[ordinal, spots] : cs.spots) + StampSpotAngles(spots, spot_geometry); + cs.done = std::async(std::launch::async, [this, &cs, &schemes, &validation, &indexer_pool, &run_indexing, + settings = config_.spot_finding, dataset, start_image, + nworkers = engines.size()] { + SpotFindingWorkers workers(reader_, *cs.experiment, *cs.mapping, *cs.mask, settings, *dataset, + start_image, config_.stride, nworkers, /*device=*/1); + const auto find = [&](const std::vector &ordinals) { + std::vector wanted; + for (const int o : ordinals) + if (!cs.spots.contains(o)) + wanted.push_back(o); + auto found = workers.Find(wanted); + for (size_t i = 0; i < wanted.size(); i++) + cs.spots.emplace(wanted[i], std::move(found[i])); + }; + // As feed_scheme, at the measured centre. + const auto gonio = cs.experiment->GetGoniometer(); + for (const auto &[name, ordinals] : schemes) { + auto ri = std::make_unique(*cs.experiment, *indexer_pool); + ri->MaxSpotsPerImage(cs.spots_per_image); + ri->VerifyMemo(config_.verify_first_pass_memo); + for (size_t i = 0; i < ordinals.size(); i++) { + if (cancelled_ || ri->AccumulationFull()) + break; + if (!cs.spots.contains(ordinals[i])) + find({ordinals.begin() + i, + ordinals.begin() + std::min(ordinals.size(), i + SPOT_PREFETCH_CHUNK)}); + std::optional angle; + if (gonio) + angle = gonio->GetAngle_deg(static_cast(ordinals[i])) + gonio->GetWedge_deg() / 2.0f; + ri->ProcessImage(ordinals[i], cs.spots.at(ordinals[i]), angle); + } + cs.ris.push_back(std::move(ri)); + } + find(validation); + std::vector rp; + for (const auto &ri : cs.ris) + rp.push_back(ri.get()); + run_indexing(rp); + }); + }; FirstPass best = pick_best(*indexer_pool, *indexer, [&] { + start_centre_speculation(); if (cancelled_ || experiment_.GetUnitCell().has_value()) return; auto &sp = short_spec.emplace(); @@ -1322,7 +1492,30 @@ bool Rugnux::FirstPassRotationIndexing(PipelineLocals &p) { { const int majority = static_cast(validation.size()) / 2; try_beam_center(measured_x, measured_y); - FirstPass alt = pick_best(*indexer_pool, *indexer); + FirstPass alt; + if (centre_spec) + centre_spec->done.get(); + if (centre_spec && centre_spec->experiment_key == ExperimentKey() + && centre_spec->key == SpotFindingKey(*spot_mapping) + && centre_spec->spots_per_image == first_pass_spots_per_image) { + // What pick_best would have found and fed here: hand the spots to the run's cache and + // its memo as prefetch_spots does, and score the schemes on the run's own indexer. + spot_cache = centre_spec->spots; + { + std::lock_guard lock(first_pass_spots_->m); + auto &entries = first_pass_spots_->entries; + if (std::none_of(entries.begin(), entries.end(), + [&](const auto &e) { return e.first == centre_spec->key; })) { + constexpr size_t MAX_SPOT_KEYS = 8; + if (entries.size() == MAX_SPOT_KEYS) + entries.erase(entries.begin()); + entries.emplace_back(centre_spec->key, centre_spec->spots); + } + } + alt = score_schemes(centre_spec->ris, *indexer); + } else + alt = pick_best(*indexer_pool, *indexer); + centre_spec.reset(); const bool header_indexes = best.result.has_value() && best.score > majority; bool measured_indexes = alt.result.has_value() && alt.score > majority; if (!header_indexes) {