From 9cc8172da585a8980d3dc0e528f0a4cc663f1bae Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 7 Oct 2026 20:13:29 +0200 Subject: [PATCH] rugnux: ingest the scaling engine's host half beside the post-refinement RotationScaleMerge::Ingest is now IngestHost (read the outcomes, sort, build every host array, smooth the geometry and recompute the partiality) followed by IngestDevice (allocate the GPU's arrays and stage them). The host half reads the outcomes and nothing else, as the post-refinement and the quality guard's merge do, so on the canonical pass it runs on its own thread beside both; the device half runs where the ingest used to, after the guard has handed the GPU back, so the two engines never hold the card at once. The device feed now stages the recomputed partiality in the same slices as every other field, straight from the ingest array, instead of in a pass of its own after the rest; the values on the device are the ones they were. 8a1a, no other load, interleaved runs: the first merge (P1) ends at 10.9 s after the image pass against 11.3 with only the guard overlapped and 12.7 on the rc175 base. p.mtz byte-identical on the gate sets (myob, thau, cytc, kdp, cytc -N 4 / -N 8, --prepass-fraction 1 on four sets, 8a1a, 8tyy, CPU build myob); --model, stills, calibration, --mode scale (GPU and CUDA_VISIBLE_DEVICES="") outputs identical; --no-merge / --mode azint / --write-process-h5 _process.h5 differ only in the datasets that differ base-vs-base (azint sums, bkgEstimate, profiling times, dates, git). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- .../scale_merge/RotationScaleMerge.cpp | 170 +++++++++--------- .../scale_merge/RotationScaleMerge.h | 8 + rugnux/RugnuxScaleMerge.cpp | 27 ++- 3 files changed, 119 insertions(+), 86 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index d3410bd0f..37518848c 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -371,6 +371,11 @@ void RotationScaleMerge::MeasureIncidentFlux(const std::vector &mean_bkg } void RotationScaleMerge::Ingest() { + IngestHost(); + IngestDevice(); +} + +void RotationScaleMerge::IngestHost() { TimingMark mark(*logger, "rsm ingest"); scale_partials_choice.reset(); n_frames = static_cast(partials_out.size()); @@ -617,23 +622,86 @@ void RotationScaleMerge::Ingest() { }); #ifdef JFJOCH_USE_CUDA - // Feed the device (probed at the top of Ingest). On the resident path the layout and every - // immutable field went up inside BuildInRangeObservations, straight from the source - // reflections; what is staged here is the partiality, which the recompute above rewrote. - if (resident_ingest) { - using F = RotationScaleMergeGPU::ObsField; - constexpr int STAGE_SLICE = 8 << 20; - const int n = n_partials_obs; - const int slice = std::min(std::max(n, 1), STAGE_SLICE); - const std::unique_ptr buf(new float[slice]); - for (int at = 0; at < n; at += slice) { - const int cnt = std::min(n - at, slice); - ParallelChunks(cnt, nthreads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) buf[i] = ingest_partiality[at + i]; + // On the resident path the ingest arrays' one host reader after the device feed is + // RockingEventFrames - which on this path sees only ingest-time values (see the header) - so its + // answer is taken now, and IngestDevice hands the arrays back once the device holds the rest. + if (resident_ingest) + rocking_event_frames_at_ingest = RockingEventFrames(&rocking_events_at_ingest); +#endif +} + +#ifdef JFJOCH_USE_CUDA +// The resident path's device feed: every per-obs field the device reads, staged straight from the +// source reflections in frame-aligned bounded slices - the partiality from the ingest array, where the +// recompute left it. +void RotationScaleMerge::StageObservations() { + gpu_->SetPartialsLayout(n_partials_obs, n_frames, frame_start.data(), frame_count.data()); + using F = RotationScaleMergeGPU::ObsField; + // Two million observations a slice: the fifteen staging arrays are then about a hundred + // megabytes together, where eight million put four hundred on top of the ingest's peak, and a + // slice still spans tens of frames for the threads to share. + constexpr int STAGE_SLICE = 2 << 20; + constexpr int NFLOAT = 12; + int max_frame_obs = 0; + for (int o = 0; o < n_frames; ++o) max_frame_obs = std::max(max_frame_obs, static_cast(frame_count[o])); + const int bufsize = std::max(1, std::max(STAGE_SLICE, max_frame_obs)); + std::unique_ptr fbuf[NFLOAT]; + for (auto &b : fbuf) b.reset(new float[bufsize]); + const std::unique_ptr frm(new int32_t[bufsize]); + const std::unique_ptr onice(new uint8_t[bufsize]), clip(new uint8_t[bufsize]); + int f0 = 0; + while (f0 < n_frames) { + int f1 = f0 + 1; + while (f1 < n_frames && frame_start[f1] + frame_count[f1] - frame_start[f0] <= bufsize) + ++f1; + const int32_t base = frame_start[f0]; + const int cnt = frame_start[f1 - 1] + frame_count[f1 - 1] - base; + if (cnt > 0) { + ParallelFor(f1 - f0, nthreads, [&](int fo) { + const int o = f0 + fo; + int32_t src = ingest_src_start[o]; + int at = frame_start[o] - base; + for (const auto &r : partials_out[o].reflections) { + if (!ingest_keep[src]) { ++src; continue; } + ++src; + fbuf[0][at] = r.I; + fbuf[1][at] = r.sigma; + fbuf[2][at] = static_cast( + r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]); + fbuf[3][at] = r.zeta; + fbuf[4][at] = r.image_scale_corr; + fbuf[5][at] = r.bkg; + fbuf[6][at] = r.var_bkg; + fbuf[7][at] = r.image_number; + fbuf[8][at] = r.d; + fbuf[9][at] = r.predicted_x; + fbuf[10][at] = r.predicted_y; + fbuf[11][at] = ingest_partiality[base + at]; + frm[at] = o; + onice[at] = r.on_ice_ring ? 1 : 0; + clip[at] = (r.clipped ? OBS_CLIPPED : 0) | (r.overloaded ? OBS_OVERLOADED : 0); + ++at; + } }); - gpu_->SetObsField(F::Partiality, at, cnt, buf.get()); + constexpr F fields[NFLOAT] = {F::I, F::Sigma, F::PrescalingCorr, F::Zeta, F::Corr0, + F::Bkg, F::VarBkg, F::ImageNumber, F::D, F::Px, F::Py, + F::Partiality}; + for (int k = 0; k < NFLOAT; ++k) + gpu_->SetObsField(fields[k], base, cnt, fbuf[k].get()); + gpu_->SetObsFrame(base, cnt, frm.get()); + gpu_->SetObsOnIce(base, cnt, onice.get()); + gpu_->SetObsClipped(base, cnt, clip.get()); } + f0 = f1; } +} +#endif + +void RotationScaleMerge::IngestDevice() { + TimingMark mark(*logger, "rsm ingest (device)"); +#ifdef JFJOCH_USE_CUDA + if (resident_ingest) + StageObservations(); if (gpu_active_ && !resident_ingest) { const int n = static_cast(partials.size()); gpu_->SetPartialsLayout(n, n_frames, frame_start.data(), frame_count.data()); @@ -682,14 +750,10 @@ void RotationScaleMerge::Ingest() { gpu_->SetFrameCellOk(frame_cell_ok.data()); logger->Info("RotationScaleMerge: GPU scaling + combine + scale-fulls + merge active"); } - // The device now holds every per-obs field the resident pipeline reads, so the ingest arrays' - // one remaining reader is RockingEventFrames - which on this path sees only ingest-time values - // (see the header). Take its answer now and hand them back. The dump path keeps the full Obs - // array: the CPU combine is what writes the dump. - if (resident_ingest) { - rocking_event_frames_at_ingest = RockingEventFrames(&rocking_events_at_ingest); + // The device now holds every per-obs field the resident pipeline reads: the ingest arrays go. + // The dump path keeps the full Obs array: the CPU combine is what writes the dump. + if (resident_ingest) partials_released = true; - } #endif ingest_delta_phi.reset(); ingest_partiality.reset(); @@ -812,10 +876,9 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr &ke ingest_image_number.reset(new float[n_keep]); #ifdef JFJOCH_USE_CUDA if (resident_ingest) { - // The resident path: fill the ingest arrays, then stage every immutable per-obs field to the - // device straight from the source reflections, in frame-aligned bounded slices. The 80-byte - // Obs is never materialised here - beside the source reflections it was another copy of the - // whole payload - and the slice buffers are bounded whatever the run. + // The resident path: fill the ingest arrays only. The 80-byte Obs is never materialised - + // beside the source reflections it would be another copy of the whole payload - and the + // device is fed straight from the source reflections by StageObservations. ingest_rock_ok.reset(new uint8_t[n_keep]); corr_ingested.resize(n_keep); ParallelFor(n_frames, nthreads, [&](int o) { @@ -840,63 +903,6 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr &ke } }); - gpu_->SetPartialsLayout(n_keep, n_frames, frame_start.data(), frame_count.data()); - using F = RotationScaleMergeGPU::ObsField; - // Two million observations a slice: the thirteen staging arrays are then under a hundred - // megabytes together, where eight million put four hundred on top of the ingest's peak, and a - // slice still spans tens of frames for the threads to share. - constexpr int STAGE_SLICE = 2 << 20; - constexpr int NFLOAT = 11; - int max_frame_obs = 0; - for (int o = 0; o < n_frames; ++o) max_frame_obs = std::max(max_frame_obs, static_cast(frame_count[o])); - const int bufsize = std::max(1, std::max(STAGE_SLICE, max_frame_obs)); - std::unique_ptr fbuf[NFLOAT]; - for (auto &b : fbuf) b.reset(new float[bufsize]); - const std::unique_ptr frm(new int32_t[bufsize]); - const std::unique_ptr onice(new uint8_t[bufsize]), clip(new uint8_t[bufsize]); - int f0 = 0; - while (f0 < n_frames) { - int f1 = f0 + 1; - while (f1 < n_frames && frame_start[f1] + frame_count[f1] - frame_start[f0] <= bufsize) - ++f1; - const int32_t base = frame_start[f0]; - const int cnt = frame_start[f1 - 1] + frame_count[f1 - 1] - base; - if (cnt > 0) { - ParallelFor(f1 - f0, nthreads, [&](int fo) { - const int o = f0 + fo; - int32_t src = src_start[o]; - int at = frame_start[o] - base; - for (const auto &r : partials_out[o].reflections) { - if (!keep[src]) { ++src; continue; } - ++src; - fbuf[0][at] = r.I; - fbuf[1][at] = r.sigma; - fbuf[2][at] = static_cast( - r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]); - fbuf[3][at] = r.zeta; - fbuf[4][at] = r.image_scale_corr; - fbuf[5][at] = r.bkg; - fbuf[6][at] = r.var_bkg; - fbuf[7][at] = r.image_number; - fbuf[8][at] = r.d; - fbuf[9][at] = r.predicted_x; - fbuf[10][at] = r.predicted_y; - frm[at] = o; - onice[at] = r.on_ice_ring ? 1 : 0; - clip[at] = (r.clipped ? OBS_CLIPPED : 0) | (r.overloaded ? OBS_OVERLOADED : 0); - ++at; - } - }); - constexpr F fields[NFLOAT] = {F::I, F::Sigma, F::PrescalingCorr, F::Zeta, F::Corr0, - F::Bkg, F::VarBkg, F::ImageNumber, F::D, F::Px, F::Py}; - for (int k = 0; k < NFLOAT; ++k) - gpu_->SetObsField(fields[k], base, cnt, fbuf[k].get()); - gpu_->SetObsFrame(base, cnt, frm.get()); - gpu_->SetObsOnIce(base, cnt, onice.get()); - gpu_->SetObsClipped(base, cnt, clip.get()); - } - f0 = f1; - } } else { #endif partials.clear(); diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index f3af5d8cb..1e9f75483 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -116,6 +116,11 @@ public: // Copy the per-frame partials into the flat buffers. Call once before the first Run(). void Ingest(); + // Ingest in its two halves, which is what Ingest() calls: the host half reads the outcomes and + // builds every host array, the device half allocates the GPU's arrays and fills them. A caller can + // run the host half beside other work and the device half once the GPU has room. + void IngestHost(); + void IngestDevice(); // Scale (per-frame G) -> smooth G -> 3D combine -> scale fulls -> merge -> error model -> statistics // for the space group currently set on the experiment, reusing the ingested buffers. @@ -506,6 +511,9 @@ private: // observation array is built, so the two never coexist - together they would be another third of // the payload on top of it. void BuildInRangeObservations(std::unique_ptr &keys, int n_obs); +#ifdef JFJOCH_USE_CUDA + void StageObservations(); +#endif // Compute the dense ASU-group id for the current space group by grouping the (pre-sorted) raw-hkl // runs by their ASU key - one gemmi ASU reduction per distinct raw hkl, not per observation. Fills diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index 3d0c3dc7b..0862e9415 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -724,6 +724,18 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { std::future guard_ahead; if (merge_first_part && !(postrefine_probe_ && !geometry_prepass && postrefine_probe_only_)) guard_ahead = std::async(std::launch::async, merge_guard); + // The scaling engine's ingest reads the outcomes and nothing else either, so its host half runs + // beside the post-refinement too; the device half waits for the guard to hand the GPU back. + Logger engine_log = Logger::Buffered(); + std::future engine_ahead; + if (is_rotation && postrefine_probe_ && !geometry_prepass && !postrefine_probe_only_ + && !experiment_.GetRefineRotationWedgeInScaling() && !rot_ss.GetRotationWedgeForScaling().has_value()) + engine_ahead = std::async(std::launch::async, [&] { + rsm.emplace(experiment_, indexer->GetIntegrationOutcome(), result.consensus_cell, + static_cast(config_.scaling_iter), + config_.nthreads, engine_log, config_.observation_dump_path); + rsm->IngestHost(); + }); if (postrefine_probe_ && !geometry_prepass && is_rotation) { post_refine_geometry(); if (speculate_geometry_probe_ && !cancelled_ && prepass_detector_geometry_ && result.post_refine @@ -774,10 +786,17 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // 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. - rsm.emplace(experiment_, indexer->GetIntegrationOutcome(), result.consensus_cell, - static_cast(config_.scaling_iter), - config_.nthreads, logger, config_.observation_dump_path); - rsm->Ingest(); + if (engine_ahead.valid()) { + engine_ahead.get(); + engine_log.ReplayInto(logger); + rsm->SetLogger(logger); + } else { + rsm.emplace(experiment_, indexer->GetIntegrationOutcome(), result.consensus_cell, + static_cast(config_.scaling_iter), + config_.nthreads, logger, config_.observation_dump_path); + rsm->IngestHost(); + } + rsm->IngestDevice(); } // The geometry pre-pass reads the per-image reflections exactly twice more: they were just