mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 67 Wiki Activity

RotationScaleMerge: chunked parallelism -> scale/merge phase ~2s

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-07-02 21:31:41 +02:00
co-authored by Claude Opus 4.8
parent 29c5410328
commit b5d9167bf4
1 changed files with 59 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/RotationScaleMerge.cpp
+59 -27
View File
@@ -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 <typename Fn>
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 <typename Fn>
void ParallelChunks(int n, size_t nthreads, Fn fn) {
if (n <= 0) return;
const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(n))));
if (nt == 1) { fn(0, n); return; }
const int chunk = (n + nt - 1) / nt;
std::vector<std::future<void>> 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<double> &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<int>(rawrun_start.size());
std::vector<uint64_t> key(n_run);
std::vector<uint8_t> 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<int32_t> 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> &obs, const std::vector<i
void RotationScaleMerge::UpdateCorr(std::vector<Obs> &obs, const std::vector<double> &g,
const std::vector<uint8_t> &frame_scaled) const {
ParallelFor(static_cast<int>(obs.size()), nthreads, [&](int i) {
auto &o = obs[i];
if (!frame_scaled[o.frame]) return;
const double denom = static_cast<double>(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<float>(o.rlp / denom);
else
o.corr = NAN;
ParallelChunks(static_cast<int>(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<double>(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<float>(o.rlp / denom);
else
o.corr = NAN;
}
});
}