From 08b318f9398d7ab82260b0936bde44ede1f8fd28 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 11:41:15 +0200 Subject: [PATCH 1/4] RotationScaleMerge: build the correction surfaces' terms on the device With the fulls resident (the GPU combine ran), ApplyCellSurface no longer builds its terms on the host and uploads them: SurfaceBuildTerms makes them from the resident fulls, this surface's cell per full and the current corr (the only two uploads) - which fulls give a term, the terms in fulls order, the group CSR over them and each subset's per-(block, cell) segments - and SurfaceShellWeights the resolution gauge's inputs: the shell edges (order statistics of 1/d^2 read off a sort), each group's shell and each (cell, shell) bin's weight. Every order the host made is a stable counting sort, which a stable radix sort of the indices reproduces; the per-bin weights are summed one bin per thread in fulls order with the host's rounding. The existing exact surface kernels read them as before. Without resident fulls (the observation dump) the host path runs. p.mtz md5 and the surfaces' log lines unchanged on myob, cytc, thau, kdp, 8a1a and 8qaw (GPU build). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- .../scale_merge/RotationScaleMerge.cpp | 87 +++-- .../scale_merge/RotationScaleMerge.h | 6 + .../scale_merge/RotationScaleMergeGPU.cu | 307 ++++++++++++++++++ .../scale_merge/RotationScaleMergeGPU.h | 15 + 4 files changed, 367 insertions(+), 48 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index bc164861c..5d4d49785 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -3677,7 +3677,26 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // Which terms each pass walks, as positions in `term`. The cross-validated halves then cost half a // pass rather than a whole one with a skip; the full fit walks the array in order. std::vector idx_all, idx_even, idx_odd; - { + // Every pass below runs on the device when the fulls are there: the terms, their group order and + // the subsets' blocks are then built there too (RotationScaleMergeGPU::SurfaceBuildTerms), from the + // resident fulls, this pass's cells and the current corr - the same arrays in the same orders. + // 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; + bool on_gpu = false; + int n_terms = 0; +#ifdef JFJOCH_USE_CUDA + on_gpu = gpu_active_ && fulls_on_device_; + if (on_gpu) { + surface_corr.resize(fulls.size()); + ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) surface_corr[i] = fulls[i].corr; + }); + n_terms = gpu_->SurfaceBuildTerms(cell.data(), surface_corr.data(), min_partiality, ncell, n_groups, + SURFACE_BLOCK); + } +#endif + if (!on_gpu) { // Count what each chunk keeps, and how much of that sits on an even frame, so the second pass // can write straight into its own span of every list. Chunk boundaries depend on nothing but // the length and the worker count, so the lists come out in fulls order on every run. @@ -3717,6 +3736,7 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int } } }); + n_terms = static_cast(term.size()); } auto subset = [&](int parity) -> const std::vector & { return parity < 0 ? idx_all : (parity ? idx_odd : idx_even); @@ -3725,10 +3745,7 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // Every pass below runs over one of these lists, ~22 times per surface, so the worker count is // set by how much data there is rather than by how many threads the machine has: a small dataset // on a large node would otherwise pay for 48 thread starts per pass to sum a few thousand terms. - const int nt = static_cast(ThreadsForWork(idx_all.size(), nthreads)); - // 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; + const int nt = static_cast(ThreadsForWork(n_terms, nthreads)); // --- The surface's RESOLUTION gauge --------------------------------------------------------- // Every symmetry equivalent of a reflection sits at one d, so a factor that is a function of d @@ -3766,14 +3783,18 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // the residual ramp inside 1.5%); 20 is fine enough to follow a curved drift and coarse enough // that each shell still holds tens of thousands of observations. constexpr int NSHELL = 20; - const int nshell = static_cast(term.size()) >= 100 * NSHELL ? NSHELL : 0; + const int nshell = n_terms >= 100 * NSHELL ? NSHELL : 0; // Weight of cell c in shell s: the statistical weight its observations carry there. Computed ONCE // from the incoming sigmas rather than per round, so the projection is a fixed linear operator and // the alternating fit below converges against it. std::vector shw(static_cast(ncell) * nshell, 0.0); std::vector shw_shell(nshell, 0.0), shw_cell(nshell > 0 ? ncell : 0, 0.0); std::vector group_shell(n_groups, 0); // the shell each ASU group sits in, for the gate - if (nshell > 0) { +#ifdef JFJOCH_USE_CUDA + if (on_gpu && nshell > 0) + gpu_->SurfaceShellWeights(n_terms, nshell, group_shell.data(), shw.data()); +#endif + if (!on_gpu && nshell > 0) { const int n_term = static_cast(term.size()); std::vector s2(n_term); ParallelChunks(n_term, nt, [&](int lo, int hi) { @@ -3798,13 +3819,13 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int if (!(sc > 0.0)) continue; shw[static_cast(t.cell) * nshell + s] += 1.0 / (sc * sc); } - for (int c = 0; c < ncell; ++c) - for (int s = 0; s < nshell; ++s) { - const double v = shw[static_cast(c) * nshell + s]; - shw_cell[c] += v; - shw_shell[s] += v; - } } + for (int c = 0; c < ncell && nshell > 0; ++c) + for (int s = 0; s < nshell; ++s) { + const double v = shw[static_cast(c) * nshell + s]; + shw_cell[c] += v; + shw_shell[s] += v; + } // The same terms once more, in ASU-GROUP order. The reference pass sums per group, so group order // is the only layout that lets a thread own a whole group's sum - and it still adds that group's @@ -3812,43 +3833,12 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector &cell, int // order). A separate copy rather than a permutation: chasing every term through an index at random // over a few hundred megabytes is the locality this pass used to stay serial to keep. struct RefTerm { float I, sigma, corr; int32_t cell; uint8_t parity; }; - std::vector g_start(n_groups + 1, 0); - for (const Term &t : term) ++g_start[t.group + 1]; - for (int g = 0; g < n_groups; ++g) g_start[g + 1] += g_start[g]; + std::vector g_start; std::vector gterm; - // With a GPU the two passes of every round run there (RotationScaleMergeGPU::Surface*), on the same - // terms in the same order: the device gets the terms with the group order as a permutation, and each - // subset cut into this fit's reduction blocks with every block's terms ordered by cell, so that one - // device thread forms what one block of the host loop sums into one cell. - bool on_gpu = false; -#ifdef JFJOCH_USE_CUDA - on_gpu = gpu_active_; - if (on_gpu) { - std::vector gperm(term.size()); - std::vector fill(g_start.begin(), g_start.end() - 1); - for (size_t k = 0; k < term.size(); ++k) - gperm[fill[term[k].group]++] = static_cast(k); - gpu_->SurfaceSetTerms(static_cast(term.size()), term.data(), term_parity.data(), - n_groups, gperm.data(), g_start.data(), ncell); - for (int parity : {0, 1, -1}) { - const std::vector &sel = subset(parity); - const int n = static_cast(sel.size()); - const int nb = ReductionBlocks(n, SURFACE_BLOCK); - std::vector perm(n), seg_start(static_cast(nb) * ncell + 1, n); - ParallelFor(nb, nt, [&](int b) { - const int lo = static_cast(static_cast(n) * b / nb); - const int hi = static_cast(static_cast(n) * (b + 1) / nb); - std::vector pos(ncell + 1, 0); - for (int k = lo; k < hi; ++k) ++pos[term[sel[k]].cell + 1]; - for (int c = 0; c < ncell; ++c) pos[c + 1] += pos[c]; - for (int c = 0; c < ncell; ++c) seg_start[static_cast(b) * ncell + c] = lo + pos[c]; - for (int k = lo; k < hi; ++k) perm[lo + pos[term[sel[k]].cell]++] = sel[k]; - }); - gpu_->SurfaceSetSubset(parity < 0 ? 2 : parity, nb, perm.data(), seg_start.data()); - } - } -#endif if (!on_gpu) { + g_start.assign(n_groups + 1, 0); + for (const Term &t : term) ++g_start[t.group + 1]; + for (int g = 0; g < n_groups; ++g) g_start[g + 1] += g_start[g]; gterm.resize(term.size()); std::vector fill(g_start.begin(), g_start.end() - 1); for (size_t k = 0; k < term.size(); ++k) { @@ -6638,6 +6628,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool scaled_fulls_on_gpu ? " + scale-fulls" : "", nf, fulls_dropped_overloaded); combined_on_gpu = true; } + fulls_on_device_ = combined_on_gpu; // The CPU combine is the one host stage left that reads corr off the partials, and it only runs for // the diagnostic observation dump - so the scatter across the fat struct happens here, where it is // wanted, instead of on every pass. diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 1e9f75483..a523ca256 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -484,6 +484,12 @@ private: // the CPU loops as the bit-parity fallback. Built in Ingest. std::unique_ptr gpu_; bool gpu_active_ = false; + // The fulls of this pass live on the device (the GPU combine ran), so the correction surfaces can + // build their terms there. Set by RunArm. + bool fulls_on_device_ = false; + // The fulls' current corr, gathered for the device term build of each surface (a member: one per + // surface, a few tens of megabytes each). + std::vector surface_corr; #endif bool write_back_per_frame_scale = true; // see SetWriteBackPerFrameScale int frame_limit = -1; // see SetFrameLimit diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index af682a1a6..004a99971 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -11,11 +11,13 @@ #include #include #include +#include #include #include "../indexing/CUDAMemHelpers.h" #include "../../common/CUDAWrapper.h" #include "../../common/JFJochException.h" +#include "../../common/ParallelFor.h" namespace { constexpr int BLK = 256; @@ -529,6 +531,118 @@ namespace { for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) p[i] = v; } + + // ===== correction-surface terms built on the device from the resident fulls ===== + // ApplyCellSurface's host build, step for step: which fulls give a term, the terms in fulls order, + // the group CSR over them, each subset's per-(block, cell) segments and the shells' weights. Every + // order the host makes is a stable counting sort, which a stable radix sort of the indices + // reproduces; the only arithmetic (1/d^2, the shell weights) is rounded as the host rounds it. + + // Per full: 1 if it gives a term (all, even frame, odd frame) - the host's `usable(o) && cell >= 0`. + // A full's partiality is 1 (the combine sets it so on both paths). + __global__ void SurfaceKeepKernel(int n, const int32_t *group, const float *corr, const uint8_t *on_ice, + const int32_t *frame, const int32_t *cell, double min_partiality, + int32_t *keep, int32_t *keep_even, int32_t *keep_odd) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const float c = corr[i]; + const bool k = group[i] >= 0 && c > 0.0f && isfinite(c) && 1.0f >= min_partiality && !on_ice[i] + && cell[i] >= 0; + const bool odd = frame[i] & 1; + keep[i] = k; keep_even[i] = k && !odd; keep_odd[i] = k && odd; + } + + // The term of each kept full at its place (the scans of the keep flags), its frame parity, and its + // place in the even / odd lists. + __global__ void SurfaceWriteTermKernel(int n, const int32_t *keep_at, const int32_t *even_at, + const int32_t *odd_at, const int32_t *last_flags, + const float *I, const float *sigma, const float *corr, const float *d, + const int32_t *cell, const int32_t *group, const int32_t *frame, + RotationScaleMergeGPU::SurfaceTerm *term, uint8_t *parity, + int32_t *idx_even, int32_t *idx_odd, uint32_t *group_key) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + // A full is kept when its exclusive-scan position steps up to the next one (the last full by + // its own flag). + const bool kept = i + 1 < n ? keep_at[i + 1] > keep_at[i] : last_flags[0] != 0; + if (!kept) return; + const int w = keep_at[i]; + term[w] = {I[i], sigma[i], corr[i], d[i], cell[i], group[i]}; + group_key[w] = static_cast(group[i]); + const uint8_t par = static_cast(frame[i] & 1); + parity[w] = par; + if (par) idx_odd[odd_at[i]] = w; + else idx_even[even_at[i]] = w; + } + + __global__ void IotaKernel(int n, int32_t *v) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) v[i] = i; + } + + // Counting into bins: integer adds, so the counts do not depend on their order. + __global__ void CountKeysKernel(int n, const uint32_t *key, int32_t *count) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) atomicAdd(&count[key[i]], 1); + } + + // The block-and-cell key of position k of a subset of n terms: block b holds positions + // [n b / nb, n (b + 1) / nb), as the host's ParallelFor over blocks cuts them. + __global__ void SurfaceSubsetKeyKernel(int n, int nb, int ncell, const int32_t *sel, + const RotationScaleMergeGPU::SurfaceTerm *term, uint32_t *key) { + const int k = blockIdx.x * blockDim.x + threadIdx.x; + if (k >= n) return; + auto lo = [n, nb](int b) { return static_cast(n) * b / nb; }; + int b = static_cast(static_cast(k) * nb / n); + while (b + 1 < nb && lo(b + 1) <= k) ++b; + while (b > 0 && lo(b) > k) --b; + key[k] = static_cast(b * ncell + term[sel[k]].cell); + } + + // 1/d^2 of each term, as the host's shells take it. + __global__ void SurfaceS2Kernel(int n, const RotationScaleMergeGPU::SurfaceTerm *term, float *s2) { + const int k = blockIdx.x * blockDim.x + threadIdx.x; + if (k >= n) return; + const float d = term[k].d; + s2[k] = d > 0.0f ? __fdiv_rn(1.0f, __fmul_rn(d, d)) : 0.0f; + } + + // The shell of each term (how many edges are at or below its s^2 - the host's upper_bound) and + // the (cell, shell) bin its weight goes to. + __global__ void SurfaceShellKernel(int n, const float *s2, const float *edge, int n_edge, int nshell, + const RotationScaleMergeGPU::SurfaceTerm *term, uint8_t *shell, + uint32_t *bin) { + const int k = blockIdx.x * blockDim.x + threadIdx.x; + if (k >= n) return; + int sh = 0; + while (sh < n_edge && edge[sh] <= s2[k]) ++sh; + shell[k] = static_cast(sh); + bin[k] = static_cast(term[k].cell * nshell + sh); + } + + // The shell of each group: that of its last term in fulls order (the host overwrites it term by term). + __global__ void SurfaceGroupShellKernel(int n_groups, const int32_t *gstart, const int32_t *gperm, + const uint8_t *shell, int32_t *group_shell) { + const int g = blockIdx.x * blockDim.x + threadIdx.x; + if (g >= n_groups) return; + group_shell[g] = gstart[g + 1] > gstart[g] ? shell[gperm[gstart[g + 1] - 1]] : 0; + } + + // One thread per (cell, shell) bin: the statistical weight of its terms, in fulls order. + __global__ void SurfaceShellWeightKernel(int n_bins, const int32_t *bin_start, const int32_t *order, + const RotationScaleMergeGPU::SurfaceTerm *term, double *shw) { + const int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= n_bins) return; + double w = 0.0; + for (int q = bin_start[b]; q < bin_start[b + 1]; ++q) { + const auto &t = term[order[q]]; + const double sc = __dmul_rn(double(t.sigma), double(t.corr)); + if (!(sc > 0.0)) continue; + w = __dadd_rn(w, 1.0 / __dmul_rn(sc, sc)); + } + shw[b] = w; + } + // ===== error-model + merge reductions over the resident, scaled fulls (mirror MergeAndStats) ===== struct MergeParams { @@ -1983,6 +2097,199 @@ void RotationScaleMergeGPU::SurfaceSetSubset(int subset, int n_blocks, const int } } +namespace { + // The positions, in order, of a stable radix sort of `value` by `key` (keys below 2^bits). + void StableSortByKey(cudaStream_t s, int n, const uint32_t *key, const int32_t *value, uint32_t *key_out, + int32_t *value_out, int bits, const char *what) { + size_t bytes = 0; + CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key, key_out, value, value_out, n, 0, bits, s), what); + CudaDevicePtr scratch(std::max(bytes, 1), ALLOC); + CudaCheck(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, key, key_out, value, value_out, n, 0, bits, s), + what); + } + + // start[b] = how many keys are below b, for b in [0, n_bins] (a counting sort's bin starts). + void BinStarts(cudaStream_t s, int n, const uint32_t *key, int n_bins, int32_t *start, const char *what) { + CudaCheck(cudaMemsetAsync(start, 0, size_t(n_bins + 1) * sizeof(int32_t), s), what); + if (n > 0) CountKeysKernel<<<(n + BLK - 1) / BLK, BLK, 0, s>>>(n, key, start); + CudaCheck(cudaGetLastError(), what); + size_t bytes = 0; + CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, start, start, n_bins + 1, s), what); + CudaDevicePtr scratch(std::max(bytes, 1), ALLOC); + CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, start, start, n_bins + 1, s), what); + } + + int BitsFor(int64_t n_keys) { + int bits = 1; + while ((int64_t(1) << bits) < n_keys) ++bits; + return bits; + } +} + +int RotationScaleMergeGPU::SurfaceBuildTerms(const int32_t *cell, const float *corr, double min_partiality, + int ncell, int n_groups, int block) { + DeviceGuard guard(impl_->device, impl_->available); + auto &d = *impl_; + const cudaStream_t st = impl_->s(); + const int nf = d.n_fulls; + d.s_ncell = ncell; + d.s_n_groups = n_groups; + for (int &nb : d.s_n_blocks) nb = 0; + for (int &ns : d.s_n_sel) ns = 0; + + // Which fulls give a term, and where each goes in the three lists. + int n_terms = 0, n_even = 0, n_odd = 0; + CudaDevicePtr idx_even, idx_odd; + CudaDevicePtr group_key; + { + CudaDevicePtr d_cell, keep, even, odd; + CudaDevicePtr d_corr; + d.Upload(d_cell, cell, nf); + d.Upload(d_corr, corr, nf); + keep = d.Alloc(std::max(1, nf)); even = d.Alloc(std::max(1, nf)); + odd = d.Alloc(std::max(1, nf)); + if (nf > 0) { + SurfaceKeepKernel<<<(nf + BLK - 1) / BLK, BLK, 0, st>>>(nf, d.f_group.get(), d_corr.get(), d.f_on_ice.get(), + d.f_frame.get(), d_cell.get(), min_partiality, keep.get(), even.get(), odd.get()); + CudaCheck(cudaGetLastError(), "surface keep launch"); + } + // The last full's flags, before the scans turn the flags into places. + int32_t last[3] = {0, 0, 0}; + auto d_last = d.Alloc(3); + if (nf > 0) { + CopyAndWait(&last[0], keep.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl keep"); + CopyAndWait(&last[1], even.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl keep"); + CopyAndWait(&last[2], odd.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl keep"); + size_t bytes = 0; + CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, keep.get(), keep.get(), nf, st), "surface scan"); + auto scratch = d.Alloc(std::max(bytes, 1)); + for (int32_t *v : {keep.get(), even.get(), odd.get()}) + CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, v, v, nf, st), "surface scan"); + int32_t at[3]; + CopyAndWait(&at[0], keep.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl count"); + CopyAndWait(&at[1], even.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl count"); + CopyAndWait(&at[2], odd.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl count"); + n_terms = at[0] + last[0]; n_even = at[1] + last[1]; n_odd = at[2] + last[2]; + } + CopyAndWait(d_last.get(), last, sizeof(last), cudaMemcpyHostToDevice, st, "ul last flags"); + d.s_term = d.Alloc(std::max(1, n_terms)); + d.s_parity = d.Alloc(std::max(1, n_terms)); + group_key = d.Alloc(std::max(1, n_terms)); + idx_even = d.Alloc(std::max(1, n_even)); + idx_odd = d.Alloc(std::max(1, n_odd)); + if (nf > 0) { + SurfaceWriteTermKernel<<<(nf + BLK - 1) / BLK, BLK, 0, st>>>(nf, keep.get(), even.get(), odd.get(), + d_last.get(), d.f_I.get(), d.f_sigma.get(), d_corr.get(), d.f_d.get(), d_cell.get(), d.f_group.get(), + d.f_frame.get(), d.s_term.get(), d.s_parity.get(), idx_even.get(), idx_odd.get(), group_key.get()); + CudaCheck(cudaGetLastError(), "surface term launch"); + } + CudaCheck(cudaStreamSynchronize(st), "surface terms sync"); + } + + // The group CSR over the terms: the terms of group g at gperm[gstart[g] .. gstart[g + 1]), in + // fulls order. + auto iota = d.Alloc(std::max(1, n_terms)); + if (n_terms > 0) IotaKernel<<<(n_terms + BLK - 1) / BLK, BLK, 0, st>>>(n_terms, iota.get()); + CudaCheck(cudaGetLastError(), "iota launch"); + d.s_gperm = d.Alloc(std::max(1, n_terms)); + d.s_gstart = d.Alloc(n_groups + 1); + { + auto key_sorted = d.Alloc(std::max(1, n_terms)); + if (n_terms > 0) + StableSortByKey(st, n_terms, group_key.get(), iota.get(), key_sorted.get(), d.s_gperm.get(), + BitsFor(n_groups), "surface group sort"); + BinStarts(st, n_terms, group_key.get(), n_groups, d.s_gstart.get(), "surface group starts"); + } + group_key = CudaDevicePtr(); + + // Each subset cut into the host's reduction blocks, every block's terms ordered by cell. + const int32_t *sel_of[3] = {idx_even.get(), idx_odd.get(), iota.get()}; + const int n_of[3] = {n_even, n_odd, n_terms}; + for (int subset = 0; subset < 3; ++subset) { + const int n = n_of[subset]; + const int nb = ReductionBlocks(n, block); + const int n_seg = nb * ncell; + d.s_n_sel[subset] = n; + d.s_n_blocks[subset] = nb; + d.s_perm[subset] = d.Alloc(std::max(1, n)); + d.s_seg_start[subset] = d.Alloc(n_seg + 1); + if (nb == 0) continue; + auto key = d.Alloc(n), key_sorted = d.Alloc(n); + SurfaceSubsetKeyKernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, nb, ncell, sel_of[subset], d.s_term.get(), + key.get()); + CudaCheck(cudaGetLastError(), "surface subset key launch"); + StableSortByKey(st, n, key.get(), sel_of[subset], key_sorted.get(), d.s_perm[subset].get(), BitsFor(n_seg), + "surface subset sort"); + BinStarts(st, n, key.get(), n_seg, d.s_seg_start[subset].get(), "surface subset starts"); + if (size_t(n_seg) > d.s_slots) { + d.s_tcross = d.Alloc(n_seg); + d.s_tref2 = d.Alloc(n_seg); + d.s_slots = n_seg; + } + } + CudaCheck(cudaStreamSynchronize(st), "surface subsets sync"); + d.s_A = d.Alloc(std::max(1, ncell)); + d.s_cross = d.Alloc(std::max(1, ncell)); + d.s_ref2 = d.Alloc(std::max(1, ncell)); + d.s_sw = d.Alloc(std::max(1, n_groups)); + d.s_swI = d.Alloc(std::max(1, n_groups)); + d.s_w_Is = d.Alloc(std::max(1, n_terms)); + d.s_w_Iref = d.Alloc(std::max(1, n_terms)); + d.s_Iref = d.Alloc(std::max(1, n_terms)); + return n_terms; +} + +void RotationScaleMergeGPU::SurfaceShellWeights(int n_terms, int nshell, int32_t *group_shell, double *shw) { + DeviceGuard guard(impl_->device, impl_->available); + auto &d = *impl_; + const cudaStream_t st = impl_->s(); + const int ncell = d.s_ncell, n_groups = d.s_n_groups; + const int n = n_terms; + // The edges: order statistics of s^2, read off a sort. + auto s2 = d.Alloc(std::max(1, n)), s2_sorted = d.Alloc(std::max(1, n)); + SurfaceS2Kernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, d.s_term.get(), s2.get()); + CudaCheck(cudaGetLastError(), "surface s2 launch"); + { + size_t bytes = 0; + CudaCheck(cub::DeviceRadixSort::SortKeys(nullptr, bytes, s2.get(), s2_sorted.get(), n, 0, 32, st), "s2 sort"); + auto scratch = d.Alloc(std::max(bytes, 1)); + CudaCheck(cub::DeviceRadixSort::SortKeys(scratch.get(), bytes, s2.get(), s2_sorted.get(), n, 0, 32, st), + "s2 sort"); + } + std::vector edge(nshell - 1); + for (int i = 1; i < nshell; ++i) + CopyAndWait(&edge[i - 1], s2_sorted.get() + size_t(n) * size_t(i) / size_t(nshell), sizeof(float), + cudaMemcpyDeviceToHost, st, "dl edge"); + s2_sorted = CudaDevicePtr(); + CudaDevicePtr d_edge; + d.Upload(d_edge, edge.data(), nshell - 1); + auto shell = d.Alloc(std::max(1, n)); + auto bin = d.Alloc(std::max(1, n)), bin_sorted = d.Alloc(std::max(1, n)); + SurfaceShellKernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, s2.get(), d_edge.get(), nshell - 1, nshell, + d.s_term.get(), shell.get(), bin.get()); + CudaCheck(cudaGetLastError(), "surface shell launch"); + auto g_shell = d.Alloc(std::max(1, n_groups)); + if (n_groups > 0) + SurfaceGroupShellKernel<<<(n_groups + BLK - 1) / BLK, BLK, 0, st>>>(n_groups, d.s_gstart.get(), + d.s_gperm.get(), shell.get(), g_shell.get()); + CudaCheck(cudaGetLastError(), "surface group shell launch"); + // The weights: the terms of each (cell, shell) bin in fulls order, summed one bin per thread. + const int n_bins = ncell * nshell; + auto iota = d.Alloc(std::max(1, n)), order = d.Alloc(std::max(1, n)); + auto bin_start = d.Alloc(n_bins + 1); + auto d_shw = d.Alloc(std::max(1, n_bins)); + IotaKernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, iota.get()); + CudaCheck(cudaGetLastError(), "iota launch"); + StableSortByKey(st, n, bin.get(), iota.get(), bin_sorted.get(), order.get(), BitsFor(n_bins), "surface bin sort"); + BinStarts(st, n, bin.get(), n_bins, bin_start.get(), "surface bin starts"); + SurfaceShellWeightKernel<<<(n_bins + BLK - 1) / BLK, BLK, 0, st>>>(n_bins, bin_start.get(), order.get(), + d.s_term.get(), d_shw.get()); + CudaCheck(cudaGetLastError(), "surface shell weight launch"); + CopyAndWait(shw, d_shw.get(), size_t(n_bins) * sizeof(double), cudaMemcpyDeviceToHost, st, "dl shw"); + CopyAndWait(group_shell, g_shell.get(), size_t(n_groups) * sizeof(int32_t), cudaMemcpyDeviceToHost, st, + "dl group shell"); +} + void RotationScaleMergeGPU::SurfaceReference(int parity, const double *A) { DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.h b/image_analysis/scale_merge/RotationScaleMergeGPU.h index 68c4de7ab..c98370e05 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -233,6 +233,21 @@ public: void SurfaceSetTerms(int n_terms, const SurfaceTerm *term, const uint8_t *parity, int n_groups, const int32_t *gperm, const int32_t *gstart, int ncell); + // The same terms, group CSR and three subsets, built on the device from the resident fulls instead of + // uploaded: cell (length n_fulls, < 0 = no cell) and corr (the host's current corr, length n_fulls) + // are the only inputs from the host. A full gives a term where it has a group, a positive finite + // corr, is off the ice rings and has a cell (ApplyCellSurface's filter); `block` is the host fit's + // reduction block. The same arrays, in the same orders, as SurfaceSetTerms + SurfaceSetSubset are + // handed by ApplyCellSurface's host build. Returns the number of terms. Requires Combine first. + int SurfaceBuildTerms(const int32_t *cell, const float *corr, double min_partiality, int ncell, int n_groups, + int block); + + // The resolution gauge's inputs for the terms SurfaceBuildTerms left: each group's shell (that of its + // last term in fulls order; length n_groups) and each (cell, shell) bin's statistical weight + // sum 1/(sigma corr)^2 in fulls order (length ncell * nshell), over nshell equal-occupancy shells in + // 1/d^2 - ApplyCellSurface's numbers to the bit. + void SurfaceShellWeights(int n_terms, int nshell, int32_t *group_shell, double *shw); + // One subset of the terms (0 = even frames, 1 = odd, 2 = all), cut into the host's n_blocks // reduction blocks and ordered within each block by cell, keeping term order inside a cell: // block b, cell c is perm[seg_start[b * ncell + c], seg_start[b * ncell + c + 1]). From 76509be7278296fc75b1a76bdfddc800a8a38fb5 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 12:37:29 +0200 Subject: [PATCH 2/4] RotationScaleMerge: MergeAndStats' per-pass arrays as members, reset on all threads The rejected flags and the observation keys and halves (per full), the merged intensities (per group) and the Wilson observation records were locals: on a large cell every one of the run's merges allocated, page-faulted and zeroed them again on one thread - hundreds of megabytes per merge on 8qaw. They are now members, only resized, and reset with ParallelChunks where a pass needs them reset. The error-model samples are no longer cleared before the device download overwrites them (the host path still clears before it appends). p.mtz md5 unchanged on myob, cytc, thau, kdp, 8a1a and 8qaw (GPU build). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- .../scale_merge/RotationScaleMerge.cpp | 29 +++++++++++++------ .../scale_merge/RotationScaleMerge.h | 10 +++++++ 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 5d4d49785..b424c990d 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -4696,7 +4696,6 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool bool error_model_b_unmeasured = false; // b had no leverage; ISa is not reported bool error_model_b_resolved = true; // b^2 stood clear of its own standard error std::vector &samples = em_samples; - samples.clear(); // A full's counting variance at intensity I_for_b. Its own sigma carries its own Poisson fluctuation: // a full that came out low has a smaller sigma, so 1/sigma^2 weights it up and the merged mean drifts // below . Rebuild the variance at the reflection's EXPECTED intensity - var_bkg + var_per_I*, @@ -4780,6 +4779,7 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool for (int g = 0; g < n_groups; ++g) if (cnt[g] >= 2 && sw[g] > 0.0) em_mean[g] = swI[g] / sw[g]; + samples.clear(); samples.reserve(fulls.size()); for (const auto &o : fulls) { if (!usable_merge(o) || cnt[o.group] < 2) continue; @@ -5079,8 +5079,10 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool // and whether it is usable at all are fixed here, and neither depends on the error model that the // merge below refits. Rank parity within the reflection - see AssignHalvesOfMembers - so the halves are // balanced and identical on the host and device paths. - std::vector obs_key(fulls.size()); - std::vector merge_half(fulls.size(), 0); + std::vector &obs_key = merge_obs_key; + std::vector &merge_half = merge_half_set; + obs_key.resize(fulls.size()); + merge_half.resize(fulls.size()); bool did_gpu_halves = false; #ifdef JFJOCH_USE_CUDA // On the device where the merge is: integer work, the same halves. @@ -5092,8 +5094,10 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool if (!did_gpu_halves) { ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) + for (int i = lo; i < hi; ++i) { obs_key[i] = ObservationKey(fulls[i].h, fulls[i].k, fulls[i].l, fulls[i].frame); + merge_half[i] = 0; + } }); // mf.group is already this array: it was filled with the same `usable_merge(o) ? o.group : -1` // above, on all threads, so a second serial walk over the fat records only rebuilds it. @@ -5139,9 +5143,9 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool // model has been refitted on the reflections that survive the cut. ---- std::vector &acc = merge_acc; Result result; - std::vector merged_I; + std::vector &merged_I = merge_merged_I; size_t reject_count = 0; - std::vector rejected_obs; // per-full outlier-rejected flag (both paths) + std::vector &rejected_obs = merge_rejected; // per-full outlier-rejected flag (both paths) // ---- The Wilson test (WilsonOutliers.h): every observation against the intensity distribution of // its shell, beside the median test. It reaches what the median cannot - a reflection measured @@ -5180,15 +5184,22 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool // vector doubles its way up to a few hundred megabytes on a large cell, copying everything it // already holds each time. result.merged.reserve(n_groups); - merged_I.assign(n_groups, NAN); + merged_I.resize(n_groups); + rejected_obs.resize(fulls.size()); + ParallelChunks(n_groups, ThreadsForWork(n_groups, nthreads), [&](int lo, int hi) { + std::fill(merged_I.begin() + lo, merged_I.begin() + hi, NAN); + }); + ParallelChunks(n_full, ThreadsForWork(fulls.size(), nthreads), [&](int lo, int hi) { + std::fill(rejected_obs.begin() + lo, rejected_obs.begin() + hi, uint8_t{0}); + }); reject_count = 0; - rejected_obs.assign(fulls.size(), 0); if (wilson_test) { // Each observation with its OWN sigma, the error model evaluated at its own intensity: the // reflection's mean, which the merge weights use, is exactly what an artefact inflates. // Which fulls take part, in order, and then each one's record on its own - on all threads. wilson_full = KeptIndices(n_full, nthreads, [&](int i) { return mf.group[i] >= 0; }); - std::vector wobs(wilson_full.size()); + std::vector &wobs = merge_wobs; + wobs.resize(wilson_full.size()); ParallelChunks(static_cast(wobs.size()), ThreadsForWork(wobs.size(), nthreads), [&](int lo, int hi) { for (int j = lo; j < hi; ++j) { diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index a523ca256..70dbf3546 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -17,6 +17,7 @@ #include "../IntegrationOutcome.h" #include "ErrorModel.h" +#include "WilsonOutliers.h" #include "Merge.h" // MergedReflection, MergeStatistics #include "SphericalHarmonicSurface.h" #ifdef JFJOCH_USE_CUDA @@ -427,6 +428,15 @@ private: }; MergeAccumStaging merge_accum; + // MergeAndStats' other per-full and per-group arrays, members for the same reason: five merges a + // run, each over up to tens of millions of fulls and millions of groups, and as locals each one + // allocated, faulted in and zeroed them again on one thread. As members they are only resized, and + // reset on all threads where a pass needs them reset. + std::vector merge_merged_I; + std::vector merge_rejected, merge_half_set; + std::vector merge_obs_key; + std::vector merge_wobs; + // Per-group scatter for the strong-reflection ISa asymptote (see MergeAndStats). A member for the // same reason: 32 bytes a group, once per merge. struct GroupScatter { double sum = 0, sum_sq = 0, sum_var = 0; int n = 0; }; From 308d1ed2580ddda551ff36680266737289bf0237 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 12:37:29 +0200 Subject: [PATCH 3/4] RotationScaleMerge: export the merged reflections on all threads Which groups are written is decided first (KeptIndices, in group order), then each written group's MergedReflection is formed on its own, on all threads, into a result sized once - where the serial loop pushed back up to millions of records, growing a vector that each merge allocated afresh. The same records in the same order. p.mtz md5 unchanged on myob, cytc, thau, kdp, 8a1a and 8qaw (GPU build). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- .../scale_merge/RotationScaleMerge.cpp | 64 ++++++++++--------- 1 file changed, 33 insertions(+), 31 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index b424c990d..838ba3c64 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -5179,11 +5179,6 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool // it with += and does need the zeroes. if (use_gpu_merge) acc.resize(n_groups); else acc.assign(n_groups, Accum{}); - result.merged.clear(); - // One reflection per group at most, and a merge runs several times per pass. Without this the - // vector doubles its way up to a few hundred megabytes on a large cell, copying everything it - // already holds each time. - result.merged.reserve(n_groups); merged_I.resize(n_groups); rejected_obs.resize(fulls.size()); ParallelChunks(n_groups, ThreadsForWork(n_groups, nthreads), [&](int lo, int hi) { @@ -5321,34 +5316,41 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool }); } // ---- Export merged reflections. ---- - for (int g = 0; g < n_groups; ++g) { - const auto &a = acc[g]; - if (a.sw <= 0.0) continue; - MergedReflection mr{}; - mr.h = group_h[g]; mr.k = group_k[g]; mr.l = group_l[g]; - mr.I = static_cast(a.swI / a.sw); - // Plain inverse-variance merged sigma. The error model's systematic term (b*I)^2 is measured - // from the scatter BETWEEN a reflection's symmetry equivalents, i.e. from exactly the part of - // the error that is NOT common to them, so it averages down over the multiplicity like the - // counting part and the merge must not hold it back. XDS behaves the same way: its merged - // I/sigma runs far above its own reported ISa in the strong low-resolution shells. - mr.sigma = static_cast(1.0 / std::sqrt(a.sw)); - mr.I_half[0] = mr.I_half[1] = NAN; - mr.sigma_half[0] = mr.sigma_half[1] = NAN; - mr.d = a.d; - mr.on_ice_ring = a.on_ice; - if (a.nh[0] + a.nh[1] > 0 && a.swh[0] > 0.0 && a.swh[1] > 0.0) { - for (int i = 0; i < 2; ++i) { - mr.I_half[i] = static_cast(a.swIh[i] / a.swh[i]); - mr.sigma_half[i] = static_cast(1.0 / std::sqrt(a.swh[i])); + // Which groups are written, in group order, then each one's record on its own - on all threads. + // A group is written when it has weight and a resolution. + const std::vector written = KeptIndices(n_groups, nthreads, [&](int g) { + return !(acc[g].sw <= 0.0) && std::isfinite(acc[g].d) && acc[g].d > 0.0f; + }); + result.merged.resize(written.size()); + ParallelChunks(static_cast(written.size()), ThreadsForWork(written.size(), nthreads), [&](int lo, int hi) { + for (int j = lo; j < hi; ++j) { + const int g = written[j]; + const auto &a = acc[g]; + MergedReflection mr{}; + mr.h = group_h[g]; mr.k = group_k[g]; mr.l = group_l[g]; + mr.I = static_cast(a.swI / a.sw); + // Plain inverse-variance merged sigma. The error model's systematic term (b*I)^2 is measured + // from the scatter BETWEEN a reflection's symmetry equivalents, i.e. from exactly the part of + // the error that is NOT common to them, so it averages down over the multiplicity like the + // counting part and the merge must not hold it back. XDS behaves the same way: its merged + // I/sigma runs far above its own reported ISa in the strong low-resolution shells. + mr.sigma = static_cast(1.0 / std::sqrt(a.sw)); + mr.I_half[0] = mr.I_half[1] = NAN; + mr.sigma_half[0] = mr.sigma_half[1] = NAN; + mr.d = a.d; + mr.on_ice_ring = a.on_ice; + if (a.nh[0] + a.nh[1] > 0 && a.swh[0] > 0.0 && a.swh[1] > 0.0) { + for (int i = 0; i < 2; ++i) { + mr.I_half[i] = static_cast(a.swIh[i] / a.swh[i]); + mr.sigma_half[i] = static_cast(1.0 / std::sqrt(a.swh[i])); + } + mr.cc_weight = static_cast((1.0 / a.swh_typ[0] + 1.0 / a.swh_typ[1]) + / (1.0 / a.swh[0] + 1.0 / a.swh[1])); } - mr.cc_weight = static_cast((1.0 / a.swh_typ[0] + 1.0 / a.swh_typ[1]) - / (1.0 / a.swh[0] + 1.0 / a.swh[1])); + merged_I[g] = mr.I; + result.merged[j] = mr; } - if (!std::isfinite(a.d) || a.d <= 0.0f) continue; - merged_I[g] = mr.I; - result.merged.push_back(mr); - } + }); }; run_merge(); From 3fa95eea29e4ba86137607be737f79fd7a4f7d78 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 12:37:29 +0200 Subject: [PATCH 4/4] RotationScaleMerge: the R_meas staging as members, its group pass on all threads The R_meas accumulators per group and the device's per-group sums were eight group-length locals plus a 56-byte record per group, value-initialised on one thread for every merge that reports statistics. They are now members, reset in parallel, and the device branch's pass that assigns each group its shell and copies its sums runs on all threads; the per-shell observation counts are integers, gathered per block. The shells' R_meas sums are still added in group order, as before. p.mtz md5 unchanged on myob, cytc, thau, kdp, 8a1a and 8qaw (GPU build). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- .../scale_merge/RotationScaleMerge.cpp | 59 +++++++++++-------- .../scale_merge/RotationScaleMerge.h | 13 ++++ 2 files changed, 47 insertions(+), 25 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 838ba3c64..f1319ccf5 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -5647,36 +5647,45 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool // n_eff = (sum v)^2 / sum v^2, which is n where the sigmas are equal, and there this is the // ordinary R_meas = sum sqrt(n/(n-1)) sum|I_i - | / sum I_i. // The unweighted sums beside them give the ordinary r_meas; the weighted ones r_meas_weighted. - struct RmeasObs { double sum_abs_dev = 0, sum_I = 0, w_abs_dev = 0, w_I = 0, sum_v = 0, sum_v2 = 0; - int n = 0, shell = -1; }; - std::vector rmeas(n_groups); + std::vector &rmeas = merge_rmeas; + rmeas.resize(n_groups); + ParallelChunks(n_groups, ThreadsForWork(n_groups, nthreads), [&](int lo, int hi) { + std::fill(rmeas.begin() + lo, rmeas.begin() + hi, RmeasObs{}); + }); bool did_gpu_rmeas = false; #ifdef JFJOCH_USE_CUDA if (use_gpu_merge) { // Per-group R_meas + usable count on the GPU; the shell is assigned per group (its fulls share d). - std::vector rabsdev(n_groups), rsumI(n_groups), rwabsdev(n_groups), rwsumI(n_groups), - rsumv(n_groups), rsumv2(n_groups); - std::vector rn(n_groups), rnusable(n_groups); - gpu_->MergeRmeas(merged_I.data(), rabsdev.data(), rsumI.data(), rwabsdev.data(), rwsumI.data(), - rsumv.data(), rsumv2.data(), rn.data(), rnusable.data()); - for (int g = 0; g < n_groups; ++g) { - if (rnusable[g] == 0 || acc[g].d < d_floor) continue; - const auto shell = shells.GetShell(acc[g].d); - if (!shell || *shell < 0 || *shell >= n_shells) continue; - // Count the MERGED population, not the R_meas one. The R_meas re-walk deliberately - // ignores the ice flag on a search pass, so its count includes observations that never - // entered `unique` - which inflates the reported multiplicity of whatever shell they - // land in. acc[g].nh is what actually went into this group's mean, and it is zero for a - // group the merge dropped entirely. (acc[g].d is NaN for such a group, so GetShell above - // already declines it; this is the same statement made where it counts.) - sa[*shell].total_obs += static_cast(acc[g].nh[0] + acc[g].nh[1]); - if (std::isfinite(merged_I[g]) && rn[g] > 0) { - auto &r = rmeas[g]; - r.sum_abs_dev = rabsdev[g]; r.sum_I = rsumI[g]; r.w_abs_dev = rwabsdev[g]; r.w_I = rwsumI[g]; - r.sum_v = rsumv[g]; r.sum_v2 = rsumv2[g]; - r.n = rn[g]; r.shell = *shell; + RmeasStaging &rs = merge_rmeas_staging; + rs.Resize(n_groups); + gpu_->MergeRmeas(merged_I.data(), rs.absdev.data(), rs.sumI.data(), rs.wabsdev.data(), rs.wsumI.data(), + rs.sumv.data(), rs.sumv2.data(), rs.n.data(), rs.nusable.data()); + // Each group on its own, on all threads; the observation counts are integers, so the blocks + // may add theirs up in any split. + const int nb = ReductionBlocks(n_groups, 4096); + std::vector> blk_obs(nb, std::vector(n_shells, 0)); + ParallelBlocks(n_groups, nthreads, [&](int blk, int g0, int g1) { + for (int g = g0; g < g1; ++g) { + if (rs.nusable[g] == 0 || acc[g].d < d_floor) continue; + const auto shell = shells.GetShell(acc[g].d); + if (!shell || *shell < 0 || *shell >= n_shells) continue; + // Count the MERGED population, not the R_meas one. The R_meas re-walk deliberately + // ignores the ice flag on a search pass, so its count includes observations that never + // entered `unique` - which inflates the reported multiplicity of whatever shell they + // land in. acc[g].nh is what actually went into this group's mean, and it is zero for a + // group the merge dropped entirely. (acc[g].d is NaN for such a group, so GetShell above + // already declines it; this is the same statement made where it counts.) + blk_obs[blk][*shell] += static_cast(acc[g].nh[0] + acc[g].nh[1]); + if (std::isfinite(merged_I[g]) && rs.n[g] > 0) { + auto &r = rmeas[g]; + r.sum_abs_dev = rs.absdev[g]; r.sum_I = rs.sumI[g]; r.w_abs_dev = rs.wabsdev[g]; + r.w_I = rs.wsumI[g]; r.sum_v = rs.sumv[g]; r.sum_v2 = rs.sumv2[g]; + r.n = rs.n[g]; r.shell = *shell; + } } - } + }); + for (const auto &c : blk_obs) + for (int sh = 0; sh < n_shells; ++sh) sa[sh].total_obs += c[sh]; did_gpu_rmeas = true; } #endif diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 70dbf3546..15f2e2725 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -436,6 +436,19 @@ private: std::vector merge_rejected, merge_half_set; std::vector merge_obs_key; std::vector merge_wobs; + // The R_meas accumulators per group (see MergeAndStats' statistics), and the device's per-group sums. + struct RmeasObs { double sum_abs_dev = 0, sum_I = 0, w_abs_dev = 0, w_I = 0, sum_v = 0, sum_v2 = 0; + int n = 0, shell = -1; }; + std::vector merge_rmeas; + struct RmeasStaging { + std::vector absdev, sumI, wabsdev, wsumI, sumv, sumv2; + std::vector n, nusable; + void Resize(int g) { + absdev.resize(g); sumI.resize(g); wabsdev.resize(g); wsumI.resize(g); sumv.resize(g); sumv2.resize(g); + n.resize(g); nusable.resize(g); + } + }; + RmeasStaging merge_rmeas_staging; // Per-group scatter for the strong-reflection ISa asymptote (see MergeAndStats). A member for the // same reason: 32 bytes a group, once per merge.