From c20cd016dc6bff6f683017f80b7af2cf6a64af3e Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 28 Sep 2026 15:17:26 +0200 Subject: [PATCH] Fixed-partition reductions: sums that do not depend on the thread count ParallelChunks cuts a range into one chunk per worker, so a floating-point sum folded per chunk and then added up rounds differently at another -N. ParallelBlocks cuts it by n alone (n / min_per_block blocks, at most 256); each block folds into its own slot and the slots are added in block order, so the sum has the same bits at any thread count. Converted: - RotationScaleMerge::ApplyCellSurface: the per-cell cross/ref2 sums of the surface fit (modulation, absorption), previously per-thread partials cut by ThreadsForWork(idx_all) threads. - PostRefine: the scale-scan cost grid (per-chunk slots cut by -N). - PostRefine joint solve: Ceres at a fixed 16 threads (as the rotation indexer's chain) instead of -N. Ceres sums cost and gradient in 4 * num_threads pieces, so the split no longer follows -N. The pieces are handed to its threads in scheduling order, which only one thread makes exact; one thread was measured at +1.0-1.4 s of post-refinement on myob and lyso (0.5 -> 1.9 s, 1.1 -> 2.1 s), so that channel is left. - IndexAndRefine supercell probe: per-frame probes kept by image number and summed in frame order, instead of added under a mutex in completion order; the primitive is the highest probed frame's instead of the last finisher's. Integer reductions and per-item passes are unchanged; the three ParallelSort callers already break ties on the index. Validation (myob, cytc, lyso, sparse; GPU and CPU builds; -N 8/16/32): p.hkl, p.mtz, p_P1.mtz and p_unmerged.mtz are md5-identical across -N, and identical to the base 351be7de0 at every -N. The base was already -N invariant on these four sets (GPU -N 1..32, CPU lyso -N 2..32), so no output moved: the converted sums differ from the old ones only in last bits that never reached a written float. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C --- common/ParallelFor.h | 23 ++++++- image_analysis/IndexAndRefine.cpp | 19 +++--- image_analysis/IndexAndRefine.h | 13 ++-- image_analysis/geom_refinement/PostRefine.cpp | 32 +++++----- .../scale_merge/RotationScaleMerge.cpp | 47 +++++++------- tests/ParallelForTest.cpp | 61 +++++++++++++++++++ 6 files changed, 143 insertions(+), 52 deletions(-) diff --git a/common/ParallelFor.h b/common/ParallelFor.h index 4e74fbaf2..3f89857cd 100644 --- a/common/ParallelFor.h +++ b/common/ParallelFor.h @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -166,7 +167,7 @@ namespace parallel_detail { // 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 -// same answer as the serial loop, bit for bit. +// same answer as the serial loop, bit for bit. A sum folded per chunk does not - use ParallelBlocks. template void ParallelChunks(int n, size_t nthreads, Fn fn) { if (n <= 0) return; @@ -196,6 +197,26 @@ void ParallelFor(int n, size_t nthreads, Fn fn) { }); } +// Fixed blocks, for a floating-point reduction. ParallelChunks cuts the range once per worker, so a sum +// folded per chunk and then added up rounds differently at another thread count - and a result sitting +// on a decision gate can then flip with -N. Here the cut depends on n alone: n / min_per_block blocks, +// at most MAX_REDUCTION_BLOCKS, block b being [n*b/nb, n*(b+1)/nb). Fold each block into its own slot +// (ReductionBlocks(n) of them), add the slots up in block order afterwards, and the sum has the same +// bits at any thread count; nthreads only says how many workers share the blocks. +constexpr int MAX_REDUCTION_BLOCKS = 256; + +inline int ReductionBlocks(int n, int min_per_block = 4096) { + if (n <= 0) return 0; + return std::clamp(n / min_per_block, 1, MAX_REDUCTION_BLOCKS); +} + +template +void ParallelBlocks(int n, size_t nthreads, Fn fn, int min_per_block = 4096) { + const int nb = ReductionBlocks(n, min_per_block); + const auto begin = [n, nb](int b) { return static_cast(static_cast(n) * b / nb); }; + ParallelFor(nb, nthreads, [&](int b) { fn(b, begin(b), begin(b + 1)); }); +} + // std::sort on several workers: the range is cut into one piece per worker, the pieces are sorted // in parallel and then merged pairwise. Only for a comparator that is a strict TOTAL order on the // elements present - no two of them equivalent - because then there is exactly one sorted sequence diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp index 8dea3f1c9..54b6eb324 100644 --- a/image_analysis/IndexAndRefine.cpp +++ b/image_analysis/IndexAndRefine.cpp @@ -639,18 +639,24 @@ SupercellProbeFit FitSupercellProbe(const SupercellProbeClass &c) { void IndexAndRefine::ProbeSupercell(std::vector image_numbers) { std::sort(image_numbers.begin(), image_numbers.end()); supercell_probe_frames_ = std::move(image_numbers); - supercell_probe_ = {}; - supercell_probe_primitive_.reset(); + supercell_probe_by_frame_.clear(); } SupercellProbe IndexAndRefine::GetSupercellProbe() { const std::unique_lock ul(supercell_probe_mutex_); - return supercell_probe_; + SupercellProbe sum{}; + for (const auto &[image, f] : supercell_probe_by_frame_) + for (int i = 0; i < 8; i++) + for (int j = 0; j < 2; j++) + sum[i][j] += f.probe[i][j]; + return sum; } std::optional IndexAndRefine::GetSupercellProbePrimitive() { const std::unique_lock ul(supercell_probe_mutex_); - return supercell_probe_primitive_; + if (supercell_probe_by_frame_.empty()) + return std::nullopt; + return supercell_probe_by_frame_.rbegin()->second.primitive; } void IndexAndRefine::ProbeSupercellFrame(const DataMessage &msg, BraggPrediction &prediction, @@ -687,10 +693,7 @@ void IndexAndRefine::ProbeSupercellFrame(const DataMessage &msg, BraggPrediction c.sum_ii += static_cast(r.I) * r.I; } const std::unique_lock ul(supercell_probe_mutex_); - for (int i = 0; i < 8; i++) - for (int j = 0; j < 2; j++) - supercell_probe_[i][j] += frame[i][j]; - supercell_probe_primitive_ = prim; + supercell_probe_by_frame_.insert_or_assign(msg.number, ProbedFrame{frame, prim}); } std::optional diff --git a/image_analysis/IndexAndRefine.h b/image_analysis/IndexAndRefine.h index f3b30055c..fd6458753 100644 --- a/image_analysis/IndexAndRefine.h +++ b/image_analysis/IndexAndRefine.h @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -120,10 +121,14 @@ class IndexAndRefine { std::vector scale_cc; std::vector > unit_cells; - // Supercell probe: the image numbers it reads, and what it has read (see ProbeSupercell). + // Supercell probe: the image numbers it reads, and what it has read (see ProbeSupercell) - per frame, + // so the sums are formed in frame order and not in the order the workers finish the frames. std::vector supercell_probe_frames_; - SupercellProbe supercell_probe_{}; - std::optional supercell_probe_primitive_; + struct ProbedFrame { + SupercellProbe probe; + CrystalLattice primitive; + }; + std::map supercell_probe_by_frame_; std::mutex supercell_probe_mutex_; void ProbeSupercellFrame(const DataMessage &msg, BraggPrediction &prediction, const BraggIntegrateFn &integrate, const IndexingOutcome &outcome, @@ -190,7 +195,7 @@ public: // The probe's integrations are the caller's to keep out of anything else it counts. void ProbeSupercell(std::vector image_numbers); [[nodiscard]] SupercellProbe GetSupercellProbe(); - // The primitive lattice the parity classes are indexed in (the last probed frame's; only its metric + // The primitive lattice the parity classes are indexed in (the highest probed frame's; only its metric // is meant to be read), so a class can be turned into the cell it would double to. [[nodiscard]] std::optional GetSupercellProbePrimitive(); // Index a single frame (no integration) with the current forced rotation lattice; used to score diff --git a/image_analysis/geom_refinement/PostRefine.cpp b/image_analysis/geom_refinement/PostRefine.cpp index 96393e738..7a08c6e94 100644 --- a/image_analysis/geom_refinement/PostRefine.cpp +++ b/image_analysis/geom_refinement/PostRefine.cpp @@ -627,8 +627,13 @@ PostRefineResult PostRefineRotationGeometry(PostRefineObservations observations, p.SetParameterLowerBound(q2, j, p2_0[j] - 0.05); p.SetParameterUpperBound(q2, j, p2_0[j] + 0.05); } + // A fixed thread count, not -N: Ceres sums the cost and gradient in 4 * num_threads + // pieces, so the steps it accepts and where it stops would otherwise move with -N. It + // hands the pieces to its threads in scheduling order, which only one thread makes + // exact - and one thread costs a second or more here - so that part is left. + constexpr int POSTREFINE_CERES_THREADS = 16; ceres::Solver::Options o; o.linear_solver_type = ceres::DENSE_QR; o.max_num_iterations = 60; - o.num_threads = std::max(1, settings.num_threads); o.logging_type = ceres::LoggingType::SILENT; + o.num_threads = POSTREFINE_CERES_THREADS; o.logging_type = ceres::LoggingType::SILENT; ceres::Solver::Summary sum; ceres::Solve(o, &p, &sum); if (distance_corr && !fix_distance && sum.IsSolutionUsable()) { std::fill(distance_corr, distance_corr + 4, std::numeric_limits::quiet_NaN()); @@ -636,7 +641,7 @@ PostRefineResult PostRefineRotationGeometry(PostRefineObservations observations, co.algorithm_type = ceres::DENSE_SVD; co.null_space_rank = -1; // a symmetry-constrained cell leaves unused length components // Threads only evaluate the Jacobian here, each residual block into its own rows. - co.num_threads = o.num_threads; + co.num_threads = std::max(1, settings.num_threads); ceres::Covariance cov(co); if (cov.Compute({{ds, ds}, {ds, q1}, {q1, q1}}, &p)) { double cdd[1], cdq[3], cqq[9]; @@ -1044,24 +1049,21 @@ PostRefineResult PostRefineRotationGeometry(PostRefineObservations observations, // Ceres minimises half the sum of the loss applied to the SQUARED residual, so that is what is // reproduced here. One pass yields the five per-fifth partial sums, which serve the all-data // fit and all five leave-a-fifth-out folds together. - // Each chunk folds into its own slot and the slots are summed in chunk order, so the - // sum does not depend on which worker finishes first: the same events always add up in - // the same sequence, and the fit is reproducible run to run. + // Each fixed block folds into its own slot and the slots are summed in block order, so the + // sum depends neither on which worker finishes first nor on how many there are: the same + // events always add up in the same sequence, and the fit is the same at any -N. const int n_terms = static_cast(terms.size()); - const int cost_nt = static_cast(std::max(1, std::min(nthreads, - static_cast(n_terms)))); - const int cost_chunk = (n_terms + cost_nt - 1) / cost_nt; + const int cost_nb = ReductionBlocks(n_terms); // Several k are always wanted at once (the grid below asks for 101), and they all sweep the // same event list, so sweep it ONCE and evaluate every k on each event while it is still in - // registers. The per-thread accumulator is one slot per (k, fifth) - 4 kB for the grid, small + // registers. The per-block accumulator is one slot per (k, fifth) - 4 kB for the grid, small // enough to stay in L1 - against re-reading the whole term array once per k. Each (k, fifth) - // still receives its events in the same order and the chunks are still summed in chunk order, + // still receives its events in the same order and the blocks are still summed in block order, // so the sums are the ones a k-at-a-time loop produced, bit for bit. const auto cost_grid = [&](const std::vector &ks) { const int nk = static_cast(ks.size()); - std::vector>> per_chunk( - cost_nt, std::vector>(nk)); - ParallelChunks(n_terms, nthreads, [&](int lo, int hi) { + std::vector>> per_block(cost_nb); + ParallelBlocks(n_terms, nthreads, [&](int b, int lo, int hi) { std::vector> acc(nk); for (int e = lo; e < hi; ++e) { const ScaleTerm &term = terms[e]; @@ -1073,10 +1075,10 @@ PostRefineResult PostRefineRotationGeometry(PostRefineObservations observations, : (2.0 * huber_delta * std::sqrt(s2) - huber_d2); } } - per_chunk[lo / cost_chunk] = std::move(acc); + per_block[b] = std::move(acc); }); std::vector> total(nk); - for (const auto &acc : per_chunk) + for (const auto &acc : per_block) for (int g = 0; g < nk; ++g) for (int j = 0; j < 5; ++j) total[g][j] += acc[g][j]; return total; diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 1178fe5eb..c32a166f0 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -3389,9 +3389,9 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // 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); }; + // Terms per block of the per-cell sums in the fit below: every block zeroes and hands back two + // ncell-long arrays, so a block has to hold many more terms than there are cells. + constexpr int SURFACE_BLOCK = 32768; // --- The surface's RESOLUTION gauge --------------------------------------------------------- // Every symmetry equivalent of a reflection sits at one d, so a factor that is a function of d @@ -3519,8 +3519,9 @@ 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)); + // Per-block cell accumulators, allocated once for the whole fit rather than per round. + const int nb = ReductionBlocks(static_cast(sel.size()), SURFACE_BLOCK); + std::vector> tcross(nb, std::vector(ncell)), tref2(nb, std::vector(ncell)); settled = false; n_clamped = 0; // What the resolution gauge takes out, accumulated over the rounds as a log per shell: the @@ -3550,28 +3551,26 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // own signal-to-noise: simulated on a flat truth, it suppressed a shell at I/sigma = 1 by // 40% over 30 rounds and left a 16x16 flat field at outer I/sigma = 0.5 with 1.9x its true // amplitude and CC 0.48 to the truth, where keeping the negatives gives 1.04x and 0.97. - // Chunked over the subset, each thread summing into its own cells: unlike the pass above + // Blocked over the subset, each block 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. + // of them, so per-block copies are cheap and the fixed blocks keep it the same at any -N. std::vector cross(ncell, 0.0), ref2(ncell, 0.0); - 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 Term &o = term[sel[k]]; - if (sw[o.group] <= 0.0) continue; - const double Iref = swI[o.group] / sw[o.group], a = A[o.cell]; - 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 || !(sc > 0.0)) continue; - const double w = 1.0 / (sc * sc); - xcross[o.cell] += w * Is * Iref; xref2[o.cell] += w * Iref * Iref; - } + ParallelBlocks(static_cast(sel.size()), nt, [&](int b, int lo, int hi) { + std::vector &xcross = tcross[b], &xref2 = tref2[b]; + std::fill(xcross.begin(), xcross.end(), 0.0); + std::fill(xref2.begin(), xref2.end(), 0.0); + for (int k = lo; k < hi; ++k) { + const Term &o = term[sel[k]]; + if (sw[o.group] <= 0.0) continue; + const double Iref = swI[o.group] / sw[o.group], a = A[o.cell]; + 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 || !(sc > 0.0)) continue; + const double w = 1.0 / (sc * sc); + xcross[o.cell] += w * Is * Iref; xref2[o.cell] += 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]; } + }, SURFACE_BLOCK); + for (int b = 0; b < nb; ++b) + for (int c = 0; c < ncell; ++c) { cross[c] += tcross[b][c]; ref2[c] += tref2[b][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]); diff --git a/tests/ParallelForTest.cpp b/tests/ParallelForTest.cpp index cec2988aa..74790440c 100644 --- a/tests/ParallelForTest.cpp +++ b/tests/ParallelForTest.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include #include @@ -104,12 +105,72 @@ TEST_CASE("ParallelFor_SplitDoesNotChangeTheResult", "[ParallelFor]") { } } +// The blocks depend on n alone: the same tiling of [0, n), block b always the same range, whatever the +// thread count. +TEST_CASE("ParallelBlocks_PartitionIgnoresTheThreadCount", "[ParallelFor]") { + for (int n: {0, 1, 7, 4095, 4096, 100000, 5000000}) { + const int nb = ReductionBlocks(n); + CHECK(nb <= MAX_REDUCTION_BLOCKS); + std::vector> reference; + for (size_t nthreads: THREAD_COUNTS) { + CAPTURE(n, nthreads); + std::vector> range(static_cast(nb), {-1, -1}); + ParallelBlocks(n, nthreads, [&](int b, int lo, int hi) { range[static_cast(b)] = {lo, hi}; }); + int expected_lo = 0; + for (const auto &[lo, hi]: range) { + CHECK(lo == expected_lo); + CHECK(hi > lo); + expected_lo = hi; + } + CHECK(expected_lo == n); + if (reference.empty()) reference = range; + CHECK(range == reference); + } + } +} + +// What the blocks are for: a floating-point sum folded per block and added up in block order has the +// same bits at any thread count. The terms span many decades, so any change of grouping would show. +TEST_CASE("ParallelBlocks_SumIsTheSameAtAnyThreadCount", "[ParallelFor]") { + constexpr int N = 1000003; + constexpr int NCELL = 37; + std::vector term(N); + for (int i = 0; i < N; i++) + term[static_cast(i)] = std::exp(std::sin(i * 0.37) * 20.0) * (i % 3 == 0 ? -1.0 : 1.0); + + auto blocked_sums = [&](size_t nthreads) { + const int nb = ReductionBlocks(N); + std::vector part(static_cast(nb), 0.0); + std::vector> cell_part(static_cast(nb), std::vector(NCELL, 0.0)); + ParallelBlocks(N, nthreads, [&](int b, int lo, int hi) { + for (int i = lo; i < hi; i++) { + part[static_cast(b)] += term[static_cast(i)]; + cell_part[static_cast(b)][static_cast(i % NCELL)] += term[static_cast(i)]; + } + }); + std::vector out(NCELL + 1, 0.0); + for (int b = 0; b < nb; b++) { + out[NCELL] += part[static_cast(b)]; + for (int c = 0; c < NCELL; c++) + out[static_cast(c)] += cell_part[static_cast(b)][static_cast(c)]; + } + return out; + }; + + const std::vector one = blocked_sums(1); + for (size_t nthreads: {size_t{3}, size_t{7}, size_t{32}}) { + CAPTURE(nthreads); + CHECK(blocked_sums(nthreads) == one); // bit for bit, not approximately + } +} + // A negative or zero count is a no-op rather than an error - callers pass a computed size. TEST_CASE("ParallelFor_DoesNothingForAnEmptyRange", "[ParallelFor]") { int calls = 0; for (int n: {0, -1, -1000}) { ParallelChunks(n, 8, [&](int, int) { calls++; }); ParallelFor(n, 8, [&](int) { calls++; }); + ParallelBlocks(n, 8, [&](int, int, int) { calls++; }); } CHECK(calls == 0); }