From eb7a65d0514f6ccc57aa51ef72d9e7f04ee006eb Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Tue, 13 Oct 2026 02:01:30 +0200 Subject: [PATCH] rotation merge: the fulls' frame and group CSRs built on the card Every merge downloaded the fulls' frame and group keys and built both CSRs by serial counting sorts on the host - four passes over the fulls, two of them scattering into an array as long as the distinct hkl (15 M on 4nwv) - and uploaded them again. They are now stable radix sorts of the resident fulls on the card, with the counts summed there; the host gets the counts per frame and the group CSR it walks itself. The same permutations (stable sorts list each frame's and each group's fulls in index order, as the counting sorts did): p.hkl identical on cytc and 4nwv. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- .../scale_merge/RotationScaleMerge.cpp | 25 +---- .../scale_merge/RotationScaleMergeGPU.cu | 102 +++++++++++++----- .../scale_merge/RotationScaleMergeGPU.h | 19 ++-- 3 files changed, 86 insertions(+), 60 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 0ec41c6f3..ab3f9f41e 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -6619,26 +6619,11 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool g_full.assign(n_frames, 1.0); if (scale_fulls && nf > 0) { - // Frame + group CSRs over the emit-ordered fulls, built by counting sort on the host (stable, - // deterministic). frame is always in [0, n_frames); group is <0 for absent/out-of-range fulls. - std::vector ff(nf), fg(nf); - gpu_->GetFullsKeys(ff.data(), fg.data()); - std::vector f_start(n_frames, 0), f_count(n_frames, 0), f_perm(nf); - for (int i = 0; i < nf; ++i) ++f_count[ff[i]]; - for (int f = 1; f < n_frames; ++f) f_start[f] = f_start[f - 1] + f_count[f - 1]; - { std::vector fill = f_start; for (int i = 0; i < nf; ++i) f_perm[fill[ff[i]]++] = i; } - gpu_->SetFullsFrameCSR(f_perm.data(), nf, f_start.data(), f_count.data()); - - std::vector g_count(n_groups, 0), g_start(n_groups, 0); - for (int i = 0; i < nf; ++i) if (fg[i] >= 0) ++g_count[fg[i]]; - int acc = 0; - for (int g = 0; g < n_groups; ++g) { g_start[g] = acc; acc += g_count[g]; } - std::vector g_perm(acc); - { std::vector fill = g_start; for (int i = 0; i < nf; ++i) if (fg[i] >= 0) g_perm[fill[fg[i]]++] = i; } - gpu_->SetFullsGroups(g_perm.data(), acc, g_start.data(), g_count.data()); - fulls_group_start = g_start; - fulls_group_start.push_back(acc); - fulls_group_perm = std::move(g_perm); + // Frame + group CSRs over the emit-ordered fulls, built on the device by stable sorts: each + // frame's and each group's fulls in index order. frame is always in [0, n_frames); group is <0 + // for absent/out-of-range fulls, which no group lists. + std::vector f_count(n_frames); + gpu_->BuildFullsCSRs(f_count.data(), fulls_group_start, fulls_group_perm); fulls_csr_built = true; gpu_->ResetFullsScale(); diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 2aeff4743..62805d660 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -2034,33 +2034,6 @@ void RotationScaleMergeGPU::GetFulls(int32_t *h, int32_t *k, int32_t *l, float * dl(clipped, dd.f_clipped.get(), n * sizeof(uint8_t)); } -void RotationScaleMergeGPU::GetFullsKeys(int32_t *frame, int32_t *group) const { - DeviceGuard guard(impl_->device, impl_->available); - const auto &d = *impl_; - if (d.n_fulls == 0) return; - const size_t bytes = size_t(d.n_fulls) * sizeof(int32_t); - CopyAndWait(frame, d.f_frame.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_frame"); - CopyAndWait(group, d.f_group.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_group"); -} - -void RotationScaleMergeGPU::SetFullsFrameCSR(const int32_t *frame_perm, int n_perm, - const int32_t *frame_start, const int32_t *frame_count) { - DeviceGuard guard(impl_->device, impl_->available); - auto &d = *impl_; - d.Upload(d.f_frame_perm, frame_perm, n_perm); - d.Upload(d.f_frame_start, frame_start, d.n_frames); - d.Upload(d.f_frame_count, frame_count, d.n_frames); -} - -void RotationScaleMergeGPU::SetFullsGroups(const int32_t *gperm, int n_gperm, - const int32_t *gstart, const int32_t *gcount) { - DeviceGuard guard(impl_->device, impl_->available); - auto &d = *impl_; - d.Upload(d.f_gperm, gperm, n_gperm); - d.Upload(d.f_gstart, gstart, d.n_groups); - d.Upload(d.f_gcount, gcount, d.n_groups); -} - void RotationScaleMergeGPU::ResetFullsScale() { DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; @@ -2611,6 +2584,19 @@ namespace { int Blocks(int n) { return std::max(1, std::min(65535, (n + BLK - 1) / BLK)); } + // --- the fulls' CSRs on the device (BuildFullsCSRs) --- + // A full's group as a sort key: the fulls in no group sort after every group. + __global__ void FullGroupKeyKernel(int n, const int32_t *__restrict__ group, int32_t *__restrict__ key) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) + key[i] = group[i] >= 0 ? group[i] : INT32_MAX; + } + + // Counts per bin (integers, so in any order); a negative bin counts nowhere. + __global__ void CountBinsKernel(int n, const int32_t *__restrict__ bin, int32_t *__restrict__ count) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) + if (bin[i] >= 0) atomicAdd(&count[bin[i]], 1); + } + // --- the ingest's geometry smoothing on the card (RotationScaleMerge::IngestOnDevice) --- // Explicit roundings throughout, so that no compiler contracts a product into a sum differently on // another card: the same bits on every GPU, and within the last bits of the host's own loops. @@ -3118,3 +3104,65 @@ void RotationScaleMergeGPU::IngestFinish(float *corr) { d.in_src_start.clear(); d.in_layout = RotationScaleMergeGPU::IngestLayout{}; } + +void RotationScaleMergeGPU::BuildFullsCSRs(int32_t *frame_count, std::vector &gstart, + std::vector &gperm) { + DeviceGuard guard(impl_->device, impl_->available); + auto &d = *impl_; + const int nf = d.n_fulls, nfr = d.n_frames, ng = d.n_groups; + const int nfa = std::max(1, nf); + const auto stable_sort = [&](const int32_t *key_in, int32_t *key_out, const int32_t *val_in, int32_t *val_out) { + size_t bytes = 0; + CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key_in, key_out, val_in, val_out, nf, 0, 32, d.s()), + "sort temp size"); + CudaDevicePtr temp = d.Alloc(std::max(1, bytes)); + CudaCheck(cub::DeviceRadixSort::SortPairs(temp.get(), bytes, key_in, key_out, val_in, val_out, nf, 0, 32, d.s()), + "sort"); + CudaCheck(cudaStreamSynchronize(d.s()), "sort"); + }; + // The counts come back for the starts, a prefix the host forms in order. + const auto count = [&](const int32_t *bin, int n_bins, CudaDevicePtr &out, std::vector &host) { + out = d.Alloc(std::max(1, n_bins)); + CudaCheck(cudaMemsetAsync(out.get(), 0, size_t(std::max(1, n_bins)) * sizeof(int32_t), d.s()), "memset counts"); + CountBinsKernel<<>>(nf, bin, out.get()); + CudaCheck(cudaGetLastError(), "CountBinsKernel launch"); + host.resize(n_bins); + if (n_bins > 0) + CopyAndWait(host.data(), out.get(), size_t(n_bins) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), + "download counts"); + }; + const auto prefix = [](const std::vector &c) { + std::vector start(c.size() + 1, 0); + for (size_t i = 0; i < c.size(); ++i) start[i + 1] = start[i] + c[i]; + return start; + }; + CudaDevicePtr idx = d.Alloc(nfa), key = d.Alloc(nfa), sorted = d.Alloc(nfa); + if (nf > 0) IotaKernel<<<(nf + BLK - 1) / BLK, BLK, 0, d.s()>>>(nf, idx.get()); + CudaCheck(cudaGetLastError(), "IotaKernel launch"); + + std::vector fcount; + count(d.f_frame.get(), nfr, d.f_frame_count, fcount); + std::copy(fcount.begin(), fcount.end(), frame_count); + const std::vector fstart = prefix(fcount); + d.Upload(d.f_frame_start, fstart.data(), nfr); + d.f_frame_perm = d.Alloc(nfa); + if (nf > 0) stable_sort(d.f_frame.get(), key.get(), idx.get(), d.f_frame_perm.get()); + + std::vector gcount; + count(d.f_group.get(), ng, d.f_gcount, gcount); + gstart = prefix(gcount); + d.Upload(d.f_gstart, gstart.data(), ng); + const int n_grouped = gstart.back(); + FullGroupKeyKernel<<>>(nf, d.f_group.get(), key.get()); + CudaCheck(cudaGetLastError(), "FullGroupKeyKernel launch"); + CudaDevicePtr by_group = d.Alloc(nfa); + if (nf > 0) stable_sort(key.get(), sorted.get(), idx.get(), by_group.get()); + d.f_gperm = d.Alloc(std::max(1, n_grouped)); + gperm.resize(n_grouped); + if (n_grouped > 0) { + CopyAndWait(d.f_gperm.get(), by_group.get(), size_t(n_grouped) * sizeof(int32_t), cudaMemcpyDeviceToDevice, + d.s(), "group perm"); + CopyAndWait(gperm.data(), by_group.get(), size_t(n_grouped) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), + "download group perm"); + } +} diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.h b/image_analysis/scale_merge/RotationScaleMergeGPU.h index 0f378d6f9..e407b84cb 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -243,18 +243,11 @@ public: // --- scale the resident fulls on the device (Unity model), no round-trip --- - // Download the fulls' frame and ASU-group keys (emit order) so the host can build the frame/group CSRs - // with a counting sort (deterministic, no GPU stable-sort) and hand them back below. - void GetFullsKeys(int32_t *frame, int32_t *group) const; - - // The fulls' per-frame CSR: frame_perm groups the emit-ordered fulls by frame (frame_start/count length - // n_frames index it), so FitPerFrameG can scale the fulls without physically reordering them. - void SetFullsFrameCSR(const int32_t *frame_perm, int n_perm, - const int32_t *frame_start, const int32_t *frame_count); - - // The fulls' per-ASU-group CSR (group-ordered permutation of the fulls with group>=0, + its CSR). - void SetFullsGroups(const int32_t *gperm, int n_gperm, - const int32_t *gstart, const int32_t *gcount); + // The fulls' per-frame CSR (each frame's fulls, for the per-frame fit) and per-ASU-group CSR, built on + // the device by stable sorts of the emit-ordered fulls, so that each frame and each group lists its + // fulls in index order as a counting sort does. Hands back the fulls per frame (length n_frames) and + // the group CSR for the host's own walks: gstart (n_groups + 1 entries, the last the total) and gperm. + void BuildFullsCSRs(int32_t *frame_count, std::vector &gstart, std::vector &gperm); // Start the Unity scaling of the resident fulls: corr, partiality, prescaling and zeta all 1, no // frame fitted yet. Once before the FitFullsScale rounds. @@ -262,7 +255,7 @@ public: // One Unity fit of the resident fulls (reduce group means -> per-frame LS G and its information, // GetG / GetInfo), leaving the fulls' working corr as it is: the host smooths the fitted scales - // and applies them with SmoothFullsCorr. Requires SetFullsFrameCSR + SetFullsGroups and + // and applies them with SmoothFullsCorr. Requires BuildFullsCSRs and // ResetFullsScale. void FitFullsScale(double min_partiality);