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