From bb55715d28d82394360ecf8132ef64e6d5fc4773 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 19 Sep 2026 23:06:36 +0200 Subject: [PATCH] RotationScaleMerge: stream the resident ingest instead of materialising a record per observation The ingest-time geometry smoothing (delta_phi from the smoothed lattice, one exact-Bragg angle per rocking event, partiality from the smoothed mosaicity) used to run over a 32-byte host record per observation, copied from the source reflections and held beside them. It now reads hkl, frame, d and zeta straight from the source reflections, frame by frame, and keeps arrays only for what it rewrites (delta_phi, partiality) and what the rocking-event walks read in raw-hkl order (image_number, and on the resident path the usability flag): 13 bytes an observation instead of 32. The full-Obs path runs the same code and copies the two rewritten fields back. Same arithmetic in the same order; merged output byte-identical. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_013nW6FNRP1bBJJ8pfHiByAT --- .../scale_merge/RotationScaleMerge.cpp | 147 +++++++++--------- .../scale_merge/RotationScaleMerge.h | 57 +++---- 2 files changed, 107 insertions(+), 97 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index c251c7fb7..ae7410a61 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -303,8 +303,8 @@ void RotationScaleMerge::Ingest() { partials_released = false; resident_ingest = false; #ifdef JFJOCH_USE_CUDA - // Probed before the observation array is built: which record the build lays down (the narrow - // resident one or the full Obs) depends on whether the device pipeline will be active. + // Probed before the observations are built: whether the build lays down the full Obs array or + // only the ingest arrays depends on whether the device pipeline will be active. gpu_ = std::make_unique(); gpu_active_ = gpu_->Available(); resident_ingest = gpu_active_ && observation_dump_path.empty(); @@ -500,6 +500,15 @@ void RotationScaleMerge::Ingest() { n_partials_obs, n_frames, rawrun_start.size()); SmoothMosaicityAndPartiality(); + if (!resident_ingest) { + ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), + [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) { + partials[i].delta_phi = ingest_delta_phi[i]; + partials[i].partiality = ingest_partiality[i]; + } + }); + } // The corr every pass starts from (see the header). The resident build filled it already, // straight from the source reflections (and the smoothing between there and here never @@ -525,7 +534,7 @@ void RotationScaleMerge::Ingest() { 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] = partials_ingest[at + i].partiality; + for (int i = lo; i < hi; ++i) buf[i] = ingest_partiality[at + i]; }); gpu_->SetObsField(F::Partiality, at, cnt, buf.get()); } @@ -577,16 +586,21 @@ 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 host record's + // 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 the record back. The dump path keeps the full - // Obs array: the CPU combine is what writes the dump. + // (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(); - std::vector().swap(partials_ingest); partials_released = true; } #endif + std::vector().swap(ingest_delta_phi); + std::vector().swap(ingest_partiality); + std::vector().swap(ingest_image_number); + std::vector().swap(ingest_rock_ok); + std::vector().swap(ingest_keep); + std::vector().swap(ingest_src_start); } // Build `partials` (and the per-frame CSR, the finiteness mask and `perm`) from the source reflections, @@ -623,14 +637,15 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector &keys) { // Where each frame's reflections sit in the key numbering, and how many of them survive. Building in // index order keeps each frame's observations contiguous and in the order they had, so the per-frame // ranges still describe them and every per-frame sum keeps its sequence of roundings. - std::vector src_start(n_frames, 0); + ingest_src_start.assign(n_frames, 0); { int32_t at = 0; for (int o = 0; o < n_frames; ++o) { - src_start[o] = at; + ingest_src_start[o] = at; at += static_cast(partials_out[o].reflections.size()); } } + const std::vector &src_start = ingest_src_start; frame_count.assign(n_frames, 0); ParallelFor(n_frames, nthreads, [&](int o) { int32_t c = 0; @@ -696,14 +711,16 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector &keys) { // finiteness (immutable) - it lets ComputeAsuGroups stamp the ASU-group id per obs from a flat // 1-byte array instead of re-reading the fat Obs struct for every space group. finite_ok.assign(n_keep, 0); + ingest_delta_phi.resize(n_keep); + ingest_partiality.resize(n_keep); + ingest_image_number.resize(n_keep); #ifdef JFJOCH_USE_CUDA if (resident_ingest) { - // The resident path: build the narrow host record, 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. - partials_ingest.clear(); - partials_ingest.resize(n_keep); + // 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. + ingest_rock_ok.resize(n_keep); corr_ingested.resize(n_keep); ParallelFor(n_frames, nthreads, [&](int o) { int32_t src = src_start[o]; @@ -711,25 +728,18 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector &keys) { for (const auto &r : partials_out[o].reflections) { if (!keep[src]) { ++src; continue; } ++src; - ObsIngest &obs = partials_ingest[at]; - obs.h = static_cast(r.h); - obs.k = static_cast(r.k); - obs.l = static_cast(r.l); - obs.frame = o; - obs.image_number = r.image_number; - obs.d = r.d; - obs.zeta = r.zeta; - obs.delta_phi = r.delta_phi_deg; - obs.partiality = r.partiality; + ingest_delta_phi[at] = r.delta_phi_deg; + ingest_partiality[at] = r.partiality; + ingest_image_number[at] = r.image_number; const float prescaling_corr = static_cast( r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]); corr_ingested[at] = r.image_scale_corr; finite_ok[at] = (std::isfinite(r.I) && std::isfinite(prescaling_corr) && prescaling_corr != 0.0f && std::isfinite(r.sigma) && r.sigma > 0.0f) ? 1 : 0; - obs.rock_ok = (std::isfinite(r.image_scale_corr) && r.image_scale_corr > 0.0f - && std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f) - ? 1 : 0; + ingest_rock_ok[at] = (std::isfinite(r.image_scale_corr) && r.image_scale_corr > 0.0f + && std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f) + ? 1 : 0; ++at; } }); @@ -816,6 +826,9 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector &keys) { obs.group = -1; finite_ok[at] = (std::isfinite(obs.I) && std::isfinite(obs.prescaling_corr) && obs.prescaling_corr != 0.0f && std::isfinite(obs.sigma) && obs.sigma > 0.0f) ? 1 : 0; + ingest_delta_phi[at] = r.delta_phi_deg; + ingest_partiality[at] = r.partiality; + ingest_image_number[at] = r.image_number; ++at; } }); @@ -824,6 +837,7 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector &keys) { #endif n_partials_obs = n_keep; + ingest_keep = std::move(keep); if (n_keep < n_obs) logger.Info("RotationScaleMerge: dropped {} of {} observations ({} of {} distinct hkl) outside the " @@ -891,7 +905,7 @@ namespace { // window would over-smooth a fast-slipping crystal and under-smooth a sparse one. The frame's own // value is excluded when scoring, so a window is only preferred if it PREDICTS that frame better, // which is exactly the noise-versus-drift trade. -template void RotationScaleMerge::SmoothGeometry(std::vector &obs) { +void RotationScaleMerge::SmoothGeometry() { const auto gon = x.GetGoniometer(); if (!gon || n_frames < 5) return; @@ -998,37 +1012,31 @@ template void RotationScaleMerge::SmoothGeometry(std::vector &obs) } std::atomic changed{0}; - ParallelChunks(static_cast(obs.size()), nthreads, [&](int lo, int hi) { + ParallelFor(n_frames, nthreads, [&](int f) { + if (!ok[f]) + return; + const std::array &rc = frame_recip[f]; int64_t local = 0; - for (int i = lo; i < hi; ++i) { - auto &o = obs[i]; - if (o.frame < 0 || o.frame >= n_frames || !ok[o.frame]) - continue; - const std::array &r = frame_recip[o.frame]; - const Coord p0 = r[0] * static_cast(o.h) + r[1] * static_cast(o.k) - + r[2] * static_cast(o.l); - const float phi_near = o.delta_phi * static_cast(PI) / 180.0f; + ForEachKeptIn(f, [&](int i, const Reflection &r) { + const Coord p0 = rc[0] * static_cast(r.h) + rc[1] * static_cast(r.k) + + rc[2] * static_cast(r.l); + const float phi_near = ingest_delta_phi[i] * static_cast(PI) / 180.0f; const auto phi = SolvePhi(p0, S0, m1, m2, m3, m2_S0, m3_S0, phi_near); if (!phi || !std::isfinite(*phi)) - continue; - o.delta_phi = *phi * 180.0f / static_cast(PI); + return; + ingest_delta_phi[i] = *phi * 180.0f / static_cast(PI); ++local; - } + }); changed += local; }); logger.Info("Smoothed per-frame geometry over +-{} frames (chosen by cross-validation); " "recomputed delta_phi for {} of {} partials", best_half, changed.load(), - obs.size()); + n_partials_obs); } void RotationScaleMerge::SmoothMosaicityAndPartiality() { - if (resident_ingest) SmoothMosaicityAndPartialityOn(partials_ingest); - else SmoothMosaicityAndPartialityOn(partials); -} - -template void RotationScaleMerge::SmoothMosaicityAndPartialityOn(std::vector &obs) { - SmoothGeometry(obs); + SmoothGeometry(); // One rocking event, one exact-Bragg angle. Every partial's delta_phi is solved from its OWN // frame's lattice - by the predictor, and again by SmoothGeometry above - so the per-frame @@ -1054,25 +1062,25 @@ template void RotationScaleMerge::SmoothMosaicityAndPartialityOn(std:: int i = lo; while (i < hi) { int j = i + 1; - while (j < hi && obs[perm[j]].image_number - - obs[perm[j - 1]].image_number <= max_frame_gap) + while (j < hi && ingest_image_number[perm[j]] + - ingest_image_number[perm[j - 1]] <= max_frame_gap) ++j; if (j - i >= 2) { - const float f0 = obs[perm[i]].image_number; + const float f0 = ingest_image_number[perm[i]]; double sum = 0.0; int n = 0; for (int m = i; m < j; ++m) { - const auto &o = obs[perm[m]]; - if (!std::isfinite(o.delta_phi)) continue; - sum += o.delta_phi + increment_deg * (o.image_number - f0); + const int p = perm[m]; + if (!std::isfinite(ingest_delta_phi[p])) continue; + sum += ingest_delta_phi[p] + increment_deg * (ingest_image_number[p] - f0); ++n; } if (n >= 2) { const double c = sum / n; for (int m = i; m < j; ++m) { - auto &o = obs[perm[m]]; - o.delta_phi = static_cast(c - increment_deg - * (o.image_number - f0)); + const int p = perm[m]; + ingest_delta_phi[p] = static_cast(c - increment_deg + * (ingest_image_number[p] - f0)); } ++l_ev; l_pa += j - i; @@ -1145,20 +1153,19 @@ template void RotationScaleMerge::SmoothMosaicityAndPartialityOn(std:: // partiality recomputed here matches the one integration was based on. sin(theta_B) = lambda/(2d). const double bandwidth_sigma = x.GetBandwidthFWHM().value_or(0.0f) / 2.3548; const double half_wavelength_A = x.GetWavelength_A() / 2.0; - ParallelChunks(static_cast(obs.size()), nthreads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - auto &o = obs[i]; - float mos = mos_smooth[o.frame]; - if (bandwidth_sigma > 0.0 && o.d > 0.0f) { - const double sin_theta = half_wavelength_A / o.d; + ParallelFor(n_frames, nthreads, [&](int f) { + ForEachKeptIn(f, [&](int i, const Reflection &r) { + float mos = mos_smooth[f]; + if (bandwidth_sigma > 0.0 && r.d > 0.0f) { + const double sin_theta = half_wavelength_A / r.d; const double sigma_bw = (180.0 / PI) * bandwidth_sigma * sin_theta / std::sqrt(1.0 - sin_theta * sin_theta); mos = static_cast(std::sqrt(mos * mos + sigma_bw * sigma_bw)); } - if (std::isfinite(mos) && mos > 1e-6f && std::isfinite(o.zeta) && o.zeta > 0.0f - && std::isfinite(o.delta_phi)) - o.partiality = RotationPartiality(o.delta_phi, o.zeta, mos, wedge); - } + if (std::isfinite(mos) && mos > 1e-6f && std::isfinite(r.zeta) && r.zeta > 0.0f + && std::isfinite(ingest_delta_phi[i])) + ingest_partiality[i] = RotationPartiality(ingest_delta_phi[i], r.zeta, mos, wedge); + }); }); logger.Info("Recomputed partiality from frame-order-smoothed mosaicity"); } @@ -2139,10 +2146,10 @@ namespace { int RotationScaleMerge::RockingEventFrames() const { if (partials_released) return rocking_event_frames_at_ingest; - if (resident_ingest) // Ingest's own call, before the narrow record is handed back + if (resident_ingest) // Ingest's own call, before the ingest arrays are handed back return RockingEventFramesOver( - [&](int i) { return partials_ingest[i].rock_ok != 0; }, - [&](int i) { return partials_ingest[i].image_number; }); + [&](int i) { return ingest_rock_ok[i] != 0; }, + [&](int i) { return ingest_image_number[i]; }); return RockingEventFramesOver( [&](int i) { const Obs &o = partials[i]; diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 81cf8164c..5fd1e8ca8 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -140,26 +140,6 @@ private: // re-run over the samples that survive the automatic resolution cutoff. See MergeAndStats. struct Sample { double s2, I2, dev2; float d; }; - // The record the ingest keeps on the host when the device pipeline is resident (a GPU and no - // observation dump): every immutable per-obs field is uploaded straight from the source - // reflections inside BuildInRangeObservations, and the host keeps only what its own remaining - // stages touch - the geometry smoothing (delta_phi, then the recomputed partiality), plus the - // rocking-event walk's inputs. Thirty-two bytes against the Obs's eighty, and on that path the fat - // record is never materialised at all: on a fine-sliced long axis it alone is many gigabytes, - // held beside the source reflections it was copied from. - struct ObsIngest { - // Sixteen bits are plenty for a Miller index - |h| <= a / d_min, in the hundreds even on the - // longest axis at atomic resolution - and on a long axis this array sets the ingest's memory - // high-water mark, so each index costs two bytes here, not four. - int16_t h, k, l; - // The rocking-event usability test (RockingEventFrames), evaluated on the ingest-time - // values - which is what that walk reads on the resident path, where the host corr is - // never refreshed (see partials_released). Beside the indices, where it fills their padding. - uint8_t rock_ok; - int32_t frame; - float image_number, d, zeta, delta_phi, partiality; - }; - // The narrow per-observation record the ingest sort orders: the raw hkl the runs are cut on, the // frame position that breaks a tie inside one, the observation's own index (which makes the order // total - see the .cpp), and the resolution the range test reads. Twenty-four bytes against the @@ -223,12 +203,36 @@ private: int report_shell_count = 9; std::optional report_range; // the second, reference-range table - // Flat buffers, allocated once by Ingest() and reused across Run() calls. Exactly one of the - // two observation arrays is built: the narrow resident record when the device pipeline is - // active and no observation dump was asked for (resident_ingest), the full Obs otherwise. + // Flat buffers, allocated once by Ingest() and reused across Run() calls. The full Obs array is + // built only when the device pipeline is NOT resident (no GPU, or an observation dump was asked + // for); on the resident path (resident_ingest) the device holds the observations and the host + // keeps only the ingest arrays below, until the end of Ingest. std::vector partials; // all per-frame partials, grouped by frame - std::vector partials_ingest; // the resident-path record (same order) bool resident_ingest = false; + // What the ingest-time geometry smoothing (SmoothMosaicityAndPartiality) needs beyond the source + // reflections, one array per field in `partials` order: the two fields it rewrites, and the two the + // rocking-event walks read in raw-hkl order, where a source reflection is out of reach. Every other + // field it reads - hkl, frame, d, zeta - is read from the source reflection itself (ForEachKeptIn), + // so on the resident path no per-observation record is built at all: on a fine-sliced long axis + // that record was gigabytes, held beside the source reflections it was copied from. Filled by + // BuildInRangeObservations and handed back at the end of Ingest. + std::vector ingest_delta_phi, ingest_partiality, ingest_image_number; + // The rocking-event usability test (RockingEventFrames) on the ingest-time values - which is what + // that walk reads on the resident path, where the host corr is never refreshed (see + // partials_released). Resident path only. + std::vector ingest_rock_ok; + // Which source reflections the resolution range kept (indexed like the sort keys: frame by frame, + // from ingest_src_start), so that the i-th kept one of a frame can be found again. + std::vector ingest_keep; + std::vector ingest_src_start; + // Call fn(observation index, source reflection) for each reflection of frame o that the ingest + // kept, in order. + template void ForEachKeptIn(int o, Fn fn) const { + int32_t src = ingest_src_start[o]; + int32_t at = frame_start[o]; + for (const auto &r : partials_out[o].reflections) + if (ingest_keep[src++]) fn(at++, r); + } // How many observations Ingest built (== partials.size() until any release below), so the // sizes survive when the array itself does not. int n_partials_obs = 0; @@ -445,9 +449,8 @@ private: // per-frame partials of one rocking event tile the curve consistently (they sum toward 1) before the // 3D combine. Deterministic (frame order); replaces the old arrival-order mosaicity moving average // that prediction applied. SG-independent, so done once in Ingest. - template void SmoothGeometry(std::vector &obs); - void SmoothMosaicityAndPartiality(); // dispatch: the resident narrow record or the full Obs - template void SmoothMosaicityAndPartialityOn(std::vector &obs); + void SmoothGeometry(); + void SmoothMosaicityAndPartiality(); void Combine(); // partials -> fulls (CPU)