From b6f2096015b8742db9a0f2aab3a7b1462276837e Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 07:07:32 +0200 Subject: [PATCH 1/8] RotationScaleMerge: the Wilson outlier test on the device in the GPU build WilsonOutliersGPU runs the test of WilsonOutliers.cpp step for step on the device and returns the same result to the bit: the two orders the host sorts into are total orders, so a stable radix sort of the indices reaches the same sequence; every sum is taken by one thread over its shell or reflection in the host's order; each product the host rounds on its own (-ffp-contract=off) is rounded with __dmul_rn; medians are read off sorts; and the logarithms (t, z, u, the tail scale) are taken on the host from downloaded counts and medians. MergeAndStats calls it where the merge itself is on the device. 8a1a final P61 merge, "rsm merge and statistics" 4.35 -> 2.82 s (the test runs twice there, once per error-model fit). p.mtz md5 unchanged on myob, cytc, 8a1a and 8qaw (GPU build). A Catch2 case checks the device result against the host one on a population that reaches every branch. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- image_analysis/scale_merge/CMakeLists.txt | 3 +- .../scale_merge/RotationScaleMerge.cpp | 12 +- .../scale_merge/WilsonOutliersGPU.cu | 396 ++++++++++++++++++ .../scale_merge/WilsonOutliersGPU.h | 17 + tests/WilsonOutliersTest.cpp | 54 +++ 5 files changed, 480 insertions(+), 2 deletions(-) create mode 100644 image_analysis/scale_merge/WilsonOutliersGPU.cu create mode 100644 image_analysis/scale_merge/WilsonOutliersGPU.h diff --git a/image_analysis/scale_merge/CMakeLists.txt b/image_analysis/scale_merge/CMakeLists.txt index be7c72f30..e52b7439d 100644 --- a/image_analysis/scale_merge/CMakeLists.txt +++ b/image_analysis/scale_merge/CMakeLists.txt @@ -41,5 +41,6 @@ TARGET_LINK_LIBRARIES(JFJochScaleMerge Ceres::ceres Eigen3::Eigen JFJochCommon f IF (JFJOCH_CUDA_AVAILABLE) TARGET_SOURCES(JFJochScaleMerge PRIVATE ../indexing/CUDAMemHelpers.h - RotationScaleMergeGPU.cu RotationScaleMergeGPU.h) + RotationScaleMergeGPU.cu RotationScaleMergeGPU.h + WilsonOutliersGPU.cu WilsonOutliersGPU.h) ENDIF() \ No newline at end of file diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 1c50c3626..2b65114e8 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -25,6 +25,7 @@ #include "FrenchWilson.h" #include "ResolutionCutoff.h" #include "WilsonOutliers.h" +#include "WilsonOutliersGPU.h" #include "OutlierBand.h" #include "../../common/CorrelationCoefficient.h" #include "../../common/CrystalLattice.h" @@ -5132,7 +5133,16 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool group_centric[g] != 0, o.clipped != 0, merge_friedel ? g : pair_of_group[g]}; } }); - wilson = WilsonOutliers(wobs, WILSON_ALPHA, nthreads); + // On the device where the merge is: the same test to the bit (WilsonOutliersGPU.h). + bool did_gpu_wilson = false; +#ifdef JFJOCH_USE_CUDA + if (use_gpu_merge) { + wilson = WilsonOutliersGPU(wobs, WILSON_ALPHA); + did_gpu_wilson = true; + } +#endif + if (!did_gpu_wilson) + wilson = WilsonOutliers(wobs, WILSON_ALPHA, nthreads); for (size_t j = 0; j < wobs.size(); ++j) if (wilson.rejected[j]) rejected_obs[wilson_full[j]] = 1; } diff --git a/image_analysis/scale_merge/WilsonOutliersGPU.cu b/image_analysis/scale_merge/WilsonOutliersGPU.cu new file mode 100644 index 000000000..fddcdfbca --- /dev/null +++ b/image_analysis/scale_merge/WilsonOutliersGPU.cu @@ -0,0 +1,396 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "WilsonOutliersGPU.h" + +#include +#include +#include + +#include + +#include "../indexing/CUDAMemHelpers.h" +#include "../../common/JFJochException.h" + +// Each step below is the step of WilsonOutliers.cpp with the same name; read that file for why. What +// this file adds is only how each one is made to give the host's bits: +// * an order the host sorts into is a total order there (d, then index; unit, then index), so a +// stable radix sort of the indices in ascending order reaches the same sequence; +// * a sum is taken by one thread over its shell or its reflection, in the host's order; +// * nvcc would fuse a product into the add that follows it, which the host build does not +// (-ffp-contract=off), so every such product is rounded on its own with __dmul_rn; +// * a median is an order statistic - a value, the same however it is found - and is read off a sort; +// * the logarithms (t, z, u, the tail scale) are taken on the host, from counts and medians +// downloaded for it, because the device's log is not the host's. +namespace { + constexpr int THREADS = 256; + constexpr size_t OBS_PER_SHELL = 2000; // WilsonOutliers.cpp + constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + + void cuda_err(cudaError_t val) { + if (val != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); + } + + unsigned Blocks(size_t n) { return static_cast(std::max(1, (n + THREADS - 1) / THREADS)); } + + __device__ __forceinline__ bool Judged(const WilsonObservation &o) { + return o.unit >= 0 && o.d > 0.0f && o.d < WILSON_OUTLIER_D_MAX && isfinite(o.I) + && o.sigma > 0.0f && isfinite(o.sigma); + } + + __global__ void JudgedKernel(int n, const WilsonObservation *obs, uint8_t *flag, int32_t *index, float *d) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + flag[i] = Judged(obs[i]) ? 1 : 0; + index[i] = i; + d[i] = obs[i].d; + } + + // lower_e2 of WilsonOutliers.cpp. + __device__ __forceinline__ double LowerE2(float I, float sigma, float epsilon, double z, double mu) { + return __dadd_rn(double(I), -__dmul_rn(z, double(sigma))) / __dmul_rn(double(epsilon), mu); + } + + struct Sorted { // the judged observations in d order, one array per field + const float *I, *sigma, *epsilon; + const uint8_t *centric; + }; + + __global__ void GatherKernel(int n, const WilsonObservation *obs, const int32_t *idx, float *I, float *sigma, + float *epsilon, uint8_t *centric) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n) return; + const WilsonObservation &o = obs[idx[j]]; + I[j] = o.I; sigma[j] = o.sigma; epsilon[j] = o.epsilon; centric[j] = o.centric ? 1 : 0; + } + + // One thread per shell: without the observations the plain bound calls improbable, and + // whether it is established; plus the shell's acentric count. + __global__ void ShellMeanKernel(int n_shells, const int64_t *shell_start, Sorted so, double z, double t, + double *mu, uint8_t *valid, int64_t *n_acentric) { + const int sh = blockIdx.x * blockDim.x + threadIdx.x; + if (sh >= n_shells) return; + const int64_t lo = shell_start[sh], hi = shell_start[sh + 1]; + double sum = 0.0; + size_t cnt = 0; + for (int64_t j = lo; j < hi; ++j) { + sum = __dadd_rn(sum, double(__fdiv_rn(so.I[j], so.epsilon[j]))); + ++cnt; + } + double sum_kept = 0.0, sum2_kept = 0.0; + size_t cnt_kept = 0; + int64_t acentric = 0; + for (int64_t j = lo; j < hi; ++j) { + const double bound = so.centric[j] ? 2.0 * t : t; + if (!so.centric[j]) ++acentric; + if (sum > 0.0 && LowerE2(so.I[j], so.sigma[j], so.epsilon[j], z, sum / double(cnt)) > bound) + continue; + const double x = __fdiv_rn(so.I[j], so.epsilon[j]); + sum_kept = __dadd_rn(sum_kept, x); + sum2_kept = __dadd_rn(sum2_kept, __dmul_rn(x, x)); + ++cnt_kept; + } + mu[sh] = 0.0; + valid[sh] = 0; + n_acentric[sh] = 0; + if (cnt_kept < 2) return; + const double nk = double(cnt_kept); + const double m = sum_kept / nk; + const double var = fmax(0.0, (sum2_kept - __dmul_rn(__dmul_rn(nk, m), m)) / (nk - 1.0)); + mu[sh] = m; + valid[sh] = m > 0.0 && m > __dmul_rn(z, sqrt(var / nk)); + n_acentric[sh] = valid[sh] ? acentric : 0; + } + + // Position j's shell, for the per-observation kernels (shell_of is the inverse of shell_start). + __global__ void ShellOfKernel(int n_shells, const int64_t *shell_start, int32_t *shell_of) { + const int sh = blockIdx.x * blockDim.x + threadIdx.x; + if (sh >= n_shells) return; + for (int64_t j = shell_start[sh]; j < shell_start[sh + 1]; ++j) shell_of[j] = sh; + } + + // sigma / (epsilon ), whose median per shell says whether the shell measures the tail. + __global__ void NoiseRatioKernel(int n, Sorted so, const int32_t *shell_of, const double *mu, double *v) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n) return; + v[j] = double(so.sigma[j]) / __dmul_rn(double(so.epsilon[j]), mu[shell_of[j]]); + } + + // The acentric excess over u in the shells that measure the tail. + __global__ void ExcessKernel(int n, Sorted so, const int32_t *shell_of, const double *mu, + const uint8_t *measures_tail, double u, double *excess, uint8_t *flag) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n) return; + const int sh = shell_of[j]; + flag[j] = 0; + if (!measures_tail[sh] || so.centric[j]) return; + const double e2 = double(so.I[j]) / __dmul_rn(double(so.epsilon[j]), mu[sh]); + if (e2 > u) { + excess[j] = e2 - u; + flag[j] = 1; + } + } + + // e2 and improbable of every observation in a valid shell, at its own index; tested[i] marks it. + __global__ void ImprobableKernel(int n, Sorted so, const int32_t *idx, const int32_t *shell_of, + const double *mu, const uint8_t *valid, double z, double t, + double tail_scale, float *e2, uint8_t *improbable, uint8_t *tested) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n) return; + const int sh = shell_of[j]; + if (!valid[sh]) return; + const int i = idx[j]; + e2[i] = float(double(so.I[j]) / __dmul_rn(double(so.epsilon[j]), mu[sh])); + const double bound = so.centric[j] ? 2.0 * t : t; + improbable[i] = LowerE2(so.I[j], so.sigma[j], so.epsilon[j], z, mu[sh]) > __dmul_rn(tail_scale, bound); + tested[i] = 1; + } + + __global__ void UnitKeyKernel(int n, const WilsonObservation *obs, const int32_t *tested, uint32_t *unit) { + const int q = blockIdx.x * blockDim.x + threadIdx.x; + if (q >= n) return; + unit[q] = static_cast(obs[tested[q]].unit); + } + + __global__ void UnitStartKernel(int n, const uint32_t *unit, uint8_t *starts, int32_t *pos) { + const int q = blockIdx.x * blockDim.x + threadIdx.x; + if (q >= n) return; + starts[q] = q == 0 || unit[q] != unit[q - 1]; + pos[q] = q; + } + + // One thread per reflection: the decision of WilsonOutliers.cpp on its observations, in its order. + __global__ void UnitDecisionKernel(int n_units, const int32_t *unit_start, int n_tested, const int32_t *tested, + const WilsonObservation *obs, const uint8_t *improbable, double z, + uint8_t *rejected) { + const int u = blockIdx.x * blockDim.x + threadIdx.x; + if (u >= n_units) return; + const int lo = unit_start[u], hi = u + 1 < n_units ? unit_start[u + 1] : n_tested; + for (int q = lo; q < hi; ++q) { + const int i = tested[q]; + if (!improbable[i]) continue; + size_t n_probable = 0; + double sw = 0.0, swI = 0.0; + for (int r = lo; r < hi; ++r) { + const WilsonObservation &o = obs[tested[r]]; + if (r == q || improbable[tested[r]] || o.clipped) continue; + const double w = 1.0 / __dmul_rn(double(o.sigma), double(o.sigma)); + sw = __dadd_rn(sw, w); + swI = __dadd_rn(swI, __dmul_rn(w, double(o.I))); + ++n_probable; + } + const size_t n_others = hi - lo - 1; + bool reject = n_others == 0; + if (2 * n_probable > n_others) { + const WilsonObservation &o = obs[i]; + reject = fabs(double(o.I) - swI / sw) + > __dmul_rn(z, sqrt(__dadd_rn(__dmul_rn(double(o.sigma), double(o.sigma)), 1.0 / sw))); + } + if (reject) rejected[i] = 1; + } + } + + template + void Download(std::vector &out, const T *src, size_t n, cudaStream_t stream) { + out.resize(n); + if (n > 0) + cuda_err(cudaMemcpyAsync(out.data(), src, n * sizeof(T), cudaMemcpyDeviceToHost, stream)); + } + + // The positions where flag is set, in ascending order (a stable compaction of `in`). + int Compact(const int32_t *in, const uint8_t *flag, int32_t *out, int n, cudaStream_t stream) { + CudaDevicePtr n_out(1, ALLOC); + size_t bytes = 0; + cuda_err(cub::DeviceSelect::Flagged(nullptr, bytes, in, flag, out, n_out.get(), n, stream)); + CudaDevicePtr scratch(std::max(bytes, 1), ALLOC); + cuda_err(cub::DeviceSelect::Flagged(scratch.get(), bytes, in, flag, out, n_out.get(), n, stream)); + int h = 0; + cuda_err(cudaMemcpyAsync(&h, n_out.get(), sizeof(int), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); + return h; + } +} + +WilsonOutlierResult WilsonOutliersGPU(const std::vector &obs, double alpha) { + WilsonOutlierResult out; + out.rejected.assign(obs.size(), 0); + out.e2.assign(obs.size(), NAN); + const int n_obs = static_cast(obs.size()); + if (n_obs == 0) + return out; + + // Observations (28 bytes) + ~13 more 4-8 byte arrays of the same length, and the sort scratch. + constexpr size_t BYTES_PER_OBS = 96; + size_t free_bytes = 0, total_bytes = 0; + cuda_err(cudaMemGetInfo(&free_bytes, &total_bytes)); + if (BYTES_PER_OBS * obs.size() > free_bytes) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, + "Wilson outlier test: " + std::to_string(obs.size()) + " observations need " + + std::to_string(BYTES_PER_OBS * obs.size() >> 20) + " MiB of GPU memory, " + + std::to_string(free_bytes >> 20) + " MiB free - run the CPU build"); + + CudaStream stream; + CudaDevicePtr d_obs(obs.size(), ALLOC); + cuda_err(cudaMemcpyAsync(d_obs, obs.data(), obs.size() * sizeof(WilsonObservation), cudaMemcpyHostToDevice, + stream)); + + // The judged observations, in index order, then sorted by d (descending; ties keep index order). + CudaDevicePtr flag(obs.size(), ALLOC); + CudaDevicePtr all_index(obs.size(), ALLOC), judged(obs.size(), ALLOC), idx(obs.size(), ALLOC); + CudaDevicePtr d_all(obs.size(), ALLOC), d_judged(obs.size(), ALLOC), d_sorted(obs.size(), ALLOC); + JudgedKernel<<>>(n_obs, d_obs, flag, all_index, d_all); + cuda_err(cudaGetLastError()); + const int n = Compact(all_index, flag, judged, n_obs, stream); + if (n == 0) + return out; + { + // d of the judged observations, in the same order as `judged`. + size_t bytes = 0; + CudaDevicePtr n_sel(1, ALLOC); + cuda_err(cub::DeviceSelect::Flagged(nullptr, bytes, d_all.get(), flag.get(), d_judged.get(), n_sel.get(), + n_obs, stream)); + CudaDevicePtr scratch(std::max(bytes, 1), ALLOC); + cuda_err(cub::DeviceSelect::Flagged(scratch.get(), bytes, d_all.get(), flag.get(), d_judged.get(), + n_sel.get(), n_obs, stream)); + bytes = 0; + cuda_err(cub::DeviceRadixSort::SortPairsDescending(nullptr, bytes, d_judged.get(), d_sorted.get(), + judged.get(), idx.get(), n, 0, 32, stream)); + CudaDevicePtr sort_scratch(std::max(bytes, 1), ALLOC); + cuda_err(cub::DeviceRadixSort::SortPairsDescending(sort_scratch.get(), bytes, d_judged.get(), d_sorted.get(), + judged.get(), idx.get(), n, 0, 32, stream)); + } + + const double p = alpha / (2.0 * static_cast(n)); + const double t = std::log(1.0 / p); + const double z = std::sqrt(2.0 * std::log(1.0 / (2.0 * p))); + + const size_t n_shells = std::max(1, static_cast(n) / OBS_PER_SHELL); + std::vector shell_start(n_shells + 1); + for (size_t sh = 0; sh <= n_shells; ++sh) + shell_start[sh] = static_cast((sh * static_cast(n) + n_shells - 1) / n_shells); + CudaDevicePtr d_shell_start(n_shells + 1, ALLOC); + cuda_err(cudaMemcpyAsync(d_shell_start, shell_start.data(), (n_shells + 1) * sizeof(int64_t), + cudaMemcpyHostToDevice, stream)); + + CudaDevicePtr s_I(n, ALLOC), s_sigma(n, ALLOC), s_eps(n, ALLOC); + CudaDevicePtr s_centric(n, ALLOC); + GatherKernel<<>>(n, d_obs, idx, s_I, s_sigma, s_eps, s_centric); + cuda_err(cudaGetLastError()); + const Sorted so{s_I, s_sigma, s_eps, s_centric}; + const int ns = static_cast(n_shells); + + CudaDevicePtr mu(n_shells, ALLOC); + CudaDevicePtr valid(n_shells, ALLOC); + CudaDevicePtr acentric(n_shells, ALLOC); + CudaDevicePtr shell_of(n, ALLOC); + ShellMeanKernel<<>>(ns, d_shell_start, so, z, t, mu, valid, acentric); + cuda_err(cudaGetLastError()); + ShellOfKernel<<>>(ns, d_shell_start, shell_of); + cuda_err(cudaGetLastError()); + + // The median noise ratio of each shell: each shell's values sorted on their own. + CudaDevicePtr v(n, ALLOC), v_sorted(n, ALLOC); + NoiseRatioKernel<<>>(n, so, shell_of, mu, v); + cuda_err(cudaGetLastError()); + { + size_t bytes = 0; + cuda_err(cub::DeviceSegmentedRadixSort::SortKeys(nullptr, bytes, v.get(), v_sorted.get(), n, ns, + d_shell_start.get(), d_shell_start.get() + 1, 0, 64, + stream)); + CudaDevicePtr scratch(std::max(bytes, 1), ALLOC); + cuda_err(cub::DeviceSegmentedRadixSort::SortKeys(scratch.get(), bytes, v.get(), v_sorted.get(), n, ns, + d_shell_start.get(), d_shell_start.get() + 1, 0, 64, + stream)); + } + std::vector h_mu, h_v_sorted; + std::vector h_valid; + std::vector h_acentric; + Download(h_mu, mu.get(), n_shells, stream); + Download(h_valid, valid.get(), n_shells, stream); + Download(h_acentric, acentric.get(), n_shells, stream); + Download(h_v_sorted, v_sorted.get(), n, stream); + cuda_err(cudaStreamSynchronize(stream)); + + size_t n_acentric = 0; + for (size_t sh = 0; sh < n_shells; ++sh) + n_acentric += static_cast(h_acentric[sh]); + const double u = std::max(0.0, std::log(static_cast(n_acentric) / 100.0)); + std::vector measures_tail(n_shells, 0); + for (size_t sh = 0; sh < n_shells; ++sh) { + const int64_t lo = shell_start[sh], hi = shell_start[sh + 1]; + if (!h_valid[sh] || hi == lo) continue; + measures_tail[sh] = z * h_v_sorted[lo + (hi - lo) / 2] <= u; + } + CudaDevicePtr d_measures_tail(n_shells, ALLOC); + cuda_err(cudaMemcpyAsync(d_measures_tail, measures_tail.data(), n_shells, cudaMemcpyHostToDevice, stream)); + + // The tail's scale: the median excess, read off the sorted excesses. + { + CudaDevicePtr excess(n, ALLOC), kept(n, ALLOC); + CudaDevicePtr eflag(n, ALLOC); + ExcessKernel<<>>(n, so, shell_of, mu, d_measures_tail, u, excess, eflag); + cuda_err(cudaGetLastError()); + CudaDevicePtr n_sel(1, ALLOC); + size_t bytes = 0; + cuda_err(cub::DeviceSelect::Flagged(nullptr, bytes, excess.get(), eflag.get(), kept.get(), n_sel.get(), n, + stream)); + CudaDevicePtr scratch(std::max(bytes, 1), ALLOC); + cuda_err(cub::DeviceSelect::Flagged(scratch.get(), bytes, excess.get(), eflag.get(), kept.get(), + n_sel.get(), n, stream)); + int n_excess = 0; + cuda_err(cudaMemcpyAsync(&n_excess, n_sel.get(), sizeof(int), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); + if (n_excess > 0) { + bytes = 0; + cuda_err(cub::DeviceRadixSort::SortKeys(nullptr, bytes, kept.get(), excess.get(), n_excess, 0, 64, + stream)); + CudaDevicePtr sort_scratch(std::max(bytes, 1), ALLOC); + cuda_err(cub::DeviceRadixSort::SortKeys(sort_scratch.get(), bytes, kept.get(), excess.get(), n_excess, + 0, 64, stream)); + double mid = 0.0; + cuda_err(cudaMemcpyAsync(&mid, excess.get() + n_excess / 2, sizeof(double), cudaMemcpyDeviceToHost, + stream)); + cuda_err(cudaStreamSynchronize(stream)); + out.tail_scale = std::max(1.0, mid / std::log(2.0)); + } + } + out.bound = out.tail_scale * t; + + // Which observations are improbable; the tested ones in index order, then by reflection. + CudaDevicePtr e2(obs.size(), ALLOC); + CudaDevicePtr improbable(obs.size(), ALLOC), tested_flag(obs.size(), ALLOC), rejected(obs.size(), ALLOC); + cuda_err(cudaMemcpyAsync(e2, out.e2.data(), obs.size() * sizeof(float), cudaMemcpyHostToDevice, stream)); + cuda_err(cudaMemsetAsync(improbable, 0, obs.size(), stream)); + cuda_err(cudaMemsetAsync(tested_flag, 0, obs.size(), stream)); + cuda_err(cudaMemsetAsync(rejected, 0, obs.size(), stream)); + ImprobableKernel<<>>(n, so, idx, shell_of, mu, valid, z, t, out.tail_scale, e2, + improbable, tested_flag); + cuda_err(cudaGetLastError()); + const int n_tested = Compact(all_index, tested_flag, judged, n_obs, stream); + out.n_tested = static_cast(n_tested); + if (n_tested > 0) { + CudaDevicePtr unit(n_tested, ALLOC), unit_sorted(n_tested, ALLOC); + UnitKeyKernel<<>>(n_tested, d_obs, judged, unit); + cuda_err(cudaGetLastError()); + size_t bytes = 0; + cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes, unit.get(), unit_sorted.get(), judged.get(), + idx.get(), n_tested, 0, 32, stream)); + CudaDevicePtr scratch(std::max(bytes, 1), ALLOC); + cuda_err(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, unit.get(), unit_sorted.get(), judged.get(), + idx.get(), n_tested, 0, 32, stream)); + // idx now holds `tested` in (unit, index) order. + UnitStartKernel<<>>(n_tested, unit_sorted, flag, all_index); + cuda_err(cudaGetLastError()); + const int n_units = Compact(all_index, flag, judged, n_tested, stream); + UnitDecisionKernel<<>>(n_units, judged, n_tested, idx, d_obs, + improbable, z, rejected); + cuda_err(cudaGetLastError()); + // all_index was overwritten with 0..n_tested-1 above; nothing reads it after this. + } + cuda_err(cudaMemcpyAsync(out.rejected.data(), rejected, obs.size(), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(out.e2.data(), e2, obs.size() * sizeof(float), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); + for (const uint8_t r : out.rejected) out.n_rejected += r; + return out; +} diff --git a/image_analysis/scale_merge/WilsonOutliersGPU.h b/image_analysis/scale_merge/WilsonOutliersGPU.h new file mode 100644 index 000000000..e60467a71 --- /dev/null +++ b/image_analysis/scale_merge/WilsonOutliersGPU.h @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +#include "WilsonOutliers.h" + +// WilsonOutliers (WilsonOutliers.h) on the GPU: the same test, step for step, with the same result to +// the last bit - the sorts are on total orders, every sum is taken by one thread in the host's order, +// every product the host rounds on its own is rounded on its own here, and the few scalars that need a +// logarithm are taken on the host. Only built with CUDA. +// +// Device memory is about 80 bytes per observation; a device that cannot hold it throws (run on the +// CPU build then). +WilsonOutlierResult WilsonOutliersGPU(const std::vector &obs, double alpha); diff --git a/tests/WilsonOutliersTest.cpp b/tests/WilsonOutliersTest.cpp index ae901362d..b13052ccc 100644 --- a/tests/WilsonOutliersTest.cpp +++ b/tests/WilsonOutliersTest.cpp @@ -133,3 +133,57 @@ TEST_CASE("WilsonOutliers: a heavier-tailed population widens the bound", "[wils CHECK(r.tail_scale > 1.5); CHECK(r.n_rejected == 0); } + +#ifdef JFJOCH_USE_CUDA +#include +#include "../common/CUDAWrapper.h" +#include "../image_analysis/scale_merge/WilsonOutliersGPU.h" + +// The device test against the host one, on a population that takes every branch: reflections measured +// one to six times, resolutions on a coarse grid (so many observations share a d and the tie order +// counts) reaching past the 4 A limit, symmetry factors 1-4, centrics, clipped mates, observations +// outside the test (no reflection, NaN, zero sigma), a heavy-tailed mixture and planted artefacts. The +// two results must be the same to the bit. +TEST_CASE("WilsonOutliers: the device test is the host test", "[wilson_outliers]") { + if (get_gpu_count() == 0) { + WARN("No CUDA GPU present. Skipping the device Wilson test"); + return; + } + for (uint32_t seed : {11u, 12u}) { + std::mt19937 rng(seed); + std::uniform_real_distribution uni(0.0, 1.0); + std::normal_distribution gauss(0.0, 1.0); + std::vector obs; + for (int u = 0; u < 120000; ++u) { + const float d = 1.5f + 0.01f * static_cast(static_cast(uni(rng) * 300)); + const double mean = 1000.0 * std::exp(-2.0 / (d * d)) * (u % 2 ? 2.0 : 0.1); + const float epsilon = static_cast(1 << static_cast(uni(rng) * 3)); + const bool centric = uni(rng) < 0.1; + const double I_true = -mean * std::log(1.0 - uni(rng)) * epsilon; + const int mult = 1 + static_cast(uni(rng) * 6); + for (int m = 0; m < mult; ++m) { + const double sigma = std::sqrt(I_true + 100.0); + double I = I_true + sigma * gauss(rng); + if (uni(rng) < 0.002) I *= 200.0; + obs.push_back({static_cast(I), static_cast(sigma), d, epsilon, centric, + uni(rng) < 0.05, u}); + } + } + obs[3].unit = -1; + obs[5].I = NAN; + obs[7].sigma = 0.0f; + + const auto host = WilsonOutliers(obs, 0.01, 8); + const auto dev = WilsonOutliersGPU(obs, 0.01); + REQUIRE(host.n_rejected > 0); + CHECK(host.tail_scale > 1.0); + CHECK(dev.n_tested == host.n_tested); + CHECK(dev.n_rejected == host.n_rejected); + CHECK(dev.tail_scale == host.tail_scale); + CHECK(dev.bound == host.bound); + CHECK(dev.rejected == host.rejected); + REQUIRE(dev.e2.size() == host.e2.size()); + CHECK(std::memcmp(dev.e2.data(), host.e2.data(), host.e2.size() * sizeof(float)) == 0); + } +} +#endif From 8f605449f25071ac97a5f85444480c0c9c753a55 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 07:14:07 +0200 Subject: [PATCH 2/8] RotationScaleMergeGPU: hand the error-model samples down compacted and complete MergeEmSamples now forms each sample whole on the device - the counting variance at the reflection's expected intensity included, every product rounded on its own as the host build rounds it - and compacts them in fulls order with a scan, so the host receives exactly the vector it used to build. That replaces the download of three full-length double arrays and a flag array, and the host pass that compacted them while reading every full's eighty-byte record. p.mtz md5 and the error-model lines 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 | 19 +-- .../scale_merge/RotationScaleMergeGPU.cu | 117 +++++++++++++----- .../scale_merge/RotationScaleMergeGPU.h | 10 +- 3 files changed, 90 insertions(+), 56 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 2b65114e8..13ecad8c6 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -4692,27 +4692,10 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool bool did_gpu = false; #ifdef JFJOCH_USE_CUDA if (use_gpu_merge) { - const int nf = static_cast(fulls.size()); - // Plain arrays: the download writes every element, and sized vectors would first zero three - // full-length arrays on one thread. - std::unique_ptr gs2(new double[nf]), gI2(new double[nf]), gdev2(new double[nf]); - std::unique_ptr gvalid(new uint8_t[nf]); gpu_->MergeEmSamples(for_search, min_partiality, merge_friedel ? obs_hand.data() : nullptr, merge_friedel ? group_has_hands.data() : nullptr, - em_mean.data(), cnt.data(), - gs2.get(), gI2.get(), gdev2.get(), gvalid.get()); - // Which fulls give a sample, in order, and then each sample on its own - on all threads. - const std::vector valid = KeptIndices(nf, nthreads, [&](int i) { return gvalid[i] != 0; }); - samples.resize(valid.size()); - ParallelChunks(static_cast(valid.size()), ThreadsForWork(valid.size(), nthreads), - [&](int lo, int hi) { - for (int j = lo; j < hi; ++j) { - const int i = valid[j]; - samples[j] = {counting_variance(fulls[i], em_mean[fulls[i].group], gs2[i]), - gI2[i], gdev2[i], mf.d[i]}; - } - }); + em_mean.data(), cnt.data(), samples); did_gpu = true; } #endif diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 564e5c023..d4830bdee 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -11,6 +11,7 @@ #include #include #include +#include #include "../indexing/CUDAMemHelpers.h" #include "../../common/CUDAWrapper.h" @@ -546,8 +547,8 @@ namespace { double *sw, *swI, *em_mean_out; double *swh, *swIh; // per (group, hand), indexed 2*g + hand int32_t *cnt, *cnth; - double *s2, *I2, *dev2; - uint8_t *valid; + int32_t *sample_at; // per full: 1 where it gives a sample, then (scanned) where it goes + ErrorModelSample *sample; // the samples, compacted in fulls order double *a_swI, *a_sw, *a_swIh0, *a_swIh1, *a_swh0, *a_swh1, *a_swht0, *a_swht1, *a_d; int32_t *a_nh0, *a_nh1, *a_rejected; uint8_t *a_on_ice; @@ -610,25 +611,54 @@ namespace { } } - // One thread per full: the leverage-corrected error-model sample, or valid=0 if dropped. - __global__ void MergeSamplesKernel(int n_obs, MergeParams p) { + // The host's counting_variance (RotationScaleMerge::MergeAndStats), to the bit: each product + // rounded on its own, as the host build rounds it, and summed in the host's order. + __device__ __forceinline__ double CountingVarianceExact(float corr, float var_bkg, float var_per_I, float capture, + double I_for_b, double own_s2) { + const double I_exp = 0.0 < I_for_b ? I_for_b : 0.0; + const double c = corr, cap = capture; + double base = __dmul_rn(__dmul_rn(c, c), double(var_bkg)); + base = __dadd_rn(base, __dmul_rn(__dmul_rn(c, double(var_per_I)), I_exp)); + base = __dadd_rn(base, __dmul_rn(__dmul_rn(__dmul_rn(cap, cap), I_exp), I_exp)); + return base > 0.0 ? base : own_s2; + } + + // 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; + const int g = p.group[i]; + if (p.cnt[g] < 2) return false; + // The hand's own mean where the hand has two of its own, the pooled pair where it has + // not - see the host obs_hand block. em_mean (the merge weights) is untouched. + const int sub = 2 * g + (p.hand ? p.hand[i] : 0); + const bool on_hand = p.swh && p.has_hands[g] && p.cnth[sub] >= 2 && p.swh[sub] > 0.0; + const double mean = on_hand ? p.swIh[sub] / p.swh[sub] : p.em_mean[g]; + if (!isfinite(mean)) return false; + 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) / (on_hand ? p.swh[sub] : p.sw[g]); + if (factor < 0.05) return false; + const double resid = double(p.I[i]) * p.corr[i] - mean; + out.s2 = CountingVarianceExact(p.corr[i], p.var_bkg[i], p.var_per_I[i], p.capture[i], p.em_mean[g], s2); + out.I2 = mean * mean; + out.dev2 = resid * resid / factor; + out.d = p.d[i]; + return true; + } + + // One thread per full: whether it gives an error-model sample (the scan turns it into where). + __global__ void MergeSampleFlagKernel(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; - // The hand's own mean where the hand has two of its own, the pooled pair where it has - // not - see the host obs_hand block. em_mean (the merge weights) is untouched. - const int sub = 2 * g + (p.hand ? p.hand[i] : 0); - const bool on_hand = p.swh && p.has_hands[g] && p.cnth[sub] >= 2 && p.swh[sub] > 0.0; - const double mean = on_hand ? p.swIh[sub] / p.swh[sub] : 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) / (on_hand ? p.swh[sub] : 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; + ErrorModelSample s; + p.sample_at[i] = EmSample(i, p, s) ? 1 : 0; + } + } + + // One thread per full: its sample, at its place among the samples in fulls order. + __global__ void MergeSampleWriteKernel(int n_obs, MergeParams p) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) { + ErrorModelSample s; + if (EmSample(i, p, s)) p.sample[p.sample_at[i]] = s; } } @@ -966,8 +996,8 @@ struct RotationScaleMergeGPU::Impl { CudaDevicePtr m_cnth; CudaDevicePtr m_hand; // per full: Bijvoet hand (0 = I(+)) CudaDevicePtr m_has_hands; // per group: acentric, and the merge pools the hands - CudaDevicePtr m_s2, m_I2, m_dev2; // per full (n_fulls) - CudaDevicePtr m_valid; + CudaDevicePtr m_sample_at; // per full (n_fulls) + CudaDevicePtr m_sample; // the compacted error-model samples CudaDevicePtr m_rejected; // per-full outlier-rejected flag (MergeAccum -> MergeRmeas) CudaDevicePtr m_half; // per-full CC1/2 half-set, assigned on the host CudaDevicePtr a_swI, a_sw, a_swIh0, a_swIh1, a_swh0, a_swh1, a_swht0, a_swht1, a_d; // merge accum per group @@ -1243,8 +1273,8 @@ void RotationScaleMergeGPU::SetFrameCellOk(const uint8_t *frame_cell_ok) { // resident+scaled fulls. Stashes the filter context for the later MergeAccum/MergeRmeas calls. void RotationScaleMergeGPU::MergeEmSamples(bool for_search, double min_partiality, const uint8_t *hand, const uint8_t *has_hands, - double *em_mean_out, int32_t *cnt_out, double *s2_out, - double *I2_out, double *dev2_out, uint8_t *valid_out) { + double *em_mean_out, int32_t *cnt_out, + std::vector &samples) { DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; const int ng = d.n_groups, nf = d.n_fulls; @@ -1264,8 +1294,7 @@ void RotationScaleMergeGPU::MergeEmSamples(bool for_search, double min_partialit d.m_cnth = CudaDevicePtr(); d.m_hand = CudaDevicePtr(); d.m_has_hands = CudaDevicePtr(); } - d.m_s2 = d.Alloc(std::max(1, nf)); d.m_I2 = d.Alloc(std::max(1, nf)); - d.m_dev2 = d.Alloc(std::max(1, nf)); d.m_valid = d.Alloc(std::max(1, nf)); + d.m_sample_at = d.Alloc(std::max(1, nf)); MergeParams p{}; p.n_groups = ng; p.min_partiality = min_partiality; @@ -1279,24 +1308,44 @@ void RotationScaleMergeGPU::MergeEmSamples(bool for_search, double min_partialit p.em_mean_out = d.m_em_mean.get(); p.cnt = d.m_cnt.get(); p.swh = d.m_swh.get(); p.swIh = d.m_swIh.get(); p.cnth = d.m_cnth.get(); p.hand = d.m_hand.get(); p.has_hands = d.m_has_hands.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(); + p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get(); p.capture = d.f_capture.get(); + p.sample_at = d.m_sample_at.get(); const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK); const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK); MergeEmStatsKernel<<s()>>>(p); - MergeSamplesKernel<<s()>>>(nf, p); + MergeSampleFlagKernel<<s()>>>(nf, p); CudaCheck(cudaGetLastError(), "merge em/samples launch"); + // Where each sample goes: the fulls that give one, in order, as the host's compaction lists them. + int n_samples = 0; + if (nf > 0) { + int32_t last_flag = 0; + CopyAndWait(&last_flag, d.m_sample_at.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, + impl_->s(), "dl last sample flag"); + size_t bytes = 0; + CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, d.m_sample_at.get(), d.m_sample_at.get(), nf, + impl_->s()), "sample scan size"); + auto scratch = d.Alloc(std::max(bytes, 1)); + CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, d.m_sample_at.get(), d.m_sample_at.get(), nf, + impl_->s()), "sample scan"); + CopyAndWait(&n_samples, d.m_sample_at.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, + impl_->s(), "dl sample count"); + n_samples += last_flag; + } + d.m_sample = d.Alloc(std::max(1, n_samples)); + p.sample = d.m_sample.get(); + MergeSampleWriteKernel<<s()>>>(nf, p); + CudaCheck(cudaGetLastError(), "merge samples write launch"); CudaCheck(cudaStreamSynchronize(impl_->s()), "merge em/samples sync"); CopyAndWait(em_mean_out, d.m_em_mean.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl em_mean"); CopyAndWait(cnt_out, d.m_cnt.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl cnt"); - if (nf > 0) { - CopyAndWait(s2_out, d.m_s2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl s2"); - CopyAndWait(I2_out, d.m_I2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl I2"); - CopyAndWait(dev2_out, d.m_dev2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl dev2"); - CopyAndWait(valid_out, d.m_valid.get(), size_t(nf) * sizeof(uint8_t), cudaMemcpyDeviceToHost, impl_->s(), "dl valid"); - } + samples.resize(n_samples); + if (n_samples > 0) + CopyAndWait(samples.data(), d.m_sample.get(), size_t(n_samples) * sizeof(ErrorModelSample), + cudaMemcpyDeviceToHost, impl_->s(), "dl samples"); + d.m_sample = CudaDevicePtr(); } void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_b, bool error_model_active, diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.h b/image_analysis/scale_merge/RotationScaleMergeGPU.h index f0d6caf4b..a53137501 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -8,6 +8,8 @@ #include #include +#include "ErrorModel.h" + // GPU engine for the RotationScaleMerge hot loops. The class keeps the per-observation data resident on // the device as a structure-of-arrays (coalesced) and runs the scaling loop there. The host keeps the // one-time raw-hkl sort and the per-space-group ASU keying (gemmi); it hands the GPU the dense group ids @@ -86,15 +88,15 @@ public: // 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. + // Per-group inv-var mean (em_mean, length n_groups) + the leverage-corrected error-model samples in + // fulls order, each with the counting variance at its reflection's expected intensity - the host's + // samples to the bit (MergeAndStats' first two error-model loops). // Stashes (for_search, min_partiality) for the MergeAccum/MergeRmeas calls that follow. // hand / has_hands are the per-full Bijvoet hand and the per-group "this group pools two hands", // which the error-model samples are formed on; both null leaves the fit on the pooled group. void MergeEmSamples(bool for_search, double min_partiality, const uint8_t *hand, const uint8_t *has_hands, - double *em_mean_out, int32_t *cnt_out, double *s2_out, double *I2_out, - double *dev2_out, uint8_t *valid_out); + double *em_mean_out, int32_t *cnt_out, std::vector &samples); // 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 From 1ad41c875f8bbefc6c18155dc9af0745e57315a5 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 07:21:47 +0200 Subject: [PATCH 3/8] RotationScaleMerge: the Bijvoet widening of the pooled outlier test measured on the device The per-group sums the widening is measured from - each hand's inverse-variance mean under the error model's sigmas, the group's d and the d range - are 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 model sigma to the bit (ModelSigmaExact: every product rounded on its own); the host keeps the shells and their medians. The host path is the same sums, now kept per group so both paths hand the shell step one array. p.mtz md5 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 | 79 +++++++++++-------- .../scale_merge/RotationScaleMergeGPU.cu | 78 ++++++++++++++++++ .../scale_merge/RotationScaleMergeGPU.h | 8 ++ 3 files changed, 133 insertions(+), 32 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 13ecad8c6..f3b8e9893 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -4932,49 +4932,64 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool // SigAno accumulator reads, with the error model's own sigmas - taking it from the previous pass // would make the written merge depend on how many merges the pipeline happened to run first. if (reject_outliers && error_model_active && merge_friedel && n_groups > 0) { - std::vector sw_h(2 * static_cast(n_groups), 0.0), swI_h(2 * static_cast(n_groups), 0.0); - // The threads split the GROUPS, each walking a group's fulls over the group CSR, so each - // (group, hand) is still summed over its fulls in index order. - std::vector t_lo(ReductionBlocks(n_groups, 512), INFINITY), t_hi(ReductionBlocks(n_groups, 512), 0.0f); - ParallelBlocks(n_groups, nthreads, [&](int t, int g0, int g1) { - for (int g = g0; g < g1; ++g) { - for (int p = fulls_group_start[g]; p < fulls_group_start[g + 1]; ++p) { - const int i = fulls_group_perm[p]; - if (mf.group[i] != g || !group_has_hands[g]) continue; - const Obs &o = fulls[i]; - const float I_corr = o.I * o.corr; - const float sigma_corr = corrected_sigma(o, I_corr, o.sigma * o.corr); - if (!(sigma_corr > 0.0f) || !std::isfinite(sigma_corr)) continue; - const double w = 1.0 / (static_cast(sigma_corr) * sigma_corr); - const size_t s = 2 * static_cast(g) + obs_hand[i]; - sw_h[s] += w; swI_h[s] += w * I_corr; - if (std::isfinite(mf.d[i]) && mf.d[i] > 0.0f) { - t_lo[t] = std::min(t_lo[t], mf.d[i]); t_hi[t] = std::max(t_hi[t], mf.d[i]); + // Per group: its d - its first usable full's (index order: the first of its CSR stretch) - and + // the excess dI^2 - var(dI) of its two hands, NaN where it has not both; and the d range of the + // fulls that entered a hand. On the device where the merge is, to the same bits. + std::vector group_d(n_groups, NAN); + std::vector group_excess(n_groups, NAN); + float d_lo = INFINITY, d_hi = 0.0f; + bool did_gpu_excess = false; +#ifdef JFJOCH_USE_CUDA + if (use_gpu_merge) { + gpu_->MergeBijvoetExcess(error_model_a, error_model_b, group_d.data(), group_excess.data(), d_lo, d_hi); + did_gpu_excess = true; + } +#endif + if (!did_gpu_excess) { + // The threads split the GROUPS, each walking a group's fulls over the group CSR, so each + // (group, hand) is still summed over its fulls in index order. + std::vector t_lo(ReductionBlocks(n_groups, 512), INFINITY), t_hi(ReductionBlocks(n_groups, 512), 0.0f); + ParallelBlocks(n_groups, nthreads, [&](int t, int g0, int g1) { + for (int g = g0; g < g1; ++g) { + double sw_h[2] = {0.0, 0.0}, swI_h[2] = {0.0, 0.0}; + bool first = true; + for (int p = fulls_group_start[g]; p < fulls_group_start[g + 1]; ++p) { + const int i = fulls_group_perm[p]; + if (mf.group[i] != g) continue; + if (first) { group_d[g] = mf.d[i]; first = false; } + if (!group_has_hands[g]) continue; + const Obs &o = fulls[i]; + const float I_corr = o.I * o.corr; + const float sigma_corr = corrected_sigma(o, I_corr, o.sigma * o.corr); + if (!(sigma_corr > 0.0f) || !std::isfinite(sigma_corr)) continue; + const double w = 1.0 / (static_cast(sigma_corr) * sigma_corr); + sw_h[obs_hand[i]] += w; swI_h[obs_hand[i]] += w * I_corr; + if (std::isfinite(mf.d[i]) && mf.d[i] > 0.0f) { + t_lo[t] = std::min(t_lo[t], mf.d[i]); t_hi[t] = std::max(t_hi[t], mf.d[i]); + } + } + if (sw_h[0] > 0.0 && sw_h[1] > 0.0 && std::isfinite(group_d[g])) { + const double di = swI_h[0] / sw_h[0] - swI_h[1] / sw_h[1]; + group_excess[g] = di * di - (1.0 / sw_h[0] + 1.0 / sw_h[1]); } } - } - }, 512); - const float d_lo = *std::min_element(t_lo.begin(), t_lo.end()); - const float d_hi = *std::max_element(t_hi.begin(), t_hi.end()); + }, 512); + d_lo = *std::min_element(t_lo.begin(), t_lo.end()); + d_hi = *std::max_element(t_hi.begin(), t_hi.end()); + } if (std::isfinite(d_lo) && d_hi > d_lo) { const ResolutionShells v_shells(d_lo * 0.999f, d_hi, report_shell_count); - // A group's d is its first usable full's (index order: the first of its CSR stretch). The - // excess of each shell is only ever asked for its median, which does not depend on the order - // the values were gathered in, so the groups gather theirs side by side. - std::vector group_d(n_groups, NAN); + // The excess of each shell is only ever asked for its median, which does not depend on the + // order the values were gathered in, so the groups gather theirs side by side. const int n_blk = ReductionBlocks(n_groups, 512); std::vector>> blk_excess( n_blk, std::vector>(report_shell_count)); ParallelBlocks(n_groups, nthreads, [&](int b, int g0, int g1) { for (int g = g0; g < g1; ++g) { - for (int p = fulls_group_start[g]; p < fulls_group_start[g + 1]; ++p) - if (mf.group[fulls_group_perm[p]] == g) { group_d[g] = mf.d[fulls_group_perm[p]]; break; } - const double swp = sw_h[2 * static_cast(g)], swm = sw_h[2 * static_cast(g) + 1]; - if (!(swp > 0.0) || !(swm > 0.0) || !std::isfinite(group_d[g])) continue; + if (std::isnan(group_excess[g])) continue; const auto sh = v_shells.GetShell(group_d[g]); if (!sh || *sh < 0 || *sh >= report_shell_count) continue; - const double di = swI_h[2 * static_cast(g)] / swp - swI_h[2 * static_cast(g) + 1] / swm; - blk_excess[b][*sh].push_back(di * di - (1.0 / swp + 1.0 / swm)); + blk_excess[b][*sh].push_back(group_excess[g]); } }, 512); std::vector> excess(report_shell_count); diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index d4830bdee..8c477d324 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -623,6 +623,51 @@ namespace { return base > 0.0 ? base : own_s2; } + // The host's model_sigma, to the bit (see CountingVarianceExact). + __device__ __forceinline__ float ModelSigmaExact(int i, const MergeParams &p, double I_for_b, float sigma_corr) { + if (!p.error_model_active) return sigma_corr; + const double a_var = CountingVarianceExact(p.corr[i], p.var_bkg[i], p.var_per_I[i], p.capture[i], I_for_b, + __dmul_rn(double(sigma_corr), double(sigma_corr))); + const double bI = __dmul_rn(p.error_model_b, I_for_b); + const double v = __dadd_rn(__dmul_rn(p.error_model_a, a_var), __dmul_rn(bI, bI)); + return v > 0.0 ? float(sqrt(v)) : sigma_corr; + } + + // One thread per group: what the pooled outlier test's widening is measured from (the host's + // reject_var_add block) - the group's d (its first usable full's), the excess dI^2 - var(dI) of its + // two hands under the error model's sigmas (NaN where a hand is missing), and the d range of the + // fulls that entered. Summed over the group's fulls in index order, as the host sums them. + __global__ void BijvoetExcessKernel(MergeParams p, float *group_d, double *excess, float *d_lo, float *d_hi) { + 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[2] = {0.0, 0.0}, swI[2] = {0.0, 0.0}; + float gd = NAN, dmin = INFINITY, dmax = 0.0f; + bool first = true; + for (int q = lo; q < hi; ++q) { + const int i = p.gperm[q]; + if (!MergeUsable(i, p)) continue; + if (first) { gd = p.d[i]; first = false; } + if (!p.has_hands[g]) continue; + const float I_corr = p.I[i] * p.corr[i]; + const double I_for_b = isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr); + const float sc = ModelSigmaExact(i, p, I_for_b, p.sigma[i] * p.corr[i]); + if (!(sc > 0.0f) || !isfinite(sc)) continue; + const double w = 1.0 / __dmul_rn(double(sc), double(sc)); + const int h = p.hand[i]; + sw[h] = __dadd_rn(sw[h], w); + swI[h] = __dadd_rn(swI[h], __dmul_rn(w, double(I_corr))); + const float d = p.d[i]; + if (isfinite(d) && d > 0.0f) { dmin = fminf(dmin, d); dmax = fmaxf(dmax, d); } + } + group_d[g] = gd; d_lo[g] = dmin; d_hi[g] = dmax; + excess[g] = NAN; + if (sw[0] > 0.0 && sw[1] > 0.0 && isfinite(gd)) { + const double di = swI[0] / sw[0] - swI[1] / sw[1]; + excess[g] = __dmul_rn(di, di) - (1.0 / sw[0] + 1.0 / sw[1]); + } + } + } + // 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; @@ -1458,6 +1503,39 @@ void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, d CopyAndWait(nusable, d.r_nusable.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl rnusable"); } +void RotationScaleMergeGPU::MergeBijvoetExcess(double error_model_a, double error_model_b, float *group_d, + double *excess, float &d_lo, float &d_hi) { + DeviceGuard guard(impl_->device, impl_->available); + auto &d = *impl_; + const int ng = d.n_groups; + auto g_d = d.Alloc(std::max(1, ng)), g_lo = d.Alloc(std::max(1, ng)), + g_hi = d.Alloc(std::max(1, ng)); + auto g_excess = 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.hand = d.m_hand.get(); p.has_hands = d.m_has_hands.get(); + p.error_model_a = error_model_a; p.error_model_b = error_model_b; p.error_model_active = 1; + 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); + BijvoetExcessKernel<<s()>>>(p, g_d.get(), g_excess.get(), g_lo.get(), g_hi.get()); + CudaCheck(cudaGetLastError(), "Bijvoet excess launch"); + std::vector lo(ng), hi(ng); + CopyAndWait(group_d, g_d.get(), size_t(ng) * sizeof(float), cudaMemcpyDeviceToHost, impl_->s(), "dl group d"); + CopyAndWait(excess, g_excess.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), + "dl excess"); + CopyAndWait(lo.data(), g_lo.get(), size_t(ng) * sizeof(float), cudaMemcpyDeviceToHost, impl_->s(), "dl d lo"); + CopyAndWait(hi.data(), g_hi.get(), size_t(ng) * sizeof(float), cudaMemcpyDeviceToHost, impl_->s(), "dl d hi"); + d_lo = INFINITY; d_hi = 0.0f; + 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::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 a53137501..1d68aa8da 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -98,6 +98,14 @@ public: const uint8_t *hand, const uint8_t *has_hands, double *em_mean_out, int32_t *cnt_out, std::vector &samples); + // What the pooled outlier test's widening by the Bijvoet variance is measured from (MergeAndStats' + // reject_var_add block), with the error model (a, b) of the first fit: per group (length n_groups) + // its d - its first usable full's - and the excess dI^2 - var(dI) of its two hands (NaN where it has + // not both); and the d range of the fulls that entered a hand. The host's numbers to the bit. + // Requires MergeEmSamples with the hands first. + void MergeBijvoetExcess(double error_model_a, double error_model_b, float *group_d, double *excess, + float &d_lo, float &d_hi); + // 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. From 8ed898d5a1dc212b4364b0acfacfb3678971b4fb Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 07:27:25 +0200 Subject: [PATCH 4/8] RotationScaleMerge: sort the outlier band's log-spread pairs on the device The multiplicative spread of the outlier band is the weighted median of (|ln(I/median)|, weight) pairs, one per usable full, and its sort was the bulk of that step. SortPairsGPU sorts them on the GPU (by the second key, then stably by the first) and leaves the sequence std::sort leaves: two pairs that compare equal are the same two numbers. The values themselves (a logarithm) and the running sum stay on the host. p.mtz md5 unchanged on myob, cytc, 8a1a and 8qaw (GPU build); a Catch2 case checks the device order against std::sort. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- image_analysis/scale_merge/CMakeLists.txt | 3 +- image_analysis/scale_merge/PairSortGPU.cu | 79 +++++++++++++++++++ image_analysis/scale_merge/PairSortGPU.h | 12 +++ .../scale_merge/RotationScaleMerge.cpp | 14 +++- tests/CMakeLists.txt | 1 + tests/PairSortGPUTest.cpp | 33 ++++++++ 6 files changed, 139 insertions(+), 3 deletions(-) create mode 100644 image_analysis/scale_merge/PairSortGPU.cu create mode 100644 image_analysis/scale_merge/PairSortGPU.h create mode 100644 tests/PairSortGPUTest.cpp diff --git a/image_analysis/scale_merge/CMakeLists.txt b/image_analysis/scale_merge/CMakeLists.txt index e52b7439d..9f5fd6c85 100644 --- a/image_analysis/scale_merge/CMakeLists.txt +++ b/image_analysis/scale_merge/CMakeLists.txt @@ -42,5 +42,6 @@ TARGET_LINK_LIBRARIES(JFJochScaleMerge Ceres::ceres Eigen3::Eigen JFJochCommon f IF (JFJOCH_CUDA_AVAILABLE) TARGET_SOURCES(JFJochScaleMerge PRIVATE ../indexing/CUDAMemHelpers.h RotationScaleMergeGPU.cu RotationScaleMergeGPU.h - WilsonOutliersGPU.cu WilsonOutliersGPU.h) + WilsonOutliersGPU.cu WilsonOutliersGPU.h + PairSortGPU.cu PairSortGPU.h) ENDIF() \ No newline at end of file diff --git a/image_analysis/scale_merge/PairSortGPU.cu b/image_analysis/scale_merge/PairSortGPU.cu new file mode 100644 index 000000000..a770cb502 --- /dev/null +++ b/image_analysis/scale_merge/PairSortGPU.cu @@ -0,0 +1,79 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "PairSortGPU.h" + +#include +#include + +#include + +#include "../indexing/CUDAMemHelpers.h" +#include "../../common/JFJochException.h" + +namespace { + constexpr int THREADS = 256; + constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + + void cuda_err(cudaError_t val) { + if (val != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); + } + + __global__ void SplitKernel(int n, const double2 *pair, double *second, int32_t *index) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + second[i] = pair[i].y; + index[i] = i; + } + + __global__ void GatherFirstKernel(int n, const double2 *pair, const int32_t *index, double *first) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + first[i] = pair[index[i]].x; + } + + __global__ void GatherPairKernel(int n, const double2 *pair, const int32_t *index, double2 *out) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + out[i] = pair[index[i]]; + } +} + +void SortPairsGPU(std::vector> &v) { + static_assert(sizeof(std::pair) == sizeof(double2)); + const int n = static_cast(v.size()); + if (n < 2) + return; + constexpr size_t BYTES_PER_PAIR = 16 + 16 + 16 + 8; // pairs in and out, keys, indices + size_t free_bytes = 0, total_bytes = 0; + cuda_err(cudaMemGetInfo(&free_bytes, &total_bytes)); + if (BYTES_PER_PAIR * v.size() > free_bytes) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, + "Sorting " + std::to_string(v.size()) + " pairs needs " + + std::to_string(BYTES_PER_PAIR * v.size() >> 20) + " MiB of GPU memory, " + + std::to_string(free_bytes >> 20) + " MiB free - run the CPU build"); + const unsigned blocks = static_cast((n + THREADS - 1) / THREADS); + CudaStream stream; + CudaDevicePtr pair(n, ALLOC), out(n, ALLOC); + CudaDevicePtr key(n, ALLOC), key_sorted(n, ALLOC); + CudaDevicePtr index(n, ALLOC), index_sorted(n, ALLOC); + cuda_err(cudaMemcpyAsync(pair, v.data(), v.size() * sizeof(double2), cudaMemcpyHostToDevice, stream)); + // By .second, then - stably - by .first: the order of the pairs. + SplitKernel<<>>(n, pair, key, index); + cuda_err(cudaGetLastError()); + size_t bytes = 0; + cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key.get(), key_sorted.get(), index.get(), + index_sorted.get(), n, 0, 64, stream)); + CudaDevicePtr scratch(std::max(bytes, 1), ALLOC); + cuda_err(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, key.get(), key_sorted.get(), index.get(), + index_sorted.get(), n, 0, 64, stream)); + GatherFirstKernel<<>>(n, pair, index_sorted, key); + cuda_err(cudaGetLastError()); + cuda_err(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, key.get(), key_sorted.get(), index_sorted.get(), + index.get(), n, 0, 64, stream)); + GatherPairKernel<<>>(n, pair, index, out); + cuda_err(cudaGetLastError()); + cuda_err(cudaMemcpyAsync(v.data(), out, v.size() * sizeof(double2), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); +} diff --git a/image_analysis/scale_merge/PairSortGPU.h b/image_analysis/scale_merge/PairSortGPU.h new file mode 100644 index 000000000..85adda070 --- /dev/null +++ b/image_analysis/scale_merge/PairSortGPU.h @@ -0,0 +1,12 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include + +// Sort pairs of finite doubles ascending, first by .first then by .second, on the GPU. Two pairs that +// compare equal are the same two numbers, so this is the sequence std::sort leaves, to the bit. Only +// built with CUDA; a device without room for about 40 bytes a pair throws (run on the CPU build then). +void SortPairsGPU(std::vector> &v); diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index f3b8e9893..1f706306e 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -26,6 +26,7 @@ #include "ResolutionCutoff.h" #include "WilsonOutliers.h" #include "WilsonOutliersGPU.h" +#include "PairSortGPU.h" #include "OutlierBand.h" #include "../../common/CorrelationCoefficient.h" #include "../../common/CrystalLattice.h" @@ -4892,8 +4893,17 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool }); if (!lv.empty()) { // Two pairs that compare equal are the same two numbers, so every sort gives the same - // sequence and the weighted median below is summed in it on any thread count. - ParallelSort(lv.begin(), lv.end(), nthreads, std::less>()); + // sequence - the device's too - and the weighted median below is summed in it on any + // thread count. + bool did_gpu_sort = false; +#ifdef JFJOCH_USE_CUDA + if (use_gpu_merge) { + SortPairsGPU(lv); + did_gpu_sort = true; + } +#endif + if (!did_gpu_sort) + ParallelSort(lv.begin(), lv.end(), nthreads, std::less>()); double total = 0.0; for (const auto &[v, w] : lv) total += w; double run = 0.0; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 887fa4221..78249a265 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -157,6 +157,7 @@ ADD_EXECUTABLE(jfjoch_test FrenchWilsonTest.cpp SphericalHarmonicSurfaceTest.cpp WilsonOutliersTest.cpp + PairSortGPUTest.cpp ReindexAmbiguityTest.cpp CrystalSettingTest.cpp LoadReferenceMtzFreeFlagsTest.cpp diff --git a/tests/PairSortGPUTest.cpp b/tests/PairSortGPUTest.cpp new file mode 100644 index 000000000..1d0812bc6 --- /dev/null +++ b/tests/PairSortGPUTest.cpp @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#ifdef JFJOCH_USE_CUDA +#include +#include +#include +#include + +#include "../common/CUDAWrapper.h" +#include "../image_analysis/scale_merge/PairSortGPU.h" + +// The device sort against std::sort, with many repeated first and second values (so both keys decide) +// and whole duplicate pairs. +TEST_CASE("SortPairsGPU: the order std::sort leaves", "[PairSortGPU]") { + if (get_gpu_count() == 0) { + WARN("No CUDA GPU present. Skipping SortPairsGPU"); + return; + } + std::mt19937 rng(5); + std::uniform_int_distribution coarse(0, 200); + std::exponential_distribution fine(1.0); + std::vector> v(500000); + for (auto &p : v) p = {coarse(rng) * 0.37, coarse(rng) % 3 ? fine(rng) : 0.5 * coarse(rng)}; + for (size_t i = 0; i < 1000; ++i) v[v.size() - 1 - i] = v[i]; + auto host = v; + std::sort(host.begin(), host.end()); + SortPairsGPU(v); + CHECK(v == host); +} +#endif From 7d3766317c4c09008fc7fd731e32565118f252b1 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 08:51:56 +0200 Subject: [PATCH 5/8] WilsonOutliersGPU: free each stage's buffers when it is done The merges of the largest sets run with 12-13.6 GB of a 16 GB device already taken by the resident observations, so the transient footprint of this pass is what decides whether it fits. Release the sort keys as soon as the d order is known and the noise ratios once their per-shell medians are read - on the device, rather than downloading every sorted value. Peak about 96 bytes per observation; the free memory is checked first and a device without it stops with a message naming the CPU build. p.mtz md5 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/WilsonOutliersGPU.cu | 19 ++++++++++++++++--- .../scale_merge/WilsonOutliersGPU.h | 4 ++-- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/image_analysis/scale_merge/WilsonOutliersGPU.cu b/image_analysis/scale_merge/WilsonOutliersGPU.cu index fddcdfbca..7287f8147 100644 --- a/image_analysis/scale_merge/WilsonOutliersGPU.cu +++ b/image_analysis/scale_merge/WilsonOutliersGPU.cu @@ -117,6 +117,13 @@ namespace { v[j] = double(so.sigma[j]) / __dmul_rn(double(so.epsilon[j]), mu[shell_of[j]]); } + __global__ void ShellMedianKernel(int n_shells, const int64_t *shell_start, const double *sorted, double *median) { + const int sh = blockIdx.x * blockDim.x + threadIdx.x; + if (sh >= n_shells) return; + const int64_t lo = shell_start[sh], hi = shell_start[sh + 1]; + median[sh] = hi > lo ? sorted[lo + (hi - lo) / 2] : 0.0; + } + // The acentric excess over u in the shells that measure the tail. __global__ void ExcessKernel(int n, Sorted so, const int32_t *shell_of, const double *mu, const uint8_t *measures_tail, double u, double *excess, uint8_t *flag) { @@ -261,6 +268,8 @@ WilsonOutlierResult WilsonOutliersGPU(const std::vector &obs, judged.get(), idx.get(), n, 0, 32, stream)); } + d_all = CudaDevicePtr(); d_judged = CudaDevicePtr(); d_sorted = CudaDevicePtr(); + const double p = alpha / (2.0 * static_cast(n)); const double t = std::log(1.0 / p); const double z = std::sqrt(2.0 * std::log(1.0 / (2.0 * p))); @@ -303,14 +312,18 @@ WilsonOutlierResult WilsonOutliersGPU(const std::vector &obs, d_shell_start.get(), d_shell_start.get() + 1, 0, 64, stream)); } - std::vector h_mu, h_v_sorted; + CudaDevicePtr noise_median(n_shells, ALLOC); + ShellMedianKernel<<>>(ns, d_shell_start, v_sorted, noise_median); + cuda_err(cudaGetLastError()); + std::vector h_mu, h_noise_median; std::vector h_valid; std::vector h_acentric; Download(h_mu, mu.get(), n_shells, stream); Download(h_valid, valid.get(), n_shells, stream); Download(h_acentric, acentric.get(), n_shells, stream); - Download(h_v_sorted, v_sorted.get(), n, stream); + Download(h_noise_median, noise_median.get(), n_shells, stream); cuda_err(cudaStreamSynchronize(stream)); + v = CudaDevicePtr(); v_sorted = CudaDevicePtr(); size_t n_acentric = 0; for (size_t sh = 0; sh < n_shells; ++sh) @@ -320,7 +333,7 @@ WilsonOutlierResult WilsonOutliersGPU(const std::vector &obs, for (size_t sh = 0; sh < n_shells; ++sh) { const int64_t lo = shell_start[sh], hi = shell_start[sh + 1]; if (!h_valid[sh] || hi == lo) continue; - measures_tail[sh] = z * h_v_sorted[lo + (hi - lo) / 2] <= u; + measures_tail[sh] = z * h_noise_median[sh] <= u; } CudaDevicePtr d_measures_tail(n_shells, ALLOC); cuda_err(cudaMemcpyAsync(d_measures_tail, measures_tail.data(), n_shells, cudaMemcpyHostToDevice, stream)); diff --git a/image_analysis/scale_merge/WilsonOutliersGPU.h b/image_analysis/scale_merge/WilsonOutliersGPU.h index e60467a71..f0309fa96 100644 --- a/image_analysis/scale_merge/WilsonOutliersGPU.h +++ b/image_analysis/scale_merge/WilsonOutliersGPU.h @@ -12,6 +12,6 @@ // every product the host rounds on its own is rounded on its own here, and the few scalars that need a // logarithm are taken on the host. Only built with CUDA. // -// Device memory is about 80 bytes per observation; a device that cannot hold it throws (run on the -// CPU build then). +// Device memory is about 96 bytes per observation at the peak; a device that cannot hold it throws +// (run on the CPU build then). WilsonOutlierResult WilsonOutliersGPU(const std::vector &obs, double alpha); From e59d2272945b00bb8faaba00d38693564e0bb327 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 08:48:46 +0200 Subject: [PATCH 6/8] 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. From 6e17ae36c82f32c5bdac953e9d9df91f44a2795d Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 08:51:37 +0200 Subject: [PATCH 7/8] RotationScaleMerge: the CC1/2 half-sets assigned on the device Every full's observation key and its half - the parity of the key's rank among its group's usable fulls - are integer work, so the device gives the host's halves exactly: one thread per full for the keys, one per group for the ranks. The host receives both (the anomalous statistics rank the same keys again within each Bijvoet cell) instead of walking every full's record for them. p.mtz md5 and p.cif 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 | 47 ++++++++------ .../scale_merge/RotationScaleMergeGPU.cu | 65 +++++++++++++++++++ .../scale_merge/RotationScaleMergeGPU.h | 5 ++ 3 files changed, 98 insertions(+), 19 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 7394dbe90..3856113b9 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -5031,25 +5031,34 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool // 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()); - ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), - [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) - obs_key[i] = ObservationKey(fulls[i].h, fulls[i].k, fulls[i].l, fulls[i].frame); - }); - std::vector merge_half; - // 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. - // A group's members are its stretch of the group CSR with the unusable fulls left out. - merge_half.assign(fulls.size(), 0); - ParallelBlocks(n_groups, nthreads, [&](int, int g0, int g1) { - std::vector member; - for (int g = g0; g < g1; ++g) { - member.clear(); - for (int p = fulls_group_start[g]; p < fulls_group_start[g + 1]; ++p) - if (mf.group[fulls_group_perm[p]] == g) member.push_back(fulls_group_perm[p]); - AssignHalvesOfMembers(obs_key, member.data(), static_cast(member.size()), merge_half); - } - }, 512); + std::vector merge_half(fulls.size(), 0); + bool did_gpu_halves = false; +#ifdef JFJOCH_USE_CUDA + // On the device where the merge is: integer work, the same halves. + if (use_gpu_merge) { + gpu_->MergeHalves(obs_key.data(), merge_half.data()); + did_gpu_halves = true; + } +#endif + if (!did_gpu_halves) { + ParallelChunks(static_cast(fulls.size()), ThreadsForWork(fulls.size(), nthreads), + [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) + obs_key[i] = ObservationKey(fulls[i].h, fulls[i].k, fulls[i].l, fulls[i].frame); + }); + // 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. + // A group's members are its stretch of the group CSR with the unusable fulls left out. + ParallelBlocks(n_groups, nthreads, [&](int, int g0, int g1) { + std::vector member; + for (int g = g0; g < g1; ++g) { + member.clear(); + for (int p = fulls_group_start[g]; p < fulls_group_start[g + 1]; ++p) + if (mf.group[fulls_group_perm[p]] == g) member.push_back(fulls_group_perm[p]); + AssignHalvesOfMembers(obs_key, member.data(), static_cast(member.size()), merge_half); + } + }, 512); + } // ---- CC1/2 weight (MergedReflection::cc_weight): the precision each half-set WOULD have had with // every observation at the run's typical frame scale, over the precision it has. An observation diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 79972d1b5..af682a1a6 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -690,6 +690,44 @@ namespace { } } + // The host's ObservationKey (RotationScaleMerge.cpp): integer arithmetic, so the same on both. + __device__ __forceinline__ uint64_t ObservationKeyDevice(int32_t h, int32_t k, int32_t l, int32_t frame) { + uint64_t z = uint64_t(uint32_t(h)) * 0x9E3779B97F4A7C15ULL + ^ uint64_t(uint32_t(k)) * 0xC2B2AE3D27D4EB4FULL + ^ uint64_t(uint32_t(l)) * 0x165667B19E3779F9ULL + ^ uint64_t(uint32_t(frame)) * 0x27D4EB2F165667C5ULL; + z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL; + z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL; + return z ^ (z >> 31); + } + + // One thread per full: its observation key, and whether the merge can use it. + __global__ void ObservationKeyKernel(int n_obs, MergeParams p, const int32_t *h, const int32_t *k, + const int32_t *l, uint64_t *key, uint8_t *usable) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) { + key[i] = ObservationKeyDevice(h[i], k[i], l[i], p.frame[i]); + usable[i] = MergeUsable(i, p) ? 1 : 0; + } + } + + // One thread per group: the CC1/2 half of each usable full, the parity of its key's rank among the + // group's usable fulls (the host's AssignHalvesOfMembers). Unusable fulls keep half 0. + __global__ void HalvesKernel(MergeParams p, const uint64_t *key, const uint8_t *usable, uint8_t *half) { + 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]; + for (int a = lo; a < hi; ++a) { + const int i = p.gperm[a]; + if (!usable[i]) continue; + int rank = 0; + for (int b = lo; b < hi; ++b) { + const int j = p.gperm[b]; + if (usable[j] && key[j] < key[i]) ++rank; + } + half[i] = uint8_t(rank & 1); + } + } + } + // 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; @@ -1585,6 +1623,33 @@ void RotationScaleMergeGPU::MergeGroupScatter(double *sum, double *sum_sq, doubl CopyAndWait(n, g_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl n"); } +void RotationScaleMergeGPU::MergeHalves(uint64_t *obs_key, uint8_t *half) { + DeviceGuard guard(impl_->device, impl_->available); + auto &d = *impl_; + const int ng = d.n_groups, nf = d.n_fulls; + if (nf == 0) return; + auto key = d.Alloc(nf); + auto usable = d.Alloc(nf), h = d.Alloc(nf); + + 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(); + + const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK); + const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK); + ObservationKeyKernel<<s()>>>(nf, p, d.f_h.get(), d.f_k.get(), d.f_l.get(), key.get(), + usable.get()); + CudaCheck(cudaGetLastError(), "observation key launch"); + CudaCheck(cudaMemsetAsync(h.get(), 0, size_t(nf), impl_->s()), "memset halves"); + HalvesKernel<<s()>>>(p, key.get(), usable.get(), h.get()); + CudaCheck(cudaGetLastError(), "halves launch"); + CopyAndWait(obs_key, key.get(), size_t(nf) * sizeof(uint64_t), cudaMemcpyDeviceToHost, impl_->s(), "dl keys"); + CopyAndWait(half, h.get(), size_t(nf), cudaMemcpyDeviceToHost, impl_->s(), "dl halves"); +} + 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 7ec3163de..68c4de7ab 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -112,6 +112,11 @@ public: // MergeEmSamples first. void MergeGroupScatter(double *sum, double *sum_sq, double *sum_var, int32_t *n); + // Every full's ObservationKey and its CC1/2 half - the parity of the key's rank among its group's + // usable fulls, 0 for an unusable one (length n_fulls each); the host's AssignHalvesOfMembers to + // the bit, as it is integer work. Requires MergeEmSamples first (the filter context). + void MergeHalves(uint64_t *obs_key, uint8_t *half); + // 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. From 87f0bdc8704d4690c4dc42ad31958c588e37fdf6 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 8 Oct 2026 08:52:04 +0200 Subject: [PATCH 8/8] RotationScaleMerge: fit the error model on the device in the GPU build (tier F) FitErrorModelGPU puts the samples in the host's rank order with four stable radix passes, so the equal-count bins hold the same samples; computes every normalised deviation with the host's rounding, so every bin median is the host's to the bit; and hands the per-bin terms to the host's own update step, now ErrorModelUpdate (a pure extraction from FitErrorModel). The one thing that differs is the order each bin's two sums are added in: a fixed tree on the device, the order nth_element happened to leave the bin in on the host. a and b^2 can therefore differ from the host fit in their last bits - measured on two synthetic pools: relative difference 2e-16 and 3e-13 in a, 5e-16 in b^2. The device fit is deterministic (same pool, same bits). p.mtz md5 nevertheless unchanged on myob, cytc, 8a1a and 8qaw (GPU build): the difference does not survive the float rounding of the merged sigmas there. A knife-edge decision elsewhere could still see it, which is why this is kept as its own commit. The host fit took 0.4-0.6 s per fit on the large merges, about half of it the initial equal-count split, with only 16-way parallelism in the iterations. Device peak about 72 bytes per sample (the rank sort; each stage's buffers are freed when it is done); the free memory is checked first and a device without it stops with a message naming the CPU build. Kept last on the branch so that it can be dropped on its own. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi --- image_analysis/scale_merge/CMakeLists.txt | 3 +- image_analysis/scale_merge/ErrorModel.cpp | 93 ++++----- image_analysis/scale_merge/ErrorModel.h | 13 ++ image_analysis/scale_merge/ErrorModelGPU.cu | 184 ++++++++++++++++++ image_analysis/scale_merge/ErrorModelGPU.h | 16 ++ .../scale_merge/RotationScaleMerge.cpp | 13 +- tests/MergeScaleTest.cpp | 29 +++ 7 files changed, 304 insertions(+), 47 deletions(-) create mode 100644 image_analysis/scale_merge/ErrorModelGPU.cu create mode 100644 image_analysis/scale_merge/ErrorModelGPU.h diff --git a/image_analysis/scale_merge/CMakeLists.txt b/image_analysis/scale_merge/CMakeLists.txt index 9f5fd6c85..9e128e029 100644 --- a/image_analysis/scale_merge/CMakeLists.txt +++ b/image_analysis/scale_merge/CMakeLists.txt @@ -43,5 +43,6 @@ IF (JFJOCH_CUDA_AVAILABLE) TARGET_SOURCES(JFJochScaleMerge PRIVATE ../indexing/CUDAMemHelpers.h RotationScaleMergeGPU.cu RotationScaleMergeGPU.h WilsonOutliersGPU.cu WilsonOutliersGPU.h - PairSortGPU.cu PairSortGPU.h) + PairSortGPU.cu PairSortGPU.h + ErrorModelGPU.cu ErrorModelGPU.h) ENDIF() \ No newline at end of file diff --git a/image_analysis/scale_merge/ErrorModel.cpp b/image_analysis/scale_merge/ErrorModel.cpp index 1204d2b01..9ae291bbc 100644 --- a/image_analysis/scale_merge/ErrorModel.cpp +++ b/image_analysis/scale_merge/ErrorModel.cpp @@ -12,7 +12,7 @@ ErrorModelFit FitErrorModel(const std::vector &pool, std::vector &smp, size_t nthreads) { - constexpr int n_bins = 16; + constexpr int n_bins = ERROR_MODEL_BINS; ErrorModelFit fit; if (pool.size() < static_cast(8 * n_bins)) return fit; @@ -63,16 +63,12 @@ ErrorModelFit FitErrorModel(const std::vector &pool, std::vect // hands b the bins' own noise instead (measured: b = 5.6, ISa 0.5, on a crystal XDS gives 6.1). Then // fit a alone and leave b at zero. Over the rotation battery the two crystals this fires on sit at // (I/sigma)^2 0.22 and 0.84 while the next is 31.7, so the threshold is not delicate. - constexpr double B_LEVER_MIN = 4.0; { const auto top_mid = smp.begin() + (bin_lo(n_bins - 1) + bin_hi(n_bins - 1)) / 2; std::nth_element(smp.begin() + bin_lo(n_bins - 1), top_mid, smp.end(), by_snr); - fit.b_measured = top_mid->snr2 > B_LEVER_MIN; + fit.b_measured = top_mid->snr2 > ERROR_MODEL_B_LEVER_MIN; } - // The median of a chi^2 with one degree of freedom: a normalised squared deviation that is right - // on average has half of its values below this. - constexpr double CHI2_1_MEDIAN = 0.454936; std::vector bs2(n_bins), bI2(n_bins), bd2(n_bins); double a = 1.0, b2 = 0.0; for (int iter = 0; iter < 30; ++iter) { @@ -93,48 +89,55 @@ ErrorModelFit FitErrorModel(const std::vector &pool, std::vect const auto mid = smp.begin() + bin_lo(bin) + n / 2; std::nth_element(smp.begin() + bin_lo(bin), mid, smp.begin() + bin_hi(bin), [](const ErrorModelBinned &p, const ErrorModelBinned &q) { return p.z2 < q.z2; }); - bs2[bin] = s; bI2[bin] = I; bd2[bin] = mid->z2 / CHI2_1_MEDIAN * static_cast(n); + bs2[bin] = s; bI2[bin] = I; bd2[bin] = mid->z2 / ERROR_MODEL_CHI2_1_MEDIAN * static_cast(n); } }); - // Every bin weighs the same in RELATIVE terms, so the weak bins still constrain a. - double Ass = 0, AsI = 0, AII = 0, Bs = 0, BI = 0; - for (int bin = 0; bin < n_bins; ++bin) { - if (!(bd2[bin] > 0.0)) continue; - const double w = 1.0 / (bd2[bin] * bd2[bin]); - Ass += w * bs2[bin] * bs2[bin]; AsI += w * bs2[bin] * bI2[bin]; AII += w * bI2[bin] * bI2[bin]; - Bs += w * bs2[bin] * bd2[bin]; BI += w * bI2[bin] * bd2[bin]; - } - const double det = Ass * AII - AsI * AsI; - double a_new, b2_new; - if (!fit.b_measured) { - if (!(Ass > 0.0)) return fit; - a_new = std::clamp(Bs / Ass, 0.25, 100.0); - b2_new = 0.0; - } else { - if (!(det > 1e-10 * Ass * AII)) return fit; - const double a_fit = (Bs * AII - BI * AsI) / det; - const double b2_fit = (Ass * BI - AsI * Bs) / det; - a_new = std::clamp(a_fit, 0.25, 100.0); - b2_new = std::max(b2_fit, 0.0); - // Leverage is not enough: b^2 can have it and still come out at zero within its own error, - // and 1/b then reads the noise in b as an I/sigma. The standard error of b^2 comes from the - // bins' own scatter about the fitted line; this only decides what is printed. - double chi = 0.0; - for (int bin = 0; bin < n_bins; ++bin) { - if (!(bd2[bin] > 0.0)) continue; - const double r = (bd2[bin] - a_fit * bs2[bin] - b2_fit * bI2[bin]) / bd2[bin]; - chi += r * r; - } - fit.b_resolved = b2_fit >= 2.0 * std::sqrt(chi / (n_bins - 2) * Ass / det); - } - const bool settled = std::fabs(a_new - a) <= 1e-4 * a && std::fabs(b2_new - b2) <= 1e-4 * b2; - a = a_new; - b2 = b2_new; - fit.active = true; - fit.a = a; - fit.b2 = b2; - if (settled) break; + if (!ErrorModelUpdate(bs2, bI2, bd2, a, b2, fit)) + break; } return fit; } + +bool ErrorModelUpdate(const std::vector &bs2, const std::vector &bI2, const std::vector &bd2, + double &a, double &b2, ErrorModelFit &fit) { + const int n_bins = static_cast(bs2.size()); + // Every bin weighs the same in RELATIVE terms, so the weak bins still constrain a. + double Ass = 0, AsI = 0, AII = 0, Bs = 0, BI = 0; + for (int bin = 0; bin < n_bins; ++bin) { + if (!(bd2[bin] > 0.0)) continue; + const double w = 1.0 / (bd2[bin] * bd2[bin]); + Ass += w * bs2[bin] * bs2[bin]; AsI += w * bs2[bin] * bI2[bin]; AII += w * bI2[bin] * bI2[bin]; + Bs += w * bs2[bin] * bd2[bin]; BI += w * bI2[bin] * bd2[bin]; + } + const double det = Ass * AII - AsI * AsI; + double a_new, b2_new; + if (!fit.b_measured) { + if (!(Ass > 0.0)) return false; + a_new = std::clamp(Bs / Ass, 0.25, 100.0); + b2_new = 0.0; + } else { + if (!(det > 1e-10 * Ass * AII)) return false; + const double a_fit = (Bs * AII - BI * AsI) / det; + const double b2_fit = (Ass * BI - AsI * Bs) / det; + a_new = std::clamp(a_fit, 0.25, 100.0); + b2_new = std::max(b2_fit, 0.0); + // Leverage is not enough: b^2 can have it and still come out at zero within its own error, + // and 1/b then reads the noise in b as an I/sigma. The standard error of b^2 comes from the + // bins' own scatter about the fitted line; this only decides what is printed. + double chi = 0.0; + for (int bin = 0; bin < n_bins; ++bin) { + if (!(bd2[bin] > 0.0)) continue; + const double r = (bd2[bin] - a_fit * bs2[bin] - b2_fit * bI2[bin]) / bd2[bin]; + chi += r * r; + } + fit.b_resolved = b2_fit >= 2.0 * std::sqrt(chi / (n_bins - 2) * Ass / det); + } + const bool settled = std::fabs(a_new - a) <= 1e-4 * a && std::fabs(b2_new - b2) <= 1e-4 * b2; + a = a_new; + b2 = b2_new; + fit.active = true; + fit.a = a; + fit.b2 = b2; + return !settled; +} diff --git a/image_analysis/scale_merge/ErrorModel.h b/image_analysis/scale_merge/ErrorModel.h index 89716bf9f..34ceb4a4b 100644 --- a/image_analysis/scale_merge/ErrorModel.h +++ b/image_analysis/scale_merge/ErrorModel.h @@ -44,5 +44,18 @@ struct ErrorModelBinned { double snr2, s2, I2, dev2, z2; }; +// The fit's equal-count bins of counting I/sigma; the strongest bin's median (I/sigma)^2 must exceed +// ERROR_MODEL_B_LEVER_MIN for b to be fitted (see FitErrorModel); and the median of a chi^2 with one +// degree of freedom, which the median normalised deviation of a bin is compared with. +constexpr int ERROR_MODEL_BINS = 16; +constexpr double ERROR_MODEL_B_LEVER_MIN = 4.0; +constexpr double ERROR_MODEL_CHI2_1_MEDIAN = 0.454936; + ErrorModelFit FitErrorModel(const std::vector &pool, std::vector &scratch, size_t nthreads); + +// One update of the fit from its per-bin terms - bs2 and bI2 the sums of s2 / var and I2 / var under +// the current (a, b^2), bd2 the bin's median normalised deviation expressed as the sum it stands for: +// the new (a, b^2) into a, b2 and fit. False when the fit stops, settled or no longer determined. +bool ErrorModelUpdate(const std::vector &bs2, const std::vector &bI2, const std::vector &bd2, + double &a, double &b2, ErrorModelFit &fit); diff --git a/image_analysis/scale_merge/ErrorModelGPU.cu b/image_analysis/scale_merge/ErrorModelGPU.cu new file mode 100644 index 000000000..0b5bbff1f --- /dev/null +++ b/image_analysis/scale_merge/ErrorModelGPU.cu @@ -0,0 +1,184 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "ErrorModelGPU.h" + +#include +#include + +#include +#include + +#include "../indexing/CUDAMemHelpers.h" +#include "../../common/JFJochException.h" + +namespace { + constexpr int THREADS = 256; + constexpr int CHUNKS = 64; // blocks per bin in the sum kernel; fixed, so the tree is too + constexpr int NB = ERROR_MODEL_BINS; + constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + + void cuda_err(cudaError_t val) { + if (val != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); + } + + unsigned Blocks(size_t n) { return static_cast(std::max(1, (n + THREADS - 1) / THREADS)); } + + // field 0..3: dev2, I2, s2, snr2 = I2 / s2 - the rank key's components, least significant first. + __global__ void KeyKernel(int n, const ErrorModelSample *pool, const int32_t *idx, int field, double *key) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n) return; + const ErrorModelSample &x = pool[idx[j]]; + key[j] = field == 0 ? x.dev2 : field == 1 ? x.I2 : field == 2 ? x.s2 : x.I2 / x.s2; + } + + __global__ void IotaKernel(int n, int32_t *idx) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j < n) idx[j] = j; + } + + __global__ void GatherKernel(int n, const ErrorModelSample *pool, const int32_t *idx, double *s2, double *I2, + double *dev2) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n) return; + const ErrorModelSample &x = pool[idx[j]]; + s2[j] = x.s2; I2[j] = x.I2; dev2[j] = x.dev2; + } + + // Block (chunk, bin): z2 of its samples, and its share of the bin's sums of s2/var and I2/var, + // reduced over a fixed tree. + __global__ void BinSumKernel(const int64_t *bin_start, const double *s2, const double *I2, const double *dev2, + double a, double b2, double *z2, double *part_s, double *part_I) { + __shared__ double sh_s[THREADS], sh_I[THREADS]; + const int bin = blockIdx.y, chunk = blockIdx.x; + const int64_t lo = bin_start[bin], n = bin_start[bin + 1] - lo; + const int64_t c0 = lo + n * chunk / CHUNKS, c1 = lo + n * (chunk + 1) / CHUNKS; + double s = 0.0, I = 0.0; + for (int64_t i = c0 + threadIdx.x; i < c1; i += THREADS) { + const double v = __dadd_rn(__dmul_rn(a, s2[i]), __dmul_rn(b2, I2[i])); + s += s2[i] / v; I += I2[i] / v; + z2[i] = dev2[i] / v; + } + sh_s[threadIdx.x] = s; sh_I[threadIdx.x] = I; + __syncthreads(); + for (int stride = THREADS / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) { + sh_s[threadIdx.x] += sh_s[threadIdx.x + stride]; + sh_I[threadIdx.x] += sh_I[threadIdx.x + stride]; + } + __syncthreads(); + } + if (threadIdx.x == 0) { + part_s[bin * CHUNKS + chunk] = sh_s[0]; + part_I[bin * CHUNKS + chunk] = sh_I[0]; + } + } + + __global__ void MedianKernel(const int64_t *bin_start, const double *z2_sorted, double *median) { + const int bin = threadIdx.x; + if (bin >= NB) return; + const int64_t lo = bin_start[bin], n = bin_start[bin + 1] - lo; + median[bin] = z2_sorted[lo + n / 2]; + } +} + +ErrorModelFit FitErrorModelGPU(const std::vector &pool) { + ErrorModelFit fit; + if (pool.size() < static_cast(8 * NB)) + return fit; + const int n = static_cast(pool.size()); + + // The rank sort is the peak: the pool, two key and two index arrays and the sort scratch. + constexpr size_t BYTES_PER_SAMPLE = 32 + 2 * 8 + 2 * 4 + 16; + size_t free_bytes = 0, total_bytes = 0; + cuda_err(cudaMemGetInfo(&free_bytes, &total_bytes)); + if (BYTES_PER_SAMPLE * pool.size() > free_bytes) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, + "Error model: " + std::to_string(pool.size()) + " samples need " + + std::to_string(BYTES_PER_SAMPLE * pool.size() >> 20) + " MiB of GPU memory, " + + std::to_string(free_bytes >> 20) + " MiB free - run the CPU build"); + + CudaStream stream; + CudaDevicePtr d_pool(pool.size(), ALLOC); + cuda_err(cudaMemcpyAsync(d_pool, pool.data(), pool.size() * sizeof(ErrorModelSample), cudaMemcpyHostToDevice, + stream)); + + // The samples in the host's rank order - (I2/s2, s2, I2, dev2), a total order - by four stable + // passes, least significant key first. The equal-count bins are then consecutive stretches. + CudaDevicePtr idx(n, ALLOC), idx_next(n, ALLOC); + CudaDevicePtr key(n, ALLOC), key_sorted(n, ALLOC); + IotaKernel<<>>(n, idx); + cuda_err(cudaGetLastError()); + size_t bytes = 0; + cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key.get(), key_sorted.get(), idx.get(), idx_next.get(), + n, 0, 64, stream)); + CudaDevicePtr scratch(std::max(bytes, 1), ALLOC); + for (int field = 0; field < 4; ++field) { + KeyKernel<<>>(n, d_pool, idx, field, key); + cuda_err(cudaGetLastError()); + cuda_err(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, key.get(), key_sorted.get(), idx.get(), + idx_next.get(), n, 0, 64, stream)); + std::swap(idx, idx_next); + } + scratch = CudaDevicePtr(); idx_next = CudaDevicePtr(); key = CudaDevicePtr(); + // After the last pass key_sorted holds I2/s2 in rank order. + CudaDevicePtr s2(n, ALLOC), I2(n, ALLOC), dev2(n, ALLOC); + GatherKernel<<>>(n, d_pool, idx, s2, I2, dev2); + cuda_err(cudaGetLastError()); + d_pool = CudaDevicePtr(); idx = CudaDevicePtr(); + + const size_t per = pool.size() / NB; + std::vector bin_start(NB + 1); + for (int b = 0; b < NB; ++b) bin_start[b] = static_cast(b * per); + bin_start[NB] = n; + CudaDevicePtr d_bin_start(NB + 1, ALLOC); + cuda_err(cudaMemcpyAsync(d_bin_start, bin_start.data(), (NB + 1) * sizeof(int64_t), cudaMemcpyHostToDevice, + stream)); + { + double top_mid = 0.0; + cuda_err(cudaMemcpyAsync(&top_mid, key_sorted.get() + (bin_start[NB - 1] + n) / 2, sizeof(double), + cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); + fit.b_measured = top_mid > ERROR_MODEL_B_LEVER_MIN; + } + + // The iterations: per bin the two sums (each chunk's partial added on the host in chunk order) and + // the median z2, read off each bin's sorted values. + CudaDevicePtr z2(n, ALLOC), z2_sorted(n, ALLOC), part_s(NB * CHUNKS, ALLOC), part_I(NB * CHUNKS, ALLOC), + median(NB, ALLOC); + size_t seg_bytes = 0; + cuda_err(cub::DeviceSegmentedRadixSort::SortKeys(nullptr, seg_bytes, z2.get(), z2_sorted.get(), n, NB, + d_bin_start.get(), d_bin_start.get() + 1, 0, 64, stream)); + CudaDevicePtr seg_scratch(std::max(seg_bytes, 1), ALLOC); + std::vector h_part_s(NB * CHUNKS), h_part_I(NB * CHUNKS), h_median(NB); + std::vector bs2(NB), bI2(NB), bd2(NB); + double a = 1.0, b2 = 0.0; + for (int iter = 0; iter < 30; ++iter) { + BinSumKernel<<>>(d_bin_start, s2, I2, dev2, a, b2, z2, part_s, part_I); + cuda_err(cudaGetLastError()); + cuda_err(cub::DeviceSegmentedRadixSort::SortKeys(seg_scratch.get(), seg_bytes, z2.get(), z2_sorted.get(), n, + NB, d_bin_start.get(), d_bin_start.get() + 1, 0, 64, + stream)); + MedianKernel<<<1, NB, 0, stream>>>(d_bin_start, z2_sorted, median); + cuda_err(cudaGetLastError()); + cuda_err(cudaMemcpyAsync(h_part_s.data(), part_s, NB * CHUNKS * sizeof(double), cudaMemcpyDeviceToHost, + stream)); + cuda_err(cudaMemcpyAsync(h_part_I.data(), part_I, NB * CHUNKS * sizeof(double), cudaMemcpyDeviceToHost, + stream)); + cuda_err(cudaMemcpyAsync(h_median.data(), median, NB * sizeof(double), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); + for (int bin = 0; bin < NB; ++bin) { + double s = 0.0, I = 0.0; + for (int c = 0; c < CHUNKS; ++c) { + s += h_part_s[bin * CHUNKS + c]; + I += h_part_I[bin * CHUNKS + c]; + } + const double nb = static_cast(bin_start[bin + 1] - bin_start[bin]); + bs2[bin] = s; bI2[bin] = I; bd2[bin] = h_median[bin] / ERROR_MODEL_CHI2_1_MEDIAN * nb; + } + if (!ErrorModelUpdate(bs2, bI2, bd2, a, b2, fit)) + break; + } + return fit; +} diff --git a/image_analysis/scale_merge/ErrorModelGPU.h b/image_analysis/scale_merge/ErrorModelGPU.h new file mode 100644 index 000000000..31ee96f1c --- /dev/null +++ b/image_analysis/scale_merge/ErrorModelGPU.h @@ -0,0 +1,16 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +#include "ErrorModel.h" + +// FitErrorModel (ErrorModel.h) on the GPU. The bins hold the same samples, every normalised deviation +// and so every bin median is the host's to the bit, and the update step is the host's own +// (ErrorModelUpdate). What differs is the order each bin's two sums are added in - a fixed tree here, +// the order nth_element left the bin in on the host - so a and b^2 can differ from the host fit in +// their last bits. Deterministic: the same pool gives the same fit on every run. Only built with CUDA; +// a device without room for about 72 bytes a sample throws (run on the CPU build then). +ErrorModelFit FitErrorModelGPU(const std::vector &pool); diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 3856113b9..23e497c76 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -27,6 +27,7 @@ #include "WilsonOutliers.h" #include "WilsonOutliersGPU.h" #include "PairSortGPU.h" +#include "ErrorModelGPU.h" #include "OutlierBand.h" #include "../../common/CorrelationCoefficient.h" #include "../../common/CrystalLattice.h" @@ -4676,8 +4677,18 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool // Fit (a, b) to a pool of samples (ErrorModel.h). A lambda because the pool changes once the cutoff // below is known. std::vector &fit_scratch = em_fit_pool; + // On the device where the merge is - the same fit to rounding (ErrorModelGPU.h). auto fit_error_model = [&](const std::vector &pool) { - const ErrorModelFit fit = FitErrorModel(pool, fit_scratch, nthreads); + ErrorModelFit fit; + bool did_gpu_fit = false; +#ifdef JFJOCH_USE_CUDA + if (use_gpu_merge) { + fit = FitErrorModelGPU(pool); + did_gpu_fit = true; + } +#endif + if (!did_gpu_fit) + fit = FitErrorModel(pool, fit_scratch, nthreads); if (!fit.active) return; error_model_a = fit.a; error_model_b = std::sqrt(fit.b2); diff --git a/tests/MergeScaleTest.cpp b/tests/MergeScaleTest.cpp index 6dcb3c6dc..ff159c451 100644 --- a/tests/MergeScaleTest.cpp +++ b/tests/MergeScaleTest.cpp @@ -502,3 +502,32 @@ TEST_CASE("ErrorModel_BNotMeasuredOnWeakData") { CHECK(fit.b2 == 0.0); CHECK(fit.a == Catch::Approx(0.9).epsilon(0.03)); } + +#ifdef JFJOCH_USE_CUDA +#include "../common/CUDAWrapper.h" +#include "../image_analysis/scale_merge/ErrorModelGPU.h" + +// The device fit against the host one. Same bins and same medians; only the order each bin's two sums +// are added in differs, so a and b^2 agree to rounding and every flag agrees exactly. The same pool +// twice on the device gives the same bits. +TEST_CASE("ErrorModel_DeviceFitIsTheHostFit") { + if (get_gpu_count() == 0) { + WARN("No CUDA GPU present. Skipping ErrorModel_DeviceFitIsTheHostFit"); + return; + } + for (const auto &pool : {SyntheticErrorModelSamples(1.3, 0.03, 300.0, 0.003), + SyntheticErrorModelSamples(0.9, 0.05, 1.5, 0.0)}) { + std::vector scratch; + const auto host = FitErrorModel(pool, scratch, 4); + const auto dev = FitErrorModelGPU(pool); + const auto again = FitErrorModelGPU(pool); + CHECK(dev.active == host.active); + CHECK(dev.b_measured == host.b_measured); + CHECK(dev.b_resolved == host.b_resolved); + CHECK(dev.a == Catch::Approx(host.a).epsilon(1e-12)); + CHECK(dev.b2 == Catch::Approx(host.b2).epsilon(1e-12).margin(1e-300)); + CHECK(again.a == dev.a); + CHECK(again.b2 == dev.b2); + } +} +#endif