From 4a8aff41c211074ec786ac06ff3b80942e3cbb54 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 19:58:30 +0200 Subject: [PATCH] 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;