From b5d9167bf4bb54ec584c1c423744b11552b65fbd Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 2 Jul 2026 21:31:41 +0200 Subject: [PATCH] RotationScaleMerge: chunked parallelism -> scale/merge phase ~2s The per-observation corr update (7.6M items) ran through a work-stealing ParallelFor that does one atomic fetch_add PER item - pure contention for trivial work (measured: update 0.60s vs reduce 0.15s / fit 0.13s in the scale-partials loop). Add ParallelChunks (one contiguous range per worker, no per-item sync) and use it for UpdateCorr, and parallelise the ASU keying (gemmi reduction per distinct raw hkl - HKLKeyGenerator is const, safe to read concurrently) and the group-stamping over disjoint raw-hkl runs. scale-partials 0.90 -> 0.28s, group-hkl 0.20 -> 0.09s, per-pass warm 0.83s, whole scale/merge phase ~3.3 -> ~2.0s. Bit-identical output (same space group, ISa, CC1/2). ParallelChunks is the CPU stand-in for a flat CUDA grid-stride kernel; ParallelFor stays for the heavy, uneven per-frame fits where the atomic amortises and work-stealing balances the load. Co-Authored-By: Claude Opus 4.8 --- .../scale_merge/RotationScaleMerge.cpp | 86 +++++++++++++------ 1 file changed, 59 insertions(+), 27 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index b91a3d6b..17285937 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -83,6 +83,9 @@ namespace { // Run fn(i) for i in [0, n) over `nthreads` workers pulling from a shared atomic counter - the same // self-load-balancing pattern the rest of the codebase uses (heavy frames don't stall light ones). + // Work-stealing per-item parallel: one atomic fetch per item. Use ONLY when the per-item work is + // heavy and uneven (e.g. per-frame fits) - the atomic amortises. For millions of tiny uniform items + // use ParallelChunks instead; a per-item atomic there is pure contention. template void ParallelFor(int n, size_t nthreads, Fn fn) { if (n <= 0) return; @@ -102,6 +105,24 @@ namespace { for (auto &f : futures) f.get(); } + // Chunked parallel: each worker gets one contiguous [lo, hi) range, no per-item synchronisation. + // Right for millions of cheap uniform items (the CPU stand-in for a flat CUDA grid-stride kernel). + template + void ParallelChunks(int n, size_t nthreads, Fn fn) { + if (n <= 0) return; + const int nt = static_cast(std::max(1, std::min(nthreads, static_cast(n)))); + if (nt == 1) { fn(0, n); return; } + const int chunk = (n + nt - 1) / nt; + std::vector> futures; + futures.reserve(nt); + for (int t = 0; t < nt; ++t) { + const int lo = t * chunk, hi = std::min(n, lo + chunk); + if (lo >= hi) break; + futures.emplace_back(std::async(std::launch::async, [&fn, lo, hi] { fn(lo, hi); })); + } + for (auto &f : futures) f.get(); + } + double median_of(std::vector &v) { std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end()); return v[v.size() / 2]; @@ -212,19 +233,25 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { // (sort the ~#distinct-hkl keys, not the millions of observations), then hand out dense ids. const int n_run = static_cast(rawrun_start.size()); std::vector key(n_run); + std::vector eligible(n_run, 0); + // The gemmi ASU reduction / absence test per raw hkl is the cost here and is independent per run + // (HKLKeyGenerator is const, so concurrent reads are safe) - compute keys in parallel chunks. + ParallelChunks(n_run, nthreads, [&](int lo, int hi) { + for (int r = lo; r < hi; ++r) { + rawrun_group[r] = -1; + if (keygen.IsSystematicallyAbsent(rawrun_h[r], rawrun_k[r], rawrun_l[r])) + continue; + const float d = rawrun_d[r]; // resolution is a per-raw-hkl property (all its partials share d) + if (!std::isfinite(d) || d <= 0.0f) continue; + if (d_min_limit && d < *d_min_limit) continue; + key[r] = keygen(rawrun_h[r], rawrun_k[r], rawrun_l[r]).pack(); + eligible[r] = 1; + } + }); std::vector idx; idx.reserve(n_run); - for (int r = 0; r < n_run; ++r) { - rawrun_group[r] = -1; - if (keygen.IsSystematicallyAbsent(rawrun_h[r], rawrun_k[r], rawrun_l[r])) - continue; - // Resolution range is a per-raw-hkl property (all its partials share d). - const float d = rawrun_d[r]; - if (!std::isfinite(d) || d <= 0.0f) continue; - if (d_min_limit && d < *d_min_limit) continue; - key[r] = keygen(rawrun_h[r], rawrun_k[r], rawrun_l[r]).pack(); - idx.push_back(r); - } + for (int r = 0; r < n_run; ++r) + if (eligible[r]) idx.push_back(r); std::sort(idx.begin(), idx.end(), [&](int32_t a, int32_t b) { return key[a] < key[b]; }); group_h.clear(); group_k.clear(); group_l.clear(); @@ -243,15 +270,18 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { // Stamp each partial's group from its raw hkl, keeping the exact per-observation AcceptReflection // finiteness (finite I, finite rlp != 0, finite sigma > 0) that grouping by raw hkl alone would miss. - for (int r = 0; r < n_run; ++r) { - const int g = rawrun_group[r]; - const int lo = rawrun_start[r], hi = rawrun_start[r] + rawrun_count[r]; - for (int p = lo; p < hi; ++p) { - auto &o = partials[perm[p]]; - o.group = (g >= 0 && std::isfinite(o.I) && std::isfinite(o.rlp) && o.rlp != 0.0f - && std::isfinite(o.sigma) && o.sigma > 0.0f) ? g : -1; + // 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) { + const int g = rawrun_group[r]; + const int lo = rawrun_start[r], hi = rawrun_start[r] + rawrun_count[r]; + for (int p = lo; p < hi; ++p) { + auto &o = partials[perm[p]]; + o.group = (g >= 0 && std::isfinite(o.I) && std::isfinite(o.rlp) && o.rlp != 0.0f + && std::isfinite(o.sigma) && o.sigma > 0.0f) ? g : -1; + } } - } + }); return n_groups; } @@ -317,14 +347,16 @@ void RotationScaleMerge::FitPerFrameG(std::vector &obs, const std::vector &obs, const std::vector &g, const std::vector &frame_scaled) const { - ParallelFor(static_cast(obs.size()), nthreads, [&](int i) { - auto &o = obs[i]; - if (!frame_scaled[o.frame]) return; - const double denom = static_cast(o.partiality) * g[o.frame]; // B_term = 1 (no B refine) - if (std::isfinite(o.rlp) && std::isfinite(denom) && denom > 0.0) - o.corr = static_cast(o.rlp / denom); - else - o.corr = NAN; + ParallelChunks(static_cast(obs.size()), nthreads, [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) { + auto &o = obs[i]; + if (!frame_scaled[o.frame]) continue; + const double denom = static_cast(o.partiality) * g[o.frame]; // B_term = 1 (no B refine) + if (std::isfinite(o.rlp) && std::isfinite(denom) && denom > 0.0) + o.corr = static_cast(o.rlp / denom); + else + o.corr = NAN; + } }); }