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) {