diff --git a/common/ParallelFor.h b/common/ParallelFor.h index 767a070f..a35d971c 100644 --- a/common/ParallelFor.h +++ b/common/ParallelFor.h @@ -12,6 +12,17 @@ // worker count from the caller rather than asking the hardware, so a run that was told how many // threads to use keeps to it. +// How many workers a pass over `n` cheap items should use: enough that each gets at least +// `min_per_thread` of them, and never more than the caller was given. Both helpers below start a +// thread per chunk, which costs tens of microseconds, so a pass that is repeated - or one whose data +// is small - can easily spend more on splitting the work than on doing it. That is not a large +// machine's problem: it is what makes the same code right on an 8-core laptop, a 16-core desktop and +// a two-socket node, none of which should be handed 48 chunks of a few thousand items. +inline size_t ThreadsForWork(size_t n, size_t nthreads, size_t min_per_thread = 32768) { + if (nthreads <= 1 || n == 0) return 1; + return std::clamp(n / min_per_thread, 1, nthreads); +} + // Chunked: each worker gets one contiguous [lo, hi) range and there is no per-item synchronisation. // Right for millions of cheap uniform items - the CPU stand-in for a flat CUDA grid-stride kernel. // The split is fixed and deterministic, so a pass whose per-element work is independent gives the diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index d92ae0b6..6470a8a0 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -1707,16 +1707,20 @@ void RotationScaleMerge::RefineAbsorption(int n_iter, int n_groups) { constexpr int NB = 8; std::vector cell(fulls.size(), -1); - for (size_t i = 0; i < fulls.size(); ++i) { - const Obs &o = fulls[i]; - if (!std::isfinite(o.px) || !std::isfinite(o.py)) - continue; - const Coord s1 = Coord((o.px - beam_x) / F, (o.py - beam_y) / F, 1.0f).Normalize(); - const Coord u = gon.GetTransformationAngle(gon.GetAngle_deg(o.image_number)).transpose() * s1; - const int ix = std::clamp(static_cast((u.x + 1.0f) * 0.5f * NB), 0, NB - 1); - const int iy = std::clamp(static_cast((u.y + 1.0f) * 0.5f * NB), 0, NB - 1); - cell[i] = ix * NB + iy; - } + // A de-rotation and a normalisation per full, each independent of the others - and a spindle + // rotation matrix costs a sine and a cosine, which is why this shows up at all. + ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) { + const Obs &o = fulls[i]; + if (!std::isfinite(o.px) || !std::isfinite(o.py)) + continue; + const Coord s1 = Coord((o.px - beam_x) / F, (o.py - beam_y) / F, 1.0f).Normalize(); + const Coord u = gon.GetTransformationAngle(gon.GetAngle_deg(o.image_number)).transpose() * s1; + const int ix = std::clamp(static_cast((u.x + 1.0f) * 0.5f * NB), 0, NB - 1); + const int iy = std::clamp(static_cast((u.y + 1.0f) * 0.5f * NB), 0, NB - 1); + cell[i] = ix * NB + iy; + } + }); ApplyCellSurface(cell, NB * NB, n_iter, n_groups, "Absorption (goniometer-frame 8x8)"); } @@ -1755,23 +1759,37 @@ void RotationScaleMerge::RefineAbsorptionTime(int n_iter, int n_groups) { if (usable(o)) { vx.push_back(o.px); vy.push_back(o.py); } if (static_cast(vx.size()) < NT * ND * ND) return; - std::sort(vx.begin(), vx.end()); - std::sort(vy.begin(), vy.end()); + // Only the ND-1 bin edges are wanted, not the order of the millions of values between them, so + // select them instead of sorting. Taken in increasing order each selection starts where the last + // one left off, on a range already partitioned below that edge, so this is a couple of passes over + // the values rather than the n log n a full sort spends. std::vector ex(ND - 1), ey(ND - 1); + size_t prev = 0; for (int i = 1; i < ND; ++i) { - ex[i - 1] = vx[vx.size() * i / ND]; - ey[i - 1] = vy[vy.size() * i / ND]; + const size_t pos = vx.size() * i / ND; + std::nth_element(vx.begin() + prev, vx.begin() + pos, vx.end()); + ex[i - 1] = vx[pos]; + prev = pos; + } + prev = 0; + for (int i = 1; i < ND; ++i) { + const size_t pos = vy.size() * i / ND; + std::nth_element(vy.begin() + prev, vy.begin() + pos, vy.end()); + ey[i - 1] = vy[pos]; + prev = pos; } const int frames_per_bin = std::max(1, n_frames / NT); std::vector cell(fulls.size(), -1); - for (size_t i = 0; i < fulls.size(); ++i) { - const Obs &o = fulls[i]; - if (!usable(o)) continue; - const int it = std::min(NT - 1, o.frame / frames_per_bin); - const int ix = static_cast(std::upper_bound(ex.begin(), ex.end(), o.px) - ex.begin()); - const int iy = static_cast(std::upper_bound(ey.begin(), ey.end(), o.py) - ey.begin()); - cell[i] = (it * ND + ix) * ND + iy; - } + ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) { + const Obs &o = fulls[i]; + if (!usable(o)) continue; + const int it = std::min(NT - 1, o.frame / frames_per_bin); + const int ix = static_cast(std::upper_bound(ex.begin(), ex.end(), o.px) - ex.begin()); + const int iy = static_cast(std::upper_bound(ey.begin(), ey.end(), o.py) - ey.begin()); + cell[i] = (it * ND + ix) * ND + iy; + } + }); ApplyCellSurface(cell, NT * ND * ND, n_iter, n_groups, "Absorption (time x detector 12x10x10)"); } @@ -1797,13 +1815,15 @@ void RotationScaleMerge::RefineModulation(int n_iter, int n_groups) { constexpr int NB = 24; const float sx = NB / (pxmax - pxmin), sy = NB / (pymax - pymin); std::vector cell(fulls.size(), -1); - for (size_t i = 0; i < fulls.size(); ++i) { - const Obs &o = fulls[i]; - if (!std::isfinite(o.px) || !std::isfinite(o.py)) continue; - const int ix = std::clamp(static_cast((o.px - pxmin) * sx), 0, NB - 1); - const int iy = std::clamp(static_cast((o.py - pymin) * sy), 0, NB - 1); - cell[i] = ix * NB + iy; - } + ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) { + const Obs &o = fulls[i]; + if (!std::isfinite(o.px) || !std::isfinite(o.py)) continue; + const int ix = std::clamp(static_cast((o.px - pxmin) * sx), 0, NB - 1); + const int iy = std::clamp(static_cast((o.py - pymin) * sy), 0, NB - 1); + cell[i] = ix * NB + iy; + } + }); ApplyCellSurface(cell, NB * NB, n_iter, n_groups, "Modulation (detector-frame 16x16)"); } @@ -1826,9 +1846,33 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int auto subset = [&](int parity) -> const std::vector & { return parity < 0 ? idx_all : (parity ? idx_odd : idx_even); }; - // Per-group reference accumulators, reused by every pass (n_groups is large, so re-allocating them - // ~22x per surface is ~22x the page faults for nothing). + + // Every pass below runs over one of these lists, ~22 times per surface, so the worker count is + // set by how much data there is rather than by how many threads the machine has: a small dataset + // on a large node would otherwise pay for 48 thread starts per pass to sum a few thousand terms. + const int nt = static_cast(ThreadsForWork(idx_all.size(), nthreads)); + // Where thread t starts in a list of n items. The boundaries depend on nothing but n and the + // thread count, so a sum split this way gives the same answer on every run. + auto part = [nt](size_t n, int t) { return n * static_cast(t) / static_cast(nt); }; + + // The per-group reference of one subset with surface A applied - the first pass of both the fit + // and the score. It stays serial: the terms of one group are spread all over the list, so the only + // way to give a thread groups of its own is to walk in group order, which trades this pass's + // near-sequential read of `fulls` for a random one over a few hundred megabytes. That is the same + // trade that made length-bucketed run ordering a loss in the combine kernel. Worth revisiting with + // per-worker accumulators (n_groups doubles each, so it costs memory rather than locality) if this + // pass ever dominates. std::vector sw(n_groups), swI(n_groups); + auto reference = [&](int parity, const std::vector &A) { + const std::vector &sel = subset(parity); + std::fill(sw.begin(), sw.end(), 0.0); + std::fill(swI.begin(), swI.end(), 0.0); + for (const int32_t i : sel) { + const Obs &o = fulls[i]; + const double a = A[cell[i]], sc = static_cast(o.sigma) * o.corr * a, w = 1.0 / (sc * sc); + sw[o.group] += w; swI[o.group] += w * static_cast(o.I) * o.corr * a; + } + }; // Fit the per-cell factor over the subset {frame&1 == parity} (parity < 0 = all fulls), n_iter // alternating rounds against that subset's own reference (Tikhonov pull to 1, gauge-fixed to a @@ -1836,14 +1880,10 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int auto fit_surface = [&](int parity) -> std::vector { const std::vector &sel = subset(parity); std::vector A(ncell, 1.0); + // Per-thread cell accumulators, allocated once for the whole fit rather than per round. + std::vector> tcross(nt, std::vector(ncell)), tref2(nt, std::vector(ncell)); for (int it = 0; it < n_iter; ++it) { - std::fill(sw.begin(), sw.end(), 0.0); - std::fill(swI.begin(), swI.end(), 0.0); - for (const int32_t i : sel) { - const Obs &o = fulls[i]; - const double a = A[cell[i]], sc = static_cast(o.sigma) * o.corr * a, w = 1.0 / (sc * sc); - sw[o.group] += w; swI[o.group] += w * static_cast(o.I) * o.corr * a; - } + reference(parity, A); // The cell's factor is the scale that carries the observation onto the reference, and it is // fitted by regressing the OBSERVATION on the reference: A = sum w Iref^2 / sum w Is Iref, // the inverse of the slope of Is against Iref. Not the other way round - a least-squares @@ -1854,16 +1894,29 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // 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. + // Chunked over the subset, each thread summing into its own cells: unlike the pass above + // there is no ordering that keeps threads off each other's bins, and there are only ncell + // of them, so per-thread copies are cheap and the fixed chunking keeps it reproducible. std::vector cross(ncell, 0.0), ref2(ncell, 0.0); - for (const int32_t i : sel) { - const Obs &o = fulls[i]; - if (sw[o.group] <= 0.0) continue; - const double Iref = swI[o.group] / sw[o.group], a = A[cell[i]]; - 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 || !(Is > 0.0) || !(sc > 0.0)) continue; - const double w = 1.0 / (sc * sc); - cross[cell[i]] += w * Is * Iref; ref2[cell[i]] += w * Iref * Iref; - } + ParallelChunks(nt, nthreads, [&](int tlo, int thi) { + for (int t = tlo; t < thi; ++t) { + std::vector &xcross = tcross[t], &xref2 = tref2[t]; + std::fill(xcross.begin(), xcross.end(), 0.0); + std::fill(xref2.begin(), xref2.end(), 0.0); + for (size_t k = part(sel.size(), t); k < part(sel.size(), t + 1); ++k) { + const int32_t i = sel[k]; + const Obs &o = fulls[i]; + if (sw[o.group] <= 0.0) continue; + const double Iref = swI[o.group] / sw[o.group], a = A[cell[i]]; + 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 || !(Is > 0.0) || !(sc > 0.0)) continue; + const double w = 1.0 / (sc * sc); + xcross[cell[i]] += w * Is * Iref; xref2[cell[i]] += w * Iref * Iref; + } + } + }); + for (int t = 0; t < nt; ++t) + for (int c = 0; c < ncell; ++c) { cross[c] += tcross[t][c]; ref2[c] += tref2[t][c]; } std::vector dsorted = cross; std::nth_element(dsorted.begin(), dsorted.begin() + dsorted.size() / 2, dsorted.end()); const double lambda = 0.1 * std::max(1e-30, dsorted[dsorted.size() / 2]); @@ -1880,27 +1933,29 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // scored against that subset's own reference (no leakage). Lower = tighter equivalents. auto score = [&](int parity, const std::vector &A) -> double { const std::vector &sel = subset(parity); - std::fill(sw.begin(), sw.end(), 0.0); - std::fill(swI.begin(), swI.end(), 0.0); - for (const int32_t i : sel) { - const Obs &o = fulls[i]; - const double a = A[cell[i]], Is = static_cast(o.I) * o.corr * a; - const double sc = static_cast(o.sigma) * o.corr * a, w = 1.0 / (sc * sc); - sw[o.group] += w; swI[o.group] += w * Is; - } + reference(parity, A); // Rmeas-like (sigma-INDEPENDENT) agreement of the held-out equivalents: sum|Is - Iref| / sum|Iref|. // Scoring on the studentized deviation instead lets a surface "improve" the held-out chi^2 by // reshaping sigma (via corr) without tightening the actual intensities - which on mis-indexed / bad // data passes cross-validation yet worsens Rmeas. A fractional metric cannot be gamed that way. + std::vector tnum(nt, 0.0), tden(nt, 0.0); + ParallelChunks(nt, nthreads, [&](int tlo, int thi) { + for (int t = tlo; t < thi; ++t) { + double num = 0.0, den = 0.0; + for (size_t k = part(sel.size(), t); k < part(sel.size(), t + 1); ++k) { + const int32_t i = sel[k]; + const Obs &o = fulls[i]; + if (sw[o.group] <= 0.0) continue; + const double a = A[cell[i]], Is = static_cast(o.I) * o.corr * a; + const double Iref = swI[o.group] / sw[o.group]; + if (!std::isfinite(Iref) || Iref <= 0.0) continue; + num += std::abs(Is - Iref); den += Iref; + } + tnum[t] = num; tden[t] = den; + } + }); double num = 0.0, den = 0.0; - for (const int32_t i : sel) { - const Obs &o = fulls[i]; - if (sw[o.group] <= 0.0) continue; - const double a = A[cell[i]], Is = static_cast(o.I) * o.corr * a; - const double Iref = swI[o.group] / sw[o.group]; - if (!std::isfinite(Iref) || Iref <= 0.0) continue; - num += std::abs(Is - Iref); den += Iref; - } + for (int t = 0; t < nt; ++t) { num += tnum[t]; den += tden[t]; } return den > 0.0 ? num / den : 0.0; }; @@ -1917,9 +1972,11 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int return; } const std::vector A = fit_surface(-1); - for (size_t i = 0; i < fulls.size(); ++i) - if (cell[i] >= 0) - fulls[i].corr = static_cast(fulls[i].corr * A[cell[i]]); + ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) + if (cell[i] >= 0) + fulls[i].corr = static_cast(fulls[i].corr * A[cell[i]]); + }); logger.Info("{} correction: cross-validated, held-out gain {:.1f}%", name, 100.0 * gain / std::max(base, 1e-30)); } diff --git a/image_analysis/scale_merge/SearchSpaceGroup.cpp b/image_analysis/scale_merge/SearchSpaceGroup.cpp index f9a61238..e58abb91 100644 --- a/image_analysis/scale_merge/SearchSpaceGroup.cpp +++ b/image_analysis/scale_merge/SearchSpaceGroup.cpp @@ -2,6 +2,7 @@ // SPDX-License-Identifier: GPL-3.0-only #include "SearchSpaceGroup.h" +#include "../../common/ParallelFor.h" #include #include @@ -315,10 +316,10 @@ SearchSpaceGroupResult SearchSpaceGroup( key_to_index.emplace(key[i], static_cast(i)); // --- Stage A: score each distinct rotation operator once --- - std::vector visited(n, 0); - uint32_t epoch = 0; - - auto score_operator = [&](const gemmi::Op& op) -> SpaceGroupOperatorScore { + // `visited` and `epoch` are scratch, taken as arguments rather than captured so that several + // operators can be scored at once - each worker below keeps its own pair. + auto score_operator = [&](const gemmi::Op& op, std::vector& visited, + uint32_t& epoch) -> SpaceGroupOperatorScore { ++epoch; std::vector x, y; for (size_t i = 0; i < n; ++i) { @@ -362,12 +363,14 @@ SearchSpaceGroupResult SearchSpaceGroup( }; std::map, SpaceGroupOperatorScore> op_cache; + std::vector visited(n, 0); + uint32_t epoch = 0; auto operator_score = [&](const gemmi::Op& op) -> const SpaceGroupOperatorScore& { const auto rk = RotKey(op); auto it = op_cache.find(rk); if (it != op_cache.end()) return it->second; - return op_cache.emplace(rk, score_operator(op)).first->second; + return op_cache.emplace(rk, score_operator(op, visited, epoch)).first->second; }; // Conjugate rotations (symmetry-equivalent within the point group) relate symmetry-equivalent @@ -423,6 +426,34 @@ SearchSpaceGroupResult SearchSpaceGroup( holohedry = HolohedryRotationSet(opt.lattice_system.value()); const auto point_groups = EnumeratePointGroups(holohedry); + // Every operator the search can ask about comes from this list, and scoring one is a pass over the + // whole merge with a hash lookup per reflection - the most expensive thing in here. They do not + // depend on each other, so score the distinct ones now and let the search below read the cache. + // One `visited` per worker, not per operator: it is as long as the merge, so allocating it per + // operator would cost more than the scoring. + { + std::vector distinct; + std::vector> keys; + for (const auto& pg : point_groups) + for (const auto& rot : pg.rotations) { + const auto rk = RotKey(rot); + if (std::find(keys.begin(), keys.end(), rk) == keys.end()) { + keys.push_back(rk); + distinct.push_back(rot); + } + } + std::vector scored(distinct.size()); + ParallelChunks(static_cast(distinct.size()), + std::min(opt.nthreads, distinct.size()), [&](int lo, int hi) { + std::vector scratch(n, 0); + uint32_t ep = 0; + for (int i = lo; i < hi; ++i) + scored[i] = score_operator(distinct[i], scratch, ep); + }); + for (size_t i = 0; i < distinct.size(); ++i) + op_cache.emplace(keys[i], scored[i]); + } + // Mapping every observation onto its symmetry representative under a candidate's rotations - one // apply_to_hkl + Canonicalize per observation per operator - is the expensive half of BOTH // quantities below, and both need exactly the same mapping. Build it once per point group. diff --git a/image_analysis/scale_merge/SearchSpaceGroup.h b/image_analysis/scale_merge/SearchSpaceGroup.h index e240f3b9..da5519e7 100644 --- a/image_analysis/scale_merge/SearchSpaceGroup.h +++ b/image_analysis/scale_merge/SearchSpaceGroup.h @@ -274,6 +274,10 @@ struct SearchSpaceGroupOptions { // in it, at p <= 2e-9 - loose enough that three well-measured dead axial reflections clear it, // tight enough that two do not. double min_screw_absence_evidence = 20.0; + + // Workers for the operator-correlation stage, which is the bulk of the search: one pass over the + // whole merge per candidate rotation, and the rotations are independent of each other. 1 = serial. + size_t nthreads = 1; }; struct SearchSpaceGroupResult { diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index a6a772cb..6bd611e8 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -2193,6 +2193,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b bool promoted_point_group = prepass_merge_sg_.has_value() && prepass_promoted_point_group_; if (!experiment_.GetGemmiSpaceGroup().has_value()) { SearchSpaceGroupOptions sg_opts; + sg_opts.nthreads = static_cast(std::max(1, config_.nthreads)); sg_opts.merge_friedel = experiment_.GetScalingSettings().GetMergeFriedel(); sg_opts.d_min_limit_A = std::max( d_min_search, experiment_.GetScalingSettings().GetHighResolutionLimit_A().value_or(0.0));