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