From bef0c6f390d095b4b1d791dbcfbc918d8d0fc2c2 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 16 Sep 2026 11:53:50 +0200 Subject: [PATCH] scale/merge: the memory peak no longer holds three copies of the partials The offline rotation pipeline held every partial observation in up to three forms at once at its peak: the per-image Reflection vectors (88 B/obs), RotationScaleMerge's flat Obs array plus its full-length GPU staging (80 + 57 B/obs), and the geometry post-refinement's own gather (2 x 32 B/obs while it bucket-sorts). On a rotation crystal with a ~640 A axis, 3600 finely-sliced images and 95.7 M first-pass partials that peaked at 28-30 GB of RSS and OOMed a 32 GB machine. Four reductions, none of which changes an output byte: * The post-refine gather scatters each partial straight into its h bucket (histogram pass first), so the gathered-order copy of the whole array - one of two full-size copies alive through the sort - no longer exists. * The gather is split from the fit (GatherPostRefineObservations), and the geometry pre-pass takes it right after Ingest, handing back each image's reflection vector as it is read: nothing later in the pre-pass reads them (its reflections are never written, and the scaling engine reads back only per-frame fields). A run with a reference MTZ keeps them, since resolving the indexing ambiguity re-reads them. * The pre-pass frees the whole RotationScaleMerge engine - about twice the partial payload again - before the post-refinement fit, which sat exactly on the pass's peak. The canonical pass keeps everything, as before. * Ingest's transients are bounded: the sort keys are consumed before the flat Obs array is built instead of coexisting with it, and the GPU upload stages bounded slices (under half a GB) instead of fourteen full-length arrays at once (5 GB on that crystal). Measured there: peak RSS 28.3 / 30.0 GB (two baseline runs) -> 20.3-22.3 GB (four runs of the fix), wall clock unchanged within run-to-run variance, and byte-identical output - merged/unmerged/P1 MTZ data blocks, .hkl, .cif, per-image table and report all equal, on this crystal and on two ordinary rotation sentinels (which also keep their wall clock and drop slightly in peak). What remains of the peak is the canonical pass, which cannot hand its reflections back (they are re-seated and written); going lower means removing the Reflection/Obs duplication itself, which this deliberately does not attempt. Co-Authored-By: Claude Opus 5 (1M context) --- docs/CHANGELOG.md | 2 + image_analysis/geom_refinement/PostRefine.cpp | 255 ++++++++++-------- image_analysis/geom_refinement/PostRefine.h | 41 ++- .../scale_merge/RotationScaleMerge.cpp | 134 +++++---- .../scale_merge/RotationScaleMerge.h | 5 +- .../scale_merge/RotationScaleMergeGPU.cu | 70 +++-- .../scale_merge/RotationScaleMergeGPU.h | 34 +-- rugnux/Rugnux.cpp | 30 ++- 8 files changed, 359 insertions(+), 212 deletions(-) diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 67948196d..a7718a4b2 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -3,6 +3,8 @@ ### 1.0.0-rc.170 +* rugnux needs about a third less peak memory to scale, merge and post-refine rotation data with tens of millions of partial observations, with identical results. +* `rugnux --model`: the placed coordinate file carries the space group its own coordinates obey, and says so when that is not the group the reflection files beside it carry. * `jfjoch_viewer`: fixes in the dataset plots (the stacked plots line up their y axes, the horizontal crosshair follows the cursor, a run in progress labels the angles of its whole sweep), in the inspector (a long dataset name wraps instead of being cut off, a live session shows its sample-head angles) and in the layout (the resize handle beside the file panel is centred). ### 1.0.0-rc.169 diff --git a/image_analysis/geom_refinement/PostRefine.cpp b/image_analysis/geom_refinement/PostRefine.cpp index b34ce7469..82a8d4027 100644 --- a/image_analysis/geom_refinement/PostRefine.cpp +++ b/image_analysis/geom_refinement/PostRefine.cpp @@ -58,16 +58,9 @@ constexpr double STEP_BOUND = 0.01; // box the solver was always given on the distance, i.e. as far as this fit has ever been trusted. constexpr int MAX_STEPS = 5; -// One integrated partial, flattened across all images. Kept as narrow as the sort and the event -// split allow: on a large cell this array is gigabytes, and the scatter and every level of the -// per-bucket sort move all of it. The goniometer angle is not stored - it is a function of the -// image number alone, and is rebuilt from it where it is needed. -struct Partial { - int h, k, l; - float img; - float I, sigma; - float obs_x, obs_y; // observed spot centroid (pixels); NAN if the box sum found no centroid -}; +// One integrated partial, flattened across all images - PostRefinePartial in the header, under its +// working name (see there for why it is kept this narrow). +using Partial = PostRefinePartial; // A rocking event: one reflection's intensity-weighted centroid over the frames it spans. struct Event { @@ -145,7 +138,133 @@ struct RotationScaleResidual { } // namespace -PostRefineResult PostRefineRotationGeometry(const std::vector &outcomes, +PostRefineObservations GatherPostRefineObservations(std::vector &outcomes, + size_t nthreads, bool release_reflections, + Logger &logger) { + PostRefineObservations out; + try { + nthreads = std::max(1, nthreads); + + // Count first, then fill. Growing one vector by push_back over tens of millions of + // reflections copies the whole thing every time it doubles - several gigabytes of pure + // copying - and the counts are cheap to take. The h range comes out of the same sweep: the + // sort below buckets by h and needs to know how many buckets that is, and this pass already + // reads every reflection. + const int n_out = static_cast(outcomes.size()); + std::vector counts(n_out, 0); + std::vector h_lo_of(n_out), h_hi_of(n_out); + ParallelChunks(n_out, nthreads, [&](int lo, int hi) { + for (int o = lo; o < hi; o++) { + size_t keep = 0; + int lmin = std::numeric_limits::max(), lmax = std::numeric_limits::min(); + for (const auto &r : outcomes[o].reflections) + if (std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f) { + keep++; + lmin = std::min(lmin, r.h); + lmax = std::max(lmax, r.h); + } + counts[o] = keep; + h_lo_of[o] = lmin; + h_hi_of[o] = lmax; + } + }); + int h_lo = std::numeric_limits::max(), h_hi = std::numeric_limits::min(); + size_t n_pts = 0; + for (int o = 0; o < n_out; o++) { + n_pts += counts[o]; + h_lo = std::min(h_lo, h_lo_of[o]); + h_hi = std::max(h_hi, h_hi_of[o]); + } + const int H = (h_lo <= h_hi) ? (h_hi - h_lo + 1) : 1; + + // The bucket histogram, in a pass of its own: the fill below scatters each partial straight + // into its bucket, so the gathered-order copy of the whole array that used to be scattered + // from - gigabytes on a large cell, alive beside the bucketed one - no longer exists at all. + // Its chunks are the outcome chunks ParallelChunks makes, so the scatter below has to be + // split the same way. + const int nt = static_cast(std::clamp(nthreads, 1, std::max(1, n_out))); + const int chunk = (n_out + nt - 1) / nt; + std::vector> hist(nt, std::vector(H, 0)); + ParallelChunks(n_out, nthreads, [&](int lo, int hi) { + std::vector &h_count = hist[lo / chunk]; + for (int o = lo; o < hi; o++) + for (const auto &r : outcomes[o].reflections) + if (std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f) + h_count[r.h - h_lo]++; + }); + + // Where each chunk writes each bucket: buckets in order, and within a bucket the chunks in + // index order, so the scatter below is stable and lands in original (per-image) order inside + // a bucket - which is what makes the tie order below reproducible at every thread count. + out.bstart.assign(H + 1, 0); + { + int32_t acc = 0; + for (int b = 0; b < H; ++b) { + out.bstart[b] = acc; + for (int t = 0; t < nt; ++t) { const int32_t c = hist[t][b]; hist[t][b] = acc; acc += c; } + } + out.bstart[H] = acc; + } + + // A vector of n partials VALUE-initialises them: on a large cell that is gigabytes of zeroing + // on one thread, and it is that one thread which first touches every page - which on a + // multi-socket machine leaves the whole array on its node, so every pass that follows runs at + // one node's memory bandwidth. new[] leaves the partials untouched, so the parallel scatter is + // the first touch and each page lands on the node of the thread that filled it. + std::unique_ptr pts(new Partial[n_pts]); + ParallelChunks(n_out, nthreads, [&](int lo, int hi) { + std::vector fill = hist[lo / chunk]; + for (int o = lo; o < hi; o++) { + for (const auto &r : outcomes[o].reflections) { + if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f) continue; + const float ox = std::isfinite(r.observed_x) ? r.observed_x : NAN; + const float oy = std::isfinite(r.observed_y) ? r.observed_y : NAN; + pts[fill[r.h - h_lo]++] = Partial{r.h, r.k, r.l, r.image_number, r.I, r.sigma, ox, oy}; + } + // The caller said it will never read this image's reflections again, so hand the + // vector back the moment it is consumed: the payload shrinks image by image as it + // moves into the (much narrower) partials instead of doubling up in full. + if (release_reflections) + std::vector().swap(outcomes[o].reflections); + } + }); + logger.Info("Post-refine: {} partials gathered", n_pts); + + // Sort each bucket, on its own thread and the largest first (the tail of this is one bucket, + // so it should be the big one). h is the leading key, so the sorted array is the buckets laid + // end to end. Sorting the whole thing in one pass moved every partial through every level of + // a comparison sort, on one thread, over tens of millions of reflections. + // + // Not a total order: two partials of one reflection on one image still tie, as they did + // before this was bucketed. What makes the result reproducible is the scatter above rather + // than the comparator - each bucket is laid out chunk by chunk and a chunk covers the images + // in index order, so a bucket reaches std::sort in global per-image order whatever the + // thread count. Ties therefore resolve the same way on every run and at every -N; they are + // simply not resolved by rank. + const auto part_less = [](const Partial &a, const Partial &b) { + if (a.h != b.h) return a.h < b.h; + if (a.k != b.k) return a.k < b.k; + if (a.l != b.l) return a.l < b.l; + return a.img < b.img; + }; + out.order.resize(H); + std::iota(out.order.begin(), out.order.end(), 0); + std::sort(out.order.begin(), out.order.end(), [&](int a, int b) { + return (out.bstart[a + 1] - out.bstart[a]) > (out.bstart[b + 1] - out.bstart[b]); + }); + ParallelFor(H, nthreads, [&](int oi) { + const int b = out.order[oi]; + std::sort(pts.get() + out.bstart[b], pts.get() + out.bstart[b + 1], part_less); + }); + out.pts = std::move(pts); + out.n_pts = n_pts; + } catch (...) { + out = PostRefineObservations{}; // empty; the fit refuses it (see the header) + } + return out; +} + +PostRefineResult PostRefineRotationGeometry(PostRefineObservations observations, const GoniometerAxis &axis, const DiffractionGeometry &nominal_geom, const CrystalLattice &reference_latt, @@ -162,116 +281,18 @@ PostRefineResult PostRefineRotationGeometry(const std::vector(outcomes.size()); - std::vector pts_offset(n_out + 1, 0); - std::vector h_lo_of(n_out), h_hi_of(n_out); - ParallelChunks(n_out, nthreads, [&](int lo, int hi) { - for (int o = lo; o < hi; o++) { - size_t keep = 0; - int lmin = std::numeric_limits::max(), lmax = std::numeric_limits::min(); - for (const auto &r : outcomes[o].reflections) - if (std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f) { - keep++; - lmin = std::min(lmin, r.h); - lmax = std::max(lmax, r.h); - } - pts_offset[o + 1] = keep; - h_lo_of[o] = lmin; - h_hi_of[o] = lmax; - } - }); - int h_lo = std::numeric_limits::max(), h_hi = std::numeric_limits::min(); - for (int o = 0; o < n_out; o++) { - pts_offset[o + 1] += pts_offset[o]; - h_lo = std::min(h_lo, h_lo_of[o]); - h_hi = std::max(h_hi, h_hi_of[o]); - } - const size_t n_pts = pts_offset[n_out]; - const int H = (h_lo <= h_hi) ? (h_hi - h_lo + 1) : 1; - - // A vector of n partials VALUE-initialises them: on a large cell that is gigabytes of zeroing - // on one thread, and it is that one thread which first touches every page - which on a - // multi-socket machine leaves the whole array on its node, so every pass that follows runs at - // one node's memory bandwidth. new[] leaves the partials untouched, so the parallel fill is - // the first touch and each page lands on the node of the thread that filled it. - std::unique_ptr pts(new Partial[n_pts]); - // The bucket histogram the sort needs is taken here rather than in a pass of its own, since - // the fill already has h in hand. Its chunks are the outcome chunks ParallelChunks makes, so - // the scatter below has to be split the same way. - const int nt = static_cast(std::clamp(nthreads, 1, std::max(1, n_out))); - const int chunk = (n_out + nt - 1) / nt; - std::vector> hist(nt, std::vector(H, 0)); - ParallelChunks(n_out, nthreads, [&](int lo, int hi) { - std::vector &h_count = hist[lo / chunk]; - for (int o = lo; o < hi; o++) { - size_t at = pts_offset[o]; - for (const auto &r : outcomes[o].reflections) { - if (!std::isfinite(r.I) || !std::isfinite(r.sigma) || r.sigma <= 0.0f) continue; - const float ox = std::isfinite(r.observed_x) ? r.observed_x : NAN; - const float oy = std::isfinite(r.observed_y) ? r.observed_y : NAN; - pts[at++] = Partial{r.h, r.k, r.l, r.image_number, r.I, r.sigma, ox, oy}; - h_count[r.h - h_lo]++; - } - } - }); - logger.Info("Post-refine: {} partials gathered", n_pts); + const size_t n_pts = observations.n_pts; if (n_pts < static_cast(settings.min_events)) return result; - - // Where each bucket starts, and the buckets largest first: the sort lays the array out this - // way and the event split below walks the same buckets. - std::vector bstart(H + 1, 0); - std::vector order(H); - - // Bucket by h, then sort the buckets. h is the leading key, so the sorted array is the - // buckets laid end to end, and each bucket sorts on its own thread. Sorting the whole thing - // in one pass moved every partial through every level of a comparison sort, on one thread, - // over tens of millions of reflections. - { - // Not a total order: two partials of one reflection on one image still tie, as they did - // before this was bucketed. What makes the result reproducible is the scatter below - // rather than the comparator - the prefix lays each bucket out chunk by chunk and a - // chunk is a contiguous span of the gathered order, so a bucket reaches std::sort in - // global gather order whatever the thread count. Ties therefore resolve the same way on - // every run and at every -N; they are simply not resolved by rank. - const auto part_less = [](const Partial &a, const Partial &b) { - if (a.h != b.h) return a.h < b.h; - if (a.k != b.k) return a.k < b.k; - if (a.l != b.l) return a.l < b.l; - return a.img < b.img; - }; - int32_t acc = 0; - for (int b = 0; b < H; ++b) { - bstart[b] = acc; - for (int t = 0; t < nt; ++t) { const int32_t c = hist[t][b]; hist[t][b] = acc; acc += c; } - } - bstart[H] = acc; - - std::unique_ptr sorted(new Partial[n_pts]); - ParallelChunks(n_out, nthreads, [&](int lo, int hi) { - std::vector fill = hist[lo / chunk]; - for (size_t i = pts_offset[lo]; i < pts_offset[hi]; ++i) - sorted[fill[pts[i].h - h_lo]++] = pts[i]; - }); - std::iota(order.begin(), order.end(), 0); - std::sort(order.begin(), order.end(), - [&](int a, int b) { return (bstart[a + 1] - bstart[a]) > (bstart[b + 1] - bstart[b]); }); - ParallelFor(H, nthreads, [&](int oi) { - const int b = order[oi]; - std::sort(sorted.get() + bstart[b], sorted.get() + bstart[b + 1], part_less); - }); - pts.swap(sorted); - } + const int H = static_cast(observations.order.size()); + const std::vector &bstart = observations.bstart; + const std::vector &order = observations.order; + const Partial *const pts = observations.pts.get(); // Split into rocking events (same raw hkl, adjacent frames). Only >=2-frame events carry an // unbiased phi_obs (a single-frame centroid is just the frame centre). diff --git a/image_analysis/geom_refinement/PostRefine.h b/image_analysis/geom_refinement/PostRefine.h index 9e67caf03..e1f3b7101 100644 --- a/image_analysis/geom_refinement/PostRefine.h +++ b/image_analysis/geom_refinement/PostRefine.h @@ -4,6 +4,8 @@ #pragma once #include +#include +#include #include #include #include @@ -75,10 +77,43 @@ struct PostRefineSettings { std::optional> measured_beam_px; }; +// One integrated partial, flattened across all images. Kept as narrow as the sort and the event +// split allow: on a large cell this array is gigabytes, and the scatter and every level of the +// per-bucket sort move all of it. The goniometer angle is not stored - it is a function of the +// image number alone, and is rebuilt from it where it is needed. +struct PostRefinePartial { + int h, k, l; + float img; + float I, sigma; + float obs_x, obs_y; // observed spot centroid (pixels); NAN if the box sum found no centroid +}; + +// The gathered form the fit runs on: every usable partial, bucketed by h (bstart = bucket starts, +// order = the buckets largest first) and sorted by (h, k, l, img) inside each bucket. Gathered +// apart from the fit itself so a caller that is finished with the per-image reflection vectors can +// hand them back before the fit's own allocations arrive - on a fine-sliced long axis the vectors +// are several times this array. +struct PostRefineObservations { + std::unique_ptr pts; + size_t n_pts = 0; + std::vector bstart; // size H+1: where each h bucket starts in pts + std::vector order; // size H: bucket indices, largest bucket first +}; + +// Gather and sort the usable partials out of the per-image outcomes. release_reflections: each +// image's reflection vector is freed the moment it has been read - only for a caller that will +// never touch them again (the rotation geometry pre-pass, whose reflections are never written). +// An allocation failure returns an empty gather (n_pts = 0), which the fit refuses - the same +// silent refusal the fit gives every other failure. +PostRefineObservations GatherPostRefineObservations(std::vector &outcomes, + size_t nthreads, bool release_reflections, + Logger &logger); + // nominal_geom / reference_latt: the current detector geometry and the phi=0 reference lattice (orientation -// + cell) from rotation indexing. outcomes: the per-image integrated reflections (observed_x/y, I, sigma, -// image_number). axis: the goniometer. Returns ok=false (geometry untouched) on failure. -PostRefineResult PostRefineRotationGeometry(const std::vector &outcomes, +// + cell) from rotation indexing. obs: the gathered per-image integrated reflections (observed_x/y, I, +// sigma, image_number), from GatherPostRefineObservations. axis: the goniometer. Returns ok=false +// (geometry untouched) on failure. +PostRefineResult PostRefineRotationGeometry(PostRefineObservations obs, const GoniometerAxis &axis, const DiffractionGeometry &nominal_geom, const CrystalLattice &reference_latt, diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index d9e19e95b..f41da9dce 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -508,30 +508,43 @@ void RotationScaleMerge::Ingest() { gpu_active_ = gpu_->Available(); if (gpu_active_) { const int n = static_cast(partials.size()); - // Plain arrays, NOT sized vectors: the sweep below writes every element of all fourteen, so the - // value-initialisation a sized vector performs is half a gigabyte of zeroes faulted in on one - // thread and overwritten on the next line. - const auto floats = [n] { return std::unique_ptr(new float[n]); }; - auto I = floats(), sigma = floats(), prescaling_corr = floats(), part = floats(), zeta = floats(), - corr = floats(), bkg = floats(), img = floats(), dd = floats(), px = floats(), - py = floats(), vbkg = floats(); - std::unique_ptr onice(new uint8_t[n]); - std::unique_ptr frm(new int32_t[n]); - // Each observation writes its own slot in each of the fourteen arrays, so this splits - // straight over the range. It is the widest pass in Ingest - it reads the whole 80-byte - // struct and writes 57 bytes of it back out - and it ran on one thread. - ParallelChunks(n, nthreads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - const auto &o = partials[i]; - I[i] = o.I; sigma[i] = o.sigma; prescaling_corr[i] = o.prescaling_corr; part[i] = o.partiality; - zeta[i] = o.zeta; onice[i] = o.on_ice; frm[i] = o.frame; corr[i] = o.corr; - bkg[i] = o.bkg; vbkg[i] = o.var_bkg; img[i] = o.image_number; dd[i] = o.d; - px[i] = o.px; py[i] = o.py; - } - }); - gpu_->SetPartials(n, n_frames, I.get(), sigma.get(), prescaling_corr.get(), part.get(), zeta.get(), - onice.get(), frm.get(), corr.get(), frame_start.data(), frame_count.data()); - gpu_->SetCombineInputs(bkg.get(), vbkg.get(), img.get(), dd.get(), px.get(), py.get()); + gpu_->SetPartialsLayout(n, n_frames, frame_start.data(), frame_count.data()); + // Staged in bounded slices: fourteen full-length arrays at once are more than half of the + // whole Obs payload live again - on a fine-sliced long axis that transient alone is + // gigabytes - where a slice of them together is under half a gigabyte whatever the run. + // One slice is still a single sweep that reads the whole 80-byte struct once and writes + // every field's slot, split straight over the range, so the traffic is the one pass the + // full-length staging paid. Plain arrays, NOT sized vectors: every element is written + // before it is read, so the value-initialisation a sized vector performs is memory faulted + // in on one thread and overwritten on the next line. + using F = RotationScaleMergeGPU::ObsField; + constexpr int STAGE_SLICE = 8 << 20; + const int slice = std::min(std::max(n, 1), STAGE_SLICE); + constexpr int NFLOAT = 12; + std::unique_ptr fbuf[NFLOAT]; + for (auto &b : fbuf) b.reset(new float[slice]); + std::unique_ptr frm(new int32_t[slice]); + std::unique_ptr onice(new uint8_t[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) { + const auto &o = partials[at + i]; + fbuf[0][i] = o.I; fbuf[1][i] = o.sigma; fbuf[2][i] = o.prescaling_corr; + fbuf[3][i] = o.partiality; fbuf[4][i] = o.zeta; fbuf[5][i] = o.corr; + fbuf[6][i] = o.bkg; fbuf[7][i] = o.var_bkg; fbuf[8][i] = o.image_number; + fbuf[9][i] = o.d; fbuf[10][i] = o.px; fbuf[11][i] = o.py; + frm[i] = o.frame; onice[i] = o.on_ice; + } + }); + constexpr F fields[NFLOAT] = {F::I, F::Sigma, F::PrescalingCorr, F::Partiality, + 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], at, cnt, fbuf[k].get()); + gpu_->SetObsFrame(at, cnt, frm.get()); + gpu_->SetObsOnIce(at, cnt, onice.get()); + } gpu_->SetRawRuns(static_cast(rawrun_start.size()), static_cast(perm.size()), perm.data(), rawrun_start.data(), rawrun_count.data(), rawrun_h.data(), rawrun_k.data(), rawrun_l.data()); @@ -549,7 +562,7 @@ void RotationScaleMerge::Ingest() { // // It has to come after the incident flux has been measured (the flux meter is the frame's mean background // over ALL its reflections) and after the sort (which is what defines a raw-hkl run in the first place). -void RotationScaleMerge::BuildInRangeObservations(const std::vector &keys) { +void RotationScaleMerge::BuildInRangeObservations(std::vector &keys) { const int n_run = static_cast(rawrun_start.size()); const int n_obs = static_cast(keys.size()); const bool limited = d_min_limit.has_value() || d_max_limit.has_value(); @@ -597,42 +610,21 @@ void RotationScaleMerge::BuildInRangeObservations(const std::vector &ke n_keep += frame_count[o]; } - // The observations themselves, one frame per thread. finite_ok is the per-obs AcceptReflection - // 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. - partials.clear(); - partials.resize(n_keep); - finite_ok.assign(n_keep, 0); + // Where each surviving observation will land: the per-frame prefix over `keep`, taken on its own + // before anything eighty bytes wide exists. The remap below is then free to consume `keys` - the + // one other reader of these indices - so the key array is gone before the observation array is + // built, instead of the two coexisting at another third of its size. std::vector new_idx(n_obs); // only written where keep[i] ParallelFor(n_frames, nthreads, [&](int o) { int32_t src = src_start[o]; int32_t at = frame_start[o]; - for (const auto &r : partials_out[o].reflections) { - if (!keep[src]) { ++src; continue; } - new_idx[src++] = at; - Obs &obs = partials[at]; - obs.h = r.h; obs.k = r.k; obs.l = r.l; - obs.I = r.I; obs.sigma = r.sigma; obs.d = r.d; - // The whole deterministic correction - Lorentz x polarization x sensor efficiency x air path - is the - // numerator of corr, so the two halves are multiplied back together here. The incident flux is - // divided out at the same time so the fits below see only what is left; it is the one term in - // the product that is measured from the data rather than computed from the geometry. - obs.prescaling_corr = static_cast(r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]); - obs.partiality = r.partiality; obs.zeta = r.zeta; obs.delta_phi = r.delta_phi_deg; obs.bkg = r.bkg; obs.var_bkg = r.var_bkg; - obs.px = r.predicted_x; obs.py = r.predicted_y; - obs.image_number = r.image_number; - obs.frame = o; - obs.on_ice = r.on_ice_ring ? 1 : 0; - obs.corr = r.image_scale_corr; - 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; - ++at; - } + const int32_t n = static_cast(partials_out[o].reflections.size()); + for (int32_t j = 0; j < n; ++j, ++src) + if (keep[src]) new_idx[src] = at++; }); // The surviving runs, in the same order, each owning the same sorted keys remapped to the - // observation indices they were just built at. + // observation indices just assigned. int n_keep_run = 0; std::vector dst_run(n_run, -1), dst_start(n_run, 0); { @@ -656,6 +648,40 @@ void RotationScaleMerge::BuildInRangeObservations(const std::vector &ke }); rawrun_start.swap(ns); rawrun_count.swap(nc); rawrun_h.swap(nh); rawrun_k.swap(nk); rawrun_l.swap(nl); rawrun_d.swap(nd); + std::vector().swap(keys); // consumed (see the header) + + // The observations themselves, one frame per thread. finite_ok is the per-obs AcceptReflection + // 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. + partials.clear(); + partials.resize(n_keep); + finite_ok.assign(n_keep, 0); + ParallelFor(n_frames, nthreads, [&](int o) { + int32_t src = src_start[o]; + int32_t at = frame_start[o]; + for (const auto &r : partials_out[o].reflections) { + if (!keep[src]) { ++src; continue; } + ++src; + Obs &obs = partials[at]; + obs.h = r.h; obs.k = r.k; obs.l = r.l; + obs.I = r.I; obs.sigma = r.sigma; obs.d = r.d; + // The whole deterministic correction - Lorentz x polarization x sensor efficiency x air path - is the + // numerator of corr, so the two halves are multiplied back together here. The incident flux is + // divided out at the same time so the fits below see only what is left; it is the one term in + // the product that is measured from the data rather than computed from the geometry. + obs.prescaling_corr = static_cast(r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]); + obs.partiality = r.partiality; obs.zeta = r.zeta; obs.delta_phi = r.delta_phi_deg; obs.bkg = r.bkg; obs.var_bkg = r.var_bkg; + obs.px = r.predicted_x; obs.py = r.predicted_y; + obs.image_number = r.image_number; + obs.frame = o; + obs.on_ice = r.on_ice_ring ? 1 : 0; + obs.corr = r.image_scale_corr; + 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; + ++at; + } + }); if (n_keep < n_obs) logger.Info("RotationScaleMerge: dropped {} of {} observations ({} of {} distinct hkl) outside the " diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 81d76dce6..9e4af0f8d 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -352,7 +352,10 @@ private: // raw-hkl RUNS are skipped, on the same per-hkl resolution ComputeAsuGroups tests, so what survives - // and the order of every sum formed over it - is exactly what it would have been had the whole array // been built and then filtered. Everything is built when no manual limit was given. - void BuildInRangeObservations(const std::vector &keys); + // `keys` is CONSUMED: it is freed as soon as the permutation has been remapped, before the + // 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::vector &keys); // 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/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 75be014a2..74920e3f9 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -727,19 +727,22 @@ RotationScaleMergeGPU::~RotationScaleMergeGPU() { bool RotationScaleMergeGPU::Available() const { return impl_->available; } -void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames, - const float *I, const float *sigma, const float *prescaling_corr, - const float *partiality, const float *zeta, const uint8_t *on_ice, - const int32_t *frame, const float *corr0, - const int32_t *frame_start, const int32_t *frame_count) { +void RotationScaleMergeGPU::SetPartialsLayout(int n_obs, int n_frames, + const int32_t *frame_start, const int32_t *frame_count) { DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; d.n_obs = n_obs; d.n_frames = n_frames; - Upload(d.I, I, n_obs); Upload(d.sigma, sigma, n_obs); Upload(d.prescaling_corr, prescaling_corr, n_obs); - Upload(d.partiality, partiality, n_obs); Upload(d.zeta, zeta, n_obs); Upload(d.on_ice, on_ice, n_obs); - Upload(d.frame, frame, n_obs); Upload(d.corr, corr0, n_obs); Upload(d.frame_start, frame_start, n_frames); Upload(d.frame_count, frame_count, n_frames); + const int n = std::max(1, n_obs); + d.I = CudaDevicePtr(n); d.sigma = CudaDevicePtr(n); + d.prescaling_corr = CudaDevicePtr(n); d.partiality = CudaDevicePtr(n); + d.zeta = CudaDevicePtr(n); d.corr = CudaDevicePtr(n); + d.bkg = CudaDevicePtr(n); d.var_bkg = CudaDevicePtr(n); + d.image_number = CudaDevicePtr(n); d.d_obs = CudaDevicePtr(n); + d.px_obs = CudaDevicePtr(n); d.py_obs = CudaDevicePtr(n); + d.frame = CudaDevicePtr(n); + d.on_ice = CudaDevicePtr(n); d.g = CudaDevicePtr(n_frames); d.scaled = CudaDevicePtr(n_frames); d.inv_sigma = CudaDevicePtr(n_obs); @@ -749,6 +752,44 @@ void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames, d.cc_n = CudaDevicePtr(std::max(1, n_frames)); } +namespace { + template + void UploadChunk(CudaDevicePtr &dst, int offset, int count, const T *v) { + if (count > 0) + CudaCheck(cudaMemcpy(dst.get() + offset, v, size_t(count) * sizeof(T), + cudaMemcpyHostToDevice), "upload chunk"); + } +} + +void RotationScaleMergeGPU::SetObsField(ObsField f, int offset, int count, const float *v) { + DeviceGuard guard(impl_->device, impl_->available); + auto &d = *impl_; + switch (f) { + case ObsField::I: UploadChunk(d.I, offset, count, v); break; + case ObsField::Sigma: UploadChunk(d.sigma, offset, count, v); break; + case ObsField::PrescalingCorr: UploadChunk(d.prescaling_corr, offset, count, v); break; + case ObsField::Partiality: UploadChunk(d.partiality, offset, count, v); break; + case ObsField::Zeta: UploadChunk(d.zeta, offset, count, v); break; + case ObsField::Corr0: UploadChunk(d.corr, offset, count, v); break; + case ObsField::Bkg: UploadChunk(d.bkg, offset, count, v); break; + case ObsField::VarBkg: UploadChunk(d.var_bkg, offset, count, v); break; + case ObsField::ImageNumber: UploadChunk(d.image_number, offset, count, v); break; + case ObsField::D: UploadChunk(d.d_obs, offset, count, v); break; + case ObsField::Px: UploadChunk(d.px_obs, offset, count, v); break; + case ObsField::Py: UploadChunk(d.py_obs, offset, count, v); break; + } +} + +void RotationScaleMergeGPU::SetObsFrame(int offset, int count, const int32_t *frame) { + DeviceGuard guard(impl_->device, impl_->available); + UploadChunk(impl_->frame, offset, count, frame); +} + +void RotationScaleMergeGPU::SetObsOnIce(int offset, int count, const uint8_t *on_ice) { + DeviceGuard guard(impl_->device, impl_->available); + UploadChunk(impl_->on_ice, offset, count, on_ice); +} + void RotationScaleMergeGPU::SetGroups(int n_groups, const int32_t *group, const int32_t *group_perm, int n_group_perm, const int32_t *group_start, const int32_t *group_count) { @@ -997,19 +1038,6 @@ void RotationScaleMergeGPU::ComputePartialCC(double min_partiality, double *cc_o cudaMemcpyDeviceToHost), "download cc_n"); } -void RotationScaleMergeGPU::SetCombineInputs(const float *bkg, const float *var_bkg, - const float *image_number, const float *d, - const float *px, const float *py) { - DeviceGuard guard(impl_->device, impl_->available); - auto &dd = *impl_; - Upload(dd.bkg, bkg, dd.n_obs); - Upload(dd.var_bkg, var_bkg, dd.n_obs); - Upload(dd.image_number, image_number, dd.n_obs); - Upload(dd.d_obs, d, dd.n_obs); - Upload(dd.px_obs, px, dd.n_obs); - Upload(dd.py_obs, py, dd.n_obs); -} - void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *perm, const int32_t *rr_start, const int32_t *rr_count, const int32_t *rr_h, const int32_t *rr_k, const int32_t *rr_l) { diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.h b/image_analysis/scale_merge/RotationScaleMergeGPU.h index c3c4c586e..cd4ac9464 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -26,13 +26,24 @@ public: // True if a GPU was found and the engine is usable. [[nodiscard]] bool Available() const; - // Upload the immutable per-observation fields (once). Arrays are length n_obs unless noted; the - // frame CSR (frame_start/frame_count) is length n_frames and indexes the obs arrays in frame order. - void SetPartials(int n_obs, int n_frames, - const float *I, const float *sigma, const float *prescaling_corr, const float *partiality, - const float *zeta, const uint8_t *on_ice, const int32_t *frame, - const float *corr0, - const int32_t *frame_start, const int32_t *frame_count); + // The observation count and the frame CSR (frame_start/frame_count, length n_frames, indexing + // the obs arrays in frame order), plus every device-side observation buffer, allocated here so + // the chunked uploads below have somewhere to land. Call once, before them. + void SetPartialsLayout(int n_obs, int n_frames, + const int32_t *frame_start, const int32_t *frame_count); + + // One slice ([offset, offset+count) of n_obs) of one immutable per-observation field, so the + // host can stage a bounded chunk at a time out of its AoS record: staged whole and all at once, + // the fourteen arrays are more than half the entire observation payload held live again. + // Corr0 seeds the resident, mutable corr (SetCorr refreshes it). Bkg/VarBkg/ImageNumber/D/Px/Py + // are the combine inputs (image-local background, fractional frame position for event + // contiguity, resolution, and the predicted detector position carried to the full for the + // absorption surface). + enum class ObsField { I, Sigma, PrescalingCorr, Partiality, Zeta, Corr0, + Bkg, VarBkg, ImageNumber, D, Px, Py }; + void SetObsField(ObsField f, int offset, int count, const float *v); + void SetObsFrame(int offset, int count, const int32_t *frame); + void SetObsOnIce(int offset, int count, const uint8_t *on_ice); // Per space group: the dense ASU-group id per obs, and a group-ordered permutation of the obs whose // group >= 0 (group_perm), with its CSR (group_start/group_count, length n_groups) - so each group's @@ -42,7 +53,7 @@ public: const int32_t *group_start, const int32_t *group_count); // Re-upload the working corr (length n_obs) before a scaling pass (the host mutates it via smooth-G - // between passes). SetPartials uploads the initial corr; this refreshes it. + // between passes). SetObsField(Corr0) uploads the initial corr; this refreshes it. void SetCorr(const float *corr); // Run `iters` of {reduce group means -> per-frame robust IRLS G -> update corr} on the device, @@ -103,13 +114,6 @@ public: // --- 3D combine (partials -> fulls), all on the device --- - // The per-obs fields the combine needs on top of the scaling inputs (image-local bkg, fractional - // frame position for event contiguity, resolution, and the predicted detector position px/py carried - // to the full for the absorption surface). Uploaded once, alongside SetPartials. - void SetCombineInputs(const float *bkg, const float *var_bkg, - const float *image_number, const float *d, - const float *px, const float *py); - // The one-time raw-hkl run layout (space-group-independent): the (raw h,k,l, image_number)-sorted // permutation of the obs, split into contiguous per-raw-hkl runs. Uploaded once in Ingest. void SetRawRuns(int n_runs, int n_perm, const int32_t *perm, diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 4b20006a0..d6077a37c 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -4877,6 +4877,20 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b rsm->Ingest(); } + // The geometry pre-pass reads the per-image reflections exactly twice more: they were just + // ingested above, and the post-refinement at the end of the pass reads a narrow + // (hkl, frame, I, sigma, observed position) view of them. Nothing else in the pre-pass + // touches them - its reflections are never written, and the scaling engine reads back only + // per-frame fields - so take the post-refine's view NOW and hand each image's vector back as + // it is read: on a fine-sliced long axis the vectors are gigabytes, held to the end of the + // pass for those 32 bytes a partial. A run with a reference MTZ keeps them: resolving the + // indexing ambiguity re-reads them (and re-ingests). + std::optional prepass_postrefine_obs; + if (geometry_prepass && rsm && config_.reference_data.empty()) + prepass_postrefine_obs = GatherPostRefineObservations( + indexer->GetIntegrationOutcome(), config_.nthreads, + /*release_reflections=*/true, logger); + auto scale_and_merge = [&](const std::string &label, bool for_search) -> ScaleMergeResult { if (rsm) { phase("Scale/combine/merge (" + label + ")"); @@ -5075,7 +5089,21 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b // were integrated at it, and step B holds the tilt fixed while it fits the beam and // the distance - so handing it the header tilt fits a beam that has to absorb a tilt // error the images do not have, and the answer describes neither geometry. - const auto pr = PostRefineRotationGeometry(outcomes, *rot->axis, + // + // On the pre-pass the observations were gathered right after Ingest (see there) and + // the scaling engine is finished - nothing after this point reads it on that pass + // (the search and the reference re-merge run only on the canonical one) - so it is + // freed BEFORE the fit: the engine holds about twice the partial payload again, + // and with it gone the fit no longer sits on the pass's memory peak. + if (prepass_postrefine_obs) + rsm.reset(); + auto pr_obs = prepass_postrefine_obs + ? std::move(*prepass_postrefine_obs) + : GatherPostRefineObservations(indexer->GetIntegrationOutcome(), + config_.nthreads, + /*release_reflections=*/false, logger); + prepass_postrefine_obs.reset(); + const auto pr = PostRefineRotationGeometry(std::move(pr_obs), *rot->axis, rot->geom, rot->lattice, prs, logger); result.post_refine = pr; if (pr.ok && pr.detector_refined) {