From 1ea0a93f41f96fe4196b9c6d9f9e76c0215932d7 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 18:26:20 +0200 Subject: [PATCH 1/3] RotationScaleMerge: the CPU merge's serial walks on all threads, large arrays first touched in parallel, output unchanged Profiled with --no-gpu on a large-cell set, the scaling-and-merging step ran under four cores for about a third of its wall time. The serial pieces, and what they are now: - The combine's tail sorted the eighty-byte fulls by frame on one thread. It now sorts each full's (frame, position) under the same frame-only comparison - std::sort's moves depend only on the outcomes of its comparisons, so the positions come out in exactly the order the records did, ties included - and gathers the records once on all threads, straight from the per-thread parts (no concatenation copy). - The fulls' group CSR (a serial counting sort) is a parallel sort on (group, index), a total order. - The host merge accumulate, the R_meas walk and the error-model sample pool were serial walks over every full scattering into per-group sums; they walk each group over the group CSR on all threads, so every group is still summed in index order, and integer counts are added up per block. - The per-pair centric test of the anomalous statistics runs on all threads; the anomalous export map is reserved once. - partials, fulls and the scaling loops' working copies (ScalingRef, ScalingFit, corr) use an allocator that default-initialises (DefaultInitAllocator / FillVector), and Obs has no member initialisers, so these arrays are first written - and their pages faulted in - by the parallel passes that fill them rather than zeroed on one thread first. On a large sweep that serial first touch was several seconds of kernel time per merge. The ingest now sets a partial's var_per_I and capture to 0 explicitly; every other producer already wrote every field. p.hkl and the MTZ data are byte-identical on three large sets (--no-gpu). Provisional, shared machine (load 20-33 on 32 threads), scaling-and-merging step: 115 -> 91 s on a 3600-frame set, 74 -> 61 s on a 500-frame P1-heavy set, 55 -> 51 s on an 1800-frame set; time under four cores 28 -> 6 s, 17 -> 3 s and 5 -> 1 s. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../scale_merge/RotationScaleMerge.cpp | 297 +++++++++++------- .../scale_merge/RotationScaleMerge.h | 49 ++- 2 files changed, 225 insertions(+), 121 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 7cfc12586..b62db5867 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -1009,6 +1009,7 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr &ke // 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.var_per_I = 0.0f; obs.capture = 0.0f; obs.px = r.predicted_x; obs.py = r.predicted_y; obs.image_number = r.image_number; obs.frame = o; @@ -1573,8 +1574,8 @@ void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vector &refs, - const std::vector &corr, + const int32_t *group_count, const FillVector &refs, + const FillVector &corr, std::vector &out_mean) const { out_mean.resize(n_groups); ParallelChunks(n_groups, ThreadsForWork(corr.size(), nthreads), [&](int glo, int ghi) { @@ -1669,7 +1670,7 @@ double RotationScaleMerge::ChooseLogScaleSmoothing(const std::vector &y, return best_lambda; } -void RotationScaleMerge::FitScalingFrames(const std::vector &fit, const std::vector &fit_start, +void RotationScaleMerge::FitScalingFrames(const FillVector &fit, const std::vector &fit_start, size_t min_obs, const std::vector &group_mean_in, std::vector &g) { const int n_fr = static_cast(fit_start.size()) - 1; @@ -4352,7 +4353,8 @@ bool RotationScaleMerge::UsableFull(const Obs &o) const { } void RotationScaleMerge::Combine() { - fulls.clear(); + // `fulls` is not cleared: SortFullsByFrame writes every record, so its old contents only save the + // resize there from constructing them again. g_full.assign(n_frames, 1.0); fulls_dropped_overloaded = 0; const int n_run = static_cast(rawrun_start.size()); @@ -4514,6 +4516,7 @@ void RotationScaleMerge::Combine() { }; size_t n_used = 0; + std::vector> part; // the fulls of consecutive runs, in run order if (!observation_dump_path.empty() || nthreads <= 1) { // Serial (diagnostic dump needs a single writer). @@ -4524,15 +4527,16 @@ void RotationScaleMerge::Combine() { // is a per-full batch (ZD) and px/py the predicted detector position (XD/YD). dump << "# h k l I sigma d n_frames captured_fraction peak_frame px py\n"; } - fulls.reserve(n_run); + part.resize(1); + part[0].reserve(n_run); for (int r = 0; r < n_run; ++r) - n_used += process_rawrun(r, fulls, dump.is_open() ? &dump : nullptr, fulls_dropped_overloaded); + n_used += process_rawrun(r, part[0], dump.is_open() ? &dump : nullptr, fulls_dropped_overloaded); } else { - // Parallel over contiguous rawrun chunks; concatenate the per-thread fulls in run order so the - // result is deterministic. + // Parallel over contiguous rawrun chunks; the per-thread fulls are taken in run order below, so + // the result is deterministic. const int nt = static_cast(std::min(nthreads, static_cast(n_run))); const int chunk = (n_run + nt - 1) / nt; - std::vector> part(nt); + part.resize(nt); std::vector used(nt, 0); std::vector overloaded(nt, 0); std::vector> futures; @@ -4547,44 +4551,74 @@ void RotationScaleMerge::Combine() { })); } for (auto &f : futures) f.get(); - size_t total = 0; for (int t = 0; t < nt; ++t) { - total += part[t].size(); n_used += used[t]; fulls_dropped_overloaded += overloaded[t]; + n_used += used[t]; fulls_dropped_overloaded += overloaded[t]; } - fulls.reserve(total); - for (int t = 0; t < nt; ++t) - fulls.insert(fulls.end(), part[t].begin(), part[t].end()); } - SortFullsByFrame(); + SortFullsByFrame(part); logger->Info("3D combine: {} fulls from {} partials, {} rocking events dropped for a saturated pixel", fulls.size(), n_used, fulls_dropped_overloaded); } void RotationScaleMerge::BuildFullsGroupCSR(int n_groups) { + // Each group's fulls in index order. (group, index) is a total order, so the sort on all threads + // returns the one sequence the serial counting sort filled. const int nf = static_cast(fulls.size()); - fulls_group_start.assign(n_groups + 1, 0); - for (int i = 0; i < nf; ++i) - if (fulls[i].group >= 0) ++fulls_group_start[fulls[i].group + 1]; - for (int g = 0; g < n_groups; ++g) fulls_group_start[g + 1] += fulls_group_start[g]; - fulls_group_perm.resize(fulls_group_start[n_groups]); - std::vector fill(fulls_group_start.begin(), fulls_group_start.end() - 1); - for (int i = 0; i < nf; ++i) - if (fulls[i].group >= 0) fulls_group_perm[fill[fulls[i].group]++] = i; + const std::vector grouped = KeptIndices(nf, nthreads, [&](int i) { return fulls[i].group >= 0; }); + const int n = static_cast(grouped.size()); + std::vector> key(n); + ParallelChunks(n, ThreadsForWork(key.size(), nthreads), [&](int lo, int hi) { + for (int j = lo; j < hi; ++j) key[j] = {fulls[grouped[j]].group, grouped[j]}; + }); + ParallelSort(key.begin(), key.end(), nthreads, std::less>()); + fulls_group_start.resize(n_groups + 1); + fulls_group_perm.resize(n); + // Group g starts where the first key of a group >= g sits, so a group's start is written by the key + // that opens it, together with the empty groups just below it. + ParallelChunks(n, ThreadsForWork(key.size(), nthreads), [&](int lo, int hi) { + for (int j = lo; j < hi; ++j) { + fulls_group_perm[j] = key[j].second; + const int g = key[j].first; + if (j > 0 && key[j - 1].first == g) continue; + for (int e = j > 0 ? key[j - 1].first + 1 : 0; e <= g; ++e) fulls_group_start[e] = j; + } + }); + for (int e = n > 0 ? key[n - 1].first + 1 : 0; e <= n_groups; ++e) fulls_group_start[e] = n; } -void RotationScaleMerge::SortFullsByFrame() { - // Sort the fulls by their (peak) frame and build per-frame CSR ranges, so the scale-fulls step can - // fit a per-frame G by slicing contiguous ranges (the same layout the partials use). - std::sort(fulls.begin(), fulls.end(), - [](const Obs &a, const Obs &b) { return a.frame < b.frame; }); +void RotationScaleMerge::SortFullsByFrame(const std::vector> &part) { + // Sort the fulls - the parts one after another - by their (peak) frame and build per-frame CSR + // ranges, so the scale-fulls step can fit a per-frame G by slicing contiguous ranges (the same layout + // the partials use). What std::sort does to a sequence depends on nothing but the outcomes of the + // comparisons it makes, so sorting each full's (frame, position) under the same frame-only test + // leaves the positions in exactly the order the eighty-byte records would have been left in - ties + // included - while moving eight bytes instead of eighty. The records are then gathered once. + struct FrameAt { int32_t frame, at; }; + std::vector part_start(part.size() + 1, 0); + for (size_t t = 0; t < part.size(); ++t) part_start[t + 1] = part_start[t] + part[t].size(); + const int nf = static_cast(part_start.back()); + std::vector order(nf); + ParallelFor(static_cast(part.size()), nthreads, [&](int t) { + for (size_t j = 0; j < part[t].size(); ++j) + order[part_start[t] + j] = {part[t][j].frame, static_cast(part_start[t] + j)}; + }); + std::sort(order.begin(), order.end(), [](const FrameAt &a, const FrameAt &b) { return a.frame < b.frame; }); + fulls.resize(nf); + ParallelChunks(nf, ThreadsForWork(order.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) { + const size_t at = order[i].at; + const size_t t = std::upper_bound(part_start.begin(), part_start.end(), at) - part_start.begin() - 1; + fulls[i] = part[t][at - part_start[t]]; + } + }); fulls_frame_start.assign(n_frames, 0); fulls_frame_count.assign(n_frames, 0); - for (int i = 0; i < static_cast(fulls.size()); ) { - const int f = fulls[i].frame; + for (int i = 0; i < nf; ) { + const int f = order[i].frame; fulls_frame_start[f] = i; int j = i; - while (j < static_cast(fulls.size()) && fulls[j].frame == f) ++j; + while (j < nf && order[j].frame == f) ++j; fulls_frame_count[f] = j - i; i = j; } @@ -4852,23 +4886,34 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool for (int g = 0; g < n_groups; ++g) if (cnt[g] >= 2 && sw[g] > 0.0) em_mean[g] = swI[g] / sw[g]; - samples.clear(); - samples.reserve(fulls.size()); - for (const auto &o : fulls) { - if (!usable_merge(o) || cnt[o.group] < 2) continue; - const size_t s = n_hand ? 2 * static_cast(o.group) + obs_hand[&o - fulls.data()] : 0; - const bool on_hand = n_hand && group_has_hands[o.group] && cnth[s] >= 2 && swh[s] > 0.0; - const double mean = on_hand ? swIh[s] / swh[s] : em_mean[o.group]; - if (!std::isfinite(mean)) continue; + // One sample per full that gives one, in index order: which fulls those are, then each one's + // sample on its own - on all threads. + auto sample_of = [&](int i, Sample &out) { + const int g = mf.group[i]; // usable_merge + if (g < 0 || cnt[g] < 2) return false; + const Obs &o = fulls[i]; + const size_t s = n_hand ? 2 * static_cast(g) + obs_hand[i] : 0; + const bool on_hand = n_hand && group_has_hands[g] && cnth[s] >= 2 && swh[s] > 0.0; + const double mean = on_hand ? swIh[s] / swh[s] : em_mean[g]; + if (!std::isfinite(mean)) return false; const double sigma_corr = static_cast(o.sigma) * o.corr; const double s2 = sigma_corr * sigma_corr; const double w = 1.0 / s2; - const double factor = 1.0 - w / (on_hand ? swh[s] : sw[o.group]); - if (factor < 0.05) continue; + const double factor = 1.0 - w / (on_hand ? swh[s] : sw[g]); + if (factor < 0.05) return false; const double resid = static_cast(o.I) * o.corr - mean; - samples.push_back({counting_variance(o, em_mean[o.group], s2), mean * mean, - resid * resid / factor, o.d}); - } + out = {counting_variance(o, em_mean[g], s2), mean * mean, resid * resid / factor, o.d}; + return true; + }; + const std::vector sampled = KeptIndices(n_full, nthreads, [&](int i) { + Sample unused; + return sample_of(i, unused); + }); + samples.resize(sampled.size()); + ParallelChunks(static_cast(sampled.size()), ThreadsForWork(sampled.size(), nthreads), + [&](int lo, int hi) { + for (int j = lo; j < hi; ++j) sample_of(sampled[j], samples[j]); + }); } // Per-group outlier-rejection median of I*corr (host both paths - a per-group median is awkward on @@ -5339,42 +5384,53 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool did_gpu_acc = true; } #endif - if (!did_gpu_acc) - for (const auto &o : fulls) { - if (!usable_merge(o)) continue; - if (rejected_obs[&o - fulls.data()]) { // by the Wilson test above - ++reject_count; - continue; - } - const float I_corr = o.I * o.corr; - float sigma_corr = o.sigma * o.corr; - sigma_corr = corrected_sigma(o, I_corr, sigma_corr); - if (reject_outliers && error_model_active && std::isfinite(reject_median[o.group])) { - // The pooled pair, widened by the shell's measured anomalous variance. - const double v_add = reject_var_add.empty() ? 0.0 : reject_var_add[o.group]; - const double I_b = std::isfinite(em_mean[o.group]) ? em_mean[o.group] : I_corr; - if (OutsideOutlierBand(I_corr, reject_median[o.group], reject_nsigma, - static_cast(sigma_corr) * sigma_corr + v_add, I_b, - error_model_b * I_b, reject_up, reject_down)) { - ++reject_count; - rejected_obs[&o - fulls.data()] = 1; - continue; + // The threads split the GROUPS, each walking a group's fulls over the group CSR, so every group + // is still summed in index order; the rejection count is an integer sum, added up per block. + if (!did_gpu_acc) { + std::vector blk_rejected(ReductionBlocks(n_groups, 512), 0); + ParallelBlocks(n_groups, nthreads, [&](int blk, int g0, int g1) { + for (int g = g0; g < g1; ++g) { + for (int p = fulls_group_start[g]; p < fulls_group_start[g + 1]; ++p) { + const int i = fulls_group_perm[p]; + if (mf.group[i] != g) continue; // usable_merge + const Obs &o = fulls[i]; + if (rejected_obs[i]) { // by the Wilson test above + ++blk_rejected[blk]; + continue; + } + const float I_corr = o.I * o.corr; + float sigma_corr = o.sigma * o.corr; + sigma_corr = corrected_sigma(o, I_corr, sigma_corr); + if (reject_outliers && error_model_active && std::isfinite(reject_median[o.group])) { + // The pooled pair, widened by the shell's measured anomalous variance. + const double v_add = reject_var_add.empty() ? 0.0 : reject_var_add[o.group]; + const double I_b = std::isfinite(em_mean[o.group]) ? em_mean[o.group] : I_corr; + if (OutsideOutlierBand(I_corr, reject_median[o.group], reject_nsigma, + static_cast(sigma_corr) * sigma_corr + v_add, I_b, + error_model_b * I_b, reject_up, reject_down)) { + ++blk_rejected[blk]; + rejected_obs[i] = 1; + continue; + } + } + const double w = 1.0 / (static_cast(sigma_corr) * sigma_corr); + const double wI = w * I_corr; + auto &a = acc[o.group]; + const int half = merge_half[i]; + a.swI += wI; a.sw += w; + a.swIh[half] += wI; a.swh[half] += w; a.nh[half]++; + a.swh_typ[half] += w * frame_cc_factor[o.frame]; + // Any contributing observation on an ice ring marks the merged reflection (as + // MergeOnTheFly does on the stills path): the intensity is still merged and written, + // the flag only lets a consumer that must not read ring signal as crystal signal - + // the CC1/2 resolution-cutoff fit - leave it out. + if (o.on_ice) a.on_ice = true; + if (!std::isfinite(a.d) && std::isfinite(o.d) && o.d > 0.0f) a.d = o.d; } } - const double w = 1.0 / (static_cast(sigma_corr) * sigma_corr); - const double wI = w * I_corr; - auto &a = acc[o.group]; - const int half = merge_half[&o - fulls.data()]; - a.swI += wI; a.sw += w; - a.swIh[half] += wI; a.swh[half] += w; a.nh[half]++; - a.swh_typ[half] += w * frame_cc_factor[o.frame]; - // Any contributing observation on an ice ring marks the merged reflection (as - // MergeOnTheFly does on the stills path): the intensity is still merged and written, - // the flag only lets a consumer that must not read ring signal as crystal signal - - // the CC1/2 resolution-cutoff fit - leave it out. - if (o.on_ice) a.on_ice = true; - if (!std::isfinite(a.d) && std::isfinite(o.d) && o.d > 0.0f) a.d = o.d; - } + }, 512); + for (size_t r : blk_rejected) reject_count += r; + } // A merged reflection's d from the cell the completeness denominator is enumerated with, not // from whichever frame first observed it: a per-frame cell puts a reflection near a shell edge // on the other side of it from its possible twin, and a shell then reads above 100 %. @@ -5765,31 +5821,45 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool did_gpu_rmeas = true; } #endif - if (!did_gpu_rmeas) - for (const auto &o : fulls) { - if (o.group < 0 || acc[o.group].d < d_floor) continue; - if (rejected_obs[&o - fulls.data()]) continue; // outlier-rejected in the merge -> also out of R_meas - if (!frame_cell_ok[o.frame]) continue; - if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) continue; - if (o.partiality < min_partiality) continue; - const float I_corr = o.I * o.corr, sigma_corr = o.sigma * o.corr; - if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) continue; - const auto shell = shells.GetShell(o.d); - if (!shell || *shell < 0 || *shell >= n_shells) continue; - // Only what the merge kept counts towards multiplicity - see the GPU branch above. The - // R_meas accumulation below keeps its own, wider filter. - if (usable_merge(o)) - sa[*shell].total_obs++; - if (std::isfinite(merged_I[o.group])) { - auto &r = rmeas[o.group]; - const double sc = corrected_sigma(o, I_corr, sigma_corr); - const double v = 1.0 / (sc * sc); - const double dev = std::fabs(static_cast(I_corr) - merged_I[o.group]); - r.sum_abs_dev += dev; r.sum_I += I_corr; - r.w_abs_dev += v * dev; r.w_I += v * I_corr; r.sum_v += v; r.sum_v2 += v * v; - r.n++; r.shell = *shell; + if (!did_gpu_rmeas) { + // The threads split the GROUPS, each walking a group's fulls over the group CSR, so every + // group is still summed in index order; the observation counts are integers, so the + // blocks may add theirs up in any split. + const int nb = ReductionBlocks(n_groups, 512); + std::vector> blk_obs(nb, std::vector(n_shells, 0)); + ParallelBlocks(n_groups, nthreads, [&](int blk, int g0, int g1) { + for (int g = g0; g < g1; ++g) { + if (acc[g].d < d_floor) continue; + for (int p = fulls_group_start[g]; p < fulls_group_start[g + 1]; ++p) { + const int i = fulls_group_perm[p]; + const Obs &o = fulls[i]; + if (rejected_obs[i]) continue; // outlier-rejected in the merge -> also out of R_meas + if (!frame_cell_ok[o.frame]) continue; + if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) continue; + if (o.partiality < min_partiality) continue; + const float I_corr = o.I * o.corr, sigma_corr = o.sigma * o.corr; + if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) continue; + const auto shell = shells.GetShell(o.d); + if (!shell || *shell < 0 || *shell >= n_shells) continue; + // Only what the merge kept counts towards multiplicity - see the GPU branch above. + // The R_meas accumulation below keeps its own, wider filter. + if (mf.group[i] == g) // usable_merge + blk_obs[blk][*shell]++; + if (std::isfinite(merged_I[g])) { + auto &r = rmeas[g]; + const double sc = corrected_sigma(o, I_corr, sigma_corr); + const double v = 1.0 / (sc * sc); + const double dev = std::fabs(static_cast(I_corr) - merged_I[g]); + r.sum_abs_dev += dev; r.sum_I += I_corr; + r.w_abs_dev += v * dev; r.w_I += v * I_corr; r.sum_v += v; r.sum_v2 += v * v; + r.n++; r.shell = *shell; + } + } } - } + }, 512); + for (const auto &c : blk_obs) + for (int sh = 0; sh < n_shells; ++sh) sa[sh].total_obs += c[sh]; + } for (const auto &r : rmeas) { if (r.n < 2 || r.shell < 0 || r.shell >= n_shells) continue; double num, den; @@ -5871,6 +5941,7 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool pair_hkl[pair_of_group[g]] = group_pair[g]; } std::vector anom(pair_hkl.size()); + if (own_table) anom_export.reserve(anom.size()); // The half-sets CCanom needs are balanced within the MATE, not merely within the reflection: // it correlates one half's I(+)-I(-) against the other's, so each hand separately needs both // of its halves filled, and a mate left entirely in one half loses the whole pair rather than @@ -5949,11 +6020,19 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool std::vector row_mean(anom.size() * 2, NAN); std::vector row_shell(anom.size() * 2, -1), row_of_cell(anom.size() * 2); for (size_t c = 0; c < row_of_cell.size(); ++c) row_of_cell[c] = static_cast(c); + // Whether each pair is centric, on its own, so on all threads: the test walks the group's + // operations, and it is the costly part of the serial pair walk below. + std::vector pair_centric(anom.size()); + ParallelChunks(n_pair, ThreadsForWork(anom.size(), nthreads, 4096), [&](int lo, int hi) { + for (int q = lo; q < hi; ++q) + pair_centric[q] = gops.is_reflection_centric( + gemmi::Op::Miller{pair_hkl[q].h, pair_hkl[q].k, pair_hkl[q].l}) ? 1 : 0; + }); for (size_t p = 0; p < anom.size(); ++p) { AnomAcc a = anom[p]; a.h = pair_hkl[p].h; a.k = pair_hkl[p].k; a.l = pair_hkl[p].l; const uint64_t fkey = pair_hkl[p].pack(); - const bool centric = gops.is_reflection_centric(gemmi::Op::Miller{a.h, a.k, a.l}); + const bool centric = pair_centric[p] != 0; { const auto hshell = shells.GetShell(a.d); if (hshell && *hshell >= 0 && *hshell < n_shells) { @@ -6424,8 +6503,8 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool std::vector scaling_index(partial_group_perm.get(), partial_group_perm.get() + n_grouped); for (int i = 0; i < n_partials; ++i) if (partials[i].group < 0) scaling_index.push_back(i); - std::vector scaling_refs(n_partials); - std::vector scaling_corr(n_partials); + FillVector scaling_refs(n_partials); + FillVector scaling_corr(n_partials); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { for (int q = lo; q < hi; ++q) { const Obs &o = partials[scaling_index[q]]; @@ -6451,7 +6530,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool fitted_obs)); }); for (int f = 0; f < n_fit_frames; ++f) fit_start[f + 1] += fit_start[f]; - std::vector scaling_fit(fit_start.back()); + FillVector scaling_fit(fit_start.back()); ParallelFor(n_fit_frames, nthreads, [&](int f) { int32_t k = fit_start[f]; for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) { @@ -6788,8 +6867,8 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool std::vector full_index(fulls_group_perm.begin(), fulls_group_perm.end()); for (int i = 0; i < nf; ++i) if (fulls[i].group < 0) full_index.push_back(i); - std::vector full_refs(nf); - std::vector full_corr(nf); + FillVector full_refs(nf); + FillVector full_corr(nf); ParallelChunks(nf, full_threads, [&](int lo, int hi) { for (int q = lo; q < hi; ++q) { const Obs &o = fulls[full_index[q]]; @@ -6812,7 +6891,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool fulls.begin() + fulls_frame_start[f] + fulls_frame_count[f], fitted_full)); }); for (int f = 0; f < n_frames; ++f) fit_start[f + 1] += fit_start[f]; - std::vector full_fit(fit_start.back()); + FillVector full_fit(fit_start.back()); ParallelFor(n_frames, nthreads, [&](int f) { int32_t k = fit_start[f]; for (int i = fulls_frame_start[f]; i < fulls_frame_start[f] + fulls_frame_count[f]; ++i) { diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 514378909..3bed6e176 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -8,7 +8,11 @@ #include #include #include +#include +#include #include +#include +#include #include #include "../../common/DiffractionExperiment.h" @@ -26,6 +30,24 @@ #include "RotationScaleMergeGPU.h" #endif +// A vector allocator that leaves each new element default-initialised - for a type without initialisers +// of its own, not written at all - instead of zeroing it. An array sized on one thread and then filled +// on all of them is otherwise first written page by page by that one thread: on a long sweep, gigabytes +// of page faults taken serially before the parallel fill starts. Everything else is std::allocator's. +template +struct DefaultInitAllocator : std::allocator { + template struct rebind { using other = DefaultInitAllocator; }; + DefaultInitAllocator() noexcept = default; + template DefaultInitAllocator(const DefaultInitAllocator &) noexcept {} + template void construct(U *p) noexcept(std::is_nothrow_default_constructible_v) { + ::new (static_cast(p)) U; + } + template void construct(U *p, Args &&... args) { + ::new (static_cast(p)) U(std::forward(args)...); + } +}; +template using FillVector = std::vector>; + // Dedicated, allocate-once scale+combine+merge for rotation data (the -P rot3d path): recompute the // per-frame partiality from the (smoothed) mosaicity, robustly fit a per-image scale G, 3D-combine each // rocking event's partials into fulls, refit a per-frame scale on the fulls (XDS order), and merge with @@ -230,17 +252,20 @@ private: // One integrated observation - a per-frame partial during scaling/combine, or a combined full during // scale-fulls/merge. Flat (not nested per image); a POD so the arrays translate straight to CUDA. + // No member initialisers, so that `partials` and `fulls` (FillVector) are first written by the + // parallel passes that fill them: whoever makes one sets every field. struct Obs { int32_t h, k, l; float I, sigma, d, prescaling_corr, partiality, zeta, delta_phi, bkg, var_bkg; // Fulls only, written by the combine: the full's variance as a function of intensity, - // var(I) = var_bkg + var_per_I * I. The merge rebuilds it at the reflection's mean. - float var_per_I = 0.0f; + // var(I) = var_bkg + var_per_I * I. The merge rebuilds it at the reflection's mean. 0 on a partial. + float var_per_I; // Fulls only: the capture uncertainty, a relative sigma - capture_uncertainty_coeff times the // part of the rocking curve the event did not catch - so the merge can add (capture * )^2 to // the variance it rebuilds at the reflection's mean, as the combine added it to the full's own. - float capture = 0.0f; - float px = NAN, py = NAN; // predicted detector position (for the absorption surface; CPU path only) + // 0 on a partial. + float capture; + float px, py; // predicted detector position (for the absorption surface; CPU path only) float image_number; // fractional frame position (for 3D-combine contiguity) int32_t frame; // index of the outcome whose per-frame scale G applies to this obs uint8_t on_ice; @@ -343,7 +368,7 @@ private: // 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 + FillVector partials; // all per-frame partials, grouped by frame bool resident_ingest = false; // What the ingest-time geometry smoothing (SmoothMosaicityAndPartiality) needs, one array per field // in `partials` order: the two fields it rewrites, the two the rocking-event walks read in raw-hkl @@ -420,7 +445,7 @@ private: std::vector rawrun_d; // representative resolution per raw hkl std::vector rawrun_group; // dense ASU-group id per raw hkl (<0 = absent/out of range) - std::vector fulls; // combined fulls (rebuilt each Run), sorted by frame + FillVector fulls; // combined fulls (rebuilt each Run), sorted by frame // The fulls by ASU group, each group's in index order: group g is fulls_group_perm[fulls_group_start[g]] // up to fulls_group_start[g + 1]. A pass over every group's fulls can then give each thread whole // groups and still add a group up in the order a serial walk over the fulls would. Rebuilt with the @@ -665,7 +690,7 @@ private: // The same again over a scaling loop's copy, refs and corr both in group order: group g's at // [group_start[g], group_start[g] + group_count[g]). void ReduceScalingGroupMeans(int n_groups, const int32_t *group_start, const int32_t *group_count, - const std::vector &refs, const std::vector &corr, + const FillVector &refs, const FillVector &corr, std::vector &out_mean) const; // Per-frame G fit against the reference group_mean_in[group], over a scaling loop's ScalingFit, frame @@ -674,7 +699,7 @@ private: // event filed under one frame, so on a sparse sweep a frame holds a handful - its fit goes in at // the information it has (min_obs 1), and the smoothing in RunFullsScalingLoop decides how far its // neighbours carry it. - void FitScalingFrames(const std::vector &fit, const std::vector &fit_start, + void FitScalingFrames(const FillVector &fit, const std::vector &fit_start, size_t min_obs, const std::vector &group_mean_in, std::vector &g); // The alternating per-frame scale fit run to convergence: `iterate` fits every frame once against @@ -710,7 +735,7 @@ private: ScalingLoopOutcome outcome; std::vector fitted; std::vector g; - std::vector corr; + FillVector corr; }; // A pass alternates its groupings (P1 full range, P1 at the search limit, the adopted group), so a // few are kept; the oldest go first. @@ -863,9 +888,9 @@ private: const char *name, bool resolution_gauge, const SphericalHarmonicBasis *sh_basis = nullptr); - // Sort `fulls` by peak frame and (re)build fulls_frame_start/count (the per-frame CSR the scale-fulls - // step slices). Shared by the CPU Combine tail and the GPU combine path. - void SortFullsByFrame(); + // Fill `fulls` with the parts' fulls sorted by peak frame and (re)build fulls_frame_start/count (the + // per-frame CSR the scale-fulls step slices). The CPU Combine's tail. + void SortFullsByFrame(const std::vector> &part); // Per-frame CC vs the partial merge reference (CPU; the GPU equivalent is gpu_->ComputePartialCC). void ComputePerFrameCC(const std::vector &partial_group_mean, From c914fa9f333047465396c117e11a4c5a8f98ea63 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 19:19:22 +0200 Subject: [PATCH 2/3] RotationScaleMerge: the CPU scaling loops move fewer bytes, output unchanged With the serial walks gone the CPU merge is bound by memory traffic over the partials - running the final merge of a large set at -N 7 takes 20 s against 18 s at -N 32 on a dual-channel machine - so these changes take passes and random accesses out of the loops rather than add threads: - The partials' corr update of one scaling iteration (UpdateCorr, then the rescale) is no longer a pass of its own over the working copy: each grouped partial is brought up to date by the next iteration's group-mean walk, which reads it anyway, and what is owed when the loop stops is applied after it (CorrUpdate). The partials in no group, which the means never read, are updated right away. Every corr goes through the same updates in the same order. - A partial's ASU group is kept in a flat per-partial array (partial_group) that the CPU passes read, instead of being stamped into the eighty-byte records at random once per grouping. - RunArm no longer restores corr onto the partials when the CPU partial scaling is about to overwrite every one of them; that loop starts its copy from corr_ingested directly. Its result goes back to the partials' order through a compact array (also the memo's) and then onto the partials in one sequential pass, instead of a random scatter into the records. - The per-frame fit prefetches the group mean it gathers, 16 observations ahead; the correction surfaces' fit reads one reference array (swI / sw, NaN where a group has no weight) instead of two, with the same prefetch. - The loops' index arrays (partials, fulls) are built on all threads. p.hkl and the MTZ data are byte-identical on three large sets with --no-gpu, and the final merge re-run (--mode scale) gives the same bytes at -N 7 and -N 32. Provisional, shared machine (load 12-27): scaling-and-merging step, against the start of this branch, 115 -> 77 s on a 3600-frame set, 74 -> 58 s on a 500-frame P1-heavy set, 55 -> 42 s on an 1800-frame set. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../scale_merge/RotationScaleMerge.cpp | 136 +++++++++++------- .../scale_merge/RotationScaleMerge.h | 32 ++++- 2 files changed, 108 insertions(+), 60 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index b62db5867..808adbc0e 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -1462,11 +1462,12 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { return n_groups; } #endif - // Without the GPU: stamp the ASU-group id per obs from its raw hkl + the precomputed finiteness, on - // the partials and in a flat array the group CSR below is counted from. + // Without the GPU: stamp the ASU-group id per obs from its raw hkl + the precomputed finiteness, in + // the flat array the CPU passes read it from (partial_group) and the group CSR below is counted from. // Plain arrays, here and for gperm below: every element is written, and a sized vector would // first zero all of them on one thread - once per merge, a sweep's worth of observations. - std::unique_ptr group_ids(new int32_t[n_partials_obs]); + partial_group.reset(new int32_t[n_partials_obs]); + int32_t *group_ids = partial_group.get(); // Parallel over raw-hkl runs: distinct runs own disjoint perm ranges, hence disjoint observations. ParallelChunks(n_run, nthreads, [&](int rlo, int rhi) { for (int r = rlo; r < rhi; ++r) { @@ -1474,9 +1475,7 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { const int lo = rawrun_start[r], hi = rawrun_start[r] + rawrun_count[r]; for (int p = lo; p < hi; ++p) { const int i = perm[p]; - const int gi = (g >= 0 && finite_ok[i]) ? g : -1; - group_ids[i] = gi; - partials[i].group = gi; + group_ids[i] = (g >= 0 && finite_ok[i]) ? g : -1; } } }); @@ -1563,7 +1562,8 @@ void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vector &refs, - const FillVector &corr, - std::vector &out_mean) const { + FillVector &corr, + std::vector &out_mean, const CorrUpdate *update) const { out_mean.resize(n_groups); ParallelChunks(n_groups, ThreadsForWork(corr.size(), nthreads), [&](int glo, int ghi) { for (int g = glo; g < ghi; ++g) { @@ -1584,6 +1584,7 @@ void RotationScaleMerge::ReduceScalingGroupMeans(int n_groups, const int32_t *gr const int lo = group_start[g], hi = lo + group_count[g]; for (int p = lo; p < hi; ++p) { const ScalingRef &r = refs[p]; + if (update) update->Apply(r, corr[p]); float I_corr = 0.0f; const double w = ReferenceWeight(g, corr[p], r.partiality, r.I, r.sigma, min_partiality, I_corr); if (w <= 0.0) continue; @@ -1674,12 +1675,17 @@ void RotationScaleMerge::FitScalingFrames(const FillVector &fit, con size_t min_obs, const std::vector &group_mean_in, std::vector &g) { const int n_fr = static_cast(fit_start.size()) - 1; + const int n_fit = static_cast(fit.size()); + constexpr int FIT_PREFETCH = 16; std::vector scaled(n_fr, 0); std::vector info(n_fr, 0.0); ParallelFor(n_fr, nthreads, [&](int f) { double num = 0.0, den = 0.0; size_t n = 0; for (int i = fit_start[f]; i < fit_start[f + 1]; ++i) { + // The group mean is a gather from an array as long as the group count - the cost of this + // loop on a P1 merge - so it is asked for FIT_PREFETCH observations ahead. + if (i + FIT_PREFETCH < n_fit) Prefetch(&group_mean_in[fit[i + FIT_PREFETCH].group]); const ScalingFit &o = fit[i]; // The reference is the group's mean WITH this observation in it. Without it the update is // an undamped Jacobi step on the sparse graph a P1 merge of a short sweep is, and it @@ -3741,6 +3747,7 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // Terms per block of the per-cell sums in the fit below: every block zeroes and hands back two // ncell-long arrays, so a block has to hold many more terms than there are cells. constexpr int SURFACE_BLOCK = 32768; + constexpr int SURFACE_PREFETCH = 16; bool on_gpu = false; int n_terms = 0; #ifdef JFJOCH_USE_CUDA @@ -3907,7 +3914,9 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // The per-group reference of one subset with surface A applied - the first pass of both the fit // and the score. - std::vector sw(n_groups), swI(n_groups); + // iref is the reference itself, swI / sw, NaN where a group has no weight: the fit below reads one + // array at random per term instead of two (CPU only). + std::vector sw(n_groups), swI(n_groups), iref(on_gpu ? 0 : n_groups); auto reference = [&](int parity, const std::vector &A) { #ifdef JFJOCH_USE_CUDA if (on_gpu) { @@ -3930,6 +3939,7 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int s_w += w; s_wI += w * Is; } sw[g] = s_w; swI[g] = s_wI; + iref[g] = s_w > 0.0 ? s_wI / s_w : std::numeric_limits::quiet_NaN(); } }); }; @@ -3992,9 +4002,11 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int std::fill(xcross.begin(), xcross.end(), 0.0); std::fill(xref2.begin(), xref2.end(), 0.0); for (int k = lo; k < hi; ++k) { + // The reference is a gather from an array as long as the group count, so it is + // asked for SURFACE_PREFETCH terms ahead. + if (k + SURFACE_PREFETCH < hi) Prefetch(&iref[term[sel[k + SURFACE_PREFETCH]].group]); const Term &o = term[sel[k]]; - if (sw[o.group] <= 0.0) continue; - const double Iref = swI[o.group] / sw[o.group], a = A[o.cell]; + const double Iref = iref[o.group], a = A[o.cell]; // NaN: the group has no weight const double Is = static_cast(o.I) * o.corr * a, sc = static_cast(o.sigma) * o.corr * a; if (!std::isfinite(Iref) || Iref <= 0.0 || !(sc > 0.0)) continue; const double w = 1.0 / (sc * sc); @@ -4635,12 +4647,13 @@ void RotationScaleMerge::ComputePerFrameCC(const std::vector &partial_gr const int lo = frame_start[f], hi = frame_start[f] + frame_count[f]; for (int i = lo; i < hi; ++i) { const auto &o = partials[i]; + const int g = partial_group[i]; if (o.on_ice) continue; - if (o.group < 0) continue; + if (g < 0) continue; if (o.partiality < min_partiality) continue; if (!std::isfinite(o.I) || !std::isfinite(o.corr) || o.corr <= 0.0f) continue; if (!std::isfinite(o.sigma) || o.sigma <= 0.0f) continue; - const double mean = partial_group_mean[o.group]; + const double mean = partial_group_mean[g]; if (!std::isfinite(mean)) continue; const double img = static_cast(o.I) * o.corr; sx += img; sy += mean; sx2 += img * img; sy2 += mean * mean; sxy += img * mean; @@ -6405,7 +6418,10 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool restored_on_gpu = true; } #endif - if (!restored_on_gpu) + // The CPU partial scaling below writes every partial's corr when it is done, from a copy it starts + // off corr_ingested itself, and nothing between here and there reads corr - so there the restore + // is not made, which saves a pass over the eighty-byte records to set four bytes of each. + if (!restored_on_gpu && !scale_partials) ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), [&](int lo, int hi) { for (int i = lo; i < hi; ++i) partials[i].corr = corr_ingested[i]; @@ -6500,16 +6516,21 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool const size_t partial_threads = ThreadsForWork(partials.size(), nthreads); const int n_grouped = partial_group_start.empty() ? 0 : partial_group_start.back() + partial_group_count.back(); - std::vector scaling_index(partial_group_perm.get(), partial_group_perm.get() + n_grouped); - for (int i = 0; i < n_partials; ++i) - if (partials[i].group < 0) scaling_index.push_back(i); + FillVector scaling_index(n_partials); + const std::vector ungrouped = KeptIndices(n_partials, nthreads, [&](int i) { + return partial_group[i] < 0; + }); + ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { + for (int q = lo; q < hi; ++q) + scaling_index[q] = q < n_grouped ? partial_group_perm[q] : ungrouped[q - n_grouped]; + }); FillVector scaling_refs(n_partials); FillVector scaling_corr(n_partials); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { for (int q = lo; q < hi; ++q) { const Obs &o = partials[scaling_index[q]]; scaling_refs[q] = {o.frame, o.I, o.sigma, o.partiality, o.prescaling_corr}; - scaling_corr[q] = o.corr; + scaling_corr[q] = corr_ingested[scaling_index[q]]; // the restore (see the top) } }); @@ -6518,68 +6539,71 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // partials are rocking-curve tails below it, measured at I/sigma ~ 0.1, and a frame's scale // fitted on them is fitted on noise that the reference does not answer to. The Rotation // model needs zeta > 0. - auto fitted_obs = [&](const Obs &o) { - return o.group >= 0 && !o.on_ice && o.partiality >= min_partiality + auto fitted_obs = [&](int i) { + const Obs &o = partials[i]; + return partial_group[i] >= 0 && !o.on_ice && o.partiality >= min_partiality && std::isfinite(o.zeta) && o.zeta > 0.0f; }; const int n_fit_frames = static_cast(frame_start.size()); std::vector fit_start(n_fit_frames + 1, 0); ParallelFor(n_fit_frames, nthreads, [&](int f) { - fit_start[f + 1] = static_cast(std::count_if( - partials.begin() + frame_start[f], partials.begin() + frame_start[f] + frame_count[f], - fitted_obs)); + int32_t n = 0; + for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) n += fitted_obs(i) ? 1 : 0; + fit_start[f + 1] = n; }); for (int f = 0; f < n_fit_frames; ++f) fit_start[f + 1] += fit_start[f]; FillVector scaling_fit(fit_start.back()); ParallelFor(n_fit_frames, nthreads, [&](int f) { int32_t k = fit_start[f]; for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) { + if (!fitted_obs(i)) continue; const Obs &o = partials[i]; - if (!fitted_obs(o)) continue; const double weight = SafeInv(o.sigma, 1.0); - scaling_fit[k++] = {o.group, o.partiality * SafeInv(o.prescaling_corr, 1.0), + scaling_fit[k++] = {partial_group[i], o.partiality * SafeInv(o.prescaling_corr, 1.0), static_cast(o.I), weight * weight}; } }); // RunScalingLoop rescales after every iteration before anything reads corr, so the iteration's - // UpdateCorr is folded into that pass over the partials: the G and the fitted frames as the - // iteration left them, then the rescale on top - the same two roundings, one pass instead of two. - std::vector g_fitted; - std::vector frame_fitted; + // UpdateCorr is folded into the rescale: the G and the fitted frames as the iteration left + // them, then the rescale on top - the same two roundings. Nor is that a pass of its own over + // the copy: the next iteration's group means are the next thing to read corr, so each partial + // in a group is brought up to date there, in the same walk, and only the partials in no group + // (which the means never read) are updated right away. Whatever is still owed when the loop + // stops is applied after it. Each corr goes through the same updates in the same order. + CorrUpdate update; + bool update_owed = false; // the grouped partials have not had `update` yet partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window, [&] { ReduceScalingGroupMeans(n_groups, partial_group_start.data(), partial_group_count.data(), - scaling_refs, scaling_corr, partial_mean); + scaling_refs, scaling_corr, partial_mean, update_owed ? &update : nullptr); + update_owed = false; FitScalingFrames(scaling_fit, fit_start, MIN_REFLECTIONS, partial_mean, g_partial); - g_fitted = g_partial; - frame_fitted = frame_scaled_scratch; + update.g = g_partial; + update.fitted = frame_scaled_scratch; }, [&](const std::vector &apply, const std::vector &ratio) { - ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int q = lo; q < hi; ++q) { - const ScalingRef &o = scaling_refs[q]; - float &corr = scaling_corr[q]; - if (frame_fitted[o.frame]) { // UpdateCorr - const double denom = static_cast(o.partiality) * g_fitted[o.frame]; - if (std::isfinite(o.prescaling_corr) && std::isfinite(denom) && denom > 0.0) - corr = static_cast(o.prescaling_corr / denom); - else - corr = NAN; - } - if (apply[o.frame] && std::isfinite(corr)) - corr = static_cast(corr * ratio[o.frame]); - } + update.apply = apply; + update.ratio = ratio; + update_owed = true; + ParallelChunks(n_partials - n_grouped, partial_threads, [&](int lo, int hi) { + for (int q = n_grouped + lo; q < n_grouped + hi; ++q) update.Apply(scaling_refs[q], scaling_corr[q]); }); }); + if (update_owed) + ParallelChunks(n_grouped, partial_threads, [&](int lo, int hi) { + for (int q = lo; q < hi; ++q) update.Apply(scaling_refs[q], scaling_corr[q]); + }); + // Back into the partials' own order - the order the memo keeps - in the compact array, and + // from there onto the partials in one sequential pass. + FillVector memo_corr(n_partials); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int q = lo; q < hi; ++q) partials[scaling_index[q]].corr = scaling_corr[q]; + for (int q = lo; q < hi; ++q) memo_corr[scaling_index[q]] = scaling_corr[q]; }); - // The memo keeps corr in the partials' own order. ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) scaling_corr[i] = partials[i].corr; + for (int i = lo; i < hi; ++i) partials[i].corr = memo_corr[i]; }); - PartialLoopMemo m{key, partial_loop, frame_scaled_scratch, g_partial, std::move(scaling_corr)}; + PartialLoopMemo m{key, partial_loop, frame_scaled_scratch, g_partial, std::move(memo_corr)}; if (partial_loop_memos.size() == MAX_PARTIAL_LOOP_MEMOS) partial_loop_memos.erase(partial_loop_memos.begin()); partial_loop_memos.push_back(std::move(m)); @@ -6864,9 +6888,13 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // full_index[q] is the full at position q, in the group CSR's order and then the fulls in no group. const int nf = static_cast(fulls.size()); const size_t full_threads = ThreadsForWork(fulls.size(), nthreads); - std::vector full_index(fulls_group_perm.begin(), fulls_group_perm.end()); - for (int i = 0; i < nf; ++i) - if (fulls[i].group < 0) full_index.push_back(i); + const int nf_grouped = static_cast(fulls_group_perm.size()); + const std::vector ungrouped = KeptIndices(nf, nthreads, [&](int i) { return fulls[i].group < 0; }); + FillVector full_index(nf); + ParallelChunks(nf, full_threads, [&](int lo, int hi) { + for (int q = lo; q < hi; ++q) + full_index[q] = q < nf_grouped ? fulls_group_perm[q] : ungrouped[q - nf_grouped]; + }); FillVector full_refs(nf); FillVector full_corr(nf); ParallelChunks(nf, full_threads, [&](int lo, int hi) { diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 3bed6e176..bb5185c97 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -649,6 +649,10 @@ private: // group over it, on all threads and in the order the serial loop would. std::vector partial_group_start, partial_group_count; NotCopied partial_group_perm; + // Without a GPU: each partial's ASU group (-1 = in none), in observation order - where the CPU + // passes read a partial's group from. A partial's own Obs::group is not stamped: written into the + // eighty-byte records at random, once per grouping, it cost more than the grouping itself. + NotCopied partial_group; // The CPU scaling loops' working copy of the partials, or of the fulls. A loop passes over them // several times per iteration and reads a handful of fields each time, so it works on these rather @@ -681,17 +685,33 @@ private: I_corr = Ic; return 1.0 / (static_cast(sigma_corr) * sigma_corr); } - static double ReferenceWeight(const Obs &o, double min_partiality, float &I_corr) { - return ReferenceWeight(o.group, o.corr, o.partiality, o.I, o.sigma, min_partiality, I_corr); - } // Inverse-variance per-group mean of I*corr over the partials (the scaling reference; no cell mask - // only the final merge applies one), using the CSR above, each group summed in index order. void ReducePartialGroupMeans(int n_groups, std::vector &out_mean) const; + // The partials' loop's corr update of one iteration - UpdateCorr from the G it fitted, then the + // rescale RunScalingLoop asked for - per frame: `fitted` and `g` as the iteration left them, `apply` + // and `ratio` the rescale's. + struct CorrUpdate { + std::vector fitted, apply; + std::vector g, ratio; + void Apply(const ScalingRef &o, float &corr) const { + if (fitted[o.frame]) { // UpdateCorr + const double denom = static_cast(o.partiality) * g[o.frame]; + if (std::isfinite(o.prescaling_corr) && std::isfinite(denom) && denom > 0.0) + corr = static_cast(o.prescaling_corr / denom); + else + corr = NAN; + } + if (apply[o.frame] && std::isfinite(corr)) + corr = static_cast(corr * ratio[o.frame]); + } + }; // The same again over a scaling loop's copy, refs and corr both in group order: group g's at - // [group_start[g], group_start[g] + group_count[g]). + // [group_start[g], group_start[g] + group_count[g]). With `update`, each corr is first brought up to + // date by it, in the same walk (see the partials' loop in RunArm). void ReduceScalingGroupMeans(int n_groups, const int32_t *group_start, const int32_t *group_count, - const FillVector &refs, const FillVector &corr, - std::vector &out_mean) const; + const FillVector &refs, FillVector &corr, + std::vector &out_mean, const CorrUpdate *update = nullptr) const; // Per-frame G fit against the reference group_mean_in[group], over a scaling loop's ScalingFit, frame // f's at [fit_start[f], fit_start[f + 1]). A frame with fewer than min_obs of them keeps its G: a From 4a8aff41c211074ec786ac06ff3b80942e3cbb54 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 19:58:30 +0200 Subject: [PATCH 3/3] RotationScaleMerge: the partials' scaling loop walks only what the means can read, smaller fit records, per-parity surface references, output unchanged Three more cuts in the bytes the CPU scaling moves per iteration: - The partials' working copy is ordered with the partials a group mean can take in first - in a group and not below the partiality floor, ReferenceWeight's fixed tests - group by group in the CSR's order, and all the others after them. The group means walk only the first part: a partial they would skip anyway adds nothing to a sum, so no sum changes, and on a fine-sliced sweep about half of the partials are rocking-curve tails below the floor. The others are read by nothing during the loop, so their per-iteration corr updates are kept (a few per-frame arrays an iteration) and applied to them once, in order, when the loop ends. - ScalingFit holds the four floats the per-frame fit reads (20 bytes, was 32) and the fit works out 1/sigma^2 and partiality/prescaling_corr from them where it uses them: the same expressions, so the same doubles. - The correction surfaces keep their group-ordered reference terms in one copy per frame parity, so a cross-validation reference pass reads only its own half; a pass over all the terms walks a group's two halves together in term order, which is the order the single copy had. Samples in the partial loop on a 3600-frame set (perf, cycles, whole scaling-and-merging step): group means 29.7k -> 15.1k, per-frame fit 22.4k -> 13.7k, the separate corr update pass 11.9k -> gone. Wall time on the shared machine (load 26-33) is within its noise of the previous commit. Validated with the battery against the start of the branch (873880bba), --no-gpu: p.hkl and MTZ data byte-identical on 4nwv, 8qaw, 9v16, lyso_x06da_ref, cytc_x10sa, myob_x10sa and yag_x10sa_20keV; the same bytes again at -N 11 on all seven; and on the GPU path the same bytes on lyso_x06da_ref and 9v16. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4 --- .../scale_merge/RotationScaleMerge.cpp | 118 ++++++++++++------ .../scale_merge/RotationScaleMerge.h | 7 +- 2 files changed, 85 insertions(+), 40 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 808adbc0e..86dc97067 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -1696,9 +1696,11 @@ void RotationScaleMerge::FitScalingFrames(const FillVector &fit, con // iterations, which is what RunScalingLoop's tolerance is for. const double mean = group_mean_in[o.group]; if (!std::isfinite(mean)) continue; - const double coeff = o.coeff_per_mean * mean; - num += o.w2 * coeff * o.Iobs; - den += o.w2 * coeff * coeff; + const double weight = SafeInv(o.sigma, 1.0), w2 = weight * weight; + const double coeff_per_mean = o.partiality * SafeInv(o.prescaling_corr, 1.0); + const double coeff = coeff_per_mean * mean; + num += w2 * coeff * static_cast(o.I); + den += w2 * coeff * coeff; ++n; } if (n < min_obs) return; @@ -3897,18 +3899,25 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // terms in fulls order, because the counting sort is stable (it walks `term`, which is in fulls // order). A separate copy rather than a permutation: chasing every term through an index at random // over a few hundred megabytes is the locality this pass used to stay serial to keep. - struct RefTerm { float I, sigma, corr; int32_t cell; uint8_t parity; }; - std::vector g_start; - std::vector gterm; + // One copy per frame parity: most reference passes are of one half (the cross-validation), and + // each then reads only its own. A pass over all the terms walks a group's two halves together in + // term order (`at`, the term's position), which is the fulls order the single copy held. + struct RefTerm { float I, sigma, corr; int32_t cell, at; }; + std::array, 2> g_start; + std::array, 2> gterm; if (!on_gpu) { - g_start.assign(n_groups + 1, 0); - for (const Term &t : term) ++g_start[t.group + 1]; - for (int g = 0; g < n_groups; ++g) g_start[g + 1] += g_start[g]; - gterm.resize(term.size()); - std::vector fill(g_start.begin(), g_start.end() - 1); + for (int par = 0; par < 2; ++par) g_start[par].assign(n_groups + 1, 0); + for (size_t k = 0; k < term.size(); ++k) ++g_start[term_parity[k]][term[k].group + 1]; + std::array, 2> fill; + for (int par = 0; par < 2; ++par) { + for (int g = 0; g < n_groups; ++g) g_start[par][g + 1] += g_start[par][g]; + gterm[par].resize(g_start[par][n_groups]); + fill[par].assign(g_start[par].begin(), g_start[par].end() - 1); + } for (size_t k = 0; k < term.size(); ++k) { const Term &t = term[k]; - gterm[fill[t.group]++] = {t.I, t.sigma, t.corr, t.cell, term_parity[k]}; + const int par = term_parity[k]; + gterm[par][fill[par][t.group]++] = {t.I, t.sigma, t.corr, t.cell, static_cast(k)}; } } @@ -3927,9 +3936,7 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int ParallelChunks(n_groups, nt, [&](int glo, int ghi) { for (int g = glo; g < ghi; ++g) { double s_w = 0.0, s_wI = 0.0; - for (int k = g_start[g]; k < g_start[g + 1]; ++k) { - const RefTerm &t = gterm[k]; - if (parity >= 0 && t.parity != parity) continue; + auto add = [&](const RefTerm &t) { // Keep the intensity's own product together before it is weighted: w * ((I*corr)*a) // is what the loops this replaced summed, and w * I * corr * a rounds differently. // The score decides whether a surface is kept, so a last-place difference here can @@ -3937,6 +3944,14 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int const double a = A[t.cell], Is = static_cast(t.I) * t.corr * a; const double sc = static_cast(t.sigma) * t.corr * a, w = 1.0 / (sc * sc); s_w += w; s_wI += w * Is; + }; + if (parity >= 0) { + for (int k = g_start[parity][g]; k < g_start[parity][g + 1]; ++k) add(gterm[parity][k]); + } else { + int k0 = g_start[0][g], k1 = g_start[1][g]; + const int e0 = g_start[0][g + 1], e1 = g_start[1][g + 1]; + while (k0 < e0 || k1 < e1) + add(k1 == e1 || (k0 < e0 && gterm[0][k0].at < gterm[1][k1].at) ? gterm[0][k0++] : gterm[1][k1++]); } sw[g] = s_w; swI[g] = s_wI; iref[g] = s_w > 0.0 ? s_wI / s_w : std::numeric_limits::quiet_NaN(); @@ -6511,18 +6526,43 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool } else { // The loop works on its own compact copy of the partials (ScalingRef and ScalingFit in the // header) and hands corr back to them when it is done. scaling_index[q] is the partial at - // position q of the copy: the CSR's group order, then the partials in no group. + // position q of the copy: first the partials a group mean can take in - in a group and not + // below the partiality floor (ReferenceWeight's fixed tests) - group by group in the CSR's + // order, then all the others in index order. The means walk only the first part: a partial + // they would skip in any case adds nothing to a sum, so leaving it out of the walk changes no + // sum, and on a fine-sliced sweep half of the partials are rocking-curve tails below the floor. const int n_partials = static_cast(partials.size()); const size_t partial_threads = ThreadsForWork(partials.size(), nthreads); - const int n_grouped = partial_group_start.empty() - ? 0 : partial_group_start.back() + partial_group_count.back(); - FillVector scaling_index(n_partials); - const std::vector ungrouped = KeptIndices(n_partials, nthreads, [&](int i) { - return partial_group[i] < 0; - }); + // Taken once, in index order, so that the walks over the CSR below gather one byte a partial + // rather than its record. + FillVector in_ref(n_partials); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int q = lo; q < hi; ++q) - scaling_index[q] = q < n_grouped ? partial_group_perm[q] : ungrouped[q - n_grouped]; + for (int i = lo; i < hi; ++i) + in_ref[i] = partial_group[i] >= 0 && !(partials[i].partiality < min_partiality) ? 1 : 0; + }); + auto in_reference = [&](int i) { return in_ref[i] != 0; }; + std::vector ref_start(n_groups), ref_count(n_groups); + ParallelChunks(n_groups, ThreadsForWork(partials.size(), nthreads), [&](int glo, int ghi) { + for (int g = glo; g < ghi; ++g) { + int32_t c = 0; + for (int p = partial_group_start[g]; p < partial_group_start[g] + partial_group_count[g]; ++p) + c += in_reference(partial_group_perm[p]) ? 1 : 0; + ref_count[g] = c; + } + }); + int n_ref = 0; + for (int g = 0; g < n_groups; ++g) { ref_start[g] = n_ref; n_ref += ref_count[g]; } + FillVector scaling_index(n_partials); + ParallelChunks(n_groups, ThreadsForWork(partials.size(), nthreads), [&](int glo, int ghi) { + for (int g = glo; g < ghi; ++g) { + int32_t q = ref_start[g]; + for (int p = partial_group_start[g]; p < partial_group_start[g] + partial_group_count[g]; ++p) + if (in_reference(partial_group_perm[p])) scaling_index[q++] = partial_group_perm[p]; + } + }); + const std::vector rest = KeptIndices(n_partials, nthreads, [&](int i) { return !in_reference(i); }); + ParallelChunks(static_cast(rest.size()), ThreadsForWork(rest.size(), nthreads), [&](int lo, int hi) { + for (int j = lo; j < hi; ++j) scaling_index[n_ref + j] = rest[j]; }); FillVector scaling_refs(n_partials); FillVector scaling_corr(n_partials); @@ -6558,9 +6598,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) { if (!fitted_obs(i)) continue; const Obs &o = partials[i]; - const double weight = SafeInv(o.sigma, 1.0); - scaling_fit[k++] = {partial_group[i], o.partiality * SafeInv(o.prescaling_corr, 1.0), - static_cast(o.I), weight * weight}; + scaling_fit[k++] = {partial_group[i], o.I, o.sigma, o.partiality, o.prescaling_corr}; } }); @@ -6568,14 +6606,17 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // UpdateCorr is folded into the rescale: the G and the fitted frames as the iteration left // them, then the rescale on top - the same two roundings. Nor is that a pass of its own over // the copy: the next iteration's group means are the next thing to read corr, so each partial - // in a group is brought up to date there, in the same walk, and only the partials in no group - // (which the means never read) are updated right away. Whatever is still owed when the loop - // stops is applied after it. Each corr goes through the same updates in the same order. + // they walk is brought up to date there, in the same walk, and whatever is still owed when the + // loop stops is applied after it. The partials the means do not walk are read by nothing until + // the loop is over, so their updates are kept - a few per-frame arrays an iteration - and + // applied to them once, in order, at the end. Each corr goes through the same updates in the + // same order. CorrUpdate update; - bool update_owed = false; // the grouped partials have not had `update` yet + bool update_owed = false; // the walked partials have not had `update` yet + std::vector rest_updates; partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window, [&] { - ReduceScalingGroupMeans(n_groups, partial_group_start.data(), partial_group_count.data(), + ReduceScalingGroupMeans(n_groups, ref_start.data(), ref_count.data(), scaling_refs, scaling_corr, partial_mean, update_owed ? &update : nullptr); update_owed = false; FitScalingFrames(scaling_fit, fit_start, MIN_REFLECTIONS, partial_mean, g_partial); @@ -6585,14 +6626,16 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool update.apply = apply; update.ratio = ratio; update_owed = true; - ParallelChunks(n_partials - n_grouped, partial_threads, [&](int lo, int hi) { - for (int q = n_grouped + lo; q < n_grouped + hi; ++q) update.Apply(scaling_refs[q], scaling_corr[q]); - }); + rest_updates.push_back(update); }); if (update_owed) - ParallelChunks(n_grouped, partial_threads, [&](int lo, int hi) { + ParallelChunks(n_ref, partial_threads, [&](int lo, int hi) { for (int q = lo; q < hi; ++q) update.Apply(scaling_refs[q], scaling_corr[q]); }); + ParallelChunks(n_partials - n_ref, partial_threads, [&](int lo, int hi) { + for (int q = n_ref + lo; q < n_ref + hi; ++q) + for (const CorrUpdate &u : rest_updates) u.Apply(scaling_refs[q], scaling_corr[q]); + }); // Back into the partials' own order - the order the memo keeps - in the compact array, and // from there onto the partials in one sequential pass. @@ -6925,8 +6968,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool for (int i = fulls_frame_start[f]; i < fulls_frame_start[f] + fulls_frame_count[f]; ++i) { const Obs &o = fulls[i]; if (!fitted_full(o)) continue; - const double weight = SafeInv(o.sigma, 1.0); - full_fit[k++] = {o.group, 1.0, static_cast(o.I), weight * weight}; + full_fit[k++] = {o.group, o.I, o.sigma, 1.0f, 1.0f}; } }); diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index bb5185c97..5c313d8d2 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -666,10 +666,13 @@ private: int32_t frame; float I, sigma, partiality, prescaling_corr; }; + // The fit's factors that do not depend on the reference - the coefficient of a mean of 1, + // partiality / prescaling_corr (1 for a full), and the weight 1/sigma^2 - are worked out from these + // floats where they are used: twenty bytes an observation read every iteration instead of thirty-two, + // and the same doubles. struct ScalingFit { int32_t group; - double coeff_per_mean; // partiality / prescaling_corr (1 for a full): the fit's coeff for a mean of 1 - double Iobs, w2; + float I, sigma, partiality, prescaling_corr; }; // The observation's inverse-variance weight in the scaling reference, and its scaled intensity;