From 4158c48dfd59d4618b85fd3eb37f6b5602b90933 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 12 Oct 2026 22:38:34 +0200 Subject: [PATCH] rugnux: post-refinement gather beside the chain; quality guard beside the engine on one card Two waits on the critical path of a canonical rotation pass go: - the post-refinement gather (2-3.5 s on a 78 M-reflection sweep) now runs in the task that fits it, beside the chain; it reads only the reflections, and the chain waits for it before the space-group search, the first step that re-indexes them; - with one GPU the quality guard's merge used to have the card to itself while the main engine waited. It now runs beside the main engine wherever both fit the card's total memory (the main engine's DeviceBytes plus its reflection share for the guard, plus the 3 GB reserve the all-observation arm uses); otherwise it runs first as before and the main engine is ingested again after it. The main engine's first merge holds its per-frame write-back until the guard has read the outcomes (as the arm on a copy does), and the guard's numbers are recorded where the pass first needs them. Same merges, same numbers; only when the guard and the gather run changes. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- rugnux/RugnuxScaleMerge.cpp | 152 +++++++++++++++++++++++++----------- 1 file changed, 107 insertions(+), 45 deletions(-) diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index 4431125b4..2b7772f86 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -736,6 +736,10 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { std::optional rsm; std::optional guard_merge; // what the quality guard reads, where not the search merge std::optional prepass_postrefine_obs; + // The canonical pass's post-refinement gather, run beside the chain (see there): set once it has + // read the reflections, which nothing may change before. + std::shared_ptr> postrefine_gathered; + std::shared_future postrefine_gathered_future; // The rotation geometry post-refinement (see its call sites below for what it is for), in two // parts. prepare_post_refine reads the pass - the outcomes, the indexer, the scaling engine - and // returns the fit with everything it needs; fit_post_refine reads only that, so the canonical pass @@ -745,7 +749,9 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { PostRefineSettings prs; PostRefineObservations obs; }; - const auto prepare_post_refine = [&](size_t max_images = SIZE_MAX) -> std::optional { + // gather = false leaves the job's observations for the caller to gather. + const auto prepare_post_refine = [&](size_t max_images = SIZE_MAX, bool gather = true) + -> std::optional { const auto &outcomes = indexer->GetIntegrationOutcome(); if (geometry_prepass) { prepass_mosaicity_.assign(outcomes.size(), NAN); @@ -804,10 +810,11 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { rsm.reset(); auto pr_obs = prepass_postrefine_obs ? std::move(*prepass_postrefine_obs) - : GatherPostRefineObservations(indexer->GetIntegrationOutcome(), - config_.nthreads, - /*release_reflections=*/false, logger, - max_images); + : gather ? GatherPostRefineObservations(indexer->GetIntegrationOutcome(), + config_.nthreads, + /*release_reflections=*/false, logger, + max_images) + : PostRefineObservations{}; prepass_postrefine_obs.reset(); return PostRefineJob{*rot, prs, std::move(pr_obs)}; } @@ -874,11 +881,13 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { const bool merge_first_part = is_rotation && quality_guard_pass1_ && prepass_images < images_to_process && !experiment_.GetRefineRotationWedgeInScaling() && !rot_ss.GetRotationWedgeForScaling().has_value(); // With two GPUs or more the guard's merge has an engine of its own on the second card, and the - // main engine does not wait for it: its ingest and first merge run beside the guard's merge. The + // main engine does not wait for it: its ingest and first merge run beside the guard's merge. With + // one, the guard's engine goes beside the main one on the card where the two fit (see below). The // guard reads the outcomes only while it ingests - the first merge writes per-frame fields back - // into them (mosaicity among them, which the ingest reads) - so that merge waits for the guard's - // ingest and nothing else. Neither merge reads anything the other writes, and on cards of one - // model a merge gives the same bits on either, so this changes only when the guard runs. + // into them (mosaicity among them, which the ingest reads) - so that merge holds its write-back + // until the guard has ingested, and nothing else waits. Neither merge reads anything the other + // writes, and on cards of one model a merge gives the same bits on either, so this changes only + // when the guard runs. const bool guard_on_second_gpu = get_gpu_count() >= 2; std::promise guard_ingested; Logger guard_log = Logger::Buffered(); @@ -903,15 +912,13 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { return ScaleMergeResult{std::move(r.merged), std::move(r.statistics), r.cc_half_before_corrections}; }; std::future guard_ahead; - if (merge_first_part && !(postrefine_probe_ && !geometry_prepass && postrefine_probe_only_)) + if (merge_first_part && guard_on_second_gpu && !(postrefine_probe_ && !geometry_prepass && postrefine_probe_only_)) guard_ahead = std::async(std::launch::async, [&] { - if (guard_on_second_gpu) - pin_gpu(1); // the card and its NUMA node, as an image-loop worker has them + pin_gpu(1); // the card and its NUMA node, as an image-loop worker has them return merge_guard(); }); // The scaling engine's ingest reads the outcomes and nothing else either, so its host half runs - // beside the post-refinement too; with one GPU the device half waits for the guard to hand the - // card back. + // beside the post-refinement too. Logger engine_log = Logger::Buffered(); std::future engine_ahead; if (is_rotation && postrefine_probe_ && !geometry_prepass && !postrefine_probe_only_ @@ -976,17 +983,28 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { if (postrefine_probe_ && !geometry_prepass && is_rotation && !postrefine_probe_only_ && !(lattice_arm_asked && lattice_arm_check_->lost)) { // Nothing in the merge chain reads the fit (RunAllPasses does, after RunPipeline has joined - // it), so only the gather runs here: it reads the outcomes, which the chain re-indexes in - // place once the space group is settled. The fit and the speculative probe it may start - // run beside the chain; the probe's copy of the run is taken now, as the chain goes on to - // change the run's members. + // it), so it runs beside the chain, and so does the gather before it: the gather reads only + // the reflections, which nothing changes until the chain re-indexes them in place once the + // space group is settled - and the chain waits for the gather (postrefine_gathered) before + // its search. The speculative probe the fit may start runs beside the chain too; the probe's + // copy of the run is taken now, as the chain goes on to change the run's members. TimingMark mark(logger, "post-refinement observations"); - if (auto job = prepare_post_refine()) { + if (auto job = prepare_post_refine(SIZE_MAX, /*gather=*/false)) { std::shared_ptr probe_run; if (speculate_geometry_probe_) probe_run = std::make_shared(*this); + postrefine_gathered = std::make_shared>(); + postrefine_gathered_future = postrefine_gathered->get_future().share(); p.postrefine_ahead = std::async(std::launch::async, - [this, &p, fit_post_refine, job = std::move(*job), probe_run]() mutable { + [this, &p, fit_post_refine, job = std::move(*job), probe_run, + gathered = postrefine_gathered]() mutable { + { + // Set however the gather ends: a gather that failed rethrows from the future, later. + struct Gathered { std::promise &p; ~Gathered() { p.set_value(); } } done{*gathered}; + TimingMark gather_mark(p.postrefine_log, "post-refinement gather"); + job.obs = GatherPostRefineObservations(p.indexer->GetIntegrationOutcome(), config_.nthreads, + /*release_reflections=*/false, p.postrefine_log); + } TimingMark fit_mark(p.postrefine_log, "post-refinement"); fit_post_refine(std::move(job), p.postrefine_log, p.postrefine_notes); fit_mark.Stop(); @@ -1044,27 +1062,63 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Rotation scaling/merging (RotationScaleMerge) does not support " "wedge refinement"); - if (merge_first_part) { + if (merge_first_part) logger.Info("Two-pass: the quality guard reads this pass's merge of the first {} images, " "the ones the pre-pass integrated", prepass_images); - if (!guard_ahead.valid() || !guard_on_second_gpu) { - guard_merge = guard_ahead.valid() ? guard_ahead.get() : merge_guard(); - guard_log.ReplayInto(logger); - } - } // A reference MTZ is allowed for rotation: it fixes the space group / cell (on the CLI) and // resolves the indexing ambiguity (below), but is NOT used to scale - the rotation merge stays // self-consistent. - if (engine_ahead.valid()) { + if (engine_ahead.valid()) engine_ahead.get(); - engine_log.ReplayInto(logger); - rsm->SetLogger(logger); - } else { + else { rsm.emplace(experiment_, indexer->GetIntegrationOutcome(), result.consensus_cell, static_cast(config_.scaling_iter), - config_.nthreads, logger, config_.observation_dump_path); + config_.nthreads, engine_log, config_.observation_dump_path); rsm->IngestHost(); } + rsm->IngestDevice(); + const int32_t cards = get_gpu_count(); + constexpr size_t CARD_BYTES_RESERVED = size_t(3) << 30; + // With one card the guard's engine goes beside this one where the two fit the card's total, + // and otherwise has the card to itself first, this engine being ingested again after it. Its + // size is reckoned as this engine's share of the reflections its frames hold: it merges those + // frames' part of these same outcomes, so it holds no more observations per reflection, and no + // more rocking events or distinct hkl, than this engine does. Judged on the total, as the + // all-observation arm below is, so that where the guard runs does not depend on the moment. + size_t guard_bytes = 0; + if (merge_first_part && !guard_ahead.valid()) { + const auto &outcomes = indexer->GetIntegrationOutcome(); + double n_first = 0.0, n_all = 0.0; + for (size_t o = 0; o < outcomes.size(); ++o) { + n_all += static_cast(outcomes[o].reflections.size()); + if (o < static_cast(prepass_images)) + n_first += static_cast(outcomes[o].reflections.size()); + } + const auto guard_estimate = static_cast(n_all > 0.0 ? rsm->DeviceBytes() * (n_first / n_all) : 0.0); + const bool fits = rsm->DeviceBytes() + guard_estimate + CARD_BYTES_RESERVED <= get_gpu_total_memory(0); + if (cards == 1) + logger.Info("Two-pass: the quality guard's merge {} this pass's engine on the card ({:.1f} GB " + "beside {:.1f} GB on a {:.1f} GB card)", fits ? "runs beside" : "runs before", + guard_estimate / 1e9, rsm->DeviceBytes() / 1e9, get_gpu_total_memory(0) / 1e9); + if (cards == 1 && fits) { + guard_bytes = guard_estimate; + guard_ahead = std::async(std::launch::async, [&] { return merge_guard(); }); + } else { + if (cards > 0) + rsm.reset(); + guard_merge = merge_guard(); + guard_log.ReplayInto(logger); + if (!rsm) { + engine_log = Logger::Buffered(); + rsm.emplace(experiment_, indexer->GetIntegrationOutcome(), result.consensus_cell, + static_cast(config_.scaling_iter), + config_.nthreads, engine_log, config_.observation_dump_path); + rsm->Ingest(); + } + } + } + engine_log.ReplayInto(logger); + rsm->SetLogger(logger); // 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 @@ -1074,13 +1128,11 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // holds more, but a merge that does not fit beside it waits for it to end. On a 16 GB card that // is the data sets up to about 50 M observations at high multiplicity; the large ones merge // their arms one after the other. - rsm->IngestDevice(); - 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); - constexpr size_t CARD_BYTES_RESERVED = size_t(3) << 30; - if (cards >= 3 || 2 * rsm->DeviceBytes() + CARD_BYTES_RESERVED <= get_gpu_total_memory(device)) { + if (cards >= 3 + || 2 * rsm->DeviceBytes() + guard_bytes + CARD_BYTES_RESERVED <= 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 @@ -1246,25 +1298,26 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { return out; }; - // First pass: P1 when searching, or directly the user-fixed space group. A guard merging on the - // second GPU (above) is still running; this merge only waits for it to have read the outcomes. - const bool guard_beside = merge_first_part && guard_on_second_gpu && guard_ahead.valid(); + // First pass: P1 when searching, or directly the user-fixed space group. A guard merging beside it + // (above) may still be reading the outcomes, so this merge holds its write-back until it has. + const bool guard_beside = merge_first_part && guard_ahead.valid(); if (guard_beside) - guard_ingested.get_future().wait(); + rsm->SetWriteBackPerFrameScale(false); const auto initial_sg = experiment_.GetGemmiSpaceGroup(); auto sm = scale_and_merge(initial_sg ? initial_sg->short_name() : "P1", !initial_sg.has_value(), /*measure_cc=*/true); if (guard_beside) { - guard_merge = guard_ahead.get(); - guard_log.ReplayInto(logger); + guard_ingested.get_future().wait(); + rsm->WriteBackPerFrameScale(); + rsm->SetWriteBackPerFrameScale(true); } // The two-pass quality guard judges the second pass against the first on THIS merge, which both // passes make in the same terms - P1 (or the group the user fixed), full range, no correction // surfaces - where their final merges need not even be in the same group. See RunAllPasses. Where - // the pre-pass integrated only part of the sweep, it is the same merge of those frames (above). - { - const ScaleMergeResult &counted = guard_merge ? *guard_merge : sm; + // the pre-pass integrated only part of the sweep, it is the same merge of those frames (above); + // a guard merging beside the chain is read where the pass first needs it (below). + const auto record_search_merge = [&](const ScaleMergeResult &counted) { const auto &o = counted.statistics.overall; result.search_merge_completeness_measured = o.possible_unique_reflections > 0; result.search_merge_completeness = result.search_merge_completeness_measured @@ -1317,7 +1370,9 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { }); result.search_merge_reflections = std::count_if(counted.merged.begin(), counted.merged.end(), centring_allows); - } + }; + if (!guard_beside) + record_search_merge(guard_merge ? *guard_merge : sm); // Rotation two-pass geometry pre-pass: this first scale/merge has now fitted a frame-order-smoothed // mosaicity and written it back onto the per-frame outcomes. Capture it for the second pass's Bragg @@ -1446,11 +1501,18 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // A canonical pass the geometry walk is going to replace stops here, where the answer is // already in (GeometryWalkTakesRound); it is asked again below, before the analyses. Not with a // _process.h5 open: that file is written whole by every pass that opens it. + if (guard_beside) { + guard_merge = guard_ahead.get(); + guard_log.ReplayInto(logger); + record_search_merge(*guard_merge); + } if (!writer && GeometryWalkTakesRound(result, p.postrefine_ahead)) { logger.Info("The geometry walk will re-integrate at the geometry the post-refinement asks for, " "so this pass stops before its space-group search"); return true; } + if (postrefine_gathered_future.valid()) + postrefine_gathered_future.wait(); // the search re-indexes the reflections the gather reads TimingMark search_mark(logger, "space-group search"); if (search_space_group) { SearchSpaceGroupOptions sg_opts;