diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index bc164861c..f1319ccf5 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) { @@ -4706,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*, @@ -4790,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; @@ -5089,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. @@ -5102,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. @@ -5149,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 @@ -5185,20 +5179,22 @@ 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.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) { @@ -5320,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(); @@ -5644,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 @@ -6638,6 +6650,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..15f2e2725 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,28 @@ 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; + // 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. struct GroupScatter { double sum = 0, sum_sq = 0, sum_var = 0; int n = 0; }; @@ -484,6 +507,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]).