diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index e7d7c2ff..fe37608c 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only +#include "../../common/ParallelFor.h" #include "RotationScaleMerge.h" #include @@ -138,48 +139,6 @@ namespace { return G; } - // Run fn(i) for i in [0, n) over `nthreads` workers pulling from a shared atomic counter - the same - // self-load-balancing pattern the rest of the codebase uses (heavy frames don't stall light ones). - // Work-stealing per-item parallel: one atomic fetch per item. Use ONLY when the per-item work is - // heavy and uneven (e.g. per-frame fits) - the atomic amortises. For millions of tiny uniform items - // use ParallelChunks instead; a per-item atomic there is pure contention. - template - void ParallelFor(int n, size_t nthreads, Fn fn) { - if (n <= 0) return; - if (nthreads <= 1 || n == 1) { - for (int i = 0; i < n; ++i) fn(i); - return; - } - const size_t local = std::min(nthreads, static_cast(n)); - std::atomic next = 0; - std::vector> futures; - futures.reserve(local); - for (size_t t = 0; t < local; ++t) - futures.emplace_back(std::async(std::launch::async, [&] { - for (int i = next.fetch_add(1); i < n; i = next.fetch_add(1)) - fn(i); - })); - for (auto &f : futures) f.get(); - } - - // Chunked parallel: each worker gets one contiguous [lo, hi) range, no per-item synchronisation. - // Right for millions of cheap uniform items (the CPU stand-in for a flat CUDA grid-stride kernel). - template - void ParallelChunks(int n, size_t nthreads, Fn fn) { - if (n <= 0) return; - const int nt = static_cast(std::max(1, std::min(nthreads, static_cast(n)))); - if (nt == 1) { fn(0, n); return; } - const int chunk = (n + nt - 1) / nt; - std::vector> futures; - futures.reserve(nt); - for (int t = 0; t < nt; ++t) { - const int lo = t * chunk, hi = std::min(n, lo + chunk); - if (lo >= hi) break; - futures.emplace_back(std::async(std::launch::async, [&fn, lo, hi] { fn(lo, hi); })); - } - for (auto &f : futures) f.get(); - } - double median_of(std::vector &v) { std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end()); return v[v.size() / 2]; @@ -319,10 +278,17 @@ void RotationScaleMerge::Ingest() { // Sort ONCE by (raw h,k,l, image_number) and split into raw-hkl runs. This is the one expensive sort; // both the 3D combine (event split) and the per-space-group ASU grouping reuse this order. - // Sorting an index array whose comparator dereferences the 72-byte Obs is a cache miss per - // comparison over a multi-GB array, so sort a packed copy of the key fields instead. Started from - // the identity order with the same comparisons in the same order, std::sort takes exactly the same - // branches and produces exactly the same permutation. + // Sorting an index array whose comparator dereferences the 80-byte Obs is a cache miss per + // comparison over a multi-GB array, so sort a packed copy of the key fields instead. + // + // The observation's own index is the last key, which makes the order TOTAL. Two observations can + // genuinely share (h,k,l) and image_number: the predictor emits both intersections of a + // reflection's rotation circle with the Ewald sphere, and near the blind region the two are close + // enough in angle that both land on the same frame. Without a final key their relative order was + // whatever the sort happened to produce, and the combine reads it - it takes on_ice from the first + // member of an event, and its per-event sums are floating point. Ordering them by arrival makes + // the result independent of the sort algorithm, which is what lets the order be reproduced by a + // faster one. struct SortKey { int32_t h, k, l; float image_number; int32_t idx; }; perm.resize(partials.size()); rawrun_start.clear(); rawrun_count.clear(); @@ -337,7 +303,8 @@ void RotationScaleMerge::Ingest() { if (a.h != b.h) return a.h < b.h; if (a.k != b.k) return a.k < b.k; if (a.l != b.l) return a.l < b.l; - return a.image_number < b.image_number; + if (a.image_number != b.image_number) return a.image_number < b.image_number; + return a.idx < b.idx; }); for (size_t i = 0; i < keys.size(); ++i) perm[i] = keys[i].idx; for (int i = 0; i < static_cast(keys.size()); ) { @@ -377,13 +344,18 @@ void RotationScaleMerge::Ingest() { std::vector onice(n); std::vector frm(n); std::vector vbkg(n); - for (int i = 0; i < n; ++i) { - const auto &o = partials[i]; - I[i] = o.I; sigma[i] = o.sigma; rlp[i] = o.rlp; part[i] = o.partiality; - zeta[i] = o.zeta; onice[i] = o.on_ice; frm[i] = o.frame; corr[i] = o.corr; - bkg[i] = o.bkg; vbkg[i] = o.var_bkg; img[i] = o.image_number; dd[i] = o.d; - px[i] = o.px; py[i] = o.py; - } + // Each observation writes its own slot in each of the fourteen arrays, so this splits + // straight over the range. It is the widest pass in Ingest - it reads the whole 80-byte + // struct and writes 57 bytes of it back out - and it ran on one thread. + ParallelChunks(n, nthreads, [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) { + const auto &o = partials[i]; + I[i] = o.I; sigma[i] = o.sigma; rlp[i] = o.rlp; part[i] = o.partiality; + zeta[i] = o.zeta; onice[i] = o.on_ice; frm[i] = o.frame; corr[i] = o.corr; + bkg[i] = o.bkg; vbkg[i] = o.var_bkg; img[i] = o.image_number; dd[i] = o.d; + px[i] = o.px; py[i] = o.py; + } + }); gpu_->SetPartials(n, n_frames, I.data(), sigma.data(), rlp.data(), part.data(), zeta.data(), onice.data(), frm.data(), corr.data(), frame_start.data(), frame_count.data()); gpu_->SetCombineInputs(bkg.data(), vbkg.data(), img.data(), dd.data(), px.data(), py.data()); diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 28ea5c52..eba7b7b2 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -70,9 +70,16 @@ namespace { const float *__restrict__ partiality, const float *__restrict__ rlp, const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice, const double *__restrict__ group_mean, + const float *__restrict__ sigma, double *__restrict__ inv_sigma, float *__restrict__ sco_coeff, uint8_t *__restrict__ sco_ok) { const int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= n_obs) return; + // sigma never changes once it is uploaded, so its reciprocal is the same in every one of the + // thirty IRLS iterations and in every scaling pass. It was being recomputed - a full 64-bit + // reciprocal, five refinement steps - for every observation of every iteration. The CPU has + // always hoisted it (ScaleObs::weight); this is the GPU catching up. Same expression on the + // same operand, so the value is the one the loop used to compute, bit for bit. + inv_sigma[i] = SafeInvD(sigma[i], 1.0); const int g = group[i]; bool ok = (g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f; double mean = 0.0; @@ -92,7 +99,7 @@ namespace { __global__ void FitPerFrameGKernel(int n_frames, double robust_k, const int32_t *__restrict__ frame_start, const int32_t *__restrict__ frame_count, - const float *__restrict__ I, const float *__restrict__ sigma, + const float *__restrict__ I, const double *__restrict__ inv_sigma, const float *__restrict__ sco_coeff, const uint8_t *__restrict__ sco_ok, const int32_t *__restrict__ perm, double *__restrict__ g, uint8_t *__restrict__ scaled) { @@ -119,7 +126,7 @@ namespace { const int a = perm ? perm[i] : i; if (!sco_ok[a]) continue; const double coeff = sco_coeff[a]; - const double w = SafeInvD(sigma[a], 1.0); + const double w = inv_sigma[a]; const double w2 = w * w; num += w2 * coeff * double(I[a]); den += w2 * coeff * coeff; @@ -141,7 +148,7 @@ namespace { const int a = perm ? perm[i] : i; if (!sco_ok[a]) continue; const double coeff = sco_coeff[a]; - const double w = SafeInvD(sigma[a], 1.0); + const double w = inv_sigma[a]; const double w2 = w * w; const double res = w * (G * coeff - double(I[a])); const double rw = 1.0 / (1.0 + res * res / k2); @@ -610,6 +617,7 @@ struct RotationScaleMergeGPU::Impl { // scratch CudaDevicePtr group_mean, g; CudaDevicePtr scaled; + CudaDevicePtr inv_sigma; // 1/sigma, hoisted out of the IRLS loop (sigma never changes) CudaDevicePtr sco_coeff; CudaDevicePtr sco_ok; CudaDevicePtr cc; // per-frame CC (diagnostic), length n_frames @@ -645,6 +653,7 @@ struct RotationScaleMergeGPU::Impl { // yield coeff=mean, plus the working corr, the per-obs scale scratch, and the fulls frame/group CSRs // (built on the host from the small f_frame/f_group key arrays, over the emit-ordered fulls). CudaDevicePtr f_corr, f_partiality, f_rlp, f_zeta, f_sco_coeff; + CudaDevicePtr f_inv_sigma; CudaDevicePtr f_sco_ok; CudaDevicePtr f_frame_perm, f_frame_start, f_frame_count; CudaDevicePtr f_gperm, f_gstart, f_gcount; @@ -706,6 +715,7 @@ void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames, Upload(d.frame_start, frame_start, n_frames); Upload(d.frame_count, frame_count, n_frames); d.g = CudaDevicePtr(n_frames); d.scaled = CudaDevicePtr(n_frames); + d.inv_sigma = CudaDevicePtr(n_obs); d.sco_coeff = CudaDevicePtr(n_obs); d.sco_ok = CudaDevicePtr(n_obs); d.cc = CudaDevicePtr(std::max(1, n_frames)); @@ -745,9 +755,10 @@ void RotationScaleMergeGPU::ScalePartials(int iters, double robust_k, double min d.group_perm.get(), d.group_start.get(), d.group_count.get(), d.I.get(), d.sigma.get(), d.partiality.get(), d.corr.get(), d.group_mean.get()); PrepScaleObsKernel<<>>(d.n_obs, d.group.get(), d.partiality.get(), d.rlp.get(), - d.zeta.get(), d.on_ice.get(), d.group_mean.get(), d.sco_coeff.get(), d.sco_ok.get()); + d.zeta.get(), d.on_ice.get(), d.group_mean.get(), d.sigma.get(), d.inv_sigma.get(), + d.sco_coeff.get(), d.sco_ok.get()); FitPerFrameGKernel<<>>(d.n_frames, robust_k, d.frame_start.get(), d.frame_count.get(), - d.I.get(), d.sigma.get(), d.sco_coeff.get(), d.sco_ok.get(), nullptr, d.g.get(), d.scaled.get()); + d.I.get(), d.inv_sigma.get(), d.sco_coeff.get(), d.sco_ok.get(), nullptr, d.g.get(), d.scaled.get()); UpdateCorrKernel<<>>(d.n_obs, d.frame.get(), d.rlp.get(), d.partiality.get(), d.g.get(), d.scaled.get(), d.corr.get()); } @@ -1006,6 +1017,7 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti d.f_on_ice = CudaDevicePtr(nf); d.f_corr = CudaDevicePtr(nf); d.f_partiality = CudaDevicePtr(nf); d.f_rlp = CudaDevicePtr(nf); d.f_zeta = CudaDevicePtr(nf); + d.f_inv_sigma = CudaDevicePtr(nf); d.f_sco_coeff = CudaDevicePtr(nf); d.f_sco_ok = CudaDevicePtr(nf); CudaCheck(cudaMemcpy(d.rr_offset.get(), offset.data(), size_t(d.n_runs) * sizeof(int32_t), cudaMemcpyHostToDevice), "upload offset"); @@ -1091,10 +1103,14 @@ void RotationScaleMergeGPU::ScaleFulls(int iters, double robust_k, double min_pa ReduceGroupMeansKernel<<>>(d.n_groups, min_partiality, d.f_gperm.get(), d.f_gstart.get(), d.f_gcount.get(), d.f_I.get(), d.f_sigma.get(), d.f_partiality.get(), d.f_corr.get(), d.group_mean.get()); - PrepScaleObsKernel<<>>(nf, d.f_group.get(), d.f_partiality.get(), d.f_rlp.get(), - d.f_zeta.get(), d.f_on_ice.get(), d.group_mean.get(), d.f_sco_coeff.get(), d.f_sco_ok.get()); + // Not grid-stride, so its grid has to cover every full - unlike the grid-stride kernels + // below, which the 65535 cap is there for. Capped, it would silently leave the tail of + // sco_coeff/sco_ok stale above 16.8M fulls. + PrepScaleObsKernel<<<(nf + BLK - 1) / BLK, BLK>>>(nf, d.f_group.get(), d.f_partiality.get(), + d.f_rlp.get(), d.f_zeta.get(), d.f_on_ice.get(), d.group_mean.get(), + d.f_sigma.get(), d.f_inv_sigma.get(), d.f_sco_coeff.get(), d.f_sco_ok.get()); FitPerFrameGKernel<<>>(d.n_frames, robust_k, - d.f_frame_start.get(), d.f_frame_count.get(), d.f_I.get(), d.f_sigma.get(), + d.f_frame_start.get(), d.f_frame_count.get(), d.f_I.get(), d.f_inv_sigma.get(), d.f_sco_coeff.get(), d.f_sco_ok.get(), d.f_frame_perm.get(), d.g.get(), d.scaled.get()); UpdateCorrKernel<<>>(nf, d.f_frame.get(), d.f_rlp.get(), d.f_partiality.get(), d.g.get(), d.scaled.get(), d.f_corr.get());