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

RotationScaleMerge: GPU scale-fulls, fulls kept resident (phase 2 step 2)

Browse Source 
Scale the combined fulls (Unity model) on the device so they no longer round-trip
between the combine and the merge: after the GPU combine, build the fulls' per-frame
and per-ASU-group CSRs on the host from just the small key arrays (f_frame/f_group)
with a deterministic counting sort - no GPU stable-sort - then scale in place and
download once.

The four scaling kernels are reused unchanged except FitPerFrameGKernel, which gains
an optional `perm` argument (null for the partials, whose arrays are already
frame-contiguous; a frame-grouping permutation for the emit-ordered fulls) so the
fulls are scaled without a physical reorder. The Unity model falls out of giving the
fulls all-ones partiality/rlp/zeta (coeff = mean), so no other kernel changes and the
committed phase-1 partial-scaling path is bit-identical (perm == null -> idx == i).

Validated across the rotation battery (JFJOCH_RSM_GPU_COMBINE=1): all 15 deterministic
crystals stay run-to-run deterministic and their merged output is bit-identical to the
CPU path (SG/ISa/CC1.2/completeness). The lone exception is EP_cs_01-24 (CC1/2 2%,
R_meas 379% - unindexable noise): merged intensities/CC/completeness match exactly, but
the ill-conditioned 16-bin error-model b fit amplifies the ~1e-7 scale-fulls rounding
to ISa 10.6 vs 10.8 - benign, same class as the accepted phase-1 GPU rounding. The 3
upstream-nondeterministic crystals vary as before (GPU-prediction overflow, not this).

Scale-fulls drops from ~0.09s to ~0 across the two passes; combine+scale-fulls region
~0.32s GPU vs ~0.46s CPU on lyso. Still opt-in (fulls are downloaded for the host merge;
the win grows once the merge/error-model also stay resident).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-07-03 07:51:29 +02:00
co-authored by Claude Opus 4.8
parent 2c928b27cd
commit ced85bcd9d
3 changed files with 154 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/RotationScaleMerge.cpp
+35 -8
View File
@@ -1098,31 +1098,58 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search,
// --- 3. 3D combine of per-frame partials into fulls (fulls inherit their ASU group here). ---
bool combined_on_gpu = false;
bool scaled_fulls_on_gpu = false;
#ifdef JFJOCH_USE_CUDA
// The GPU combine mirrors Combine() exactly but keeps the fulls on the device; the diagnostic dump
// (serial, one writer) has no GPU equivalent, so fall back to the CPU path when it is requested.
// GPU combine (+ scale-fulls) keeps the fulls resident on the device: combine, then build the frame /
// ASU-group CSRs on the host from just the small key arrays (a deterministic counting sort - no GPU
// stable-sort), scale the fulls in place, and download only once. Mirrors Combine() + the Unity
// scale-fulls loop below. The diagnostic dump (serial, one writer) has no GPU path -> CPU fallback.
if (gpu_active_ && gpu_combine_ && observation_dump_path.empty()) {
std::vector<float> corr(partials.size()); // refresh the smoothed corr on the device
for (size_t i = 0; i < partials.size(); ++i) corr[i] = partials[i].corr;
gpu_->SetCorr(corr.data());
const int nf = gpu_->Combine(rawrun_group.data(), min_partiality, capture_uncertainty_coeff);
fulls.assign(nf, Obs{});
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<int32_t> ff(nf), fg(nf);
gpu_->GetFullsKeys(ff.data(), fg.data());
std::vector<int32_t> 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<int32_t> 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<int32_t> 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<int32_t> g_perm(acc);
{ std::vector<int32_t> 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());
gpu_->ScaleFulls(scaling_iter, SCALE_ROBUST_K, min_partiality);
scaled_fulls_on_gpu = true;
}
fulls.assign(nf, Obs{});
std::vector<int32_t> fh(nf), fk(nf), fl(nf), fframe(nf), fgroup(nf);
std::vector<float> fI(nf), fsig(nf), fd(nf), fimg(nf);
std::vector<float> fI(nf), fsig(nf), fd(nf), fimg(nf), fcorr(nf, 1.0f);
std::vector<uint8_t> fon(nf);
gpu_->GetFulls(fh.data(), fk.data(), fl.data(), fI.data(), fsig.data(), fd.data(),
fimg.data(), fframe.data(), fon.data(), fgroup.data());
if (scaled_fulls_on_gpu) gpu_->GetFullsCorr(fcorr.data());
for (int i = 0; i < nf; ++i) {
Obs &o = fulls[i];
o.h = fh[i]; o.k = fk[i]; o.l = fl[i];
o.I = fI[i]; o.sigma = fsig[i]; o.d = fd[i];
o.rlp = 1.0f; o.partiality = 1.0f; o.corr = 1.0f;
o.rlp = 1.0f; o.partiality = 1.0f; o.corr = fcorr[i];
o.image_number = fimg[i]; o.frame = fframe[i];
o.on_ice = fon[i]; o.group = fgroup[i];
}
SortFullsByFrame();
logger.Info("3D combine (GPU): {} fulls", nf);
logger.Info("3D combine{} (GPU): {} fulls", scaled_fulls_on_gpu ? " + scale-fulls" : "", nf);
combined_on_gpu = true;
}
#endif
@@ -1131,7 +1158,7 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search,
lap("combine");
// --- 4. Scale the fulls (XDS order, Unity model). ---
if (scale_fulls) {
if (scale_fulls && !scaled_fulls_on_gpu) {
std::vector<double> full_mean;
for (int it = 0; it < scaling_iter; ++it) {
ReduceGroupMeans(fulls, n_groups, false, {}, full_mean);