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;