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

Merge branch 'fixed-partition' into rc173 (float reductions on a partition fixed by the data, not the thread count)

Browse Source 
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C
This commit is contained in:
leonarski_f
2026-09-28 18:58:26 +02:00
co-authored by Claude Opus 5.5
parent 2a9ecc0abf c20cd016dc
commit 82e6583e00
6 changed files with 143 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/ParallelFor.h
+22 -1
View File
@@ -5,6 +5,7 @@
#include <algorithm>
#include <atomic>
#include <cstdint>
#include <condition_variable>
#include <deque>
#include <exception>
@@ -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 <typename Fn>
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 <typename Fn>
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<int>(static_cast<int64_t>(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
image_analysis/IndexAndRefine.cpp
+11 -8
View File
@@ -639,18 +639,24 @@ SupercellProbeFit FitSupercellProbe(const SupercellProbeClass &c) {
void IndexAndRefine::ProbeSupercell(std::vector<int64_t> 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<CrystalLattice> 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<double>(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<IndexAndRefine::IndexingOutcome>
image_analysis/IndexAndRefine.h
+9 -4
View File
@@ -5,6 +5,7 @@
#include <array>
#include <atomic>
#include <map>
#include <vector>
#include <mutex>
#include <functional>
@@ -120,10 +121,14 @@ class IndexAndRefine {
std::vector<float> scale_cc;
std::vector<std::optional<UnitCell> > 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<int64_t> supercell_probe_frames_;
SupercellProbe supercell_probe_{};
std::optional<CrystalLattice> supercell_probe_primitive_;
struct ProbedFrame {
SupercellProbe probe;
CrystalLattice primitive;
};
std::map<int64_t, ProbedFrame> 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<int64_t> 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<CrystalLattice> GetSupercellProbePrimitive();
// Index a single frame (no integration) with the current forced rotation lattice; used to score
image_analysis/geom_refinement/PostRefine.cpp
+17 -15
View File
@@ -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<double>::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<int>(terms.size());
const int cost_nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads,
static_cast<size_t>(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<double> &ks) {
const int nk = static_cast<int>(ks.size());
std::vector<std::vector<std::array<double, 5>>> per_chunk(
cost_nt, std::vector<std::array<double, 5>>(nk));
ParallelChunks(n_terms, nthreads, [&](int lo, int hi) {
std::vector<std::vector<std::array<double, 5>>> per_block(cost_nb);
ParallelBlocks(n_terms, nthreads, [&](int b, int lo, int hi) {
std::vector<std::array<double, 5>> 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<std::array<double, 5>> 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;
image_analysis/scale_merge/RotationScaleMerge.cpp
+23 -24
View File
@@ -3389,9 +3389,9 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &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<int>(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<size_t>(t) / static_cast<size_t>(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<int32_t> &cell, int
auto fit_surface = [&](int parity) -> std::vector<double> {
const std::vector<int32_t> &sel = subset(parity);
std::vector<double> A(ncell, 1.0);
// Per-thread cell accumulators, allocated once for the whole fit rather than per round.
std::vector<std::vector<double>> tcross(nt, std::vector<double>(ncell)), tref2(nt, std::vector<double>(ncell));
// Per-block cell accumulators, allocated once for the whole fit rather than per round.
const int nb = ReductionBlocks(static_cast<int>(sel.size()), SURFACE_BLOCK);
std::vector<std::vector<double>> tcross(nb, std::vector<double>(ncell)), tref2(nb, std::vector<double>(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<int32_t> &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<double> cross(ncell, 0.0), ref2(ncell, 0.0);
ParallelChunks(nt, nthreads, [&](int tlo, int thi) {
for (int t = tlo; t < thi; ++t) {
std::vector<double> &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<double>(o.I) * o.corr * a, sc = static_cast<double>(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<int>(sel.size()), nt, [&](int b, int lo, int hi) {
std::vector<double> &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<double>(o.I) * o.corr * a, sc = static_cast<double>(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<double> 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]);
tests/ParallelForTest.cpp
+61
View File
@@ -4,6 +4,7 @@
#include <catch2/catch_all.hpp>
#include <algorithm>
#include <cmath>
#include <mutex>
#include <numeric>
#include <stdexcept>
@@ -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<std::pair<int, int>> reference;
for (size_t nthreads: THREAD_COUNTS) {
CAPTURE(n, nthreads);
std::vector<std::pair<int, int>> range(static_cast<size_t>(nb), {-1, -1});
ParallelBlocks(n, nthreads, [&](int b, int lo, int hi) { range[static_cast<size_t>(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<double> term(N);
for (int i = 0; i < N; i++)
term[static_cast<size_t>(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<double> part(static_cast<size_t>(nb), 0.0);
std::vector<std::vector<double>> cell_part(static_cast<size_t>(nb), std::vector<double>(NCELL, 0.0));
ParallelBlocks(N, nthreads, [&](int b, int lo, int hi) {
for (int i = lo; i < hi; i++) {
part[static_cast<size_t>(b)] += term[static_cast<size_t>(i)];
cell_part[static_cast<size_t>(b)][static_cast<size_t>(i % NCELL)] += term[static_cast<size_t>(i)];
}
});
std::vector<double> out(NCELL + 1, 0.0);
for (int b = 0; b < nb; b++) {
out[NCELL] += part[static_cast<size_t>(b)];
for (int c = 0; c < NCELL; c++)
out[static_cast<size_t>(c)] += cell_part[static_cast<size_t>(b)][static_cast<size_t>(c)];
}
return out;
};
const std::vector<double> 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);
}