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);