From c8692e320c968680535a0209c1f664b766742ab4 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 28 Sep 2026 02:19:31 +0200 Subject: [PATCH 1/5] CPU pixel loops: ring sums held in registers, vectorised vertical window, shared packed mask AdaptiveSpotFinderCPU::AccumulateRings: consecutive pixels mostly share a ring, so the ring's sum/sum2/count and the fused azint sums are held in locals while they do and stored when the ring changes - the same additions in the same order (az_sum2 is still contracted to the same FMA), without a store-and-reload chain through memory on every pixel. ImageSpotFinderCPU::DetectPass: the vertical-sum update (add the entering row, take out the leaving one) is one branch-free loop over the raw image that GCC vectorises (int64 lanes). A pixel strong in the previous pass used to be substituted per pixel through a bit test, which kept the loop scalar; it is now added with its row and taken out again from the few set bits of prev_strong. Integer sums, so the same totals. (A first, fully branch-free version that kept the per-pixel bit test did not vectorise on the prev_strong path and was measured slower; this is its replacement.) ImagePreprocessorCPU: the per-engine std::vector built bit by bit from the 32-bit mask (~10 core-s per cytc run, one per worker per pass) is replaced by 32-pixel mask words that PixelMask derives once beside its binary mask; each engine copies 2 MB. A branch-free rewrite of the Analyze loop was measured and dropped: the loop is bound by reading the decompressed image (330 vs 328 core-s on cytc), so only the mask test changed. Measured (perf, 499 Hz, CPU-only build, cytc, first version of this change): AccumulateRings 591 -> 539 core-s. Byte-identical p.hkl, p.mtz, p_P1.mtz, p_unmerged.mtz on myob, cytc, lyso, sparse (CPU) and myob, lyso (GPU). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C --- common/PixelMask.cpp | 7 ++ common/PixelMask.h | 4 + .../ImagePreprocessorCPU.cpp | 7 +- .../ImagePreprocessorCPU.h | 4 +- .../spot_finding/AdaptiveSpotFinderCPU.cpp | 46 ++++++++-- .../spot_finding/ImageSpotFinderCPU.cpp | 85 ++++++++++++------- 6 files changed, 107 insertions(+), 46 deletions(-) diff --git a/common/PixelMask.cpp b/common/PixelMask.cpp index a5338f4ee..8e455195e 100644 --- a/common/PixelMask.cpp +++ b/common/PixelMask.cpp @@ -59,6 +59,9 @@ void PixelMask::UpdateBinaryMask() { for (size_t i = 0; i < mask.size(); i++) binary_mask[i] = (mask[i] != 0); binary_mask_checksum = TableChecksum(binary_mask.data(), binary_mask.size()); + packed_mask.assign((mask.size() + 31) / 32, 0); + for (size_t i = 0; i < mask.size(); i++) + packed_mask[i / 32] |= static_cast(binary_mask[i]) << (i % 32); } void PixelMask::CalcEdgePixels_i(const DiffractionExperiment &experiment) { @@ -127,6 +130,10 @@ const std::vector &PixelMask::GetBinaryMask() const { return binary_mask; } +const std::vector &PixelMask::GetPackedMask() const { + return packed_mask; +} + uint64_t PixelMask::GetBinaryMaskChecksum() const { return binary_mask_checksum; } diff --git a/common/PixelMask.h b/common/PixelMask.h index 5408b0789..ddcf50be2 100644 --- a/common/PixelMask.h +++ b/common/PixelMask.h @@ -26,6 +26,9 @@ class PixelMask { // derived here once instead of in every one of those engines. std::vector binary_mask; uint64_t binary_mask_checksum = 0; + // The same flag packed 32 pixels to a word (pixel i is bit i % 32 of word i / 32) - the form the CPU + // image preprocessor reads, derived here for the same reason. + std::vector packed_mask; uint32_t LoadMask(const std::vector& mask, uint8_t bit); // Everything that follows from the mask, recomputed wherever the mask changes. @@ -68,6 +71,7 @@ public: [[nodiscard]] const std::vector &GetMask(const DiffractionExperiment& experiment) const; [[nodiscard]] const std::vector &GetMask() const; [[nodiscard]] const std::vector &GetBinaryMask() const; + [[nodiscard]] const std::vector &GetPackedMask() const; [[nodiscard]] uint64_t GetBinaryMaskChecksum() const; [[nodiscard]] std::vector GetUserMask(const DiffractionExperiment& experiment) const; diff --git a/image_analysis/image_preprocessing/ImagePreprocessorCPU.cpp b/image_analysis/image_preprocessing/ImagePreprocessorCPU.cpp index 8803a73db..bb3a98931 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorCPU.cpp +++ b/image_analysis/image_preprocessing/ImagePreprocessorCPU.cpp @@ -5,10 +5,7 @@ ImagePreprocessorCPU::ImagePreprocessorCPU(const DiffractionExperiment &experiment, const PixelMask &mask) : ImagePreprocessor(experiment), - mask_1bit(npixels, false) { - for (int i = 0; i < npixels; i++) - mask_1bit[i] = (mask.GetMask().at(i) != 0); -} + mask_bits(mask.GetPackedMask()) {} ImageStatistics ImagePreprocessorCPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *image_ptr, CompressedImageMode image_mode) { switch (image_mode) { @@ -43,7 +40,7 @@ ImageStatistics ImagePreprocessorCPU::Analyze(ImagePreprocessorBuffer &processed sat_pixel_val = static_cast(saturation_limit); for (int i = 0; i < npixels; i++) { - if (mask_1bit[i] != 0) { + if ((mask_bits[i / 32] >> (i % 32)) & 1U) { processed_image[i] = INT32_MIN; ++ret.masked_pixel_count; } else if (image[i] == err_pixel_val) { diff --git a/image_analysis/image_preprocessing/ImagePreprocessorCPU.h b/image_analysis/image_preprocessing/ImagePreprocessorCPU.h index aa326d77b..2b07bd54a 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorCPU.h +++ b/image_analysis/image_preprocessing/ImagePreprocessorCPU.h @@ -7,7 +7,9 @@ class ImagePreprocessorCPU : public ImagePreprocessor { - std::vector mask_1bit; + // The pixel mask one bit per pixel, 32 to a word - PixelMask derives it once, so each engine only + // copies it. + std::vector mask_bits; template ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *input, T err_value, T sat_value); public: diff --git a/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp b/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp index 93b7b9597..ef66199a3 100644 --- a/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp +++ b/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp @@ -55,26 +55,58 @@ void AdaptiveSpotFinderCPU::AccumulateRings(const ImagePreprocessorBuffer &image std::fill(azint_sum2.begin(), azint_sum2.end(), 0.0f); std::fill(azint_count.begin(), azint_count.end(), 0); } + // Consecutive pixels mostly share a ring, so the ring's sums are held in locals while they do + // and written back when the ring changes: the same additions in the same order, without a + // store and a reload of the same address on every pixel. + uint32_t *hist = ring_hist.data(); + size_t cur = nbins; // the ring held in the locals below; nbins = none + int64_t sum = 0, cnt = 0; + uint64_t sum2 = 0; + float az_sum = 0.0f, az_sum2 = 0.0f; + uint32_t az_cnt = 0; + const auto store = [&] { + if (cur == nbins) + return; + ring_sum[cur] = sum; + ring_sum2[cur] = sum2; + ring_cnt[cur] = cnt; + if (fuse_azint) { + azint_sum[cur] = az_sum; + azint_sum2[cur] = az_sum2; + azint_count[cur] = az_cnt; + } + }; for (size_t pxl = 0; pxl < npix; ++pxl) { const int32_t v = image[pxl]; if (v == INT32_MIN || v == INT32_MAX) continue; // bad / saturated const uint16_t b = pixel_to_bin[pxl]; if (b >= nbins) continue; // masked / out of range (UINT16_MAX) + if (b != cur) { + store(); + cur = b; + sum = ring_sum[b]; + sum2 = ring_sum2[b]; + cnt = ring_cnt[b]; + az_sum = azint_sum[b]; + az_sum2 = azint_sum2[b]; + az_cnt = azint_count[b]; + } if (fuse_azint) { const float val = static_cast(v) * corrections[pxl]; const float val_sq = val * val; - azint_sum[b] += val; - azint_sum2[b] += val_sq; - ++azint_count[b]; + az_sum += val; + az_sum2 += val_sq; + ++az_cnt; } - ring_sum[b] += v; - ring_sum2[b] += static_cast(static_cast(v) * v); - ring_cnt[b] += 1; + sum += v; + sum2 += static_cast(static_cast(v) * v); + cnt += 1; if (v >= 0 && v < HIST_VALUES) - ring_hist[b * HIST_VALUES + v] += 1; + hist[b * HIST_VALUES + v] += 1; else ring_overflow.emplace_back(b, v); } + store(); } else { // A sigma-clip pass over the plain pass's values: each distinct value of a ring meets the same // test the pixels holding it would, and its pixels are added as a count. diff --git a/image_analysis/spot_finding/ImageSpotFinderCPU.cpp b/image_analysis/spot_finding/ImageSpotFinderCPU.cpp index e2a41bc15..b69993fa6 100644 --- a/image_analysis/spot_finding/ImageSpotFinderCPU.cpp +++ b/image_analysis/spot_finding/ImageSpotFinderCPU.cpp @@ -146,42 +146,61 @@ void ImageSpotFinderCPU::DetectPass(const ImagePreprocessorBuffer &image, std::vector sum2_vert(width, 0); std::vector valid_vert(width, 0); - for (int line = 0; line < NBX; line++) { - for (int col = 0; col < width; col++) { - auto pxl = line * width + col; - - int64_t tmp = value_at(pxl); - if (tmp != INT32_MAX && tmp != INT32_MIN) { - sum_vert[col] += tmp; - sum2_vert[col] += tmp * tmp; - valid_vert[col] += 1; - } + // Add the row entering the window (line_in) and take out the one leaving it (line_out); -1 is + // no row. A bad or saturated pixel adds 0. Written without branches, so the loop vectorises. + // A pixel strong in the previous pass is not counted either; it is added here with the rest + // of its row and taken out again below - there are few, and the sums are integers, so the + // totals are the same as never adding it. + const int32_t *img = image.data(); + const int32_t w = width; + int64_t *sv = sum_vert.data(); + int64_t *sv2 = sum2_vert.data(); + uint16_t *vv = valid_vert.data(); + auto slide_vert = [&](int line_in, int line_out) { + for (int32_t col = 0; col < w; col++) { + const int64_t in = line_in >= 0 ? img[line_in * w + col] : INT32_MIN; + const int64_t out = line_out >= 0 ? img[line_out * w + col] : INT32_MIN; + const bool in_ok = in != INT32_MAX && in != INT32_MIN; + const bool out_ok = out != INT32_MAX && out != INT32_MIN; + const int64_t in_v = in_ok ? in : 0; + const int64_t out_v = out_ok ? out : 0; + sv[col] += in_v - out_v; + sv2[col] += in_v * in_v - out_v * out_v; + vv[col] += static_cast(in_ok - out_ok); } - } + if (!prev_strong) + return; + // sign = +1 takes out what line_in added, -1 puts back what line_out took out. + const auto uncount_strong = [&](int line, int64_t sign) { + const int32_t first = line * w, last = first + w - 1; + for (int32_t word = first / 32; word <= last / 32; word++) { + uint32_t bits = prev_strong[word]; + if (word == first / 32) + bits &= UINT32_MAX << (first % 32); + if (word == last / 32 && last % 32 != 31) + bits &= (1U << (last % 32 + 1)) - 1; + for (; bits; bits &= bits - 1) { + const int32_t pxl = word * 32 + std::countr_zero(bits); + const int64_t v = img[pxl]; + if (v == INT32_MAX || v == INT32_MIN) + continue; + sv[pxl - first] -= sign * v; + sv2[pxl - first] -= sign * v * v; + vv[pxl - first] -= static_cast(sign); + } + } + }; + if (line_in >= 0) + uncount_strong(line_in, 1); + if (line_out >= 0) + uncount_strong(line_out, -1); + }; + + for (int line = 0; line < NBX; line++) + slide_vert(line, -1); for (int line = 0; line < height; line++) { - for (int col = 0; col < width; col++) { - if (line < height - NBX) { - auto pxl = (line + NBX) * width + col; - - int64_t tmp = value_at(pxl); - if (tmp != INT32_MAX && tmp != INT32_MIN) { - sum_vert[col] += tmp; - sum2_vert[col] += tmp * tmp; - valid_vert[col] += 1; - } - } - - if (line >= NBX + 1) { - auto pxl = (line - (NBX + 1)) * width + col; - int64_t tmp = value_at(pxl); - if (tmp != INT32_MAX && tmp != INT32_MIN) { - sum_vert[col] -= tmp; - sum2_vert[col] -= tmp * tmp; - valid_vert[col] -= 1; - } - } - } + slide_vert(line < height - NBX ? line + NBX : -1, line >= NBX + 1 ? line - (NBX + 1) : -1); if (candidates) { // Only the runs of 32-column blocks whose result is read (first_pass_needed; every From 59c4d6fe41c96f4547665436b26d9bf7793a7e4c Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 28 Sep 2026 02:19:51 +0200 Subject: [PATCH 2/5] BraggPredictionRot: reject far-off-frame solutions before the atan2 Most rotation solutions are far from the frame, and the existing phi prefilter rejected them only after computing phi = atan2(sin, cos). The same test is now asked of the (cos, sin) pair first: |phi| > phi_limit is cos(phi) < cos(phi_limit), i.e. c < cos(phi_limit) * |(c, s)| in double. phi_limit is the existing test's threshold at the largest rocking width any reflection can have (bandwidth term at the resolution limit) plus a 1e-3 rad margin, far above float rounding, so every solution rejected here is rejected by the existing test too, and the survivors go through it unchanged. One cos per Calc; min_zeta <= 0 disables it, as it disables the existing test. Measured (perf, 499 Hz, CPU-only build, cytc): atanf32 + atan2f32 + __atan2f_finite 196 -> 12 core-s, Calc itself 161 -> 199 core-s (the sqrt of the new test), net ~ -150 core-s. Byte-identical p.hkl, p.mtz, p_P1.mtz, p_unmerged.mtz on myob, cytc, lyso, sparse (CPU) and myob, lyso (GPU); Bragg* tests pass. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C --- .../bragg_prediction/BraggPredictionRot.cpp | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/image_analysis/bragg_prediction/BraggPredictionRot.cpp b/image_analysis/bragg_prediction/BraggPredictionRot.cpp index dbcd90fd4..4a516262f 100644 --- a/image_analysis/bragg_prediction/BraggPredictionRot.cpp +++ b/image_analysis/bragg_prediction/BraggPredictionRot.cpp @@ -77,6 +77,28 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys experiment.GetBraggIntegrationSettings().GetFlightPath(), geom.GetDetectorDistance_mm(), geom.GetWavelength_A()); + // The phi test after the atan2 below, taken before it: a solution with |phi| above phi_limit fails + // that test whatever its own bandwidth term, since mos_eff_rad is largest at the resolution limit. + // |phi| > phi_limit is cos(phi) < cos(phi_limit), which the (cos, sin) pair answers without the + // angle. phi_limit carries a 1e-3 rad margin, far above float rounding, so every solution rejected + // here is rejected by that test too. -1 = nothing is rejected here. + double cos_phi_limit = -1.0; + if (settings.min_zeta > 0.0f) { + double mos_max_rad = mos_angle_rad; + if (bandwidth_sigma > 0.0f) { + const double sin_theta = half_wavelength_A * static_cast(one_over_dmax); + const double dphi_bw = sin_theta < 1.0 + ? bandwidth_sigma * sin_theta / std::sqrt(1.0 - sin_theta * sin_theta) + : INFINITY; + mos_max_rad = std::sqrt(mos_max_rad * mos_max_rad + dphi_bw * dphi_bw); + } + const double phi_limit = half_wedge_angle_rad + + (settings.mosaicity_multiplier * mos_max_rad + 1e-5) / settings.min_zeta + + 1e-3; + if (phi_limit < PI) + cos_phi_limit = std::cos(phi_limit); + } + for (int h = -settings.max_h; h <= settings.max_h; h++) { // Precompute A* h contribution @@ -127,6 +149,9 @@ int BraggPredictionRot::Calc(const DiffractionExperiment &experiment, const Crys const float cosphi = (p_m1 * p0_m1 + p_m3 * p0_m3) / rho_sq; const float sinphi = (p_m1 * p0_m3 - p_m3 * p0_m1) / rho_sq; + const double c = cosphi, s = sinphi; + if (c < cos_phi_limit * std::sqrt(c * c + s * s)) + continue; float phi = -1.0f * std::atan2(sinphi, cosphi); // Most solutions are far from this frame. zeta >= min_zeta, so the rocking-curve test From 2bad92ddd6c86a90c20977924f80ec79c11a46a6 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 28 Sep 2026 02:19:51 +0200 Subject: [PATCH 3/5] RotationScaleMerge: the CPU partial scaling loop works on a compact copy of the partials Each iteration of the partial scaling loop passes over all partials three times (group means, per-frame G fit, corr update + rescale), each time reading a few fields of the 80-byte Obs, and the group means gathered whole Obs records through the group permutation. The loop now works on ScalingObs (frame, group, I, sigma, partiality, prescaling_corr, zeta, on_ice; 36 bytes, partials' order) with corr in its own float array, and the group means read ScalingRef (I, sigma, partiality) laid out in the CSR's group order, so only corr is gathered. corr is written back to the partials (and kept by the memo) when the loop ends. ReferenceWeight takes the fields, FitPerFrameG is a template over Obs/ScalingObs; the same values, so the same arithmetic in the same order. Transient memory: 52 bytes per partial during the loop (~0.4 GB on 8M partials). Measured (perf, 499 Hz, CPU-only build, cytc): update/rescale lambda 104 -> 35, group means 87 -> 24, FitPerFrameG 63 -> 27 core-s (-167 core-s, of which the copy costs ~10). CPU-only path; the GPU scaling path is untouched. Byte-identical p.hkl, p.mtz, p_P1.mtz, p_unmerged.mtz on myob, cytc, lyso, sparse (CPU) and myob, lyso (GPU); rotation scale/merge tests pass. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C --- .../scale_merge/RotationScaleMerge.cpp | 88 ++++++++++++++----- .../scale_merge/RotationScaleMerge.h | 27 +++++- 2 files changed, 92 insertions(+), 23 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 10c5829d6..1178fe5eb 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -1336,12 +1336,13 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { return n_groups; } -double RotationScaleMerge::ReferenceWeight(const Obs &o, double min_partiality, float &I_corr) { - if (o.group < 0) return 0.0; - if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) return 0.0; - if (o.partiality < min_partiality) return 0.0; - const float Ic = o.I * o.corr; - const float sigma_corr = o.sigma * o.corr; +double RotationScaleMerge::ReferenceWeight(int32_t group, float corr, float partiality, float I, float sigma, + double min_partiality, float &I_corr) { + if (group < 0) return 0.0; + if (!(corr > 0.0f) || !std::isfinite(corr)) return 0.0; + if (partiality < min_partiality) return 0.0; + const float Ic = I * corr; + const float sigma_corr = sigma * corr; if (!std::isfinite(Ic) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0; I_corr = Ic; return 1.0 / (static_cast(sigma_corr) * sigma_corr); @@ -1383,7 +1384,30 @@ void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vector &obs, const std::vector &fstart, +void RotationScaleMerge::ReduceScalingGroupMeans(int n_groups, const std::vector &refs, + const std::vector &corr, + std::vector &out_mean) const { + out_mean.assign(n_groups, NAN); + ParallelChunks(n_groups, ThreadsForWork(corr.size(), nthreads), [&](int glo, int ghi) { + for (int g = glo; g < ghi; ++g) { + double sw = 0.0, swI = 0.0; + const int lo = partial_group_start[g], hi = lo + partial_group_count[g]; + for (int p = lo; p < hi; ++p) { + const ScalingRef &r = refs[p]; + float I_corr = 0.0f; + const double w = ReferenceWeight(g, corr[partial_group_perm[p]], r.partiality, r.I, r.sigma, + min_partiality, I_corr); + if (w <= 0.0) continue; + sw += w; + swI += w * I_corr; + } + if (sw > 0.0) out_mean[g] = swI / sw; + } + }); +} + +template +void RotationScaleMerge::FitPerFrameG(const std::vector &obs, const std::vector &fstart, const std::vector &fcount, const std::vector &group_mean_in, bool unity, std::vector &g) { @@ -5560,6 +5584,30 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st logger.Info("Per-frame scaling (partials): the same grouping and settings as an earlier merge, " "so its {} iterations stand", partial_loop.iterations); } else { + // The loop works on its own compact copy of the partials (ScalingObs in the header) and hands + // corr back to them when it is done. + const int n_partials = static_cast(partials.size()); + const size_t partial_threads = ThreadsForWork(partials.size(), nthreads); + std::vector scaling_obs(n_partials); + std::vector scaling_corr(n_partials); + ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) { + const Obs &o = partials[i]; + scaling_obs[i] = {o.frame, o.group, o.I, o.sigma, o.partiality, o.prescaling_corr, o.zeta, + o.on_ice}; + scaling_corr[i] = o.corr; + } + }); + std::vector scaling_refs(partial_group_start.empty() + ? 0 : partial_group_start.back() + partial_group_count.back()); + ParallelChunks(n_groups, partial_threads, [&](int glo, int ghi) { + for (int g = glo; g < ghi; ++g) + for (int p = partial_group_start[g]; p < partial_group_start[g] + partial_group_count[g]; ++p) { + const ScalingObs &o = scaling_obs[partial_group_perm[p]]; + scaling_refs[p] = {o.I, o.sigma, o.partiality}; + } + }); + // RunScalingLoop rescales after every iteration before anything reads corr, so the iteration's // UpdateCorr is folded into that pass over the partials: the G and the fitted frames as the // iteration left them, then the rescale on top - the same two roundings, one pass instead of two. @@ -5567,34 +5615,32 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st std::vector frame_fitted; partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window, /*release_restraint=*/false, [&] { - ReducePartialGroupMeans(n_groups, partial_mean); - FitPerFrameG(partials, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial); + ReduceScalingGroupMeans(n_groups, scaling_refs, scaling_corr, partial_mean); + FitPerFrameG(scaling_obs, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial); g_fitted = g_partial; frame_fitted = frame_scaled_scratch; }, [&](const std::vector &apply, const std::vector &ratio) { - ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), - [&](int lo, int hi) { + ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { for (int i = lo; i < hi; ++i) { - Obs &o = partials[i]; + const ScalingObs &o = scaling_obs[i]; + float &corr = scaling_corr[i]; if (frame_fitted[o.frame]) { // UpdateCorr const double denom = static_cast(o.partiality) * g_fitted[o.frame]; if (std::isfinite(o.prescaling_corr) && std::isfinite(denom) && denom > 0.0) - o.corr = static_cast(o.prescaling_corr / denom); + corr = static_cast(o.prescaling_corr / denom); else - o.corr = NAN; + corr = NAN; } - if (apply[o.frame] && std::isfinite(o.corr)) - o.corr = static_cast(o.corr * ratio[o.frame]); + if (apply[o.frame] && std::isfinite(corr)) + corr = static_cast(corr * ratio[o.frame]); } }); }); - PartialLoopMemo m{key, partial_loop, frame_scaled_scratch, g_partial, {}}; - m.corr.resize(partials.size()); - ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), - [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) m.corr[i] = partials[i].corr; + ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) partials[i].corr = scaling_corr[i]; }); + PartialLoopMemo m{key, partial_loop, frame_scaled_scratch, g_partial, std::move(scaling_corr)}; if (partial_loop_memos.size() == MAX_PARTIAL_LOOP_MEMOS) partial_loop_memos.erase(partial_loop_memos.begin()); partial_loop_memos.push_back(std::move(m)); diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index ebef3de15..ea1947c8f 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -446,17 +446,40 @@ private: std::vector partial_group_start, partial_group_count; std::unique_ptr partial_group_perm; + // The CPU partial scaling loop's working copy of the partials. It passes over them several times + // per iteration and reads a handful of fields each time, so it works on these rather than on the + // eighty-byte Obs: ScalingObs holds, in the partials' own order, what FitPerFrameG and the corr + // update read; ScalingRef holds, in the CSR's group order, what the group means read besides corr, + // so that only corr is gathered. The same values, so the same arithmetic in the same order. + struct ScalingObs { + int32_t frame, group; + float I, sigma, partiality, prescaling_corr, zeta; + uint8_t on_ice; + }; + struct ScalingRef { + float I, sigma, partiality; + }; + // The observation's inverse-variance weight in the scaling reference, and its scaled intensity; // 0 when the observation is not in the reference at all. - static double ReferenceWeight(const Obs &o, double min_partiality, float &I_corr); + static double ReferenceWeight(int32_t group, float corr, float partiality, float I, float sigma, + double min_partiality, float &I_corr); + static double ReferenceWeight(const Obs &o, double min_partiality, float &I_corr) { + return ReferenceWeight(o.group, o.corr, o.partiality, o.I, o.sigma, min_partiality, I_corr); + } // Inverse-variance per-group mean of I*corr over `obs` (the merge reference). void ReduceGroupMeans(const std::vector &obs, int n_groups, std::vector &out_mean) const; // The same over the partials, using the CSR above - the same sums in the same order. void ReducePartialGroupMeans(int n_groups, std::vector &out_mean) const; + // The same again over the scaling loop's copy: refs in group order, corr in the partials' order. + void ReduceScalingGroupMeans(int n_groups, const std::vector &refs, + const std::vector &corr, std::vector &out_mean) const; // Robust per-frame G fit (IRLS, Cauchy k=3), unity=false uses the rotation partiality, unity=true the // scale-fulls (partiality already folded in). Reads out_mean[group] as the reference intensity. - void FitPerFrameG(std::vector &obs, const std::vector &fstart, + // O is Obs or ScalingObs. + template + void FitPerFrameG(const std::vector &obs, const std::vector &fstart, const std::vector &fcount, const std::vector &group_mean_in, bool unity, std::vector &g); From 61d88cdefe7b9afc257d9075c5c6ca550bdd945d Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 28 Sep 2026 02:19:51 +0200 Subject: [PATCH 4/5] BraggIntegrationEngineCPU: reflection mask and owner map kept as 16x16 tiles refl_mask (uint8) and owner (uint32) were full-frame arrays per engine, ~90 MB at 18 Mpixel, of which a call writes only the neighbourhoods of its reflections; over a sweep every page of both got touched in every worker, and the dirty-rectangle clears were row-strided memsets. They are now a TiledFrame: a per-tile directory plus a pool of 16x16 tiles allocated on first write and dropped by Clear(); reading a tile nothing wrote gives the empty value (0 / BRAGG_OWNER_NONE). Writes and reads are exactly those of the flat arrays, so the same ownership decisions. Measured (perf, 499 Hz, CPU-only build, cytc): RunImpl 1633 -> 1557 core-s, memset 88 -> 32 core-s. Peak RSS and page faults did not move measurably (myob 15.3 GB; the peak is set elsewhere). Byte-identical p.hkl, p.mtz, p_P1.mtz, p_unmerged.mtz on myob, cytc, lyso, sparse (CPU) and myob, lyso (GPU); Bragg* tests pass. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C --- .../BraggIntegrationEngineCPU.cpp | 32 ++++--------- .../BraggIntegrationEngineCPU.h | 47 +++++++++++++++---- 2 files changed, 48 insertions(+), 31 deletions(-) diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp index 64a864868..edbe47b1b 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp @@ -44,7 +44,9 @@ struct ImageSampler { } // namespace BraggIntegrationEngineCPU::BraggIntegrationEngineCPU(const DiffractionExperiment &experiment) - : BraggIntegrationEngine(experiment) {} + : BraggIntegrationEngine(experiment), + refl_mask(static_cast(xpixel), static_cast(ypixel), 0), + owner(static_cast(xpixel), static_cast(ypixel), BRAGG_OWNER_NONE) {} template std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, @@ -71,17 +73,8 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, // signal is excluded from this reflection's r2..r3 background ring. The region is the INNER // stencil ellipse, taken in each neighbour's OWN frame - an elongated neighbour whose streak // is still marked as a disk would leak its tails into this reflection's ring. --- - // Back to empty what the previous call wrote (see the header). - for (const Rect &b : refl_mask_dirty) - for (int y = b.y0; y <= b.y1; ++y) - std::fill(refl_mask.begin() + y * W + b.x0, refl_mask.begin() + y * W + b.x1 + 1, 0); - refl_mask_dirty.clear(); - for (const Rect &b : owner_dirty) - for (int y = b.y0; y <= b.y1; ++y) - std::fill(owner.begin() + y * W + b.x0, owner.begin() + y * W + b.x1 + 1, BRAGG_OWNER_NONE); - owner_dirty.clear(); - - refl_mask.resize(npixel, 0); + refl_mask.Clear(); + owner.Clear(); for (size_t i = 0; i < npredicted; ++i) { const auto &r = predicted[i]; const BraggStencil st = MakeBraggStencil(r.predicted_x, r.predicted_y, stencil); @@ -89,12 +82,10 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, const int x1 = std::min(W - 1, static_cast(std::ceil(r.predicted_x + st.ex_in + 1.0f))); const int y0 = std::max(0, static_cast(std::floor(r.predicted_y - st.ey_in - 1.0f))); const int y1 = std::min(H - 1, static_cast(std::ceil(r.predicted_y + st.ey_in + 1.0f))); - if (x0 <= x1 && y0 <= y1) - refl_mask_dirty.push_back({x0, x1, y0, y1}); for (int y = y0; y <= y1; ++y) for (int x = x0; x <= x1; ++x) { const auto d = BraggStencilDistances(st, x - r.predicted_x, y - r.predicted_y); - if (d.inner < r2_sq) refl_mask[y * W + x] = 1; + if (d.inner < r2_sq) refl_mask.At(x, y) = 1; } } @@ -106,7 +97,6 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, const bool exclude = overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum; const bool use_owner = overlap != OverlapMode::Off; if (use_owner) { - owner.resize(npixel, BRAGG_OWNER_NONE); const float claim_sq = claim * claim; for (size_t i = 0; i < npredicted; ++i) { const auto &r = predicted[i]; @@ -114,21 +104,19 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, const int x1 = std::min(W - 1, static_cast(std::ceil(r.predicted_x + claim))); const int y0 = std::max(0, static_cast(std::floor(r.predicted_y - claim))); const int y1 = std::min(H - 1, static_cast(std::ceil(r.predicted_y + claim))); - if (x0 <= x1 && y0 <= y1) - owner_dirty.push_back({x0, x1, y0, y1}); for (int y = y0; y <= y1; ++y) for (int x = x0; x <= x1; ++x) { const float dx = x - r.predicted_x, dy = y - r.predicted_y; const float d2 = dx * dx + dy * dy; if (d2 >= claim_sq) continue; - uint32_t &o = owner[y * W + x]; + uint32_t &o = owner.At(x, y); o = std::min(o, BraggOwnerKey(std::sqrt(d2), inv_claim, static_cast(i))); } } } auto clean = [&](int x, int y, size_t i) { return !use_owner - || BraggOwnedBy(owner[static_cast(y) * W + x], static_cast(i)); + || BraggOwnedBy(owner.Get(x, y), static_cast(i)); }; // --- Pass A: box-sum every reflection (rough I, background, centroid, strong flag). --- @@ -213,7 +201,7 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, y_sum += y; ++n_inner_valid; } else if (d.inner >= r2_sq && d.outer < r3_sq) { - if (refl_mask[y * W + x]) { ++n_bkg_neighbour; continue; } + if (refl_mask.Get(x, y)) { ++n_bkg_neighbour; continue; } if (!valid(px)) continue; bkg_sum += static_cast(px); if (bkg_trim_frac > 0.0) bkg_vals.push_back(px); @@ -267,7 +255,7 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, for (int x = x0; x <= x1; ++x) { const auto d = BraggStencilDistances(st, x - r.predicted_x, y - r.predicted_y); if (!(d.inner >= r2_sq && d.outer < r3_sq)) continue; - if (refl_mask[y * W + x]) continue; + if (refl_mask.Get(x, y)) continue; const int32_t px = img[y * W + x]; if (!valid(px)) continue; if (static_cast(px) <= thr) { diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h index ed1861599..f125f391b 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h @@ -14,15 +14,44 @@ class BraggIntegrationEngineCPU : public BraggIntegrationEngine { // Core integrator, templated on a pixel sampler so it reads either the preprocessed int32 buffer // or a raw CompressedImage of any pixel type - both presented per-pixel in the INT32_MIN(masked)/ // INT32_MAX(saturated) convention - without ever materialising a second full-image copy. - // Full-frame scratch of RunImpl: the reflection mask and the signal-region owner map. Kept from - // call to call rather than allocated per image - at 18 Mpixel that was 90 MB, a fresh mmap and a - // page fault per 4 kB every call - and put back to their empty state one rectangle at a time: the - // rectangles a call writes are recorded, and the next call clears exactly those before it starts. - struct Rect { int x0, x1, y0, y1; }; - std::vector refl_mask; - std::vector owner; - std::vector refl_mask_dirty; - std::vector owner_dirty; + // Full-frame scratch of RunImpl: the reflection mask and the signal-region owner map. A call writes + // only around its reflections, so these keep only the 16x16-pixel tiles written since the last + // Clear(): a frame-sized array was mostly never read, yet over a sweep every page of it got + // touched, in every worker. Reading a tile nothing wrote gives `empty`. + template + class TiledFrame { + static constexpr int TILE = 16; + int tiles_x; + T empty; + std::vector tile_start; // per tile: where it starts in `pixels`, -1 = not written + std::vector written; // the tiles written, for Clear() + std::vector pixels; + public: + TiledFrame(int width, int height, T empty) + : tiles_x((width + TILE - 1) / TILE), empty(empty), + tile_start(static_cast(tiles_x) * ((height + TILE - 1) / TILE), -1) {} + T Get(int x, int y) const { + const int32_t start = tile_start[(y / TILE) * tiles_x + x / TILE]; + return start < 0 ? empty : pixels[start + (y % TILE) * TILE + x % TILE]; + } + T &At(int x, int y) { + const int t = (y / TILE) * tiles_x + x / TILE; + if (tile_start[t] < 0) { + tile_start[t] = static_cast(pixels.size()); + pixels.resize(pixels.size() + TILE * TILE, empty); + written.push_back(t); + } + return pixels[tile_start[t] + (y % TILE) * TILE + x % TILE]; + } + void Clear() { + for (int t : written) + tile_start[t] = -1; + written.clear(); + pixels.clear(); + } + }; + TiledFrame refl_mask; + TiledFrame owner; template std::vector RunImpl(const Sampler &img, const std::vector &predicted, From 8e6ecedd5c5762ce7cd6f30767555d65810bd629 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 28 Sep 2026 03:57:10 +0200 Subject: [PATCH 5/5] CPU pixel loops: keep the ring sums and the vertical slide out of memory for real Two fixes found in the post-merge profile, where the previous commit's loops did not do what they were written to do: - AccumulateRingsBlock: the ring_overflow.emplace_back call inside the loop left no callee-saved registers for the locals, so GCC kept sum/sum2/count on the stack and the store-reload chain was still there (29% of the function's samples on one stack add). The azint sums get a loop of their own over the block, and the out-of-histogram values are listed by a second loop run only when the block has any, in the same pixel order; now the sums live in registers. Same additions in the same order. - DetectPass: inlined into DetectPass's main loop the vertical slide was not vectorised (the vectorised copies GCC made were for the other call sites). It is now a free function, SlideVertical, which GCC vectorises on its own. Measured (perf, 75 s of myob, CPU-only, relative to the unchanged FlagRow/AnalyzeBlock): AccumulateRingsBlock -4..-11%, DetectPass + SlideVertical -7..-14%. Byte-identical p.hkl, p.mtz, p_P1.mtz, p_unmerged.mtz on myob, cytc, lyso, sparse (CPU) and myob (GPU); ImageSpotFinderCPU*, AdaptiveSpotFinder, SpotFinding, AzimuthalIntegration and portable tests pass. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C --- .../spot_finding/AdaptiveSpotFinderCPU.cpp | 89 +++++++++++++------ .../spot_finding/ImageSpotFinderCPU.cpp | 38 ++++---- 2 files changed, 83 insertions(+), 44 deletions(-) diff --git a/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp b/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp index 4a8107b22..62474e7e3 100644 --- a/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp +++ b/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp @@ -60,48 +60,65 @@ void AdaptiveSpotFinderCPU::AccumulateRingsBlock(const ImagePreprocessorBuffer & const size_t nbins = ring_sum.size(); const float *corrections = mapping.Corrections().data(); - // Consecutive pixels mostly share a ring, so the ring's sums are held in locals while they do and + // Consecutive pixels mostly share a ring, so a ring's sums are held in locals while they do and // written back when the ring changes: the same additions in the same order, without a store and a - // reload of the same address on every pixel. - uint32_t *hist = ring_hist.data(); - size_t cur = nbins; // the ring held in the locals below; nbins = none - int64_t sum = 0, cnt = 0; - uint64_t sum2 = 0; - float az_sum = 0.0f, az_sum2 = 0.0f; - uint32_t az_cnt = 0; - const auto store = [&] { - if (cur == nbins) - return; - ring_sum[cur] = sum; - ring_sum2[cur] = sum2; - ring_cnt[cur] = cnt; - if (fuse_azint) { + // reload of the same address on every pixel. The azimuthal-integration sums get a loop of their own + // over the block, so that neither loop runs out of registers for its sums. + if (fuse_azint) { + size_t cur = nbins; // the ring held in the locals below; nbins = none + float az_sum = 0.0f, az_sum2 = 0.0f; + uint32_t az_cnt = 0; + for (size_t pxl = first; pxl < first + n; ++pxl) { + const int32_t v = image[pxl]; + if (v == INT32_MIN || v == INT32_MAX) continue; // bad / saturated + const uint16_t b = pixel_to_bin[pxl]; + if (b >= nbins) continue; // masked / out of range (UINT16_MAX) + if (b != cur) { + if (cur != nbins) { + azint_sum[cur] = az_sum; + azint_sum2[cur] = az_sum2; + azint_count[cur] = az_cnt; + } + cur = b; + az_sum = azint_sum[b]; + az_sum2 = azint_sum2[b]; + az_cnt = azint_count[b]; + } + const float val = static_cast(v) * corrections[pxl]; + const float val_sq = val * val; + az_sum += val; + az_sum2 += val_sq; + ++az_cnt; + } + if (cur != nbins) { azint_sum[cur] = az_sum; azint_sum2[cur] = az_sum2; azint_count[cur] = az_cnt; } - }; + } + + // Values outside the histogram are listed by a second loop, run only when the block has any: a + // call in this loop would leave the sums in memory again. + uint32_t *hist = ring_hist.data(); + bool overflow = false; + size_t cur = nbins; + int64_t sum = 0, cnt = 0; + uint64_t sum2 = 0; for (size_t pxl = first; pxl < first + n; ++pxl) { const int32_t v = image[pxl]; if (v == INT32_MIN || v == INT32_MAX) continue; // bad / saturated const uint16_t b = pixel_to_bin[pxl]; if (b >= nbins) continue; // masked / out of range (UINT16_MAX) if (b != cur) { - store(); + if (cur != nbins) { + ring_sum[cur] = sum; + ring_sum2[cur] = sum2; + ring_cnt[cur] = cnt; + } cur = b; sum = ring_sum[b]; sum2 = ring_sum2[b]; cnt = ring_cnt[b]; - az_sum = azint_sum[b]; - az_sum2 = azint_sum2[b]; - az_cnt = azint_count[b]; - } - if (fuse_azint) { - const float val = static_cast(v) * corrections[pxl]; - const float val_sq = val * val; - az_sum += val; - az_sum2 += val_sq; - ++az_cnt; } sum += v; sum2 += static_cast(static_cast(v) * v); @@ -109,9 +126,23 @@ void AdaptiveSpotFinderCPU::AccumulateRingsBlock(const ImagePreprocessorBuffer & if (v >= 0 && v < HIST_VALUES) hist[b * HIST_VALUES + v] += 1; else - ring_overflow.emplace_back(b, v); + overflow = true; } - store(); + if (cur != nbins) { + ring_sum[cur] = sum; + ring_sum2[cur] = sum2; + ring_cnt[cur] = cnt; + } + + if (overflow) + for (size_t pxl = first; pxl < first + n; ++pxl) { + const int32_t v = image[pxl]; + if (v == INT32_MIN || v == INT32_MAX) continue; + const uint16_t b = pixel_to_bin[pxl]; + if (b >= nbins) continue; + if (v < 0 || v >= HIST_VALUES) + ring_overflow.emplace_back(b, v); + } } // A sigma-clip pass over the plain pass's values: each distinct value of a ring meets the same test the diff --git a/image_analysis/spot_finding/ImageSpotFinderCPU.cpp b/image_analysis/spot_finding/ImageSpotFinderCPU.cpp index ea7549f2f..1cb6773f9 100644 --- a/image_analysis/spot_finding/ImageSpotFinderCPU.cpp +++ b/image_analysis/spot_finding/ImageSpotFinderCPU.cpp @@ -45,6 +45,24 @@ bool StrongInWindow(int64_t pxl_val, int64_t sum, int64_t sum2, int64_t valid, // pixel is above SNR threshold } +// Adds row `in` to the vertical sums and takes row `out` out of them; nullptr is no row. A bad or +// saturated pixel adds 0. Written without branches, and as a function of its own so that the loop +// vectorises whatever it is inlined into; integer sums, so the same totals as pixel by pixel. +void SlideVertical(const int32_t *in, const int32_t *out, int32_t width, + int64_t *sum, int64_t *sum2, uint16_t *valid) { + for (int32_t col = 0; col < width; col++) { + const int64_t a = in ? in[col] : INT32_MIN; + const int64_t r = out ? out[col] : INT32_MIN; + const bool a_ok = a != INT32_MAX && a != INT32_MIN; + const bool r_ok = r != INT32_MAX && r != INT32_MIN; + const int64_t a_v = a_ok ? a : 0; + const int64_t r_v = r_ok ? r : 0; + sum[col] += a_v - r_v; + sum2[col] += a_v * a_v - r_v * r_v; + valid[col] += static_cast(a_ok - r_ok); + } +} + } // namespace void ImageSpotFinderCPU::DetectAt(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings, @@ -164,27 +182,17 @@ void ImageSpotFinderCPU::DetectPass(const ImagePreprocessorBuffer &image, std::vector valid_vert(width, 0); // Add the row entering the window (line_in) and take out the one leaving it (line_out); -1 is - // no row. A bad or saturated pixel adds 0. Written without branches, so the loop vectorises. - // A pixel strong in the previous pass is not counted either; it is added here with the rest - // of its row and taken out again below - there are few, and the sums are integers, so the - // totals are the same as never adding it. + // no row. A pixel strong in the previous pass is not counted either; SlideVertical adds it with + // the rest of its row and it is taken out again below - there are few, and the sums are + // integers, so the totals are the same as never adding it. const int32_t *img = image.data(); const int32_t w = width; int64_t *sv = sum_vert.data(); int64_t *sv2 = sum2_vert.data(); uint16_t *vv = valid_vert.data(); auto slide_vert = [&](int line_in, int line_out) { - for (int32_t col = 0; col < w; col++) { - const int64_t in = line_in >= 0 ? img[line_in * w + col] : INT32_MIN; - const int64_t out = line_out >= 0 ? img[line_out * w + col] : INT32_MIN; - const bool in_ok = in != INT32_MAX && in != INT32_MIN; - const bool out_ok = out != INT32_MAX && out != INT32_MIN; - const int64_t in_v = in_ok ? in : 0; - const int64_t out_v = out_ok ? out : 0; - sv[col] += in_v - out_v; - sv2[col] += in_v * in_v - out_v * out_v; - vv[col] += static_cast(in_ok - out_ok); - } + SlideVertical(line_in >= 0 ? img + static_cast(line_in) * w : nullptr, + line_out >= 0 ? img + static_cast(line_out) * w : nullptr, w, sv, sv2, vv); if (!prev_strong) return; // sign = +1 takes out what line_in added, -1 puts back what line_out took out.