From bd27d802b76f17fe8dbb59f2c273ad00bacd94af Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 09:28:20 +0200 Subject: [PATCH] RotationScaleMerge: CPU fulls scaling loop on a working copy, combine prefetches, output unchanged The fulls' scaling loop now works as the partials' loop does since 4e177bdb2: a compact copy of the fulls in the group CSR's order (ScalingRef, corr beside it), so the group means read sequentially instead of gathering eighty-byte Obs through fulls_group_perm, and the per-frame fit over a ScalingFit list built once with the Unity model's filters (FitScalingFrames, which takes the minimum count as a parameter). FitPerFrameG, SolveScale and ReduceFullGroupMeans had no callers left and are gone; their rationale comments moved to where the filters and the fit now are. Combine prefetches the partials of the raw hkl 8 runs ahead: a run's partials lie a frame apart in memory and the work per rocking event kept the CPU from reaching the next runs' loads, so it waited on a cache miss for nearly every partial. Same values, same operations in the same order: p.hkl and the merged MTZ data are byte-identical on lysozyme, cytochrome C, myoglobin, a small-molecule set and 4nwv. Measured on 4nwv (Apple M6, two interleaved runs each): the combine-and-scale-fulls stages 12.6 / 12.7 s -> 5.6 / 8.9 s; the scaling steps (5 + 9) 54.6 / 53.0 s -> 52.0 / 51.0 s; total wall time is within run-to-run noise (127-131 s). Cytochrome C scaling and merging 3.35 / 3.43 s -> 3.19 / 3.19 s. Co-Authored-By: Claude Opus 5.5 --- .../scale_merge/RotationScaleMerge.cpp | 204 +++++++++--------- .../scale_merge/RotationScaleMerge.h | 42 ++-- .../scale_merge/RotationScaleMergeGPU.cu | 6 +- 3 files changed, 128 insertions(+), 124 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 08bbd1b51..2424c2d2b 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -19,6 +19,10 @@ #include #include +#if defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)) +#include +#endif + #include #include @@ -264,7 +268,14 @@ namespace { for (size_t j = 0; j < to[k].size(); ++j) to[k][j] += from[k][j]; } - struct ScaleObs { double coeff, Iobs, weight; }; + // Ask for the cache line at p ahead of its use. Only a hint: it changes no result. + inline void Prefetch(const void *p) { +#if defined(_MSC_VER) && (defined(_M_X64) || defined(_M_IX86)) + _mm_prefetch(static_cast(p), _MM_HINT_T0); +#elif defined(__GNUC__) || defined(__clang__) + __builtin_prefetch(p); +#endif + } // Per-frame scale, linear in G: the weighted least-squares slope of the observations on their // reference, G = sum w^2 c I / sum w^2 c^2. PLAIN, not the Cauchy-reweighted fit ScaleOnTheFly @@ -281,16 +292,6 @@ namespace { return std::isfinite(G) ? std::max(0.0, G) : 1.0; } - double SolveScale(const std::vector &obs, double &info) { - double num = 0.0, den = 0.0; - for (const auto &o : obs) { - const double w2 = o.weight * o.weight; - num += w2 * o.coeff * o.Iobs; - den += w2 * o.coeff * o.coeff; - } - return ScaleFromSums(num, den, info); - } - double median_of(std::vector &v) { std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end()); return v[v.size() / 2]; @@ -1516,26 +1517,6 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { return n_groups; } -void RotationScaleMerge::ReduceFullGroupMeans(int n_groups, std::vector &out_mean) const { - // Inverse-variance per-group mean of I*corr = the merge reference (no cell mask: the scaling - // reference is built without a reference cell - only the final merge applies it). A group per task - // over the group CSR, each summed over its fulls in index order, as the serial walk summed it. - out_mean.assign(n_groups, NAN); - ParallelBlocks(n_groups, nthreads, [&](int, int glo, int ghi) { - for (int g = glo; g < ghi; ++g) { - double sw = 0.0, swI = 0.0; - for (int p = fulls_group_start[g]; p < fulls_group_start[g + 1]; ++p) { - float I_corr = 0.0f; - const double w = ReferenceWeight(fulls[fulls_group_perm[p]], min_partiality, I_corr); - if (w <= 0.0) continue; - sw += w; - swI += w * I_corr; - } - if (sw > 0.0) out_mean[g] = swI / sw; - } - }, 512); -} - void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vector &out_mean) const { out_mean.assign(n_groups, NAN); ParallelChunks(n_groups, ThreadsForWork(partials.size(), nthreads), [&](int glo, int ghi) { @@ -1554,14 +1535,15 @@ void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vector &refs, +void RotationScaleMerge::ReduceScalingGroupMeans(int n_groups, const int32_t *group_start, + const int32_t *group_count, const std::vector &refs, const std::vector &corr, std::vector &out_mean) const { out_mean.resize(n_groups); ParallelChunks(n_groups, ThreadsForWork(corr.size(), nthreads), [&](int glo, int ghi) { for (int g = glo; g < ghi; ++g) { double sw = 0.0, swI = 0.0; - const int lo = partial_group_start[g], hi = lo + partial_group_count[g]; + const int lo = group_start[g], hi = lo + group_count[g]; for (int p = lo; p < hi; ++p) { const ScalingRef &r = refs[p]; float I_corr = 0.0f; @@ -1650,25 +1632,17 @@ double RotationScaleMerge::ChooseLogScaleSmoothing(const std::vector &y, return best_lambda; } -template -void RotationScaleMerge::FitPerFrameG(const std::vector &obs, const std::vector &fstart, - const std::vector &fcount, - const std::vector &group_mean_in, - bool unity, std::vector &g) { - std::vector scaled(fstart.size(), 0); - std::vector info(fstart.size(), 0.0); - const int n_fr = static_cast(fstart.size()); +void RotationScaleMerge::FitScalingFrames(const std::vector &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; + std::vector scaled(n_fr, 0); + std::vector info(n_fr, 0.0); ParallelFor(n_fr, nthreads, [&](int f) { - std::vector so; - for (int i = fstart[f]; i < fstart[f] + fcount[f]; ++i) { - const auto &o = obs[i]; - if (o.group < 0) continue; - if (o.on_ice) continue; - // Only the observations the reference is built from (ReduceGroupMeans' partiality floor): - // on a fine-sliced sweep half of the 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. - if (o.partiality < min_partiality) continue; + double num = 0.0, den = 0.0; + size_t n = 0; + for (int i = fit_start[f]; i < fit_start[f + 1]; ++i) { + 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 // oscillates (measured: the gauge alternating 0.56 and 1.39 every other iteration on the @@ -1678,46 +1652,12 @@ void RotationScaleMerge::FitPerFrameG(const std::vector &obs, const std::vect // iterations, which is what RunScalingLoop's tolerance is for. const double mean = group_mean_in[o.group]; if (!std::isfinite(mean)) continue; - double coeff; - if (unity) { - coeff = mean; // partiality already folded into the full, prescaling_corr = 1 - } else { - if (!(std::isfinite(o.zeta) && o.zeta > 0.0f)) continue; // Rotation model needs zeta > 0 - coeff = o.partiality * SafeInv(o.prescaling_corr, 1.0) * mean; - } - so.push_back({coeff, static_cast(o.I), SafeInv(o.sigma, 1.0)}); - } - // A partial frame needs MIN_REFLECTIONS of its own (corr is left untouched otherwise, as - // ScaleOnTheFly does). A full is a whole rocking event filed under one frame, so on a sparse - // sweep a frame holds a handful; its fit goes in at the information it has, and the smoothing - // in RunFullsScalingLoop decides how far its neighbours carry it. - if (so.empty() || (!unity && so.size() < MIN_REFLECTIONS)) return; - g[f] = SolveScale(so, info[f]); - scaled[f] = 1; - }); - // Remember which frames were fitted this call (so the caller updates corr only there). - frame_scaled_scratch = std::move(scaled); - frame_info_scratch = std::move(info); -} - -void RotationScaleMerge::FitScalingFrames(const std::vector &fit, const std::vector &fit_start, - const std::vector &group_mean_in, std::vector &g) { - const int n_fr = static_cast(fit_start.size()) - 1; - 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) { - const ScalingFit &o = fit[i]; - 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; ++n; } - if (n < MIN_REFLECTIONS) return; + if (n < min_obs) return; g[f] = ScaleFromSums(num, den, info[f]); scaled[f] = 1; }); @@ -3990,7 +3930,7 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // sits on the detector, so with Is as the regressor the bias is a smooth radial ramp - which // the gauge fix spreads over the whole surface and the n_iter rounds compound, and which no // even/odd cross-validation can see because both halves carry it equally. The per-frame - // scale (FitPerFrameG) already regresses this way round. + // scale (FitScalingFrames) already regresses this way round. // A NEGATIVE observation is a legitimate measurement of a weak reflection and enters both // sums as it stands. Both are linear in Is - there is no logarithm and no division by it // here, unlike the relative-B fits above, which regress log(Iref/Is) and do need Is > 0 - @@ -4378,11 +4318,27 @@ void RotationScaleMerge::Combine() { fulls.clear(); g_full.assign(n_frames, 1.0); fulls_dropped_overloaded = 0; + const int n_run = static_cast(rawrun_start.size()); + + // A raw hkl's partials lie a frame apart in memory, and the work on one rocking event is long + // enough that the CPU does not reach the next runs' loads on its own - the combine waited on a + // cache miss for nearly every partial. So each run asks for the partials of the run + // COMBINE_PREFETCH_RUNS ahead. + constexpr int COMBINE_PREFETCH_RUNS = 8; + auto prefetch_run = [&](int r) { + if (r >= n_run) return; + for (int p = rawrun_start[r]; p < rawrun_start[r] + rawrun_count[r]; ++p) { + const char *o = reinterpret_cast(&partials[perm[p]]); + Prefetch(o); + Prefetch(o + sizeof(Obs) - 1); + } + }; // Combine one raw-hkl run into fulls, appended to `out`. Independent per run (the events of one hkl // touch no shared state), so runs parallelise cleanly. Returns the number of usable partials seen. // `dump` (serial path only) writes each emitted full for the diagnostic observation dump. auto process_rawrun = [&](int r, std::vector &out, std::ofstream *dump, int64_t &n_overloaded) -> size_t { + prefetch_run(r + COMBINE_PREFETCH_RUNS); const int lo = rawrun_start[r], hi = rawrun_start[r] + rawrun_count[r]; std::vector ev; // usable perm-indices of this raw hkl, in image-number order ev.reserve(hi - lo); @@ -4520,7 +4476,6 @@ void RotationScaleMerge::Combine() { return ev.size(); }; - const int n_run = static_cast(rawrun_start.size()); size_t n_used = 0; if (!observation_dump_path.empty() || nthreads <= 1) { @@ -6441,7 +6396,11 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool } }); - // FitPerFrameG's filters, all but the reference mean's, frame by frame. + // The observations a frame's scale is fitted on, frame by frame. Only those the reference + // is built from (ReferenceWeight's partiality floor): on a fine-sliced sweep half of the + // 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 && std::isfinite(o.zeta) && o.zeta > 0.0f; @@ -6473,8 +6432,9 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool std::vector frame_fitted; partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window, [&] { - ReduceScalingGroupMeans(n_groups, scaling_refs, scaling_corr, partial_mean); - FitScalingFrames(scaling_fit, fit_start, partial_mean, g_partial); + ReduceScalingGroupMeans(n_groups, partial_group_start.data(), partial_group_count.data(), + scaling_refs, scaling_corr, partial_mean); + FitScalingFrames(scaling_fit, fit_start, MIN_REFLECTIONS, partial_mean, g_partial); g_fitted = g_partial; frame_fitted = frame_scaled_scratch; }, [&](const std::vector &apply, const std::vector &ratio) { @@ -6781,20 +6741,66 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // --- 4. Scale the fulls (XDS order, Unity model). --- if (scale_fulls && !scaled_fulls_on_gpu) { + // On a working copy of the fulls, as the partials' loop (ScalingRef and ScalingFit in the header): + // 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); + std::vector full_refs(nf); + std::vector 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]]; + full_refs[q] = {o.frame, o.I, o.sigma, o.partiality, o.prescaling_corr}; + full_corr[q] = o.corr; + } + }); + std::vector full_group_count(n_groups); + for (int g = 0; g < n_groups; ++g) + full_group_count[g] = fulls_group_start[g + 1] - fulls_group_start[g]; + // The partials' filters for the Unity model, which needs no zeta, and whose coefficient is the + // mean itself: the partiality is folded into the full, and prescaling_corr is 1. + auto fitted_full = [&](const Obs &o) { + return o.group >= 0 && !o.on_ice && o.partiality >= min_partiality; + }; + std::vector fit_start(n_frames + 1, 0); + ParallelFor(n_frames, nthreads, [&](int f) { + fit_start[f + 1] = static_cast(std::count_if( + fulls.begin() + fulls_frame_start[f], + 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()); + 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) { + 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}; + } + }); + std::vector full_mean, g_fit(n_frames); fulls_loop = RunFullsScalingLoop(g_full, g_fit, fulls_frame_count, [&] { - ReduceFullGroupMeans(n_groups, full_mean); - FitPerFrameG(fulls, fulls_frame_start, fulls_frame_count, full_mean, /*unity=*/true, g_fit); + ReduceScalingGroupMeans(n_groups, fulls_group_start.data(), full_group_count.data(), full_refs, + full_corr, full_mean); + FitScalingFrames(full_fit, fit_start, 1, full_mean, g_fit); }, [&](const std::vector &apply, const std::vector &ratio) { - ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), - [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - Obs &o = fulls[i]; - if (apply[o.frame] && std::isfinite(o.corr)) - o.corr = static_cast(o.corr * ratio[o.frame]); + ParallelChunks(nf, full_threads, [&](int lo, int hi) { + for (int q = lo; q < hi; ++q) { + float &corr = full_corr[q]; + const int f = full_refs[q].frame; + if (apply[f] && std::isfinite(corr)) + corr = static_cast(corr * ratio[f]); } }); }); + ParallelChunks(nf, full_threads, [&](int lo, int hi) { + for (int q = lo; q < hi; ++q) fulls[full_index[q]].corr = full_corr[q]; + }); logger->Info("Scaled fulls (XDS order, Unity model)"); } // The total scale each frame's fulls were amplified by - partial scale, incident flux and the fulls' diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index cb77e956b..9ab8e0481 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -514,7 +514,7 @@ private: std::vector em_fit_pool; std::vector em_chi2; - // Set by FitPerFrameG: which frames were fitted this call (so corr/G is updated only there), and + // Set by FitScalingFrames: which frames were fitted this call (so corr/G is updated only there), and // the information each fit had on G, sum w^2 c^2 (0 where the frame was not fitted). std::vector frame_scaled_scratch; std::vector frame_info_scratch; @@ -614,12 +614,12 @@ private: std::vector partial_group_start, partial_group_count; NotCopied partial_group_perm; - // The CPU partial scaling loop's working copy of the partials. It passes over them several times - // per iteration and reads a handful of fields each time, so it works on these rather than on the - // eighty-byte Obs. ScalingRef holds what the group means and the corr update read, in the CSR's + // 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 + // than on the eighty-byte Obs. ScalingRef holds what the group means and the corr update read, in the CSR's // group order with the partials in no group after them, and corr is kept beside it in the same // order - so neither pass gathers. ScalingFit holds, frame by frame, the observations - // FitScalingFrames fits: FitPerFrameG's filters, which do not change over the loop, applied once, + // FitScalingFrames fits: the fit's filters, which do not change over the loop, applied once, // and the factors of its sums that do not depend on the reference worked out once. The same // values, so the same arithmetic in the same order. struct ScalingRef { @@ -628,7 +628,7 @@ private: }; struct ScalingFit { int32_t group; - double coeff_per_mean; // partiality / prescaling_corr: FitPerFrameG's coeff for a mean of 1 + double coeff_per_mean; // partiality / prescaling_corr (1 for a full): the fit's coeff for a mean of 1 double Iobs, w2; }; @@ -648,25 +648,23 @@ private: 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 fulls (the merge reference), by the fulls' - // group CSR, so each group is summed in index order. - void ReduceFullGroupMeans(int n_groups, std::vector &out_mean) const; - // The same over the partials, using the CSR above - the same sums in the same order. + // 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 same again over the scaling loop's copy, refs and corr both in group order. - void ReduceScalingGroupMeans(int n_groups, const std::vector &refs, - const std::vector &corr, std::vector &out_mean) const; + // 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, + std::vector &out_mean) const; - // Robust per-frame G fit (IRLS, Cauchy k=3), unity=false uses the rotation partiality, unity=true the - // scale-fulls (partiality already folded in). Reads out_mean[group] as the reference intensity. - template - void FitPerFrameG(const std::vector &obs, const std::vector &fstart, - const std::vector &fcount, const std::vector &group_mean_in, - bool unity, std::vector &g); - // FitPerFrameG(unity = false) over the scaling loop's ScalingFit, frame f's at - // [fit_start[f], fit_start[f + 1]). + // 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 + // partial frame needs MIN_REFLECTIONS of its own (as ScaleOnTheFly), while a full is a whole rocking + // 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, - const std::vector &group_mean_in, std::vector &g); + 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 // the current reference (leaving G in `g` and the frames it fitted in frame_scaled_scratch), diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 2582589b5..d9a0fb567 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -102,7 +102,7 @@ namespace { // The scale-fit coefficient of an observation (rotation model), recomputed each scaling iteration // once the group means are known: coeff = partiality * (1/prescaling_corr) * mean[group], or false - // where the observation is not one the reference is built from (the host FitPerFrameG says why). + // where the observation is not one the reference is built from (the host RunArm says why, where it fills ScalingFit). __device__ __forceinline__ bool ScaleCoeff(int i, double min_partiality, const int32_t *__restrict__ group, const ObsRecord &o, const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice, @@ -117,13 +117,13 @@ namespace { } // One block per frame: the per-frame scale G as the weighted least-squares slope over the frame's - // contiguous obs - the same plain fit as the host SolveScale (which says why it is not robust). + // contiguous obs - the same plain fit as the host ScaleFromSums (which says why it is not robust). // Leaves g/scaled untouched for frames with fewer than min_obs usable observations. `perm` (null // for the partials, whose arrays are already frame-contiguous) maps a position in the frame's // [lo,hi) range to the obs index, so the same kernel scales the fulls (emit-ordered) through a // frame-grouping permutation without physically reordering the fulls arrays. info[f] receives the // fit's information on G, sum w^2 c^2 (0 for a frame left alone). The weight is 1/sigma, as the - // CPU's ScaleObs::weight. + // CPU's ScalingFit::w2 is its square. __global__ void FitPerFrameGKernel(int n_frames, const int32_t *__restrict__ frame_start, const int32_t *__restrict__ frame_count,