From e59d2272945b00bb8faaba00d38693564e0bb327 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 08:48:46 +0200 Subject: [PATCH] RotationScaleMerge: the ISa asymptote's group sums on the device The strong-reflection asymptote is read from per-group sums of I*corr, its square and the counting variance at the group mean - a walk over every full's record on the host. In the GPU build a kernel takes them, one thread per group in index order with the host's rounding (CountingVarianceExact); the host keeps the strong-group selection and the median. p.mtz md5, p.cif and the error-model lines (which print the asymptote) unchanged on myob, cytc, 8a1a and 8qaw (GPU build). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- .../scale_merge/RotationScaleMerge.cpp | 41 +++++++++++----- .../scale_merge/RotationScaleMergeGPU.cu | 49 +++++++++++++++++++ .../scale_merge/RotationScaleMergeGPU.h | 6 +++ 3 files changed, 83 insertions(+), 13 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 1f706306e..7394dbe90 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -5352,21 +5352,36 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool if (full_stats && error_model_active) { std::vector &gs = asymptote_scatter; gs.assign(n_groups, GroupScatter{}); + bool did_gpu_scatter = false; +#ifdef JFJOCH_USE_CUDA + // On the device where the merge is, to the same bits. + if (use_gpu_merge) { + std::vector sum(n_groups), sum_sq(n_groups), sum_var(n_groups); + std::vector n(n_groups); + gpu_->MergeGroupScatter(sum.data(), sum_sq.data(), sum_var.data(), n.data()); + ParallelChunks(n_groups, ThreadsForWork(n_groups, nthreads), [&](int lo, int hi) { + for (int g = lo; g < hi; ++g) gs[g] = {sum[g], sum_sq[g], sum_var[g], n[g]}; + }); + did_gpu_scatter = true; + } +#endif // Split by groups, as the Bijvoet sums above: every group is still summed in index order. - ParallelBlocks(n_groups, nthreads, [&](int, int g0, int g1) { - for (int gi = g0; gi < g1; ++gi) { - for (int p = fulls_group_start[gi]; p < fulls_group_start[gi + 1]; ++p) { - const int i = fulls_group_perm[p]; - if (mf.group[i] != gi) continue; - const double I_corr = static_cast(mf.I[i]) * mf.corr[i]; - const double sigma_corr = static_cast(mf.sigma[i]) * mf.corr[i]; - auto &g = gs[gi]; - g.sum += I_corr; g.sum_sq += I_corr * I_corr; - g.sum_var += counting_variance(fulls[i], em_mean[gi], sigma_corr * sigma_corr); - ++g.n; + if (!did_gpu_scatter) { + ParallelBlocks(n_groups, nthreads, [&](int, int g0, int g1) { + for (int gi = g0; gi < g1; ++gi) { + for (int p = fulls_group_start[gi]; p < fulls_group_start[gi + 1]; ++p) { + const int i = fulls_group_perm[p]; + if (mf.group[i] != gi) continue; + const double I_corr = static_cast(mf.I[i]) * mf.corr[i]; + const double sigma_corr = static_cast(mf.sigma[i]) * mf.corr[i]; + auto &g = gs[gi]; + g.sum += I_corr; g.sum_sq += I_corr * I_corr; + g.sum_var += counting_variance(fulls[i], em_mean[gi], sigma_corr * sigma_corr); + ++g.n; + } } - } - }, 512); + }, 512); + } // Per-group counting-subtracted fractional systematic variance, paired with the group's I/sigma. // Two conventions decide whether this measures anything: // * the counting term is the error model's OWN, a*sigma^2. Subtracting the raw sigma^2 while the diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 8c477d324..79972d1b5 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -668,6 +668,28 @@ namespace { } } + // One thread per group: the sums the strong-reflection ISa asymptote is read from (MergeAndStats), + // over the group's usable fulls in index order, with the host's rounding. + __global__ void GroupScatterKernel(MergeParams p, double *sum, double *sum_sq, double *sum_var, int32_t *n) { + for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) { + const int lo = p.gstart[g], hi = lo + p.gcount[g]; + double s = 0.0, s2 = 0.0, sv = 0.0; + int cnt = 0; + for (int q = lo; q < hi; ++q) { + const int i = p.gperm[q]; + if (!MergeUsable(i, p)) continue; + const double I_corr = __dmul_rn(double(p.I[i]), double(p.corr[i])); + const double sigma_corr = __dmul_rn(double(p.sigma[i]), double(p.corr[i])); + s = __dadd_rn(s, I_corr); + s2 = __dadd_rn(s2, __dmul_rn(I_corr, I_corr)); + sv = __dadd_rn(sv, CountingVarianceExact(p.corr[i], p.var_bkg[i], p.var_per_I[i], p.capture[i], + p.em_mean[g], __dmul_rn(sigma_corr, sigma_corr))); + ++cnt; + } + sum[g] = s; sum_sq[g] = s2; sum_var[g] = sv; n[g] = cnt; + } + } + // The leverage-corrected error-model sample of full i; false where the full gives none. __device__ __forceinline__ bool EmSample(int i, const MergeParams &p, ErrorModelSample &out) { if (!MergeUsable(i, p)) return false; @@ -1536,6 +1558,33 @@ void RotationScaleMergeGPU::MergeBijvoetExcess(double error_model_a, double erro for (int g = 0; g < ng; ++g) { d_lo = std::min(d_lo, lo[g]); d_hi = std::max(d_hi, hi[g]); } } +void RotationScaleMergeGPU::MergeGroupScatter(double *sum, double *sum_sq, double *sum_var, int32_t *n) { + DeviceGuard guard(impl_->device, impl_->available); + auto &d = *impl_; + const int ng = d.n_groups; + auto g_sum = d.Alloc(std::max(1, ng)), g_sum_sq = d.Alloc(std::max(1, ng)), + g_sum_var = d.Alloc(std::max(1, ng)); + auto g_n = d.Alloc(std::max(1, ng)); + + MergeParams p{}; + p.n_groups = ng; p.min_partiality = d.merge_min_part; p.for_search = d.merge_for_search; + p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get(); p.partiality = d.f_partiality.get(); + p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get(); + p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get(); + p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get(); + p.em_mean = d.m_em_mean.get(); p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get(); + p.capture = d.f_capture.get(); + + const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK); + GroupScatterKernel<<s()>>>(p, g_sum.get(), g_sum_sq.get(), g_sum_var.get(), g_n.get()); + CudaCheck(cudaGetLastError(), "group scatter launch"); + CopyAndWait(sum, g_sum.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sum"); + CopyAndWait(sum_sq, g_sum_sq.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sum_sq"); + CopyAndWait(sum_var, g_sum_var.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), + "dl sum_var"); + CopyAndWait(n, g_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl n"); +} + void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio) { 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 1d68aa8da..7ec3163de 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -106,6 +106,12 @@ public: void MergeBijvoetExcess(double error_model_a, double error_model_b, float *group_d, double *excess, float &d_lo, float &d_hi); + // The sums the strong-reflection ISa asymptote is read from (MergeAndStats), per group (length + // n_groups): sum and sum of squares of I*corr over the usable fulls, the sum of their counting + // variances at the group's em_mean, and their count. The host's numbers to the bit. Requires + // MergeEmSamples first. + void MergeGroupScatter(double *sum, double *sum_sq, double *sum_var, int32_t *n); + // Per-group merge accumulators (inv-var sums + deterministic half-sets, error-model-corrected sigma // from a/b). `half` is the per-full CC1/2 half-set (length n_fulls) assigned on the host, so this // kernel and the host merge loop cannot disagree about it.