From 34b3c3c4e7c79097cee75c834e08467a084f8cba Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Fri, 3 Jul 2026 10:38:36 +0200 Subject: [PATCH] RotationScaleMerge: GPU merge + error-model reductions over the resident fulls Port the four fulls-walking reductions of MergeAndStats to the GPU, over the fulls group-CSR already resident from scale-fulls: the per-group inv-var mean + leverage- corrected error-model samples, the merge accumulate (inv-var sums + deterministic half-sets, error-model-corrected sigma, with outlier rejection), and R_meas + the per-shell usable count. The host keeps the parts that don't parallelise cleanly or are tiny: the I2-sort + 16-bin (a,b) median fit, the per-group reject median (a per-group median is awkward on the GPU - cheap on the host from the GPU cnt), the merged export, the shells and the gemmi completeness. Only per-group arrays (~55k) + the samples (~n_fulls, for the fit) come back - the fulls are not re-walked on the host. Device HalfForImage (splitmix64) + IceRingIndex mirror the host; the corrected-sigma uses (b*I_for_b)^2 (not b^2*I^2) to match the host rounding; the R_meas usable count requires finite d (the host counts only fulls with a valid shell, and a group's fulls share d, so the shell is assigned per group). Gated on fulls_resident (GPU combine+scale-fulls active); reject is fully supported so it runs for the default rot3d command. merge+stats ~0.49 -> ~0.37s, taking RSM on lyso to ~0.78s (was ~0.91). Validated across the battery: 15/15 deterministic crystals bit-identical to the CPU path (SG / ISa / CC1.2 / completeness / total-obs, and the exact outlier-reject count), only EP_cs_01-24 noise wobbles. The em-sort + a,b fit are the remaining host floor. Non-CUDA build unaffected (use_gpu_merge is always false there). Co-Authored-By: Claude Opus 4.8 --- .../scale_merge/RotationScaleMerge.cpp | 194 ++++++++---- .../scale_merge/RotationScaleMerge.h | 5 +- .../scale_merge/RotationScaleMergeGPU.cu | 298 ++++++++++++++++++ .../scale_merge/RotationScaleMergeGPU.h | 26 ++ 4 files changed, 459 insertions(+), 64 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 7dfee378..fca148f4 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -271,6 +271,7 @@ void RotationScaleMerge::Ingest() { gpu_->SetRawRuns(static_cast(rawrun_start.size()), static_cast(perm.size()), perm.data(), rawrun_start.data(), rawrun_count.data(), rawrun_h.data(), rawrun_k.data(), rawrun_l.data()); + gpu_->SetFrameCellOk(frame_cell_ok.data()); gpu_combine_ = std::getenv("JFJOCH_RSM_CPU_COMBINE") == nullptr; logger.Info("RotationScaleMerge: GPU partial-scaling{} active", gpu_combine_ ? " + combine + scale-fulls" : ""); @@ -810,7 +811,8 @@ namespace { } RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool for_search, - const std::vector &masked) { + const std::vector &masked, + bool fulls_resident) { // A full is usable for the merge / error model if it passes AddImage's filters (with the current // ice/masked-ring context). group >= 0 already encodes "not absent and passes AcceptReflection". auto masked_ring = [&](const Obs &o) { @@ -829,24 +831,68 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool return std::isfinite(I_corr) && std::isfinite(sigma_corr) && sigma_corr > 0.0f; }; + // The em-stats / samples / merge-accumulate / R_meas reductions run on the resident, scaled fulls + // (their group CSR is still on the device from scale-fulls) when fulls_resident; the host keeps the + // I2-sort, the (a,b) fit, the export and the statistics. reject_outliers is excluded upstream. + bool use_gpu_merge = false; +#ifdef JFJOCH_USE_CUDA + use_gpu_merge = fulls_resident && !fulls.empty(); + std::vector gpu_masked(masked.begin(), masked.end()); +#endif + // ---- Error model: fit dev2 = a*sigma^2 + b^2*^2 from symmetry-equivalent scatter. ---- std::vector em_mean(n_groups, NAN); std::vector reject_median(n_groups, NAN); double error_model_a = 1.0, error_model_b = 0.0, error_model_chi2 = 0.0; bool error_model_active = false; { - // Per-group inverse-variance mean over usable fulls (>=2 obs), and the leverage-corrected samples. - std::vector sw(n_groups, 0.0), swI(n_groups, 0.0); - std::vector cnt(n_groups, 0); - for (const auto &o : fulls) { - if (!usable_merge(o)) continue; - const double sigma_corr = static_cast(o.sigma) * o.corr; - const double w = 1.0 / (sigma_corr * sigma_corr); - sw[o.group] += w; swI[o.group] += w * (static_cast(o.I) * o.corr); cnt[o.group]++; + struct Sample { double s2, I2, dev2; }; + std::vector samples; + std::vector cnt(n_groups, 0); // per-group usable count (both paths; feeds reject-median) + bool did_gpu = false; +#ifdef JFJOCH_USE_CUDA + if (use_gpu_merge) { + const int nf = static_cast(fulls.size()); + std::vector gs2(nf), gI2(nf), gdev2(nf); + std::vector gvalid(nf); + gpu_->MergeEmSamples(for_search, gpu_masked.data(), static_cast(gpu_masked.size()), + ice_half_width_q, min_partiality, em_mean.data(), cnt.data(), + gs2.data(), gI2.data(), gdev2.data(), gvalid.data()); + samples.reserve(nf); + for (int i = 0; i < nf; ++i) + if (gvalid[i]) samples.push_back({gs2[i], gI2[i], gdev2[i]}); + did_gpu = true; } - for (int g = 0; g < n_groups; ++g) - if (cnt[g] >= 2 && sw[g] > 0.0) em_mean[g] = swI[g] / sw[g]; +#endif + if (!did_gpu) { + // Per-group inverse-variance mean over usable fulls (>=2 obs), and the leverage-corrected samples. + std::vector sw(n_groups, 0.0), swI(n_groups, 0.0); + for (const auto &o : fulls) { + if (!usable_merge(o)) continue; + const double sigma_corr = static_cast(o.sigma) * o.corr; + const double w = 1.0 / (sigma_corr * sigma_corr); + sw[o.group] += w; swI[o.group] += w * (static_cast(o.I) * o.corr); cnt[o.group]++; + } + for (int g = 0; g < n_groups; ++g) + if (cnt[g] >= 2 && sw[g] > 0.0) em_mean[g] = swI[g] / sw[g]; + samples.reserve(fulls.size()); + for (const auto &o : fulls) { + if (!usable_merge(o) || cnt[o.group] < 2) continue; + const double mean = em_mean[o.group]; + if (!std::isfinite(mean)) continue; + const double sigma_corr = static_cast(o.sigma) * o.corr; + const double s2 = sigma_corr * sigma_corr; + const double w = 1.0 / s2; + const double factor = 1.0 - w / sw[o.group]; + if (factor < 0.05) continue; + const double resid = static_cast(o.I) * o.corr - mean; + samples.push_back({s2, mean * mean, resid * resid / factor}); + } + } + + // Per-group outlier-rejection median of I*corr (host both paths - a per-group median is awkward on + // the GPU; cheap here, cnt >= 2 filter from the em pass). Fed to the merge accumulate. if (reject_outliers) { std::vector> iv(n_groups); for (const auto &o : fulls) @@ -858,23 +904,6 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool reject_median[g] = iv[g][iv[g].size() / 2]; } } - - struct Sample { double s2, I2, dev2; }; - std::vector samples; - samples.reserve(fulls.size()); - for (const auto &o : fulls) { - if (!usable_merge(o) || cnt[o.group] < 2) continue; - const double mean = em_mean[o.group]; - if (!std::isfinite(mean)) continue; - const double sigma_corr = static_cast(o.sigma) * o.corr; - const double s2 = sigma_corr * sigma_corr; - const double w = 1.0 / s2; - const double factor = 1.0 - w / sw[o.group]; - if (factor < 0.05) continue; - const double resid = static_cast(o.I) * o.corr - mean; - samples.push_back({s2, mean * mean, resid * resid / factor}); - } - constexpr int n_bins = 16; if (samples.size() >= static_cast(8 * n_bins)) { std::sort(samples.begin(), samples.end(), @@ -929,30 +958,50 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool + (error_model_b * I_for_b) * (error_model_b * I_for_b); return v > 0.0 ? static_cast(std::sqrt(v)) : sigma_corr; }; - // ---- Merge: per-group inverse-variance sums with corrected sigma + deterministic half sets. ---- struct Accum { double swI = 0, sw = 0, swIh[2] = {0, 0}, swh[2] = {0, 0}; size_t nh[2] = {0, 0}; float d = NAN; }; std::vector acc(n_groups); size_t reject_count = 0; - for (const auto &o : fulls) { - if (!usable_merge(o)) continue; - const float I_corr = o.I * o.corr; - float sigma_corr = o.sigma * o.corr; - sigma_corr = corrected_sigma(I_corr, sigma_corr, o.group); - if (reject_outliers && error_model_active && std::isfinite(reject_median[o.group]) - && std::fabs(I_corr - reject_median[o.group]) > reject_nsigma * sigma_corr) { - ++reject_count; - continue; + bool did_gpu_acc = false; +#ifdef JFJOCH_USE_CUDA + if (use_gpu_merge) { + std::vector aswI(n_groups), asw(n_groups), aswIh0(n_groups), aswIh1(n_groups), + aswh0(n_groups), aswh1(n_groups), ad(n_groups); + std::vector anh0(n_groups), anh1(n_groups), arej(n_groups); + gpu_->MergeAccum(error_model_a, error_model_b, error_model_active, + reject_outliers, reject_nsigma, reject_median.data(), + aswI.data(), asw.data(), aswIh0.data(), aswIh1.data(), + aswh0.data(), aswh1.data(), anh0.data(), anh1.data(), ad.data(), arej.data()); + for (int g = 0; g < n_groups; ++g) { + Accum &a = acc[g]; + a.swI = aswI[g]; a.sw = asw[g]; a.swIh[0] = aswIh0[g]; a.swIh[1] = aswIh1[g]; + a.swh[0] = aswh0[g]; a.swh[1] = aswh1[g]; + a.nh[0] = static_cast(anh0[g]); a.nh[1] = static_cast(anh1[g]); + a.d = static_cast(ad[g]); + reject_count += static_cast(arej[g]); } - const double w = 1.0 / (static_cast(sigma_corr) * sigma_corr); - const double wI = w * I_corr; - const int half = HalfForImage(o.frame); - auto &a = acc[o.group]; - a.swI += wI; a.sw += w; - a.swIh[half] += wI; a.swh[half] += w; a.nh[half]++; - if (!std::isfinite(a.d) && std::isfinite(o.d) && o.d > 0.0f) a.d = o.d; + did_gpu_acc = true; } - +#endif + if (!did_gpu_acc) + for (const auto &o : fulls) { + if (!usable_merge(o)) continue; + const float I_corr = o.I * o.corr; + float sigma_corr = o.sigma * o.corr; + sigma_corr = corrected_sigma(I_corr, sigma_corr, o.group); + if (reject_outliers && error_model_active && std::isfinite(reject_median[o.group]) + && std::fabs(I_corr - reject_median[o.group]) > reject_nsigma * sigma_corr) { + ++reject_count; + continue; + } + const double w = 1.0 / (static_cast(sigma_corr) * sigma_corr); + const double wI = w * I_corr; + const int half = HalfForImage(o.frame); + auto &a = acc[o.group]; + a.swI += wI; a.sw += w; + a.swIh[half] += wI; a.swh[half] += w; a.nh[half]++; + if (!std::isfinite(a.d) && std::isfinite(o.d) && o.d > 0.0f) a.d = o.d; + } // ---- Export merged reflections (+ resolution-shell R-free flags). ---- Result result; result.isa = error_model_b > 0 ? 1.0 / error_model_b : 0.0; @@ -997,7 +1046,6 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool if (reject_count > 0) logger.Info("Merge outlier rejection: dropped {} observations", reject_count); - // ---- Statistics (10 shells): completeness, multiplicity, , R_meas, CC1/2. ---- constexpr int n_shells = 10; float sd_min = std::numeric_limits::max(), sd_max = 0.0f; @@ -1045,22 +1093,43 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool // |I_i - | per reflection. struct RmeasObs { double sum_abs_dev = 0, sum_I = 0; int n = 0, shell = -1; }; std::vector rmeas(n_groups); - for (const auto &o : fulls) { - if (o.group < 0) continue; - if (!frame_cell_ok[o.frame]) continue; - if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) continue; - if (o.partiality < min_partiality) continue; - const float I_corr = o.I * o.corr, sigma_corr = o.sigma * o.corr; - if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) continue; - const auto shell = shells.GetShell(o.d); - if (!shell || *shell < 0 || *shell >= n_shells) continue; - sa[*shell].total_obs++; - if (std::isfinite(merged_I[o.group])) { - auto &r = rmeas[o.group]; - r.sum_abs_dev += std::fabs(static_cast(I_corr) - merged_I[o.group]); - r.sum_I += I_corr; r.n++; r.shell = *shell; + 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); + std::vector rn(n_groups), rnusable(n_groups); + gpu_->MergeRmeas(merged_I.data(), rabsdev.data(), rsumI.data(), rn.data(), rnusable.data()); + for (int g = 0; g < n_groups; ++g) { + if (rnusable[g] == 0) continue; + const auto shell = shells.GetShell(acc[g].d); + if (!shell || *shell < 0 || *shell >= n_shells) continue; + sa[*shell].total_obs += rnusable[g]; + 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.n = rn[g]; r.shell = *shell; + } } + did_gpu_rmeas = true; } +#endif + if (!did_gpu_rmeas) + for (const auto &o : fulls) { + if (o.group < 0) continue; + if (!frame_cell_ok[o.frame]) continue; + if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) continue; + if (o.partiality < min_partiality) continue; + const float I_corr = o.I * o.corr, sigma_corr = o.sigma * o.corr; + if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) continue; + const auto shell = shells.GetShell(o.d); + if (!shell || *shell < 0 || *shell >= n_shells) continue; + sa[*shell].total_obs++; + if (std::isfinite(merged_I[o.group])) { + auto &r = rmeas[o.group]; + r.sum_abs_dev += std::fabs(static_cast(I_corr) - merged_I[o.group]); + r.sum_I += I_corr; r.n++; r.shell = *shell; + } + } std::vector rmeas_num(n_shells, 0.0), rmeas_den(n_shells, 0.0); double rmeas_num_all = 0.0, rmeas_den_all = 0.0; for (const auto &r : rmeas) { @@ -1098,7 +1167,6 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool overall.cc_half = cc_half_overall.GetCC(); overall.cc_ref = NAN; overall.r_meas = rmeas_den_all > 0.0 ? rmeas_num_all / rmeas_den_all : NAN; - logger.Info("Merge complete ({} unique reflections)", result.merged.size()); return result; } @@ -1274,7 +1342,7 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, lap("scale fulls"); // --- 5. Error model + merge + statistics. --- - auto r = MergeAndStats(n_groups, for_search, masked_ice_rings); + auto r = MergeAndStats(n_groups, for_search, masked_ice_rings, combined_on_gpu && scaled_fulls_on_gpu); lap("merge+stats"); return r; } diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index cf2284f8..0803c084 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -192,5 +192,8 @@ private: const std::vector &frame_scaled); // Error model + merge + statistics over the fulls (the last stage). n_groups is the fulls group count. - Result MergeAndStats(int n_groups, bool for_search, const std::vector &masked_ice_rings); + // fulls_resident: the (scaled) fulls + their group CSR are still on the GPU, so the em-stats / samples + // / merge-accumulate / R_meas reductions run there (only per-group + samples come back). + Result MergeAndStats(int n_groups, bool for_search, const std::vector &masked_ice_rings, + bool fulls_resident); }; diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 955c79f4..8e3e74ba 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -406,6 +406,172 @@ namespace { for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) p[i] = v; } + // ===== error-model + merge reductions over the resident, scaled fulls (mirror MergeAndStats) ===== + + // Device copy of common/Definitions.h IceRingIndex (only consulted when the merge has masked rings). + __device__ __forceinline__ int IceRingIndexDev(float d, float hw) { + if (!(d > 0.0f)) return -1; + const float two_pi = 6.283185307f; + const float q = two_pi / d; + const float rings[11] = {3.895f, 3.661f, 3.438f, 2.667f, 2.249f, 2.068f, + 1.947f, 1.916f, 1.882f, 1.719f, 1.522f}; + for (int i = 0; i < 11; ++i) + if (fabsf(q - two_pi / rings[i]) < hw) return i; + return -1; + } + + // Deterministic CC1/2 half from the frame index (splitmix64), matching Merge.cpp / RSM HalfForImage. + __device__ __forceinline__ int HalfForImageDev(long long image_id) { + unsigned long long z = (unsigned long long)image_id + 0x9e3779b97f4a7c15ULL; + z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL; + z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL; + z = z ^ (z >> 31); + return (int)(z & 1ULL); + } + + struct MergeParams { + int n_groups; + double min_partiality, error_model_a, error_model_b, reject_nsigma; + float ice_hw; + int for_search, n_masked, error_model_active, reject_outliers; + const float *I, *sigma, *corr, *partiality, *d, *reject_median; + const int32_t *group, *frame; + const uint8_t *on_ice, *frame_cell_ok, *masked; + const int32_t *gperm, *gstart, *gcount; + const double *em_mean, *merged_I; + // outputs + double *sw, *swI, *em_mean_out; + int32_t *cnt; + double *s2, *I2, *dev2; + uint8_t *valid; + double *a_swI, *a_sw, *a_swIh0, *a_swIh1, *a_swh0, *a_swh1, *a_d; + int32_t *a_nh0, *a_nh1, *a_rejected; + double *r_absdev, *r_sumI; + int32_t *r_n, *r_nusable; + }; + + // A full passes the merge / error-model filter (mirrors MergeAndStats::usable_merge). + __device__ __forceinline__ bool MergeUsable(int i, const MergeParams &p) { + const int g = p.group[i]; + if (g < 0) return false; + if (!p.frame_cell_ok[p.frame[i]]) return false; + const float c = p.corr[i]; + if (!(c > 0.0f) || !isfinite(c)) return false; + if (p.for_search && p.on_ice[i]) return false; + if (p.n_masked > 0) { + const int ring = IceRingIndexDev(p.d[i], p.ice_hw); + if (ring >= 0 && ring < p.n_masked && p.masked[ring]) return false; + } + if (p.partiality[i] < p.min_partiality) return false; + const float I_corr = p.I[i] * c, sigma_corr = p.sigma[i] * c; + return isfinite(I_corr) && isfinite(sigma_corr) && sigma_corr > 0.0f; + } + + // The looser R_meas filter (no ice / masked-ring / for_search - Mask = cell only). + __device__ __forceinline__ bool RmeasUsable(int i, const MergeParams &p) { + const int g = p.group[i]; + if (g < 0) return false; + if (!p.frame_cell_ok[p.frame[i]]) return false; + const float c = p.corr[i]; + if (!(c > 0.0f) || !isfinite(c)) return false; + if (p.partiality[i] < p.min_partiality) return false; + const float I_corr = p.I[i] * c, sigma_corr = p.sigma[i] * c; + return isfinite(I_corr) && isfinite(sigma_corr) && sigma_corr > 0.0f; + } + + // One thread per group: inverse-variance sums over the group's usable fulls + the group mean (>=2). + __global__ void MergeEmStatsKernel(MergeParams p) { + 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 sw = 0.0, swI = 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 sigma_corr = double(p.sigma[i]) * p.corr[i]; + const double w = 1.0 / (sigma_corr * sigma_corr); + sw += w; swI += w * (double(p.I[i]) * p.corr[i]); ++cnt; + } + p.sw[g] = sw; p.swI[g] = swI; p.cnt[g] = cnt; + p.em_mean_out[g] = (cnt >= 2 && sw > 0.0) ? swI / sw : NAN; + } + } + + // One thread per full: the leverage-corrected error-model sample, or valid=0 if dropped. + __global__ void MergeSamplesKernel(int n_obs, MergeParams p) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) { + p.valid[i] = 0; + if (!MergeUsable(i, p)) continue; + const int g = p.group[i]; + if (p.cnt[g] < 2) continue; + const double mean = p.em_mean[g]; + if (!isfinite(mean)) continue; + const double sigma_corr = double(p.sigma[i]) * p.corr[i]; + const double s2 = sigma_corr * sigma_corr; + const double factor = 1.0 - (1.0 / s2) / p.sw[g]; + if (factor < 0.05) continue; + const double resid = double(p.I[i]) * p.corr[i] - mean; + p.s2[i] = s2; p.I2[i] = mean * mean; p.dev2[i] = resid * resid / factor; p.valid[i] = 1; + } + } + + // One thread per group: the merge accumulators (inv-var sums + deterministic half-sets), with the + // error-model-corrected sigma. Mirrors MergeAndStats' merge loop (reject path stays on the host). + __global__ void MergeAccumKernel(MergeParams p) { + 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 swI = 0, sw = 0, swIh0 = 0, swIh1 = 0, swh0 = 0, swh1 = 0, dd = NAN; + int nh0 = 0, nh1 = 0, rejected = 0; + const float rmed = p.reject_median[g]; + for (int q = lo; q < hi; ++q) { + const int i = p.gperm[q]; + if (!MergeUsable(i, p)) continue; + const float I_corr = p.I[i] * p.corr[i]; + float sigma_corr = p.sigma[i] * p.corr[i]; + if (p.error_model_active) { + const double I_for_b = isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr); + const double bi = p.error_model_b * I_for_b; + const double v = p.error_model_a * double(sigma_corr) * sigma_corr + bi * bi; + if (v > 0.0) sigma_corr = float(sqrt(v)); + } + if (p.reject_outliers && p.error_model_active && isfinite(rmed) + && fabsf(I_corr - rmed) > p.reject_nsigma * sigma_corr) { ++rejected; continue; } + const double w = 1.0 / (double(sigma_corr) * sigma_corr); + const double wI = w * I_corr; + const int half = HalfForImageDev(p.frame[i]); + swI += wI; sw += w; + if (half) { swIh1 += wI; swh1 += w; ++nh1; } else { swIh0 += wI; swh0 += w; ++nh0; } + if (!isfinite(dd) && isfinite(p.d[i]) && p.d[i] > 0.0f) dd = p.d[i]; + } + p.a_swI[g] = swI; p.a_sw[g] = sw; p.a_swIh0[g] = swIh0; p.a_swIh1[g] = swIh1; + p.a_swh0[g] = swh0; p.a_swh1[g] = swh1; p.a_nh0[g] = nh0; p.a_nh1[g] = nh1; p.a_d[g] = dd; + p.a_rejected[g] = rejected; + } + } + + // One thread per group: R_meas accumulators (sum|I_corr - merged_I|, sum_I, n) + usable count for the + // per-shell total_observations. Mirrors MergeAndStats' R_meas re-walk (looser, cell-only filter). + __global__ void MergeRmeasKernel(MergeParams p) { + 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]; + const double mI = p.merged_I[g]; + const bool have = isfinite(mI); + double absdev = 0, sumI = 0; + int n = 0, nusable = 0; + for (int q = lo; q < hi; ++q) { + const int i = p.gperm[q]; + if (!RmeasUsable(i, p)) continue; + if (!isfinite(p.d[i]) || !(p.d[i] > 0.0f)) continue; // host counts only fulls with a shell + ++nusable; + if (have) { + const double I_corr = double(p.I[i]) * p.corr[i]; + absdev += fabs(I_corr - mI); sumI += I_corr; ++n; + } + } + p.r_absdev[g] = absdev; p.r_sumI[g] = sumI; p.r_n[g] = n; p.r_nusable[g] = nusable; + } + } + void CudaCheck(cudaError_t e, const char *what) { if (e != cudaSuccess) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, @@ -441,6 +607,20 @@ struct RotationScaleMergeGPU::Impl { CudaDevicePtr cc_n; CudaDevicePtr smooth_apply; // per-frame smooth-G apply flag + ratio, length n_frames CudaDevicePtr smooth_ratio; + // merge / error-model reductions over the resident fulls (reuse the fulls group CSR f_gperm/...) + CudaDevicePtr frame_cell_ok, merge_masked; + CudaDevicePtr m_sw, m_swI, m_em_mean; // per group (n_groups) + CudaDevicePtr m_cnt; + CudaDevicePtr m_s2, m_I2, m_dev2; // per full (n_fulls) + CudaDevicePtr m_valid; + CudaDevicePtr a_swI, a_sw, a_swIh0, a_swIh1, a_swh0, a_swh1, a_d; // merge accum per group + CudaDevicePtr a_nh0, a_nh1, a_rejected; + CudaDevicePtr reject_median; + CudaDevicePtr merged_I, r_absdev, r_sumI; // R_meas per group (merged_I uploaded) + CudaDevicePtr r_n, r_nusable; + int merge_for_search = 0, merge_n_masked = 0; // filter context for one MergeAndStats call + float merge_ice_hw = 0.0f; + double merge_min_part = 0.0; // combine: extra per-obs inputs + the one-time raw-hkl run layout CudaDevicePtr bkg, image_number, d_obs; @@ -544,6 +724,124 @@ void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const { cudaMemcpyDeviceToHost), "download scaled"); } +void RotationScaleMergeGPU::SetFrameCellOk(const uint8_t *frame_cell_ok) { + Upload(impl_->frame_cell_ok, frame_cell_ok, impl_->n_frames); +} + +// The per-group inv-var mean (em_mean) + the per-full leverage-corrected error-model samples over the +// resident+scaled fulls. Stashes the filter context for the later MergeAccum/MergeRmeas calls. +void RotationScaleMergeGPU::MergeEmSamples(bool for_search, const uint8_t *masked, int n_masked, + double ice_half_width_q, double min_partiality, + double *em_mean_out, int32_t *cnt_out, double *s2_out, + double *I2_out, double *dev2_out, uint8_t *valid_out) { + auto &d = *impl_; + const int ng = d.n_groups, nf = d.n_fulls; + d.merge_for_search = for_search ? 1 : 0; d.merge_n_masked = n_masked; + d.merge_ice_hw = float(ice_half_width_q); d.merge_min_part = min_partiality; + d.m_sw = CudaDevicePtr(std::max(1, ng)); d.m_swI = CudaDevicePtr(std::max(1, ng)); + d.m_em_mean = CudaDevicePtr(std::max(1, ng)); d.m_cnt = CudaDevicePtr(std::max(1, ng)); + d.m_s2 = CudaDevicePtr(std::max(1, nf)); d.m_I2 = CudaDevicePtr(std::max(1, nf)); + d.m_dev2 = CudaDevicePtr(std::max(1, nf)); d.m_valid = CudaDevicePtr(std::max(1, nf)); + Upload(d.merge_masked, masked, n_masked); + + MergeParams p{}; + p.n_groups = ng; p.min_partiality = min_partiality; p.ice_hw = d.merge_ice_hw; + p.for_search = d.merge_for_search; p.n_masked = n_masked; + 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.masked = d.merge_masked.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.sw = d.m_sw.get(); p.swI = d.m_swI.get(); + p.em_mean_out = d.m_em_mean.get(); p.cnt = d.m_cnt.get(); + p.s2 = d.m_s2.get(); p.I2 = d.m_I2.get(); p.dev2 = d.m_dev2.get(); p.valid = d.m_valid.get(); + + const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK); + const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK); + MergeEmStatsKernel<<>>(p); + MergeSamplesKernel<<>>(nf, p); + CudaCheck(cudaGetLastError(), "merge em/samples launch"); + CudaCheck(cudaDeviceSynchronize(), "merge em/samples sync"); + CudaCheck(cudaMemcpy(em_mean_out, d.m_em_mean.get(), size_t(ng) * sizeof(double), + cudaMemcpyDeviceToHost), "dl em_mean"); + CudaCheck(cudaMemcpy(cnt_out, d.m_cnt.get(), size_t(ng) * sizeof(int32_t), + cudaMemcpyDeviceToHost), "dl cnt"); + if (nf > 0) { + CudaCheck(cudaMemcpy(s2_out, d.m_s2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl s2"); + CudaCheck(cudaMemcpy(I2_out, d.m_I2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl I2"); + CudaCheck(cudaMemcpy(dev2_out, d.m_dev2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl dev2"); + CudaCheck(cudaMemcpy(valid_out, d.m_valid.get(), size_t(nf) * sizeof(uint8_t), cudaMemcpyDeviceToHost), "dl valid"); + } +} + +void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_b, bool error_model_active, + bool reject_outliers, double reject_nsigma, const float *reject_median, + double *swI, double *sw, double *swIh0, double *swIh1, + double *swh0, double *swh1, int32_t *nh0, int32_t *nh1, double *d_out, + int32_t *rejected) { + auto &d = *impl_; + const int ng = d.n_groups; + d.a_swI = CudaDevicePtr(std::max(1, ng)); d.a_sw = CudaDevicePtr(std::max(1, ng)); + d.a_swIh0 = CudaDevicePtr(std::max(1, ng)); d.a_swIh1 = CudaDevicePtr(std::max(1, ng)); + d.a_swh0 = CudaDevicePtr(std::max(1, ng)); d.a_swh1 = CudaDevicePtr(std::max(1, ng)); + d.a_nh0 = CudaDevicePtr(std::max(1, ng)); d.a_nh1 = CudaDevicePtr(std::max(1, ng)); + d.a_d = CudaDevicePtr(std::max(1, ng)); d.a_rejected = CudaDevicePtr(std::max(1, ng)); + Upload(d.reject_median, reject_median, ng); + + MergeParams p{}; + p.n_groups = ng; p.min_partiality = d.merge_min_part; p.ice_hw = d.merge_ice_hw; + p.for_search = d.merge_for_search; p.n_masked = d.merge_n_masked; + p.error_model_a = error_model_a; p.error_model_b = error_model_b; + p.error_model_active = error_model_active ? 1 : 0; + p.reject_outliers = reject_outliers ? 1 : 0; p.reject_nsigma = reject_nsigma; + p.reject_median = d.reject_median.get(); + 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.masked = d.merge_masked.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.a_swI = d.a_swI.get(); p.a_sw = d.a_sw.get(); p.a_swIh0 = d.a_swIh0.get(); p.a_swIh1 = d.a_swIh1.get(); + p.a_swh0 = d.a_swh0.get(); p.a_swh1 = d.a_swh1.get(); p.a_nh0 = d.a_nh0.get(); p.a_nh1 = d.a_nh1.get(); + p.a_d = d.a_d.get(); p.a_rejected = d.a_rejected.get(); + + const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK); + MergeAccumKernel<<>>(p); + CudaCheck(cudaGetLastError(), "merge accum launch"); + CudaCheck(cudaDeviceSynchronize(), "merge accum sync"); + auto dl = [&](void *h, const CudaDevicePtr &s) { + CudaCheck(cudaMemcpy(h, s.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl accum"); }; + dl(swI, d.a_swI); dl(sw, d.a_sw); dl(swIh0, d.a_swIh0); dl(swIh1, d.a_swIh1); + dl(swh0, d.a_swh0); dl(swh1, d.a_swh1); dl(d_out, d.a_d); + CudaCheck(cudaMemcpy(nh0, d.a_nh0.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl nh0"); + CudaCheck(cudaMemcpy(nh1, d.a_nh1.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl nh1"); + CudaCheck(cudaMemcpy(rejected, d.a_rejected.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rej"); +} + +void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, double *sumI, + int32_t *n, int32_t *nusable) { + auto &d = *impl_; + const int ng = d.n_groups; + Upload(d.merged_I, merged_I, ng); + d.r_absdev = CudaDevicePtr(std::max(1, ng)); d.r_sumI = CudaDevicePtr(std::max(1, ng)); + d.r_n = CudaDevicePtr(std::max(1, ng)); d.r_nusable = CudaDevicePtr(std::max(1, ng)); + + MergeParams p{}; + p.n_groups = ng; p.min_partiality = d.merge_min_part; + 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.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.merged_I = d.merged_I.get(); + p.r_absdev = d.r_absdev.get(); p.r_sumI = d.r_sumI.get(); p.r_n = d.r_n.get(); p.r_nusable = d.r_nusable.get(); + + const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK); + MergeRmeasKernel<<>>(p); + CudaCheck(cudaGetLastError(), "merge rmeas launch"); + CudaCheck(cudaDeviceSynchronize(), "merge rmeas sync"); + CudaCheck(cudaMemcpy(absdev, d.r_absdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl absdev"); + CudaCheck(cudaMemcpy(sumI, d.r_sumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl sumI"); + CudaCheck(cudaMemcpy(n, d.r_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rn"); + CudaCheck(cudaMemcpy(nusable, d.r_nusable.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rnusable"); +} + void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio) { auto &d = *impl_; Upload(d.smooth_apply, apply, d.n_frames); diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.h b/image_analysis/scale_merge/RotationScaleMergeGPU.h index 725b16a0..e8704300 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -58,6 +58,32 @@ public: // apply[f] (both length n_frames), matching CPU SmoothG. Keeps corr resident (no round-trip). void SmoothCorr(const uint8_t *apply, const double *ratio); + // --- merge + error-model reductions over the resident, scaled fulls (reuse the fulls group CSR) --- + + // The per-frame cell-consistency mask (length n_frames) used by the merge filter. Uploaded once. + void SetFrameCellOk(const uint8_t *frame_cell_ok); + + // Per-group inv-var mean (em_mean, length n_groups) + per-full leverage-corrected error-model samples + // (s2/I2/dev2 + valid flag, length n_fulls), mirroring MergeAndStats' first two error-model loops. + // Stashes (for_search, masked, ice, min_partiality) for the MergeAccum/MergeRmeas calls that follow. + void MergeEmSamples(bool for_search, const uint8_t *masked, int n_masked, + double ice_half_width_q, double min_partiality, + double *em_mean_out, int32_t *cnt_out, double *s2_out, double *I2_out, + double *dev2_out, uint8_t *valid_out); + + // Per-group merge accumulators (inv-var sums + deterministic half-sets, error-model-corrected sigma + // from a/b). Outputs length n_groups; rejected[g] counts outliers dropped (reject_median uploaded, NAN + // where none). Requires MergeEmSamples first (em_mean resident). + void MergeAccum(double error_model_a, double error_model_b, bool error_model_active, + bool reject_outliers, double reject_nsigma, const float *reject_median, + double *swI, double *sw, double *swIh0, double *swIh1, + double *swh0, double *swh1, int32_t *nh0, int32_t *nh1, double *d_out, + int32_t *rejected); + + // Per-group R_meas accumulators (sum|I_corr-merged_I|, sum_I, n, and the usable count for per-shell + // total_observations); merged_I is uploaded. All arrays length n_groups. + void MergeRmeas(const double *merged_I, double *absdev, double *sumI, int32_t *n, int32_t *nusable); + // Post-smooth per-frame diagnostic CC: recompute the group means from the resident (smoothed) corr // and the Pearson CC of each frame's I*corr vs its group mean, downloading only the per-frame cc / // cc_n (length n_frames). Mirrors ReduceGroupMeans(partials) + FinalizePerFrameScale's CC loop.