diff --git a/image_analysis/scale_merge/CMakeLists.txt b/image_analysis/scale_merge/CMakeLists.txt index be7c72f30..9e128e029 100644 --- a/image_analysis/scale_merge/CMakeLists.txt +++ b/image_analysis/scale_merge/CMakeLists.txt @@ -41,5 +41,8 @@ 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 + 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/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 05cb80e42..bc164861c 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -26,6 +26,9 @@ #include "FrenchWilson.h" #include "ResolutionCutoff.h" #include "WilsonOutliers.h" +#include "WilsonOutliersGPU.h" +#include "PairSortGPU.h" +#include "ErrorModelGPU.h" #include "OutlierBand.h" #include "../../common/CorrelationCoefficient.h" #include "../../common/CrystalLattice.h" @@ -4722,8 +4725,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); @@ -4739,27 +4752,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 @@ -4956,8 +4952,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; @@ -4996,49 +5001,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); @@ -5070,25 +5090,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 @@ -5180,7 +5209,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; } @@ -5382,21 +5420,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 564e5c023..af682a1a6 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,159 @@ 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 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]); + } + } + } + + // 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 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) { - 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; + 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; + 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) { + 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 +1101,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 +1378,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 +1399,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 +1413,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, @@ -1409,6 +1563,93 @@ 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::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::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 f0d6caf4b..68c4de7ab 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,34 @@ 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); + + // 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); + + // 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); + + // 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 diff --git a/image_analysis/scale_merge/WilsonOutliersGPU.cu b/image_analysis/scale_merge/WilsonOutliersGPU.cu new file mode 100644 index 000000000..7287f8147 --- /dev/null +++ b/image_analysis/scale_merge/WilsonOutliersGPU.cu @@ -0,0 +1,409 @@ +// 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]]); + } + + __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) { + 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)); + } + + 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))); + + 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)); + } + 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_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) + 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_noise_median[sh] <= 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..f0309fa96 --- /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 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); 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/MergeScaleTest.cpp b/tests/MergeScaleTest.cpp index 74c3819f0..187c2cea7 100644 --- a/tests/MergeScaleTest.cpp +++ b/tests/MergeScaleTest.cpp @@ -526,3 +526,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 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 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