mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 81 Wiki Activity

Merge branch 'gpu-tail2' into rc175

Browse Source
This commit is contained in:
leonarski_f
2026-10-08 13:33:04 +02:00
parent 49abb4a079 3fa95eea29
commit 09a8c4e5f3
4 changed files with 477 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/RotationScaleMerge.cpp
+126 -113
View File
@@ -3677,7 +3677,26 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &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<int32_t> 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<int>(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<int32_t> &cell, int
}
}
});
n_terms = static_cast<int>(term.size());
}
auto subset = [&](int parity) -> const std::vector<int32_t> & {
return parity < 0 ? idx_all : (parity ? idx_odd : idx_even);
@@ -3725,10 +3745,7 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &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<int>(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<int>(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<int32_t> &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<int>(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<double> shw(static_cast<size_t>(ncell) * nshell, 0.0);
std::vector<double> shw_shell(nshell, 0.0), shw_cell(nshell > 0 ? ncell : 0, 0.0);
std::vector<int32_t> 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<int>(term.size());
std::vector<float> s2(n_term);
ParallelChunks(n_term, nt, [&](int lo, int hi) {
@@ -3798,13 +3819,13 @@ void RotationScaleMerge::ApplyCellSurface(const std::vector<int32_t> &cell, int
if (!(sc > 0.0)) continue;
shw[static_cast<size_t>(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<size_t>(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<size_t>(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<int32_t> &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<int32_t> 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<int32_t> g_start;
std::vector<RefTerm> 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<int32_t> gperm(term.size());
std::vector<int32_t> fill(g_start.begin(), g_start.end() - 1);
for (size_t k = 0; k < term.size(); ++k)
gperm[fill[term[k].group]++] = static_cast<int32_t>(k);
gpu_->SurfaceSetTerms(static_cast<int>(term.size()), term.data(), term_parity.data(),
n_groups, gperm.data(), g_start.data(), ncell);
for (int parity : {0, 1, -1}) {
const std::vector<int32_t> &sel = subset(parity);
const int n = static_cast<int>(sel.size());
const int nb = ReductionBlocks(n, SURFACE_BLOCK);
std::vector<int32_t> perm(n), seg_start(static_cast<size_t>(nb) * ncell + 1, n);
ParallelFor(nb, nt, [&](int b) {
const int lo = static_cast<int>(static_cast<int64_t>(n) * b / nb);
const int hi = static_cast<int>(static_cast<int64_t>(n) * (b + 1) / nb);
std::vector<int32_t> 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<size_t>(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<int32_t> 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<Sample> &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 <I>. Rebuild the variance at the reflection's EXPECTED intensity - var_bkg + var_per_I*<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<uint64_t> obs_key(fulls.size());
std::vector<uint8_t> merge_half(fulls.size(), 0);
std::vector<uint64_t> &obs_key = merge_obs_key;
std::vector<uint8_t> &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<int>(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<Accum> &acc = merge_acc;
Result result;
std::vector<double> merged_I;
std::vector<double> &merged_I = merge_merged_I;
size_t reject_count = 0;
std::vector<uint8_t> rejected_obs; // per-full outlier-rejected flag (both paths)
std::vector<uint8_t> &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<WilsonObservation> wobs(wilson_full.size());
std::vector<WilsonObservation> &wobs = merge_wobs;
wobs.resize(wilson_full.size());
ParallelChunks(static_cast<int>(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<float>(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<float>(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<float>(a.swIh[i] / a.swh[i]);
mr.sigma_half[i] = static_cast<float>(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<int32_t> 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<int>(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<float>(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<float>(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<float>(a.swIh[i] / a.swh[i]);
mr.sigma_half[i] = static_cast<float>(1.0 / std::sqrt(a.swh[i]));
}
mr.cc_weight = static_cast<float>((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<float>((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 - <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<RmeasObs> rmeas(n_groups);
std::vector<RmeasObs> &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<double> rabsdev(n_groups), rsumI(n_groups), rwabsdev(n_groups), rwsumI(n_groups),
rsumv(n_groups), rsumv2(n_groups);
std::vector<int32_t> 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<int>(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<std::vector<int>> blk_obs(nb, std::vector<int>(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<int>(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.
image_analysis/scale_merge/RotationScaleMerge.h
+29
View File
@@ -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<double> merge_merged_I;
std::vector<uint8_t> merge_rejected, merge_half_set;
std::vector<uint64_t> merge_obs_key;
std::vector<WilsonObservation> 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<RmeasObs> merge_rmeas;
struct RmeasStaging {
std::vector<double> absdev, sumI, wabsdev, wsumI, sumv, sumv2;
std::vector<int32_t> 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<RotationScaleMergeGPU> 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<float> surface_corr;
#endif
bool write_back_per_frame_scale = true; // see SetWriteBackPerFrameScale
int frame_limit = -1; // see SetFrameLimit
image_analysis/scale_merge/RotationScaleMergeGPU.cu
+307
View File
@@ -11,11 +11,13 @@
#include <string>
#include <vector>
#include <cuda_runtime.h>
#include <cub/device/device_radix_sort.cuh>
#include <cub/device/device_scan.cuh>
#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<uint32_t>(group[i]);
const uint8_t par = static_cast<uint8_t>(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<int64_t>(n) * b / nb; };
int b = static_cast<int>(static_cast<int64_t>(k) * nb / n);
while (b + 1 < nb && lo(b + 1) <= k) ++b;
while (b > 0 && lo(b) > k) --b;
key[k] = static_cast<uint32_t>(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<uint8_t>(sh);
bin[k] = static_cast<uint32_t>(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<uint8_t> scratch(std::max<size_t>(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<uint8_t> scratch(std::max<size_t>(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<int32_t> idx_even, idx_odd;
CudaDevicePtr<uint32_t> group_key;
{
CudaDevicePtr<int32_t> d_cell, keep, even, odd;
CudaDevicePtr<float> d_corr;
d.Upload(d_cell, cell, nf);
d.Upload(d_corr, corr, nf);
keep = d.Alloc<int32_t>(std::max(1, nf)); even = d.Alloc<int32_t>(std::max(1, nf));
odd = d.Alloc<int32_t>(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<int32_t>(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<uint8_t>(std::max<size_t>(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<SurfaceTerm>(std::max(1, n_terms));
d.s_parity = d.Alloc<uint8_t>(std::max(1, n_terms));
group_key = d.Alloc<uint32_t>(std::max(1, n_terms));
idx_even = d.Alloc<int32_t>(std::max(1, n_even));
idx_odd = d.Alloc<int32_t>(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<int32_t>(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<int32_t>(std::max(1, n_terms));
d.s_gstart = d.Alloc<int32_t>(n_groups + 1);
{
auto key_sorted = d.Alloc<uint32_t>(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<uint32_t>();
// 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<int32_t>(std::max(1, n));
d.s_seg_start[subset] = d.Alloc<int32_t>(n_seg + 1);
if (nb == 0) continue;
auto key = d.Alloc<uint32_t>(n), key_sorted = d.Alloc<uint32_t>(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<double>(n_seg);
d.s_tref2 = d.Alloc<double>(n_seg);
d.s_slots = n_seg;
}
}
CudaCheck(cudaStreamSynchronize(st), "surface subsets sync");
d.s_A = d.Alloc<double>(std::max(1, ncell));
d.s_cross = d.Alloc<double>(std::max(1, ncell));
d.s_ref2 = d.Alloc<double>(std::max(1, ncell));
d.s_sw = d.Alloc<double>(std::max(1, n_groups));
d.s_swI = d.Alloc<double>(std::max(1, n_groups));
d.s_w_Is = d.Alloc<double>(std::max(1, n_terms));
d.s_w_Iref = d.Alloc<double>(std::max(1, n_terms));
d.s_Iref = d.Alloc<double>(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<float>(std::max(1, n)), s2_sorted = d.Alloc<float>(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<uint8_t>(std::max<size_t>(bytes, 1));
CudaCheck(cub::DeviceRadixSort::SortKeys(scratch.get(), bytes, s2.get(), s2_sorted.get(), n, 0, 32, st),
"s2 sort");
}
std::vector<float> 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<float>();
CudaDevicePtr<float> d_edge;
d.Upload(d_edge, edge.data(), nshell - 1);
auto shell = d.Alloc<uint8_t>(std::max(1, n));
auto bin = d.Alloc<uint32_t>(std::max(1, n)), bin_sorted = d.Alloc<uint32_t>(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<int32_t>(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<int32_t>(std::max(1, n)), order = d.Alloc<int32_t>(std::max(1, n));
auto bin_start = d.Alloc<int32_t>(n_bins + 1);
auto d_shw = d.Alloc<double>(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_;
image_analysis/scale_merge/RotationScaleMergeGPU.h
+15
View File
@@ -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]).