From f849e2d1bea1faf5a78f5a92cd79bb0563d16506 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 3 Oct 2026 12:19:05 +0200 Subject: [PATCH 1/2] Pre-scan on the GPU: background beam-centre walk and beam-stop mask The two pre-scan steps that were still CPU-bound in a GPU build now run where the projection already is. - FindBeamCenterFromBackground: the per-iteration binning pass and the two clip rounds run on the device (BeamCenterBackgroundGPU); the fit itself stays on the host. Each cell is summed in the host's order (pixel order within the host's row blocks, blocks in order), and the per-pixel cell/derivative formula is shared (BackgroundBand.h). The angles come from BackgroundAtan2 (IEEE ops only) instead of atan2f, and both translation units are compiled without FMA contraction, so host and device give the same bits: 0 of 6.5 M pixels in a different cell, identical walks on the three in-house rotation sets. With glibc/CUDA atan2f and default contraction ~30 pixels per 16 Mpx sweep changed cell and the fitted centre moved by up to 0.05 px. - ShadowFinder::GetMask: the whole mask (pooling, ring medians, components, morphology, hole fill, arm search) runs on the device from ShadowAccumulatorGPU's projection (ShadowMaskGPU), so the 360 MB projection no longer comes back; the mean projection is divided on the device too (same bits). The two small fits over rings and sectors (BlockedOutTo, HarmonicFit) are shared with the host path in ShadowFinderInternal.h. Integers, comparisons, sorts and components are exact; the polarization trig, the Poisson log and the arm-search azimuth are not, so a pixel at a threshold can differ. The one-time change against the previous CPU arithmetic (BackgroundAtan2, no contraction), measured on the myoglobin, cytochrome C and thaumatin rotation sets: ring centre moves 0.002-0.045 px (fit sigma 0.75-1.2 px), beam-centre capture 0.01-0.04 px; beam-stop mask differs on 31 / 144 / 53 pixels of 259k / 144k / 198k (25 of the myoglobin ones are GPU-vs-CPU arithmetic in the mask, the rest follow the centre); hot-pixel mask identical. Spot width, integration radii, bandwidth, beam-centre arbitration, indexing, space group, cell, resolution and the merged statistics table are identical; only the error model moves in its 4th digit. CPU build: the same centres and decisions. Timing (GPU, box at load 30-38): ring walk 0.54 -> 0.23-0.27 s, mask 1.24-1.44 -> 0.18-0.22 s, beam-centre capture walk 1.1-1.3 -> 0.31-0.35 s. Tests: ShadowFinder_DeviceMaskMatchesHost, BeamCenterFromBackground_DeviceMatchesHost (bit-exact), plus [ShadowFinder], [BeamCenter], [HotPixelFinder]. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01K5K8jvPPbmCrbqnWkddTuB --- image_analysis/CMakeLists.txt | 3 + .../beam_stop/ShadowAccumulatorGPU.cu | 10 + .../beam_stop/ShadowAccumulatorGPU.h | 6 + image_analysis/beam_stop/ShadowFinder.cpp | 242 +++---- .../beam_stop/ShadowFinderInternal.h | 90 +++ image_analysis/beam_stop/ShadowMaskGPU.cu | 636 ++++++++++++++++++ image_analysis/beam_stop/ShadowMaskGPU.h | 39 ++ .../geom_refinement/BackgroundBand.h | 107 +++ .../BeamCenterBackgroundGPU.cu | 210 ++++++ .../geom_refinement/BeamCenterBackgroundGPU.h | 40 ++ .../BeamCenterFromBackground.cpp | 126 ++-- .../BeamCenterFromBackground.h | 6 +- image_analysis/geom_refinement/CMakeLists.txt | 14 +- tests/BeamCenterFromBackgroundTest.cpp | 26 + tests/ShadowFinderTest.cpp | 88 +++ 15 files changed, 1468 insertions(+), 175 deletions(-) create mode 100644 image_analysis/beam_stop/ShadowFinderInternal.h create mode 100644 image_analysis/beam_stop/ShadowMaskGPU.cu create mode 100644 image_analysis/beam_stop/ShadowMaskGPU.h create mode 100644 image_analysis/geom_refinement/BackgroundBand.h create mode 100644 image_analysis/geom_refinement/BeamCenterBackgroundGPU.cu create mode 100644 image_analysis/geom_refinement/BeamCenterBackgroundGPU.h diff --git a/image_analysis/CMakeLists.txt b/image_analysis/CMakeLists.txt index 6abc665d9..d9a50ec73 100644 --- a/image_analysis/CMakeLists.txt +++ b/image_analysis/CMakeLists.txt @@ -99,7 +99,10 @@ ADD_LIBRARY(JFJochImageAnalysis STATIC beam_stop/ShadowFinder.cpp beam_stop/ShadowFinder.h $<$:beam_stop/ShadowAccumulatorGPU.cu> + $<$:beam_stop/ShadowMaskGPU.cu> beam_stop/ShadowAccumulatorGPU.h + beam_stop/ShadowFinderInternal.h + beam_stop/ShadowMaskGPU.h rotation_indexer/RotationIndexer.cpp rotation_indexer/RotationIndexer.h WriteReflections.cpp diff --git a/image_analysis/beam_stop/ShadowAccumulatorGPU.cu b/image_analysis/beam_stop/ShadowAccumulatorGPU.cu index 240c83d86..1499f1fec 100644 --- a/image_analysis/beam_stop/ShadowAccumulatorGPU.cu +++ b/image_analysis/beam_stop/ShadowAccumulatorGPU.cu @@ -190,3 +190,13 @@ void ShadowAccumulatorGPU::Download(std::vector &max_value, std::vector cudaMemcpyDeviceToHost, *stream)); cuda_err(cudaStreamSynchronize(*stream)); } + +std::vector ShadowAccumulatorGPU::Mask(const ShadowMaskSetup &setup, const std::vector &pixel_mask) { + FoldPending(); + return ShadowMaskOnDevice(setup, pixel_mask, gpu_max, gpu_sum, gpu_count, frames, *stream); +} + +std::vector ShadowAccumulatorGPU::MeanProjection(const std::vector &pixel_mask) { + FoldPending(); + return MeanProjectionOnDevice(pixel_mask, gpu_sum, gpu_count, npixels, *stream); +} diff --git a/image_analysis/beam_stop/ShadowAccumulatorGPU.h b/image_analysis/beam_stop/ShadowAccumulatorGPU.h index c45a424c7..71907a0be 100644 --- a/image_analysis/beam_stop/ShadowAccumulatorGPU.h +++ b/image_analysis/beam_stop/ShadowAccumulatorGPU.h @@ -10,6 +10,7 @@ #include "../../common/CompressedImage.h" #include "../image_preprocessing/BSLZ4DecoderGPU.h" #include "../indexing/CUDAMemHelpers.h" +#include "ShadowMaskGPU.h" // The beam-stop projection accumulated on the device: only the compressed chunk crosses PCIe, and // both the decode and the per-pixel maximum / sum / count run on the GPU. The projection comes back @@ -59,6 +60,11 @@ public: [[nodiscard]] uint32_t GetFrameCount() const { return frames; } + // The beam-stop mask and the mean projection, made where the projection is (ShadowMaskGPU.h), so + // that it does not have to come back at all. + std::vector Mask(const ShadowMaskSetup &setup, const std::vector &pixel_mask); + std::vector MeanProjection(const std::vector &pixel_mask); + // Bring the projection back to the host, folding in whatever the last batch still holds. Cheap // to call once; it moves 20 bytes per pixel. void Download(std::vector &max_value, std::vector &sum_value, diff --git a/image_analysis/beam_stop/ShadowFinder.cpp b/image_analysis/beam_stop/ShadowFinder.cpp index 95846c533..3d62d6d6d 100644 --- a/image_analysis/beam_stop/ShadowFinder.cpp +++ b/image_analysis/beam_stop/ShadowFinder.cpp @@ -2,6 +2,7 @@ // SPDX-License-Identifier: GPL-3.0-only #include "ShadowFinder.h" +#include "ShadowFinderInternal.h" #include #include @@ -18,64 +19,9 @@ #include "../../common/ParallelFor.h" #include "../../common/JFJochException.h" -// A pixel is shadow when its background is below this fraction of the background it is -// compared against. -constexpr float SHADOW_RATIO = 0.50f; +using namespace shadow_finder; -// The boundary grows outward into partially shadowed pixels down to this fraction, but no -// further than PENUMBRA_MAX_PX from the core. A pin or a loop casts a wide half-shadow, so the -// reach is a good deal more than the beam stop's own edge needs. -constexpr float PENUMBRA_RATIO = 0.75f; -constexpr int PENUMBRA_MAX_PX = 30; - -// Bridge small breaks along the holder arm. A module gap wider than this is bridged separately for -// the arm search (see bridge_gaps). -constexpr int BRIDGE_PX = 6; - -// A pixel whose maximum reaches this recorded a real reflection and is never masked - a -// beam stop cannot block a reflection that was measured. -constexpr int64_t MIN_REFLECTION = 25; - -// How far below the background it is compared against a pixel must sit before the dip is -// believed, in standard deviations of the counts that back it. The counts are photons, so their -// scatter is Poisson and the deficit is measured against it rather than against a fixed number: -// on a well-exposed sweep a third of the background missing is overwhelming, and on a handful of -// low-background frames the same third is noise. Without this a six-frame pre-scan of a -// low-background sweep masks three quarters of the detector. -constexpr double MIN_DEFICIT_SIGMA = 6.0; - -// Smallest region the per-pixel test may return. A shadow is cast by something physical and is -// correspondingly large; an isolated patch this small is the background wandering, not hardware. -// This is what keeps the test specific now that a shadow no longer has to touch the direct beam. -constexpr int MIN_SHADOW_PIXELS = 2000; - -// ... and of those, how many must be deep (below SHADOW_RATIO) for a region of merely DIM pixels - -// hardware that lets part of the beam through - to count. A thin holder arm is dim along most of -// its length and deep in places; the background drifting over a detector's edge is dim everywhere -// and deep nowhere. -constexpr int MIN_CORE_PIXELS = 200; - -// The arm search compares a pixel with its ring as the ring actually varies around the beam. A ring -// of background is not flat once divided by the polarization factor when that factor is not the -// beam's: the remainder is a second harmonic in azimuth, cos 2phi, which reaches tens of percent at -// high angle. It is measured over radial bands of HARMONIC_BAND_PX, from the median of each of -// HARMONIC_SECTORS sectors that holds at least MIN_SECTOR_PIXELS pixels. -constexpr int HARMONIC_BAND_PX = 64; -constexpr int HARMONIC_SECTORS = 24; -constexpr int MIN_SECTOR_PIXELS = 200; - -// Side of the box the background is pooled over before testing. Its area is how many pixels back -// a ring's countability test, which decides where an azimuthal comparison is possible at all. -constexpr int POOL_PX = 5; -constexpr double MEAN_POOLED_PIXELS = POOL_PX * POOL_PX; - -// A ring with fewer valid pixels than this says nothing about whether it was counted. -constexpr int MIN_RING_PIXELS = 32; - -// A ring lies wholly inside the stop when its background is below this fraction of the background -// further out. This asks about a whole ring rather than about a pixel, so it keeps a threshold of -// its own and does not follow SHADOW_RATIO. -constexpr float BLOCKED_RING_RATIO = 0.35f; +static_assert(ShadowFinder::SHADOW == MASK_SHADOW && ShadowFinder::TRANSMITTING == MASK_TRANSMITTING); // Binary-image helpers on a width*height frame stored row-major as char (0/1). All run once, // at GetMask() time, and all are O(pixels) rather than O(pixels * radius). @@ -635,6 +581,10 @@ void ShadowFinder::ReleaseProjection() { std::vector ShadowFinder::GetMeanProjection() const { std::unique_lock ul(m); +#ifdef JFJOCH_USE_CUDA + if (Gpu() && gpu->GetFrameCount() > 0 && host.frames == 0) + return gpu->MeanProjection(pixel_mask); +#endif const Projection &p = Reduced(); const auto &sum_value = p.sum_value; const auto &valid_count = p.valid_count; @@ -652,6 +602,28 @@ std::vector ShadowFinder::GetMask(size_t nthreads) const { std::unique_lock ul(m); if (nthreads == 0) nthreads = std::max(1u, std::thread::hardware_concurrency()); +#ifdef JFJOCH_USE_CUDA + // Where every frame went to the device the mask is made there too, from the projection as it lies. + if (Gpu() && gpu->GetFrameCount() > 0 && host.frames == 0) { + const float diag = std::hypot(static_cast(width), static_cast(height)); + if (!std::isfinite(beam_x) || !std::isfinite(beam_y) + || std::fabs(beam_x - width * 0.5f) > 4.0f * diag || std::fabs(beam_y - height * 0.5f) > 4.0f * diag) + return std::vector(static_cast(width) * height, 0); + ShadowMaskSetup setup; + setup.width = width; + setup.height = height; + setup.beam_x = beam_x; + setup.beam_y = beam_y; + const auto rot = geometry.GetDetectorMatrix().arr(); + for (int k = 0; k < 9; k++) + setup.det_matrix[k] = rot[k]; + setup.pixel_size_mm = geometry.GetPixelSize_mm(); + setup.distance_mm = geometry.GetDetectorDistance_mm(); + setup.has_polarization = polarization.has_value(); + setup.polarization = polarization.value_or(0.0f); + return gpu->Mask(setup, pixel_mask); + } +#endif const Projection &p = Reduced(); const auto &max_value = p.max_value; const auto &sum_value = p.sum_value; @@ -804,25 +776,7 @@ std::vector ShadowFinder::GetMask(size_t nthreads) const { // disk it masked. The median over the rings this walk is willing to judge is what "typically" // means, and it is robust from both sides - to a bright ring, and to a corner ring of a handful // of pixels whose median is one pixel's mean. - std::vector judgeable; - for (int rad = 0; rad <= max_radius; rad++) - if (ring_pixels[rad] >= MIN_RING_PIXELS) - judgeable.push_back(baseline[rad]); - float typical_background = 0.0f; - if (!judgeable.empty()) { - const auto middle = judgeable.begin() + judgeable.size() / 2; - std::nth_element(judgeable.begin(), middle, judgeable.end()); - typical_background = *middle; - } - - int blocked_out_to = -1; - for (int rad = 0; rad <= max_radius; rad++) { - if (ring_pixels[rad] < MIN_RING_PIXELS) - continue; - if (baseline[rad] >= BLOCKED_RING_RATIO * typical_background) - break; - blocked_out_to = rad; - } + const int blocked_out_to = BlockedOutTo(baseline, ring_pixels); // The counts a pixel's pooled background is made of, and the counts the ring says it should // have had. The test is on the deficit between them, in units of its own Poisson scatter. @@ -975,51 +929,17 @@ std::vector ShadowFinder::GetMask(size_t nthreads) const { sector_values[k].insert(sector_values[k].end(), values[k].begin(), values[k].end()); }); block_values.clear(); - std::vector harm_c(n_bands, 0.0f), harm_s(n_bands, 0.0f); - ParallelFor(n_bands, nthreads, [&](int band) { - std::vector med(HARMONIC_SECTORS, -1.0), c(HARMONIC_SECTORS), s(HARMONIC_SECTORS); - for (int k = 0; k < HARMONIC_SECTORS; k++) { - auto &v = sector_values[static_cast(band) * HARMONIC_SECTORS + k]; - if (v.size() >= MIN_SECTOR_PIXELS) { - std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end()); - med[k] = v[v.size() / 2]; - } - const double phi = (k + 0.5) * 2.0 * std::numbers::pi / HARMONIC_SECTORS - std::numbers::pi; - c[k] = std::cos(2 * phi); - s[k] = std::sin(2 * phi); - } - // Least squares of med = m + p cos 2phi + q sin 2phi over the sectors that are not themselves - // dim against the fit, three times over as the baseline is; the model is relative, 1 + (p cos + q sin)/m. - std::vector use(HARMONIC_SECTORS); - for (int k = 0; k < HARMONIC_SECTORS; k++) - use[k] = med[k] >= 0; - for (int iter = 0; iter < 3; iter++) { - double n = 0, sc = 0, ss = 0, scc = 0, sss = 0, scs = 0, y = 0, yc = 0, ys = 0; - for (int k = 0; k < HARMONIC_SECTORS; k++) { - if (!use[k]) continue; - n++; sc += c[k]; ss += s[k]; scc += c[k] * c[k]; sss += s[k] * s[k]; scs += c[k] * s[k]; - y += med[k]; yc += med[k] * c[k]; ys += med[k] * s[k]; - } - // Sectors crowded into a narrow arc cannot tell a harmonic from a level; the determinant of - // the normal matrix, per sector cubed, is 1/4 on a full ring. - const double det = n * (scc * sss - scs * scs) - sc * (sc * sss - scs * ss) + ss * (sc * scs - scc * ss); - if (n < 6 || det < 0.01 * n * n * n) { - harm_c[band] = harm_s[band] = 0.0f; - return; - } - const double m = (y * (scc * sss - scs * scs) - sc * (yc * sss - scs * ys) + ss * (yc * scs - scc * ys)) / det; - const double p = (n * (yc * sss - ys * scs) - y * (sc * sss - scs * ss) + ss * (sc * ys - yc * ss)) / det; - const double q = (n * (scc * ys - scs * yc) - sc * (sc * ys - yc * ss) + y * (sc * scs - scc * ss)) / det; - if (m <= 0) { - harm_c[band] = harm_s[band] = 0.0f; - return; - } - harm_c[band] = static_cast(p / m); - harm_s[band] = static_cast(q / m); - for (int k = 0; k < HARMONIC_SECTORS; k++) - use[k] = med[k] >= 0 && med[k] >= PENUMBRA_RATIO * (1.0 + harm_c[band] * c[k] + harm_s[band] * s[k]); + // The median of each sector with enough pixels to have one. + std::vector sector_median(n_sectors, -1.0); + ParallelFor(static_cast(n_sectors), nthreads, [&](int k) { + auto &v = sector_values[k]; + if (v.size() >= MIN_SECTOR_PIXELS) { + std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end()); + sector_median[k] = v[v.size() / 2]; } }); + std::vector harm_c, harm_s; + HarmonicFit(sector_median, n_bands, harm_c, harm_s); Plane dim = filled_plane(n_pixels, 0, nthreads); ParallelChunks(n_pixels, nthreads, [&](int lo, int hi) { @@ -1062,3 +982,87 @@ std::vector ShadowFinder::GetMask(size_t nthreads) const { }); return mask; } + +namespace shadow_finder { + +int BlockedOutTo(const std::vector &baseline, const std::vector &ring_pixels) { + const int max_radius = static_cast(baseline.size()) - 1; + // A ring lies inside the stop when its background is a fraction of what this detector's + // background typically is. Counting statistics cannot decide this: on a bright dataset the + // shadow is still well counted. The comparison used to be against the LARGEST background of any + // ring further out, and that reads a sample whose background peaks in a strong ring away from + // the beam - a powder standard, a strong solvent ring - as a beam stop the size of that ring: + // the ordinary background inside it is legitimately below a third of the peak. On one corpus + // dataset it declared 16 % of the detector to be stop, with diffraction rings visible inside the + // disk it masked. The median over the rings this walk is willing to judge is what "typically" + // means, and it is robust from both sides - to a bright ring, and to a corner ring of a handful + // of pixels whose median is one pixel's mean. + std::vector judgeable; + for (int rad = 0; rad <= max_radius; rad++) + if (ring_pixels[rad] >= MIN_RING_PIXELS) + judgeable.push_back(baseline[rad]); + float typical_background = 0.0f; + if (!judgeable.empty()) { + const auto middle = judgeable.begin() + judgeable.size() / 2; + std::nth_element(judgeable.begin(), middle, judgeable.end()); + typical_background = *middle; + } + + int blocked_out_to = -1; + for (int rad = 0; rad <= max_radius; rad++) { + if (ring_pixels[rad] < MIN_RING_PIXELS) + continue; + if (baseline[rad] >= BLOCKED_RING_RATIO * typical_background) + break; + blocked_out_to = rad; + } + return blocked_out_to; +} + +void HarmonicFit(const std::vector §or_median, int n_bands, + std::vector &harm_c, std::vector &harm_s) { + harm_c.assign(n_bands, 0.0f); + harm_s.assign(n_bands, 0.0f); + for (int band = 0; band < n_bands; band++) { + std::vector med(HARMONIC_SECTORS), c(HARMONIC_SECTORS), s(HARMONIC_SECTORS); + for (int k = 0; k < HARMONIC_SECTORS; k++) { + med[k] = sector_median[static_cast(band) * HARMONIC_SECTORS + k]; + const double phi = (k + 0.5) * 2.0 * std::numbers::pi / HARMONIC_SECTORS - std::numbers::pi; + c[k] = std::cos(2 * phi); + s[k] = std::sin(2 * phi); + } + // Least squares of med = m + p cos 2phi + q sin 2phi over the sectors that are not themselves + // dim against the fit, three times over as the baseline is; the model is relative, 1 + (p cos + q sin)/m. + std::vector use(HARMONIC_SECTORS); + for (int k = 0; k < HARMONIC_SECTORS; k++) + use[k] = med[k] >= 0; + for (int iter = 0; iter < 3; iter++) { + double n = 0, sc = 0, ss = 0, scc = 0, sss = 0, scs = 0, y = 0, yc = 0, ys = 0; + for (int k = 0; k < HARMONIC_SECTORS; k++) { + if (!use[k]) continue; + n++; sc += c[k]; ss += s[k]; scc += c[k] * c[k]; sss += s[k] * s[k]; scs += c[k] * s[k]; + y += med[k]; yc += med[k] * c[k]; ys += med[k] * s[k]; + } + // Sectors crowded into a narrow arc cannot tell a harmonic from a level; the determinant of + // the normal matrix, per sector cubed, is 1/4 on a full ring. + const double det = n * (scc * sss - scs * scs) - sc * (sc * sss - scs * ss) + ss * (sc * scs - scc * ss); + if (n < 6 || det < 0.01 * n * n * n) { + harm_c[band] = harm_s[band] = 0.0f; + break; + } + const double m = (y * (scc * sss - scs * scs) - sc * (yc * sss - scs * ys) + ss * (yc * scs - scc * ys)) / det; + const double p = (n * (yc * sss - ys * scs) - y * (sc * sss - scs * ss) + ss * (sc * ys - yc * ss)) / det; + const double q = (n * (scc * ys - scs * yc) - sc * (sc * ys - yc * ss) + y * (sc * scs - scc * ss)) / det; + if (m <= 0) { + harm_c[band] = harm_s[band] = 0.0f; + break; + } + harm_c[band] = static_cast(p / m); + harm_s[band] = static_cast(q / m); + for (int k = 0; k < HARMONIC_SECTORS; k++) + use[k] = med[k] >= 0 && med[k] >= PENUMBRA_RATIO * (1.0 + harm_c[band] * c[k] + harm_s[band] * s[k]); + } + } +} + +} // namespace shadow_finder diff --git a/image_analysis/beam_stop/ShadowFinderInternal.h b/image_analysis/beam_stop/ShadowFinderInternal.h new file mode 100644 index 000000000..63e292a70 --- /dev/null +++ b/image_analysis/beam_stop/ShadowFinderInternal.h @@ -0,0 +1,90 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// What ShadowFinder::GetMask shares with its device twin (ShadowMaskGPU): the constants it tests +// against, and the two small fits made over whole rings and sectors, which stay on the host on both +// paths. + +#include +#include + +namespace shadow_finder { + +// The values of ShadowFinder::SHADOW and ShadowFinder::TRANSMITTING, for the device code, which does +// not include ShadowFinder.h. +inline constexpr uint32_t MASK_SHADOW = 1; +inline constexpr uint32_t MASK_TRANSMITTING = 2; + +// A pixel is shadow when its background is below this fraction of the background it is +// compared against. +inline constexpr float SHADOW_RATIO = 0.50f; + +// The boundary grows outward into partially shadowed pixels down to this fraction, but no +// further than PENUMBRA_MAX_PX from the core. A pin or a loop casts a wide half-shadow, so the +// reach is a good deal more than the beam stop's own edge needs. +inline constexpr float PENUMBRA_RATIO = 0.75f; +inline constexpr int PENUMBRA_MAX_PX = 30; + +// Bridge small breaks along the holder arm. A module gap wider than this is bridged separately for +// the arm search (see bridge_gaps). +inline constexpr int BRIDGE_PX = 6; + +// A pixel whose maximum reaches this recorded a real reflection and is never masked - a +// beam stop cannot block a reflection that was measured. +inline constexpr int64_t MIN_REFLECTION = 25; + +// How far below the background it is compared against a pixel must sit before the dip is +// believed, in standard deviations of the counts that back it. The counts are photons, so their +// scatter is Poisson and the deficit is measured against it rather than against a fixed number: +// on a well-exposed sweep a third of the background missing is overwhelming, and on a handful of +// low-background frames the same third is noise. Without this a six-frame pre-scan of a +// low-background sweep masks three quarters of the detector. +inline constexpr double MIN_DEFICIT_SIGMA = 6.0; + +// Smallest region the per-pixel test may return. A shadow is cast by something physical and is +// correspondingly large; an isolated patch this small is the background wandering, not hardware. +// This is what keeps the test specific now that a shadow no longer has to touch the direct beam. +inline constexpr int MIN_SHADOW_PIXELS = 2000; + +// ... and of those, how many must be deep (below SHADOW_RATIO) for a region of merely DIM pixels - +// hardware that lets part of the beam through - to count. A thin holder arm is dim along most of +// its length and deep in places; the background drifting over a detector's edge is dim everywhere +// and deep nowhere. +inline constexpr int MIN_CORE_PIXELS = 200; + +// The arm search compares a pixel with its ring as the ring actually varies around the beam. A ring +// of background is not flat once divided by the polarization factor when that factor is not the +// beam's: the remainder is a second harmonic in azimuth, cos 2phi, which reaches tens of percent at +// high angle. It is measured over radial bands of HARMONIC_BAND_PX, from the median of each of +// HARMONIC_SECTORS sectors that holds at least MIN_SECTOR_PIXELS pixels. +inline constexpr int HARMONIC_BAND_PX = 64; +inline constexpr int HARMONIC_SECTORS = 24; +inline constexpr int MIN_SECTOR_PIXELS = 200; + +// Side of the box the background is pooled over before testing. Its area is how many pixels back +// a ring's countability test, which decides where an azimuthal comparison is possible at all. +inline constexpr int POOL_PX = 5; +inline constexpr double MEAN_POOLED_PIXELS = POOL_PX * POOL_PX; + +// A ring with fewer valid pixels than this says nothing about whether it was counted. +inline constexpr int MIN_RING_PIXELS = 32; + +// A ring lies wholly inside the stop when its background is below this fraction of the background +// further out. This asks about a whole ring rather than about a pixel, so it keeps a threshold of +// its own and does not follow SHADOW_RATIO. +inline constexpr float BLOCKED_RING_RATIO = 0.35f; + +// The rings that lie wholly inside the stop: walking outward, every judgeable ring (at least +// MIN_RING_PIXELS pixels) before the first whose baseline reaches BLOCKED_RING_RATIO of the typical +// background. -1 when there is none. See GetMask. +int BlockedOutTo(const std::vector &baseline, const std::vector &ring_pixels); + +// The second harmonic in azimuth of each radial band, relative to its level, fitted to the medians of +// its HARMONIC_SECTORS sectors (sector_median[band * HARMONIC_SECTORS + k], negative where the sector +// has too few pixels). See GetMask. +void HarmonicFit(const std::vector §or_median, int n_bands, + std::vector &harm_c, std::vector &harm_s); + +} // namespace shadow_finder diff --git a/image_analysis/beam_stop/ShadowMaskGPU.cu b/image_analysis/beam_stop/ShadowMaskGPU.cu new file mode 100644 index 000000000..5961535c2 --- /dev/null +++ b/image_analysis/beam_stop/ShadowMaskGPU.cu @@ -0,0 +1,636 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "ShadowMaskGPU.h" + +#include + +#include "ShadowFinderInternal.h" +#include "../../common/JFJochMath.h" +#include "../indexing/CUDAMemHelpers.h" +#include "../../common/JFJochException.h" + +using namespace shadow_finder; + +namespace { + +constexpr int THREADS = 256; + +void check(cudaError_t err, const char *what) { + if (err != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, + std::string("Beam stop mask: ") + what + ": " + cudaGetErrorString(err)); +} + +unsigned grid(size_t n) { + return static_cast((n + THREADS - 1) / THREADS); +} + +// A float as an unsigned integer that sorts the same way, and back. +__device__ uint32_t float_key(float f) { + const uint32_t u = __float_as_uint(f); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} +__device__ float key_float(uint32_t k) { + return __uint_as_float((k & 0x80000000u) ? (k & 0x7fffffffu) : ~k); +} + +// The first index in a sorted key array whose key is not below `value`. +__device__ size_t lower_bound(const uint64_t *keys, size_t n, uint64_t value) { + size_t lo = 0, hi = n; + while (lo < hi) { + const size_t mid = (lo + hi) / 2; + if (keys[mid] < value) lo = mid + 1; + else hi = mid; + } + return lo; +} + +__device__ double poisson_deficit_sigma(double observed, double expected) { + if (expected <= 0.0 || observed >= expected) + return 0.0; + const double ll = 2.0 * (expected - observed + (observed > 0.0 ? observed * log(observed / expected) : 0.0)); + return ll > 0.0 ? sqrt(ll) : 0.0; +} + +// DiffractionGeometry::CalcAzIntPolarizationCorr about the centre the rings are drawn about. +__device__ float polarization_factor(const ShadowMaskSetup &s, float x, float y) { + const float u = (x - s.beam_x) * s.pixel_size_mm; + const float v = (y - s.beam_y) * s.pixel_size_mm; + const float *m = s.det_matrix; + const float lx = m[0] * u + m[1] * v + m[2] * s.distance_mm; + const float ly = m[3] * u + m[4] * v + m[5] * s.distance_mm; + const float lz = m[6] * u + m[7] * v + m[8] * s.distance_mm; + const float two_theta = atan2f(sqrtf(lx * lx + ly * ly), lz); + float phi = atan2f(ly, lx); + if (phi < 0) + phi += 2.0f * PI; + const float cos_2theta = cosf(two_theta); + const float cos_2theta_2 = cos_2theta * cos_2theta; + const float cos_2phi = cosf(2.0f * phi); + return 0.5f * (1.0f + cos_2theta_2 - s.polarization * cos_2phi * (1.0f - cos_2theta_2)); +} + +__global__ void setup_kernel(ShadowMaskSetup s, const uint32_t *__restrict__ pixel_mask, + const int64_t *__restrict__ sum_value, const uint32_t *__restrict__ valid_count, + float *__restrict__ pol, char *__restrict__ valid, int *__restrict__ radius, + double *__restrict__ num, int32_t *__restrict__ den, int *__restrict__ max_radius) { + const size_t n = static_cast(s.width) * s.height; + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) + return; + const int x = static_cast(i % s.width), y = static_cast(i / s.width); + const float dx = x - s.beam_x, dy = y - s.beam_y; + const float p = s.has_polarization ? polarization_factor(s, static_cast(x), static_cast(y)) : 1.0f; + pol[i] = p; + float mean = 0.0f; + char v = 0; + if (valid_count[i] > 0 && pixel_mask[i] == 0 && p > 0.0f) { + mean = static_cast(static_cast(sum_value[i]) / valid_count[i] / p); + v = 1; + } + valid[i] = v; + num[i] = v ? mean : 0.0; + den[i] = v ? 1 : 0; + const int r = static_cast(lroundf(sqrtf(dx * dx + dy * dy))); + radius[i] = r; + atomicMax(max_radius, r); +} + +// Sum over the k x k box about each pixel, zero outside the frame: one running sum per row, then one +// per column, each with exactly the terms and order of the host's (box_sum in ShadowFinder.cpp). +template +__global__ void box_rows(const T *__restrict__ in, T *__restrict__ out, int W, int H, int half) { + const int y = blockIdx.x * blockDim.x + threadIdx.x; + if (y >= H) return; + const T *src = in + static_cast(y) * W; + T *dst = out + static_cast(y) * W; + T s = 0; + for (int x = 0; x <= min(half, W - 1); x++) + s += src[x]; + for (int x = 0; x < W; x++) { + dst[x] = s; + if (x + half + 1 < W) s += src[x + half + 1]; + if (x - half >= 0) s -= src[x - half]; + } +} + +template +__global__ void box_columns(const T *__restrict__ in, T *__restrict__ out, int W, int H, int half) { + const int x = blockIdx.x * blockDim.x + threadIdx.x; + if (x >= W) return; + T s = 0; + for (int y = 0; y <= min(half, H - 1); y++) + s += in[static_cast(y) * W + x]; + for (int y = 0; y < H; y++) { + out[static_cast(y) * W + x] = s; + if (y + half + 1 < H) s += in[static_cast(y + half + 1) * W + x]; + if (y - half >= 0) s -= in[static_cast(y - half) * W + x]; + } +} + +// Dilation of a 0/1 plane by the (2r+1) square clipped to the frame, as a count over a sliding window +// along rows and then along columns (dilate in ShadowFinder.cpp). +__global__ void dilate_rows(const char *__restrict__ in, char *__restrict__ out, int W, int H, int r) { + const int y = blockIdx.x * blockDim.x + threadIdx.x; + if (y >= H) return; + const char *src = in + static_cast(y) * W; + char *dst = out + static_cast(y) * W; + int count = 0; + for (int x = 0; x <= min(r, W - 1); x++) + count += src[x]; + for (int x = 0; x < W; x++) { + dst[x] = count > 0; + if (x + r + 1 < W) count += src[x + r + 1]; + if (x - r >= 0) count -= src[x - r]; + } +} + +__global__ void dilate_columns(const char *__restrict__ in, char *__restrict__ out, int W, int H, int r) { + const int x = blockIdx.x * blockDim.x + threadIdx.x; + if (x >= W) return; + int count = 0; + for (int y = 0; y <= min(r, H - 1); y++) + count += in[static_cast(y) * W + x]; + for (int y = 0; y < H; y++) { + out[static_cast(y) * W + x] = count > 0; + if (y + r + 1 < H) count += in[static_cast(y + r + 1) * W + x]; + if (y - r >= 0) count -= in[static_cast(y - r) * W + x]; + } +} + +__global__ void pooled_kernel(size_t n, const double *__restrict__ pooled_sum, const int32_t *__restrict__ pooled_count, + const char *__restrict__ valid, const int *__restrict__ radius, + float *__restrict__ pooled, uint64_t *__restrict__ ring_key) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + const float p = pooled_count[i] > 0 ? static_cast(pooled_sum[i] / pooled_count[i]) : 0.0f; + pooled[i] = p; + ring_key[i] = valid[i] ? (static_cast(radius[i]) << 32) | float_key(p) : UINT64_MAX; +} + +// Where each of the keys' leading 32-bit groups (ring or sector) starts in a sorted key array. +__global__ void group_offsets(const uint64_t *__restrict__ keys, size_t n, int groups, int *__restrict__ offset) { + const int g = blockIdx.x * blockDim.x + threadIdx.x; + if (g > groups) return; + offset[g] = static_cast(lower_bound(keys, n, static_cast(g) << 32)); +} + +// The ring's baseline, three iterations of an order statistic over its sorted values (GetMask). +__global__ void baseline_kernel(int rings, const uint64_t *__restrict__ keys, const int *__restrict__ offset, + float *__restrict__ baseline) { + const int r = blockIdx.x * blockDim.x + threadIdx.x; + if (r >= rings) return; + const int lo = offset[r], n = offset[r + 1] - offset[r]; + int excluded = 0; + float b = 0.0f; + for (int iter = 0; iter < 3; iter++) { + const int avail = n - excluded; + b = (avail <= 0) ? 0.0f : key_float(static_cast(keys[lo + excluded + avail / 2])); + const float d = fmaxf(b, 1e-6f); + int excl = 0; + while (excl < n && key_float(static_cast(keys[lo + excl])) / d < SHADOW_RATIO) + excl++; + excluded = excl; + } + baseline[r] = b; +} + +__global__ void low_kernel(size_t n, uint32_t frames, const char *__restrict__ valid, const float *__restrict__ pooled, + const int32_t *__restrict__ pooled_count, const float *__restrict__ pol, + const int *__restrict__ radius, const float *__restrict__ baseline, + float *__restrict__ ratio, float *__restrict__ deficit, char *__restrict__ low) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + const float base = baseline[radius[i]]; + const float rt = valid[i] ? pooled[i] / fmaxf(base, 1e-6f) : 1.0f; + ratio[i] = rt; + float df = 0.0f; + char l = 0; + if (valid[i]) { + const double counted = static_cast(frames) * pooled_count[i] * pol[i]; + df = static_cast(poisson_deficit_sigma(pooled[i] * counted, base * counted)); + l = rt < SHADOW_RATIO && df > MIN_DEFICIT_SIGMA; + } + deficit[i] = df; + low[i] = l; +} + +// 8-connected components by union-find: every component ends up named by its smallest pixel index, +// whatever order the unions ran in. +__device__ int find_root(const int *parent, int x) { + while (parent[x] != x) + x = parent[x]; + return x; +} + +__global__ void cc_init(size_t n, const char *__restrict__ member, int *__restrict__ parent) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + parent[i] = member[i] ? static_cast(i) : -1; +} + +__device__ void cc_unite(int *parent, int a, int b) { + while (true) { + a = find_root(parent, a); + b = find_root(parent, b); + if (a == b) return; + if (a < b) { const int t = a; a = b; b = t; } + if (atomicCAS(&parent[a], a, b) == a) return; + } +} + +__global__ void cc_union(int W, int H, const char *__restrict__ member, int *parent) { + const size_t n = static_cast(W) * H; + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n || !member[i]) return; + const int x = static_cast(i % W), y = static_cast(i / W); + // The four neighbours before this pixel; the other four see it from their side. + if (x > 0 && member[i - 1]) cc_unite(parent, static_cast(i), static_cast(i - 1)); + if (y > 0) { + const size_t up = i - W; + if (member[up]) cc_unite(parent, static_cast(i), static_cast(up)); + if (x > 0 && member[up - 1]) cc_unite(parent, static_cast(i), static_cast(up - 1)); + if (x + 1 < W && member[up + 1]) cc_unite(parent, static_cast(i), static_cast(up + 1)); + } +} + +__global__ void cc_flatten(size_t n, int *parent) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n || parent[i] < 0) return; + parent[i] = find_root(parent, static_cast(i)); +} + +// Per component (by root): how many of its pixels are `a`, and how many are both `a` and `b`. +__global__ void cc_count(size_t n, const int *__restrict__ root, const char *__restrict__ a, const char *__restrict__ b, + int *__restrict__ count_a, int *__restrict__ count_ab) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n || root[i] < 0 || !a[i]) return; + atomicAdd(&count_a[root[i]], 1); + if (count_ab && b[i]) + atomicAdd(&count_ab[root[i]], 1); +} + +__global__ void region_kernel(size_t n, const int *__restrict__ root, const int *__restrict__ n_low, + const char *__restrict__ low, const char *__restrict__ valid, const int *__restrict__ radius, + int blocked_out_to, char *__restrict__ region) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + char r = root[i] >= 0 && n_low[root[i]] >= MIN_SHADOW_PIXELS ? low[i] : 0; + if (valid[i] && radius[i] <= blocked_out_to) + r = 1; + region[i] = r; +} + +__global__ void lit_kernel(size_t n, const uint32_t *__restrict__ valid_count, const int64_t *__restrict__ max_value, + char *__restrict__ lit) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + lit[i] = (valid_count[i] > 0) && (max_value[i] >= MIN_REFLECTION); +} + +__global__ void reflection_kernel(int W, int H, const char *__restrict__ lit, char *__restrict__ reflection) { + const size_t n = static_cast(W) * H; + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + const int x = static_cast(i % W), y = static_cast(i / W); + char r = 0; + if (lit[i]) { + int neighbours = 0; + for (int dy = -1; dy <= 1; dy++) + for (int dx = -1; dx <= 1; dx++) { + const int yy = y + dy, xx = x + dx; + if ((dx || dy) && yy >= 0 && yy < H && xx >= 0 && xx < W && lit[static_cast(yy) * W + xx]) + neighbours++; + } + r = neighbours >= 2; + } + reflection[i] = r; +} + +__global__ void penumbra_kernel(size_t n, const char *__restrict__ penumbra, const char *__restrict__ valid, + const float *__restrict__ ratio, const float *__restrict__ deficit, char *__restrict__ region) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + if (penumbra[i] && valid[i] && ratio[i] < PENUMBRA_RATIO && deficit[i] > MIN_DEFICIT_SIGMA) + region[i] = 1; +} + +__global__ void invert_kernel(size_t n, const char *__restrict__ in, char *__restrict__ out) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + out[i] = !in[i]; +} + +// Background components that touch the frame's edge are outside; the rest are holes. +__global__ void outside_kernel(int W, int H, const int *__restrict__ root, int *__restrict__ outside) { + const int k = blockIdx.x * blockDim.x + threadIdx.x; + const int perimeter = 2 * W + 2 * H; + if (k >= perimeter) return; + int x, y; + if (k < W) { x = k; y = 0; } + else if (k < 2 * W) { x = k - W; y = H - 1; } + else if (k < 2 * W + H) { x = 0; y = k - 2 * W; } + else { x = W - 1; y = k - 2 * W - H; } + const int r = root[static_cast(y) * W + x]; + if (r >= 0) outside[r] = 1; +} + +__global__ void final_kernel(size_t n, const int *__restrict__ background_root, const int *__restrict__ outside, + const char *__restrict__ reflection_grown, char *__restrict__ region, + uint32_t *__restrict__ mask) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + char r = region[i]; + if (background_root[i] >= 0 && !outside[background_root[i]]) + r = 1; // a hole + if (reflection_grown[i]) + r = 0; + region[i] = r; + mask[i] = r ? MASK_SHADOW : 0; +} + +__global__ void sector_key_kernel(ShadowMaskSetup s, const char *__restrict__ valid, const char *__restrict__ region, + const int *__restrict__ radius, const float *__restrict__ ratio, + uint64_t *__restrict__ key) { + const size_t n = static_cast(s.width) * s.height; + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + if (!valid[i] || region[i]) { + key[i] = UINT64_MAX; + return; + } + const float dx = static_cast(i % s.width) - s.beam_x, dy = static_cast(i / s.width) - s.beam_y; + const double phi = atan2f(dy, dx) + PI; + const int sector = min(HARMONIC_SECTORS - 1, static_cast(phi / (2.0 * PI) * HARMONIC_SECTORS)); + const uint64_t k = static_cast(radius[i] / HARMONIC_BAND_PX) * HARMONIC_SECTORS + sector; + key[i] = (k << 32) | float_key(ratio[i]); +} + +__global__ void sector_median_kernel(int sectors, const uint64_t *__restrict__ keys, const int *__restrict__ offset, + double *__restrict__ median) { + const int k = blockIdx.x * blockDim.x + threadIdx.x; + if (k >= sectors) return; + const int n = offset[k + 1] - offset[k]; + median[k] = n >= MIN_SECTOR_PIXELS ? key_float(static_cast(keys[offset[k] + n / 2])) : -1.0; +} + +__global__ void dim_kernel(ShadowMaskSetup s, uint32_t frames, const char *__restrict__ valid, + const char *__restrict__ region, const float *__restrict__ ratio, + const float *__restrict__ deficit, const int *__restrict__ radius, + const float *__restrict__ harm_c, const float *__restrict__ harm_s, + const int32_t *__restrict__ pooled_count, const float *__restrict__ pol, + const float *__restrict__ pooled, const float *__restrict__ baseline, + char *__restrict__ dim) { + const size_t n = static_cast(s.width) * s.height; + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + char d = 0; + if (valid[i] && !region[i] && ratio[i] < PENUMBRA_RATIO) { + // cos 2phi and sin 2phi from the offset to the beam. + const float dx = static_cast(i % s.width) - s.beam_x, dy = static_cast(i / s.width) - s.beam_y; + const float r2 = fmaxf(dx * dx + dy * dy, 1e-6f); + const int band = radius[i] / HARMONIC_BAND_PX; + const float model = fminf(1.0f, 1.0f + harm_c[band] * (dx * dx - dy * dy) / r2 + + harm_s[band] * 2.0f * dx * dy / r2); + if (model >= 1.0f) { + d = deficit[i] > MIN_DEFICIT_SIGMA; + } else { + const double counted = static_cast(frames) * pooled_count[i] * pol[i]; + d = ratio[i] < PENUMBRA_RATIO * model + && poisson_deficit_sigma(pooled[i] * counted, baseline[radius[i]] * model * counted) > MIN_DEFICIT_SIGMA; + } + } + dim[i] = d; +} + +// Join a region across the module gaps it crosses, one line per thread (bridge_gaps in +// ShadowFinder.cpp). Both directions read `region` and only ever set pixels of `out` to 1. +__global__ void bridge_rows(int W, int H, const char *__restrict__ region, const char *__restrict__ valid, + char *out) { + const int y = blockIdx.x * blockDim.x + threadIdx.x; + if (y >= H) return; + const size_t row = static_cast(y) * W; + int k = 0; + while (k < W) { + if (valid[row + k]) { k++; continue; } + const int start = k; + while (k < W && !valid[row + k]) k++; + if (start > 0 && k < W && region[row + start - 1] && region[row + k]) + for (int j = start; j < k; j++) out[row + j] = 1; + } +} + +__global__ void bridge_columns(int W, int H, const char *__restrict__ region, const char *__restrict__ valid, + char *out) { + const int x = blockIdx.x * blockDim.x + threadIdx.x; + if (x >= W) return; + const auto at = [&](int y) { return static_cast(y) * W + x; }; + int k = 0; + while (k < H) { + if (valid[at(k)]) { k++; continue; } + const int start = k; + while (k < H && !valid[at(k)]) k++; + if (start > 0 && k < H && region[at(start - 1)] && region[at(k)]) + for (int j = start; j < k; j++) out[at(j)] = 1; + } +} + +__global__ void transmitting_kernel(size_t n, const int *__restrict__ root, const char *__restrict__ dim, + const int *__restrict__ n_dim, const int *__restrict__ n_low, + uint32_t *__restrict__ mask) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + const int c = root[i]; + if (c >= 0 && dim[i] && n_dim[c] >= MIN_SHADOW_PIXELS && n_low[c] >= MIN_CORE_PIXELS) + mask[i] = MASK_TRANSMITTING; +} + +__global__ void mean_kernel(size_t n, const uint32_t *__restrict__ pixel_mask, const int64_t *__restrict__ sum_value, + const uint32_t *__restrict__ valid_count, float *__restrict__ mean) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + mean[i] = valid_count[i] > 0 && pixel_mask[i] == 0 + ? static_cast(static_cast(sum_value[i]) / valid_count[i]) : NAN; +} + +// The device side of one GetMask: planes, sort scratch and the stream to run on. +class MaskEngine { +public: + const ShadowMaskSetup s; + const int W, H; + const size_t n; + cudaStream_t stream; + + MaskEngine(const ShadowMaskSetup &setup, cudaStream_t st) + : s(setup), W(setup.width), H(setup.height), n(static_cast(setup.width) * setup.height), stream(st) {} + + void Check(const char *what) const { + check(cudaGetLastError(), what); + } + + void Dilate(const char *in, char *out, char *scratch, int r) const { + dilate_rows<<>>(in, scratch, W, H, r); + dilate_columns<<>>(scratch, out, W, H, r); + Check("dilate"); + } + + // Components of `member`, each pixel's root in `root` (-1 outside every component). + void Label(const char *member, int *root) const { + cc_init<<>>(n, member, root); + cc_union<<>>(W, H, member, root); + cc_flatten<<>>(n, root); + Check("components"); + } + + void SortKeys(uint64_t *keys, uint64_t *sorted) const { + size_t bytes = 0; + check(cub::DeviceRadixSort::SortKeys(nullptr, bytes, keys, sorted, n, 0, 64, stream), "sort size"); + CudaDevicePtr scratch(bytes); + check(cub::DeviceRadixSort::SortKeys(scratch.get(), bytes, keys, sorted, n, 0, 64, stream), "sort"); + // The scratch is freed on the allocation stream, which knows nothing of this one. + check(cudaStreamSynchronize(stream), "sort"); + } + + template + std::vector Download(const T *device, size_t count) const { + std::vector host(count); + check(cudaMemcpyAsync(host.data(), device, count * sizeof(T), cudaMemcpyDeviceToHost, stream), "download"); + check(cudaStreamSynchronize(stream), "download"); + return host; + } + + template + void Upload(T *device, const std::vector &host) const { + check(cudaMemcpyAsync(device, host.data(), host.size() * sizeof(T), cudaMemcpyHostToDevice, stream), "upload"); + } +}; + +} // namespace + +std::vector ShadowMaskOnDevice(const ShadowMaskSetup &setup, const std::vector &pixel_mask, + const int64_t *max_value, const int64_t *sum_value, + const uint32_t *valid_count, uint32_t frames, cudaStream_t stream) { + const MaskEngine e(setup, stream); + const size_t n = e.n; + const int W = e.W, H = e.H; + + CudaDevicePtr d_pixel_mask(n), mask(n); + e.Upload(d_pixel_mask.get(), pixel_mask); + + // Mean projection over the polarization factor, usable pixels and radius from the beam centre. + CudaDevicePtr pol(n), pooled(n), ratio(n), deficit(n); + CudaDevicePtr valid(n), low(n), region(n), a(n), b(n), c(n); + CudaDevicePtr radius(n), root(n), count_a(n), count_b(n), max_radius(1); + CudaDevicePtr num(n), dsum(n); + CudaDevicePtr den(n), dcount(n); + check(cudaMemsetAsync(max_radius.get(), 0, sizeof(int), stream), "memset"); + setup_kernel<<>>(setup, d_pixel_mask, sum_value, valid_count, pol, valid, radius, + num, den, max_radius); + e.Check("setup"); + const int max_r = e.Download(max_radius.get(), 1)[0]; + const int rings = max_r + 1; + + // The background pooled over a small box. + box_rows<<>>(num, dsum, W, H, POOL_PX / 2); + box_columns<<>>(dsum, num, W, H, POOL_PX / 2); + box_rows<<>>(den, dcount, W, H, POOL_PX / 2); + box_columns<<>>(dcount, den, W, H, POOL_PX / 2); + e.Check("pooling"); + const double *pooled_sum = num; + const int32_t *pooled_count = den; + + // The rings, each sorted once; the baseline is an order statistic of them. + CudaDevicePtr keys(n), sorted(n); + pooled_kernel<<>>(n, pooled_sum, pooled_count, valid, radius, pooled, keys); + e.Check("pooled"); + e.SortKeys(keys, sorted); + CudaDevicePtr ring_offset(rings + 1); + group_offsets<<>>(sorted, n, rings, ring_offset); + CudaDevicePtr baseline(rings); + baseline_kernel<<>>(rings, sorted, ring_offset, baseline); + e.Check("baseline"); + const auto host_baseline = e.Download(baseline.get(), rings); + const auto offsets = e.Download(ring_offset.get(), rings + 1); + std::vector ring_pixels(rings); + for (int r = 0; r < rings; r++) + ring_pixels[r] = offsets[r + 1] - offsets[r]; + const int blocked_out_to = BlockedOutTo(host_baseline, ring_pixels); + + // Low pixels, and the regions of them large enough to be hardware. + low_kernel<<>>(n, frames, valid, pooled, pooled_count, pol, radius, baseline, + ratio, deficit, low); + e.Check("low"); + e.Dilate(low, a, c, BRIDGE_PX); + e.Label(a, root); + check(cudaMemsetAsync(count_a.get(), 0, n * sizeof(int), stream), "memset"); + cc_count<<>>(n, root, low, low, count_a, nullptr); + region_kernel<<>>(n, root, count_a, low, valid, radius, blocked_out_to, region); + e.Check("region"); + + // Recorded reflections; `b` holds them until they are given back at the end. + lit_kernel<<>>(n, valid_count, max_value, a); + reflection_kernel<<>>(W, H, a, b); + e.Check("reflections"); + + // Penumbra, round and fill. + e.Dilate(region, a, c, PENUMBRA_MAX_PX); + penumbra_kernel<<>>(n, a, valid, ratio, deficit, region); + e.Dilate(region, a, c, 2); + invert_kernel<<>>(n, a, region); + e.Dilate(region, a, c, 2); + invert_kernel<<>>(n, a, region); // region = erode(dilate(region)) + invert_kernel<<>>(n, region, a); // the background + e.Label(a, root); + check(cudaMemsetAsync(count_a.get(), 0, n * sizeof(int), stream), "memset"); + outside_kernel<<>>(W, H, root, count_a); + e.Dilate(b, c, a, 1); // the reflections, grown by one + final_kernel<<>>(n, root, count_a, c, region, mask); + e.Check("fill"); + + // The arm search: sector medians of the ratio, the harmonic of each band, the dim pixels. + const int n_bands = max_r / HARMONIC_BAND_PX + 1; + const int n_sectors = n_bands * HARMONIC_SECTORS; + sector_key_kernel<<>>(setup, valid, region, radius, ratio, keys); + e.Check("sectors"); + e.SortKeys(keys, sorted); + CudaDevicePtr sector_offset(n_sectors + 1); + group_offsets<<>>(sorted, n, n_sectors, sector_offset); + CudaDevicePtr sector_median(n_sectors); + sector_median_kernel<<>>(n_sectors, sorted, sector_offset, sector_median); + e.Check("sector medians"); + std::vector harm_c, harm_s; + HarmonicFit(e.Download(sector_median.get(), n_sectors), n_bands, harm_c, harm_s); + CudaDevicePtr d_harm_c(n_bands), d_harm_s(n_bands); + e.Upload(d_harm_c.get(), harm_c); + e.Upload(d_harm_s.get(), harm_s); + dim_kernel<<>>(setup, frames, valid, region, ratio, deficit, radius, d_harm_c, d_harm_s, + pooled_count, pol, pooled, baseline, a); + e.Check("dim"); + e.Dilate(a, b, c, BRIDGE_PX); + check(cudaMemcpyAsync(c.get(), b.get(), n, cudaMemcpyDeviceToDevice, stream), "copy"); + bridge_rows<<>>(W, H, b, valid, c); + bridge_columns<<>>(W, H, b, valid, c); + e.Label(c, root); + check(cudaMemsetAsync(count_a.get(), 0, n * sizeof(int), stream), "memset"); + check(cudaMemsetAsync(count_b.get(), 0, n * sizeof(int), stream), "memset"); + cc_count<<>>(n, root, a, low, count_a, count_b); + transmitting_kernel<<>>(n, root, a, count_a, count_b, mask); + e.Check("transmitting"); + + return e.Download(mask.get(), n); +} + +std::vector MeanProjectionOnDevice(const std::vector &pixel_mask, const int64_t *sum_value, + const uint32_t *valid_count, size_t npixels, cudaStream_t stream) { + CudaDevicePtr d_pixel_mask(npixels); + CudaDevicePtr mean(npixels); + check(cudaMemcpyAsync(d_pixel_mask.get(), pixel_mask.data(), npixels * sizeof(uint32_t), cudaMemcpyHostToDevice, + stream), "upload"); + mean_kernel<<>>(npixels, d_pixel_mask, sum_value, valid_count, mean); + check(cudaGetLastError(), "mean"); + std::vector host(npixels); + check(cudaMemcpyAsync(host.data(), mean.get(), npixels * sizeof(float), cudaMemcpyDeviceToHost, stream), "download"); + check(cudaStreamSynchronize(stream), "mean"); + return host; +} diff --git a/image_analysis/beam_stop/ShadowMaskGPU.h b/image_analysis/beam_stop/ShadowMaskGPU.h new file mode 100644 index 000000000..baa25a00f --- /dev/null +++ b/image_analysis/beam_stop/ShadowMaskGPU.h @@ -0,0 +1,39 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// Included only under JFJOCH_USE_CUDA. + +#include +#include + +#include + +// What the mask is drawn about: the detector, the centre the rings are drawn about, and the geometry +// the polarization factor is read off (ShadowFinder keeps it as a DiffractionGeometry; the device +// takes it as numbers). +struct ShadowMaskSetup { + int width = 0, height = 0; + float beam_x = 0.0f, beam_y = 0.0f; + float det_matrix[9] = {}; // row major + float pixel_size_mm = 0.0f; + float distance_mm = 0.0f; + bool has_polarization = false; + float polarization = 0.0f; +}; + +// ShadowFinder::GetMask on the device, from the projection ShadowAccumulatorGPU holds there. Step for +// step the host's algorithm - the same pooling, ring medians, components, morphology and arm search - +// and the same answer wherever the arithmetic is exact: every integer, comparison, sort and component +// is. What is not is the floating point the two compilers evaluate differently - the polarization +// factor's trigonometry, the Poisson test's logarithm and the azimuth of the arm search - so a pixel +// within a rounding of one of those thresholds can come out the other way. +std::vector ShadowMaskOnDevice(const ShadowMaskSetup &setup, const std::vector &pixel_mask, + const int64_t *max_value, const int64_t *sum_value, + const uint32_t *valid_count, uint32_t frames, cudaStream_t stream); + +// The mean projection the host's GetMeanProjection makes, computed where the sums are: the same +// division, so the same bits, and a quarter of the bytes to bring back. +std::vector MeanProjectionOnDevice(const std::vector &pixel_mask, const int64_t *sum_value, + const uint32_t *valid_count, size_t npixels, cudaStream_t stream); diff --git a/image_analysis/geom_refinement/BackgroundBand.h b/image_analysis/geom_refinement/BackgroundBand.h new file mode 100644 index 000000000..394ab3e84 --- /dev/null +++ b/image_analysis/geom_refinement/BackgroundBand.h @@ -0,0 +1,107 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// Where one pixel falls in the background beam-centre fit (FindBeamCenterFromBackground): which +// radial-bin x sector cell of the fitted band it lands in about a trial centre, and the derivative +// of its 2theta with respect to that centre. Written once and compiled both for the host fit and for +// its device twin (BeamCenterBackgroundGPU), so the two read the same formula. +// +// The same formula is the same number only when both sides evaluate it the same way. The library +// atan2f differs between glibc and CUDA in the last bit, and the compilers fuse multiply-adds each +// in their own way; either moves a pixel within a rounding of a cell edge into the neighbouring cell +// (measured on 16 Mpx sweeps: ~30 of 6.5 million pixels, enough to move the fitted centre by up to +// 0.05 px). So the angles come from BackgroundAtan2 below and both translation units are compiled +// without contraction (geom_refinement/CMakeLists.txt), and the host and the device then agree to +// the bit. + +#include +#include + +#include "../../common/JFJochMath.h" + +#ifdef __CUDACC__ +#define BACKGROUND_BAND_HD __host__ __device__ inline +#else +#define BACKGROUND_BAND_HD inline +#endif + +struct BackgroundBand { + static constexpr int SECTORS = 36; + static constexpr int RADIAL_BINS = 120; + static constexpr int CELLS = RADIAL_BINS * SECTORS; + + float rot[9]; // detector matrix, row major + float pixel_size; // mm + float distance; // mm + float tt_lo, tt_hi; // the band in 2theta + float d_tt; // width of one radial bin + double tan_lo, tan_hi; // the band in tan(2theta), widened for the quick rejection +}; + +// atan2(y, x) from IEEE operations alone - add, multiply, divide, square root, all correctly rounded on +// the host and on the device - so that, compiled without contraction (see CMakeLists.txt), the host +// and the device return the same bits. The library atan2f does not: glibc's and CUDA's differ in the +// last place. Two half-angle reductions take the argument below tan(pi/16), where eleven terms of the +// series leave an error under 1e-17. +BACKGROUND_BAND_HD double BackgroundAtan2(double y, double x) { + const double ax = x < 0 ? -x : x, ay = y < 0 ? -y : y; + if (ax == 0.0 && ay == 0.0) + return 0.0; + const bool swap = ay > ax; + double t = swap ? ax / ay : ay / ax; + t = t / (1.0 + sqrt(1.0 + t * t)); + t = t / (1.0 + sqrt(1.0 + t * t)); + const double t2 = t * t; + double series = 1.0 / 23.0; + for (int k = 21; k >= 1; k -= 2) + series = 1.0 / k - t2 * series; + double a = 4.0 * t * series; + if (swap) a = PI / 2 - a; + if (x < 0) a = PI - a; + return y < 0 ? -a : a; +} + +// The cell of pixel (x, y) about (beam_x, beam_y), or -1 when it is outside the band; for a pixel in +// it, also the two components of the derivative of its 2theta with respect to the centre. +BACKGROUND_BAND_HD int BackgroundBandCell(const BackgroundBand &b, int x, int y, float beam_x, float beam_y, + float &jac_x, float &jac_y) { + const float *rot = b.rot; + const float u = (x - beam_x) * b.pixel_size; + const float v = (y - beam_y) * b.pixel_size; + const float lx = rot[0] * u + rot[1] * v + rot[2] * b.distance; + const float ly = rot[3] * u + rot[4] * v + rot[5] * b.distance; + const float lz = rot[6] * u + rot[7] * v + rot[8] * b.distance; + const float rho_sq = lx * lx + ly * ly; + // Most pixels lie outside the band. Those clearly outside it in tan(2theta) = rho / lz - by a + // margin far above float rounding - skip the square root and both atan2 below; every pixel the + // exact test would keep still reaches it. + if (lz > 0.0f) { + const double lz_sq = static_cast(lz) * lz; + if (rho_sq < b.tan_lo * b.tan_lo * lz_sq || rho_sq > b.tan_hi * b.tan_hi * lz_sq) + return -1; + } + const float rho = sqrtf(rho_sq); + const float two_theta = static_cast(BackgroundAtan2(rho, lz)); + if (two_theta < b.tt_lo || two_theta >= b.tt_hi || rho == 0.0f) + return -1; + + const float phi = static_cast(BackgroundAtan2(ly, lx)); + // Both bins are clamped: a pixel one float ulp below the top of the band divides to exactly + // RADIAL_BINS, which is one cell past the end of every accumulator. + int r_bin = static_cast((two_theta - b.tt_lo) / b.d_tt); + r_bin = r_bin < 0 ? 0 : (r_bin > BackgroundBand::RADIAL_BINS - 1 ? BackgroundBand::RADIAL_BINS - 1 : r_bin); + int s_bin = static_cast((phi + PI) / (2 * PI) * BackgroundBand::SECTORS); + s_bin = s_bin < 0 ? 0 : (s_bin > BackgroundBand::SECTORS - 1 ? BackgroundBand::SECTORS - 1 : s_bin); + + // d(2theta)/d(beam), through the lab coordinate: the detector coordinate depends on the centre + // only as (x - beam_x), so moving the centre is moving the pixel. + const float denominator = rho * rho + lz * lz; + const float g_x = lz * lx / (rho * denominator); + const float g_y = lz * ly / (rho * denominator); + const float g_z = -rho / denominator; + jac_x = -b.pixel_size * (g_x * rot[0] + g_y * rot[3] + g_z * rot[6]); + jac_y = -b.pixel_size * (g_x * rot[1] + g_y * rot[4] + g_z * rot[7]); + return r_bin * BackgroundBand::SECTORS + s_bin; +} diff --git a/image_analysis/geom_refinement/BeamCenterBackgroundGPU.cu b/image_analysis/geom_refinement/BeamCenterBackgroundGPU.cu new file mode 100644 index 000000000..550593f1d --- /dev/null +++ b/image_analysis/geom_refinement/BeamCenterBackgroundGPU.cu @@ -0,0 +1,210 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "BeamCenterBackgroundGPU.h" + +#include + +#include "../indexing/CUDAMemHelpers.h" + +namespace { + +constexpr int THREADS = 256; + +void check(cudaError_t err, const char *what) { + if (err != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, + std::string("Beam centre from background: ") + what + ": " + cudaGetErrorString(err)); +} + +// Every pixel's cell (BackgroundBand::CELLS where it is outside the band or unusable), its index, +// and its derivatives. +__global__ void bin_kernel(BackgroundBand band, int width, size_t npixels, float beam_x, float beam_y, + const char *__restrict__ usable, int32_t *__restrict__ key, + int32_t *__restrict__ index, float *__restrict__ jac_x, + float *__restrict__ jac_y) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= npixels) + return; + int cell = BackgroundBand::CELLS; + float jx = 0.0f, jy = 0.0f; + if (usable[i]) { + const int c = BackgroundBandCell(band, static_cast(i % width), static_cast(i / width), + beam_x, beam_y, jx, jy); + if (c >= 0) + cell = c; + } + key[i] = cell; + index[i] = static_cast(i); + jac_x[i] = jx; + jac_y[i] = jy; +} + +// Where each cell's pixels start in the sorted keys; offset[CELLS] is the number of binned pixels. +__global__ void offset_kernel(const int32_t *__restrict__ sorted_key, size_t n, int32_t *__restrict__ offset) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c > BackgroundBand::CELLS) + return; + size_t lo = 0, hi = n; + while (lo < hi) { + const size_t mid = (lo + hi) / 2; + if (sorted_key[mid] < c) lo = mid + 1; + else hi = mid; + } + offset[c] = static_cast(lo); +} + +// One thread per cell, walking its pixels in pixel order: a partial sum per host row block, added to +// the cell's total when the block changes - the host's arithmetic, step for step. A cell whose pixels +// are clipped (clip_limit != nullptr) skips the ones above its limit, as the host's clip rounds do. +__global__ void sum_kernel(int width, const int32_t *__restrict__ offset, const int32_t *__restrict__ index, + const int32_t *__restrict__ row_block, const float *__restrict__ mean, + const float *__restrict__ jac_x, const float *__restrict__ jac_y, + const float *__restrict__ clip_limit, + double *__restrict__ sum, double *__restrict__ sum_sq, + double *__restrict__ sum_jx, double *__restrict__ sum_jy, + int32_t *__restrict__ count) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= BackgroundBand::CELLS) + return; + const bool clipped = clip_limit != nullptr; + const float limit = clipped ? clip_limit[c] : 0.0f; + + double s = 0, ss = 0, jx = 0, jy = 0; + double bs = 0, bss = 0, bjx = 0, bjy = 0; + int32_t n = 0; + int block = -1; + for (int32_t p = offset[c]; p < offset[c + 1]; p++) { + const int32_t i = index[p]; + const float value = mean[i]; + if (clipped && (limit < 0.0f || value > limit)) + continue; + const int b = row_block[i / width]; + if (b != block) { + s += bs; ss += bss; jx += bjx; jy += bjy; + bs = bss = bjx = bjy = 0; + block = b; + } + n++; + bs += value; + bss += static_cast(value) * value; + if (!clipped) { + bjx += jac_x[i]; + bjy += jac_y[i]; + } + } + s += bs; ss += bss; jx += bjx; jy += bjy; + sum[c] = s; + sum_sq[c] = ss; + count[c] = n; + if (!clipped) { + sum_jx[c] = jx; + sum_jy[c] = jy; + } +} + +} // namespace + +struct BeamCenterBackgroundGPU::Impl { + int width, height; + size_t npixels; + CudaStream stream; + CudaDevicePtr usable; + CudaDevicePtr mean; + CudaDevicePtr row_block; + CudaDevicePtr key, sorted_key, index, sorted_index; + CudaDevicePtr jac_x, jac_y; + CudaDevicePtr offset; + CudaDevicePtr clip_limit; + CudaDevicePtr sum, sum_sq, sum_jx, sum_jy; + CudaDevicePtr count; + CudaDevicePtr sort_scratch; + size_t sort_scratch_bytes = 0; + + Impl(int w, int h) + : width(w), height(h), npixels(static_cast(w) * h), + usable(npixels), mean(npixels), row_block(h), + key(npixels), sorted_key(npixels), index(npixels), sorted_index(npixels), + jac_x(npixels), jac_y(npixels), offset(BackgroundBand::CELLS + 1), + clip_limit(BackgroundBand::CELLS), + sum(BackgroundBand::CELLS), sum_sq(BackgroundBand::CELLS), + sum_jx(BackgroundBand::CELLS), sum_jy(BackgroundBand::CELLS), + count(BackgroundBand::CELLS) { + // Keys run to CELLS inclusive, the bin of everything outside the band. + check(cub::DeviceRadixSort::SortPairs(nullptr, sort_scratch_bytes, key.get(), sorted_key.get(), + index.get(), sorted_index.get(), npixels, 0, end_bit(), stream), + "sort size"); + sort_scratch = CudaDevicePtr(sort_scratch_bytes); + } + + static int end_bit() { + int bits = 0; + while ((1 << bits) <= BackgroundBand::CELLS) bits++; + return bits; + } + + void Download(std::vector &s, std::vector &ss, std::vector *jx, + std::vector *jy, std::vector &n) { + const size_t cells = BackgroundBand::CELLS; + check(cudaMemcpyAsync(s.data(), sum.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy"); + check(cudaMemcpyAsync(ss.data(), sum_sq.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy"); + if (jx) check(cudaMemcpyAsync(jx->data(), sum_jx.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy"); + if (jy) check(cudaMemcpyAsync(jy->data(), sum_jy.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy"); + check(cudaMemcpyAsync(n.data(), count.get(), cells * sizeof(int32_t), cudaMemcpyDeviceToHost, stream), "copy"); + check(cudaStreamSynchronize(stream), "sums"); + } +}; + +BeamCenterBackgroundGPU::BeamCenterBackgroundGPU(int width, int height, const std::vector &block_row, + const char *usable, const float *mean) + : impl(std::make_unique(width, height)) { + std::vector row_block(height); + for (size_t b = 0; b + 1 < block_row.size(); b++) + for (int y = block_row[b]; y < block_row[b + 1]; y++) + row_block[y] = static_cast(b); + check(cudaMemcpyAsync(impl->row_block.get(), row_block.data(), height * sizeof(int32_t), + cudaMemcpyHostToDevice, impl->stream), "upload"); + check(cudaMemcpyAsync(impl->usable.get(), usable, impl->npixels, cudaMemcpyHostToDevice, impl->stream), "upload"); + check(cudaMemcpyAsync(impl->mean.get(), mean, impl->npixels * sizeof(float), cudaMemcpyHostToDevice, + impl->stream), "upload"); + check(cudaStreamSynchronize(impl->stream), "upload"); +} + +BeamCenterBackgroundGPU::~BeamCenterBackgroundGPU() = default; + +void BeamCenterBackgroundGPU::Bin(const BackgroundBand &band, float beam_x, float beam_y, + std::vector &sum, std::vector &sum_sq, + std::vector &sum_jx, std::vector &sum_jy, + std::vector &count) { + Impl &d = *impl; + const auto blocks = static_cast((d.npixels + THREADS - 1) / THREADS); + bin_kernel<<>>(band, d.width, d.npixels, beam_x, beam_y, d.usable, + d.key, d.index, d.jac_x, d.jac_y); + check(cudaGetLastError(), "bin"); + // A radix sort is stable, so each cell's pixels come out in pixel order. + check(cub::DeviceRadixSort::SortPairs(d.sort_scratch.get(), d.sort_scratch_bytes, d.key.get(), + d.sorted_key.get(), d.index.get(), d.sorted_index.get(), + d.npixels, 0, Impl::end_bit(), d.stream), "sort"); + constexpr int cell_blocks = (BackgroundBand::CELLS + 1 + THREADS - 1) / THREADS; + offset_kernel<<>>(d.sorted_key, d.npixels, d.offset); + check(cudaGetLastError(), "offsets"); + sum_kernel<<>>(d.width, d.offset, d.sorted_index, d.row_block, d.mean, + d.jac_x, d.jac_y, nullptr, d.sum, d.sum_sq, + d.sum_jx, d.sum_jy, d.count); + check(cudaGetLastError(), "sums"); + d.Download(sum, sum_sq, &sum_jx, &sum_jy, count); +} + +void BeamCenterBackgroundGPU::Clip(const std::vector &clip_limit, + std::vector &sum, std::vector &sum_sq, + std::vector &count) { + Impl &d = *impl; + check(cudaMemcpyAsync(d.clip_limit.get(), clip_limit.data(), BackgroundBand::CELLS * sizeof(float), + cudaMemcpyHostToDevice, d.stream), "upload"); + constexpr int cell_blocks = (BackgroundBand::CELLS + THREADS - 1) / THREADS; + sum_kernel<<>>(d.width, d.offset, d.sorted_index, d.row_block, d.mean, + d.jac_x, d.jac_y, d.clip_limit, d.sum, d.sum_sq, + d.sum_jx, d.sum_jy, d.count); + check(cudaGetLastError(), "clip"); + d.Download(sum, sum_sq, nullptr, nullptr, count); +} diff --git a/image_analysis/geom_refinement/BeamCenterBackgroundGPU.h b/image_analysis/geom_refinement/BeamCenterBackgroundGPU.h new file mode 100644 index 000000000..72d47b266 --- /dev/null +++ b/image_analysis/geom_refinement/BeamCenterBackgroundGPU.h @@ -0,0 +1,40 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// Included only under JFJOCH_USE_CUDA. Free of CUDA headers, so the host fit can hold one. + +#include +#include +#include + +#include "BackgroundBand.h" + +// The two passes over the pixels of the background beam-centre fit (FindBeamCenterFromBackground), +// on the device: binning every usable pixel into its cell about a trial centre, and summing the +// binned pixels again under a clip. They are all of the fit's cost; the fit itself stays on the host. +// +// Each cell is summed in the order the host sums it - pixel order within each of the host's row +// blocks, the blocks then added in block order - so a pixel that lands in the same cell on both +// sides adds the same rounding on both. Whatever differs comes from BackgroundBandCell (see there). +class BeamCenterBackgroundGPU { + struct Impl; + std::unique_ptr impl; +public: + // block_row: the first row of each of the host's row blocks, and one past the last row at the end. + BeamCenterBackgroundGPU(int width, int height, const std::vector &block_row, + const char *usable, const float *mean); + ~BeamCenterBackgroundGPU(); + + // Bin the band about (beam_x, beam_y) and sum each cell: the values, their squares, the two + // derivatives, and the count. + void Bin(const BackgroundBand &band, float beam_x, float beam_y, + std::vector &sum, std::vector &sum_sq, + std::vector &sum_jx, std::vector &sum_jy, std::vector &count); + + // Sum the pixels the last Bin put in each cell again, leaving out those above the cell's + // clip_limit and every pixel of a cell whose limit is negative. + void Clip(const std::vector &clip_limit, + std::vector &sum, std::vector &sum_sq, std::vector &count); +}; diff --git a/image_analysis/geom_refinement/BeamCenterFromBackground.cpp b/image_analysis/geom_refinement/BeamCenterFromBackground.cpp index 19f468a65..0b99a030c 100644 --- a/image_analysis/geom_refinement/BeamCenterFromBackground.cpp +++ b/image_analysis/geom_refinement/BeamCenterFromBackground.cpp @@ -2,6 +2,7 @@ // SPDX-License-Identifier: GPL-3.0-only #include "BeamCenterFromBackground.h" +#include "BackgroundBand.h" #include #include @@ -10,6 +11,10 @@ #include "../../common/CompressedImage.h" #include "../../common/JFJochMath.h" #include "../../common/ParallelFor.h" +#ifdef JFJOCH_USE_CUDA +#include "../../common/CUDAWrapper.h" +#include "BeamCenterBackgroundGPU.h" +#endif namespace { @@ -18,8 +23,8 @@ namespace { constexpr float BAND_LOW_RES_A = 12.0f; constexpr float BAND_HIGH_RES_A = 2.2f; -constexpr int SECTORS = 36; -constexpr int RADIAL_BINS = 120; +constexpr int SECTORS = BackgroundBand::SECTORS; +constexpr int RADIAL_BINS = BackgroundBand::RADIAL_BINS; // A cell with fewer pixels than this has no usable mean. constexpr int MIN_PIXELS_PER_CELL = 20; @@ -84,7 +89,7 @@ float median_of(std::vector &v) { std::optional FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const PixelMask &mask, const std::vector &mean, size_t nthreads, - std::optional> start) { + std::optional> start, bool allow_device) { if (nthreads == 0) nthreads = std::max(1u, std::thread::hardware_concurrency()); const auto W = static_cast(experiment.GetXPixelsNumConv()); @@ -165,10 +170,31 @@ FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const Pixe const double tan_lo = std::tan(static_cast(tt_lo)) * (1.0 - 1e-3); const double tan_hi = tt_hi < PI / 2 ? std::tan(static_cast(tt_hi)) * (1.0 + 1e-3) : INFINITY; + BackgroundBand band{}; + for (int k = 0; k < 9; k++) + band.rot[k] = rot[k]; + band.pixel_size = pixel_size; + band.distance = distance; + band.tt_lo = tt_lo; + band.tt_hi = tt_hi; + band.d_tt = d_tt; + band.tan_lo = tan_lo; + band.tan_hi = tan_hi; + +#ifndef JFJOCH_USE_CUDA + (void) allow_device; +#else + // With a GPU the two passes over the pixels run there, and only the cells come back. + std::unique_ptr gpu; + if (allow_device && get_gpu_count() > 0) + gpu = std::make_unique(W, H, block_row, usable.data(), mean.data()); +#endif + float step_x = 0.0f, step_y = 0.0f, sigma_x = 0.0f, sigma_y = 0.0f; float previous_x = 0.0f, previous_y = 0.0f; int reversals = 0; - for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++) { + // The binning pass about the current centre, then the clip rounds over the pixels it binned. + const auto bin_cpu = [&] { ParallelFor(BLOCKS, nthreads, [&](int b) { double *b_sum = block_sum.data() + static_cast(b) * n_cells; double *b_sum_sq = block_sum_sq.data() + static_cast(b) * n_cells; @@ -188,48 +214,54 @@ FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const Pixe const size_t i = static_cast(y) * W + x; if (!usable[i]) continue; - const float u = (x - beam_x) * pixel_size; - const float v = (y - beam_y) * pixel_size; - const float lx = rot[0] * u + rot[1] * v + rot[2] * distance; - const float ly = rot[3] * u + rot[4] * v + rot[5] * distance; - const float lz = rot[6] * u + rot[7] * v + rot[8] * distance; - const float rho_sq = lx * lx + ly * ly; - if (lz > 0.0f) { - const double lz_sq = static_cast(lz) * lz; - if (rho_sq < tan_lo * tan_lo * lz_sq || rho_sq > tan_hi * tan_hi * lz_sq) - continue; - } - const float rho = std::sqrt(rho_sq); - const float two_theta = std::atan2(rho, lz); - if (two_theta < tt_lo || two_theta >= tt_hi || rho == 0.0f) + float jac_x, jac_y; + const int cell = BackgroundBandCell(band, x, y, beam_x, beam_y, jac_x, jac_y); + if (cell < 0) continue; - - const float phi = std::atan2(ly, lx); - // Both bins are clamped: a pixel one float ulp below the top of the band divides - // to exactly RADIAL_BINS, which is one cell past the end of every accumulator. - const int r_bin = std::clamp(static_cast((two_theta - tt_lo) / d_tt), 0, RADIAL_BINS - 1); - const int s_bin = std::clamp(static_cast((phi + PI) / (2 * PI) * SECTORS), 0, SECTORS - 1); - const int cell = r_bin * SECTORS + s_bin; - - // d(2theta)/d(beam), through the lab coordinate: the detector coordinate depends - // on the centre only as (x - beam_x), so moving the centre is moving the pixel. - const float denominator = rho * rho + lz * lz; - const float g_x = lz * lx / (rho * denominator); - const float g_y = lz * ly / (rho * denominator); - const float g_z = -rho / denominator; cells[n_band] = cell; values[n_band] = mean[i]; n_band++; b_count[cell]++; b_sum[cell] += mean[i]; b_sum_sq[cell] += static_cast(mean[i]) * mean[i]; - b_jx[cell] += -pixel_size * (g_x * rot[0] + g_y * rot[3] + g_z * rot[6]); - b_jy[cell] += -pixel_size * (g_x * rot[1] + g_y * rot[4] + g_z * rot[7]); + b_jx[cell] += jac_x; + b_jy[cell] += jac_y; } } band_pixels[b] = n_band; }); fold(true); + }; + const auto clip_cpu = [&] { + ParallelFor(BLOCKS, nthreads, [&](int b) { + double *b_sum = block_sum.data() + static_cast(b) * n_cells; + double *b_sum_sq = block_sum_sq.data() + static_cast(b) * n_cells; + int32_t *b_count = block_count.data() + static_cast(b) * n_cells; + std::fill(b_sum, b_sum + n_cells, 0.0); + std::fill(b_sum_sq, b_sum_sq + n_cells, 0.0); + std::fill(b_count, b_count + n_cells, 0); + const int32_t *cells = band_cell.data() + static_cast(block_row[b]) * W; + const float *values = band_value.data() + static_cast(block_row[b]) * W; + for (size_t j = 0; j < band_pixels[b]; j++) { + const int32_t c = cells[j]; + const float value = values[j]; + if (clip_limit[c] < 0.0f || value > clip_limit[c]) + continue; + b_count[c]++; + b_sum[c] += value; + b_sum_sq[c] += static_cast(value) * value; + } + }); + fold(false); + }; + + for (int iteration = 0; iteration < MAX_ITERATIONS; iteration++) { +#ifdef JFJOCH_USE_CUDA + if (gpu) + gpu->Bin(band, beam_x, beam_y, sum, sum_sq, sum_jx, sum_jy, count); + else +#endif + bin_cpu(); count_all = count; // the Jacobian sums belong to the unclipped pixel set @@ -240,26 +272,12 @@ FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const Pixe const double variance = std::max(sum_sq[c] / count[c] - m * m, 0.0); clip_limit[c] = static_cast(m + CLIP_SIGMA * std::sqrt(variance)); } - ParallelFor(BLOCKS, nthreads, [&](int b) { - double *b_sum = block_sum.data() + static_cast(b) * n_cells; - double *b_sum_sq = block_sum_sq.data() + static_cast(b) * n_cells; - int32_t *b_count = block_count.data() + static_cast(b) * n_cells; - std::fill(b_sum, b_sum + n_cells, 0.0); - std::fill(b_sum_sq, b_sum_sq + n_cells, 0.0); - std::fill(b_count, b_count + n_cells, 0); - const int32_t *cells = band_cell.data() + static_cast(block_row[b]) * W; - const float *values = band_value.data() + static_cast(block_row[b]) * W; - for (size_t j = 0; j < band_pixels[b]; j++) { - const int32_t c = cells[j]; - const float value = values[j]; - if (clip_limit[c] < 0.0f || value > clip_limit[c]) - continue; - b_count[c]++; - b_sum[c] += value; - b_sum_sq[c] += static_cast(value) * value; - } - }); - fold(false); +#ifdef JFJOCH_USE_CUDA + if (gpu) + gpu->Clip(clip_limit, sum, sum_sq, count); + else +#endif + clip_cpu(); } // Radial profile: the median over the sectors that have a mean, on rings that are diff --git a/image_analysis/geom_refinement/BeamCenterFromBackground.h b/image_analysis/geom_refinement/BeamCenterFromBackground.h index 350cb6ddd..7f1df9efb 100644 --- a/image_analysis/geom_refinement/BeamCenterFromBackground.h +++ b/image_analysis/geom_refinement/BeamCenterFromBackground.h @@ -40,10 +40,14 @@ struct BeamCenterEstimate { // `start` is where the walk begins; the centre in the file when it is not given. The walk advances // by a bounded distance per iteration, so where it starts decides how much of its budget is spent // travelling and - on a surface with more than one basin - which fixed point it can reach at all. +// +// With a GPU the passes over the pixels run on it (BeamCenterBackgroundGPU); allow_device = false +// keeps them on the host, which is what the parity test compares against. std::optional FindBeamCenterFromBackground(const DiffractionExperiment &experiment, const PixelMask &mask, const std::vector &mean, size_t nthreads = 0, - std::optional> start = {}); + std::optional> start = {}, + bool allow_device = true); // The precision of a centre that is the FFT capture alone, with no walk behind it. The capture is // a half-pixel grid position read off a surface, measured over 75 rotation datasets at a median diff --git a/image_analysis/geom_refinement/CMakeLists.txt b/image_analysis/geom_refinement/CMakeLists.txt index b004e8d4e..f431ddbb8 100644 --- a/image_analysis/geom_refinement/CMakeLists.txt +++ b/image_analysis/geom_refinement/CMakeLists.txt @@ -24,6 +24,7 @@ ADD_LIBRARY(JFJochGeomRefinement STATIC XtalOptimizer.cpp XtalOptimizer.h XtalResidual.h + BackgroundBand.h PostRefine.cpp PostRefine.h GeometryRefiner.cpp @@ -37,7 +38,8 @@ ADD_LIBRARY(JFJochGeomRefinement STATIC TARGET_LINK_LIBRARIES(JFJochGeomRefinement Ceres::ceres Eigen3::Eigen JFJochCommon fftw3f) IF (JFJOCH_CUDA_AVAILABLE) - TARGET_SOURCES(JFJochGeomRefinement PRIVATE BeamCenterFFTGPU.cu BeamCenterFFTGPU.h) + TARGET_SOURCES(JFJochGeomRefinement PRIVATE BeamCenterFFTGPU.cu BeamCenterFFTGPU.h + BeamCenterBackgroundGPU.cu BeamCenterBackgroundGPU.h) # Same static/dynamic cuFFT choice as the FFT indexer, and for the same reasons - see the long # note in image_analysis/indexing/CMakeLists.txt. IF (JFJOCH_PORTABLE_ONLY AND TARGET CUDA::cufft_static) @@ -46,3 +48,13 @@ IF (JFJOCH_CUDA_AVAILABLE) TARGET_LINK_LIBRARIES(JFJochGeomRefinement CUDA::cufft) ENDIF() ENDIF() + +# The background beam-centre fit bins every pixel with BackgroundBand.h on the host and on the device +# and must get the same bits on both (see there). That needs no multiply-add contracted on either side: +# GCC and Clang contract by default and nvcc does too, while MSVC does not without /fp:contract. +IF (JFJOCH_CUDA_AVAILABLE) + SET_SOURCE_FILES_PROPERTIES(BeamCenterBackgroundGPU.cu PROPERTIES COMPILE_OPTIONS "--fmad=false") +ENDIF() +IF (CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + SET_SOURCE_FILES_PROPERTIES(BeamCenterFromBackground.cpp PROPERTIES COMPILE_OPTIONS "-ffp-contract=off") +ENDIF() diff --git a/tests/BeamCenterFromBackgroundTest.cpp b/tests/BeamCenterFromBackgroundTest.cpp index 73a547e92..d38db14e2 100644 --- a/tests/BeamCenterFromBackgroundTest.cpp +++ b/tests/BeamCenterFromBackgroundTest.cpp @@ -307,3 +307,29 @@ TEST_CASE("FindBeamCenter_BlanksTheBeamStopOutOfTheCapture", "[BeamCenter]") { CHECK(masked_error < 12.0f); CHECK(masked_error <= unmasked_error); } + +#ifdef JFJOCH_USE_CUDA +#include "../common/CUDAWrapper.h" + +// With a GPU the passes over the pixels run on it. The cell a pixel lands in is computed from IEEE +// operations alone and without contraction on both sides, and each cell is summed in the host's order, +// so the walk ends on the same bits - a tilted detector and an offset centre included. +TEST_CASE("BeamCenterFromBackground_DeviceMatchesHost", "[BeamCenter]") { + if (get_gpu_count() == 0) + SKIP("no GPU"); + DiffractionExperiment x = TestExperiment(); + x.PoniRot1_rad(0.005f).PoniRot2_rad(-0.003f); + PixelMask pixel_mask(x); + + const DiffractionGeometry geom_true = OffsetBy(x.GetDiffractionGeometry(), 7.0f, -4.5f); + const auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 0.5f); + const auto host = FindBeamCenterFromBackground(x, pixel_mask, projection, 0, {}, /*allow_device=*/false); + const auto device = FindBeamCenterFromBackground(x, pixel_mask, projection, 0, {}, /*allow_device=*/true); + + REQUIRE(host.has_value()); + REQUIRE(device.has_value()); + CHECK(device->beam_x_pxl == host->beam_x_pxl); + CHECK(device->beam_y_pxl == host->beam_y_pxl); + CHECK(device->sigma_pxl == host->sigma_pxl); +} +#endif diff --git a/tests/ShadowFinderTest.cpp b/tests/ShadowFinderTest.cpp index cc89345f4..f01e2c0b4 100644 --- a/tests/ShadowFinderTest.cpp +++ b/tests/ShadowFinderTest.cpp @@ -471,3 +471,91 @@ TEST_CASE("ShadowFinder_ABrightRingIsNotABeamStop", "[ShadowFinder]") { // Anything much beyond the stop, its arm and their penumbra means the walk ran away. CHECK(std::count(mask.begin(), mask.end(), 1u) < 6000); } + +#ifdef JFJOCH_USE_CUDA +#include "../common/CUDAWrapper.h" +#include "../compression/JFJochCompressor.h" + +namespace { + // The mask and the mean projection of the same frames, once from the host projection (the frames + // handed over uncompressed) and once from the device's (handed over as bitshuffle+LZ4, which is + // what sends them to the GPU). + struct HostAndDevice { + std::vector host_mask, device_mask; + std::vector host_mean, device_mean; + }; + + HostAndDevice MaskBothWays(const DiffractionExperiment &x, const std::vector> &frames) { + const PixelMask pixel_mask(x); + HostAndDevice out; + std::vector buffer; + { + ShadowFinder finder(x, pixel_mask); + for (const auto &frame : frames) { + DataMessage msg{}; + msg.image = CompressedImage(frame, W, H); + finder.AddImage(msg, buffer); + } + out.host_mask = finder.GetMask(); + out.host_mean = finder.GetMeanProjection(); + } + { + ShadowFinder finder(x, pixel_mask); + JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4); + std::vector> compressed; + for (const auto &frame : frames) { + compressed.push_back(compressor.Compress(frame)); + DataMessage msg{}; + msg.image = CompressedImage(compressed.back().data(), compressed.back().size(), W, H, + CompressedImageMode::Int32, CompressionAlgorithm::BSHUF_LZ4); + finder.AddImage(msg, buffer); + } + out.device_mask = finder.GetMask(); + out.device_mean = finder.GetMeanProjection(); + } + return out; + } +} + +// The mask is made on the GPU wherever the projection is there. On scenes where no pixel sits within a +// rounding of a threshold the two must agree to the pixel: every step but the polarization factor, +// the Poisson test's logarithm and the arm search's azimuth is exact on both, and these scenes are +// built so that none of the three decides anything at an edge. +TEST_CASE("ShadowFinder_DeviceMaskMatchesHost", "[ShadowFinder]") { + if (get_gpu_count() == 0) + SKIP("no GPU"); + + SECTION("a beam stop with a reflection behind it") { + std::vector> frames; + for (int f = 0; f < NFRAMES; f++) + frames.push_back(Scene(false, f == 0)); + const auto r = MaskBothWays(TestExperiment(), frames); + CHECK(std::count(r.host_mask.begin(), r.host_mask.end(), 1u) == 2612); + CHECK(r.device_mask == r.host_mask); + CHECK(std::memcmp(r.device_mean.data(), r.host_mean.data(), r.host_mean.size() * sizeof(float)) == 0); + } + + SECTION("an arm that lets part of the beam through, across a module gap") { + constexpr int32_t BRIGHT = 50; + constexpr int ARM_HALF_WIDE = 15, GAP_X0 = 200, GAP_X1 = 216, OPAQUE_FROM_X = 232; + std::vector> frames; + for (int f = 0; f < NFRAMES; f++) { + frames.emplace_back(static_cast(W) * H, BRIGHT); + auto &frame = frames.back(); + for (int y = 0; y < H; y++) + for (int xi = 0; xi < W; xi++) { + const int dx = xi - C, dy = y - C; + if (dx * dx + dy * dy <= STOP_R * STOP_R) + frame[I(xi, y)] = 0; + else if (dx >= 0 && std::abs(dy) <= ARM_HALF_WIDE) + frame[I(xi, y)] = xi >= OPAQUE_FROM_X ? 0 : BRIGHT * 6 / 10; + if (xi >= GAP_X0 && xi <= GAP_X1) + frame[I(xi, y)] = INT32_MIN; + } + } + const auto r = MaskBothWays(TestExperiment(), frames); + CHECK(std::count(r.host_mask.begin(), r.host_mask.end(), ShadowFinder::TRANSMITTING) > 0); + CHECK(r.device_mask == r.host_mask); + } +} +#endif From d4f4bab158f6f97575359fea4340a77776a11fab Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 3 Oct 2026 13:33:59 +0200 Subject: [PATCH 2/2] Pre-scan defective pixels: keep the GPU finder's work on the device The hot-pixel step of the pre-scan (MaskDefectivePixels) on a GPU build: - The device half is built once, before the workers start (HotPixelFinder::PrepareDevice), instead of by the first worker's frame while the others waited. The unmasked pixels grouped by key are sorted on the device (stable radix sort: the same order the host fill gave) instead of scattered on the host and uploaded. - No per-frame host round trip: the ring-sector levels and lit thresholds are made on the device from the order statistics, and the per-key frame and level sums are kept there too; frames queue on their workers' streams and the per-pixel accumulation is ordered by an event instead of a host synchronisation. - The mask: the chance rate's per-ring counts are summed on the device, and only the pixels the tests can pass (error value on most frames, or lit on at least min(max(2, k_chance), valid frames)) come back with their sums - not the five per-pixel arrays (470 MB pageable on a 16 Mpx detector). The host tests run on them unchanged. The threshold is written as fma(nsigma, noise, level) + offset on the host - what GCC already contracted it to - and the device takes the same two roundings, so the levels are bit-identical (and no longer depend on whether a compiler contracts). Exact: hot-pixel mask and p.mtz byte-identical to f849e2d1b on myoglobin, cytochrome C and thaumatin, GPU and CPU builds. Hot-pixel step (GPU, box at load 18-23, interleaved A/B, two pairs each): 0.94-1.32 s -> 0.65-0.81 s; frames 0.37-0.46 -> 0.18-0.23 s, mask 0.21-0.34 -> 0.08-0.18 s. Device memory of the finder: 543 MB as before, plus a ~0.36 GB transient for the sort while it is built. Tests: [HotPixelFinder] (HotPixelFinder_DeviceMatchesHost bit-exact), [ShadowFinder], [BeamCenter]. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01K5K8jvPPbmCrbqnWkddTuB --- rugnux/HotPixels.cpp | 120 ++++++++++---------- rugnux/HotPixels.h | 8 +- rugnux/HotPixelsGPU.cu | 245 +++++++++++++++++++++++++++++++++-------- rugnux/HotPixelsGPU.h | 77 +++++++++---- rugnux/Rugnux.cpp | 5 + 5 files changed, 327 insertions(+), 128 deletions(-) diff --git a/rugnux/HotPixels.cpp b/rugnux/HotPixels.cpp index 5b898f8ca..88fdadc9e 100644 --- a/rugnux/HotPixels.cpp +++ b/rugnux/HotPixels.cpp @@ -216,7 +216,9 @@ void HotPixelFinder::AddLevels(const std::vector §or_level, const s const int r = static_cast(k / SECTORS); level[k] = std::max(ring_level[r], sector_level[k]); const float noise = std::max(std::sqrt(static_cast(std::max(level[k], 0))), ring_spread[r]); - threshold[k] = static_cast(level[k]) + LIT_NSIGMA * noise + LIT_OFFSET; + // One rounding for level + nsigma * noise and one for the offset, written out so that every + // compiler takes the same two - the device takes them too (HotPixelsGPU.cu). + threshold[k] = std::fma(LIT_NSIGMA, noise, static_cast(level[k])) + LIT_OFFSET; } // Every sum is an integer, so the result does not depend on the order the frames arrive in. @@ -230,50 +232,24 @@ void HotPixelFinder::AddLevels(const std::vector §or_level, const s } #ifdef JFJOCH_USE_CUDA +void HotPixelFinder::PrepareDevice() { + std::lock_guard lock(m); + if (!gpu) + gpu = std::make_unique(key.get(), width * height, key_begin, nrings, SECTORS, + HotPixelLevelRules{MIN_SECTOR_PIXELS, MIN_RING_PIXELS, + LIT_NSIGMA, LIT_OFFSET}); +} + void HotPixelFinder::AddDeviceImage(const int32_t *device_image, HotPixelFinderGPU::Frame &frame) { - HotPixelFinderGPU *device; - { - std::lock_guard lock(m); - if (!gpu) - gpu = std::make_unique(key.get(), width * height, key_begin, nrings, SECTORS); - device = gpu.get(); - } - std::vector count; - std::vector sector_median, ring_median, ring_mad; - device->Statistics(device_image, frame, count, sector_median, ring_median, ring_mad); - - // The same levels AddImage takes off its scratch buffer, and under the same pixel minima. - const size_t nkeys = static_cast(nrings) * SECTORS; - std::vector sector_level(nkeys, 0); - for (size_t k = 0; k < nkeys; k++) - if (count[k] >= MIN_SECTOR_PIXELS) - sector_level[k] = sector_median[k]; - std::vector ring_level(nrings, 0); - std::vector ring_spread(nrings, 0.0f); - std::vector ring_ok(nrings, 0); - for (int r = 0; r < nrings; r++) { - size_t n = 0; - for (int s = 0; s < SECTORS; s++) - n += count[r * SECTORS + s]; - if (n < MIN_RING_PIXELS) continue; - ring_ok[r] = 1; - ring_level[r] = ring_median[r]; - ring_spread[r] = 1.4826f * static_cast(ring_mad[r]); - } - - std::vector level; - std::vector threshold; - AddLevels(sector_level, ring_level, ring_spread, ring_ok, level, threshold); - device->Accumulate(device_image, frame, level, threshold, ring_ok); + PrepareDevice(); + gpu->Add(device_image, frame); + std::lock_guard lock(m); + frames++; } #endif HotPixelFinder::Result HotPixelFinder::GetMask(double oscillation_deg, double spacing_deg, size_t nthreads) { std::lock_guard lock(m); -#ifdef JFJOCH_USE_CUDA - if (gpu) - gpu->Download(n_lit.get(), n_error.get(), sum_value.get(), n_error_ring_ok.get(), error_level_sum.get()); -#endif Result ret; ret.frames = frames; ret.mask.assign(width * height, 0); @@ -282,25 +258,37 @@ HotPixelFinder::Result HotPixelFinder::GetMask(double oscillation_deg, double sp // The chance rate per ring, from the pixels lit on no more than half of their frames: whatever // lights those - reflections, zingers, noise above the bound - lights a defect-free pixel too. // Counted in integers by blocks of rows in parallel, so the totals do not depend on the split. - std::vector> block_lit(BANDS), block_seen(BANDS); - const size_t rows_per_band = (height + BANDS - 1) / BANDS; - ParallelFor(static_cast(BANDS), nthreads, [&](int b) { - block_lit[b].assign(nrings, 0); - block_seen[b].assign(nrings, 0); - const size_t begin = std::min(width * height, b * rows_per_band * width); - const size_t end = std::min(width * height, (b + 1) * rows_per_band * width); - for (size_t i = begin; i < end; i++) - if (key[i] >= 0 && n_valid(i) > 0 && 2 * n_lit[i] <= n_valid(i)) { - block_lit[b][key[i] / SECTORS] += n_lit[i]; - block_seen[b][key[i] / SECTORS] += n_valid(i); - } - }); std::vector lit(nrings, 0.0), seen(nrings, 0.0); - for (int r = 0; r < nrings; r++) - for (size_t b = 0; b < BANDS; b++) { - lit[r] += static_cast(block_lit[b][r]); - seen[r] += static_cast(block_seen[b][r]); +#ifdef JFJOCH_USE_CUDA + if (gpu) { + std::vector device_lit, device_seen; + gpu->ChanceCounts(device_lit, device_seen); + for (int r = 0; r < nrings; r++) { + lit[r] = static_cast(device_lit[r]); + seen[r] = static_cast(device_seen[r]); } + } else +#endif + { + std::vector> block_lit(BANDS), block_seen(BANDS); + const size_t rows_per_band = (height + BANDS - 1) / BANDS; + ParallelFor(static_cast(BANDS), nthreads, [&](int b) { + block_lit[b].assign(nrings, 0); + block_seen[b].assign(nrings, 0); + const size_t begin = std::min(width * height, b * rows_per_band * width); + const size_t end = std::min(width * height, (b + 1) * rows_per_band * width); + for (size_t i = begin; i < end; i++) + if (key[i] >= 0 && n_valid(i) > 0 && 2 * n_lit[i] <= n_valid(i)) { + block_lit[b][key[i] / SECTORS] += n_lit[i]; + block_seen[b][key[i] / SECTORS] += n_valid(i); + } + }); + for (int r = 0; r < nrings; r++) + for (size_t b = 0; b < BANDS; b++) { + lit[r] += static_cast(block_lit[b][r]); + seen[r] += static_cast(block_seen[b][r]); + } + } std::vector k_chance(nrings, n + 1); for (int r = 0; r < nrings; r++) if (seen[r] > 0.0) @@ -309,6 +297,26 @@ HotPixelFinder::Result HotPixelFinder::GetMask(double oscillation_deg, double sp // Persistent: lit on more frames than one reflection or chance explains. const int min_valid = std::max(10, n / 2); +#ifdef JFJOCH_USE_CUDA + // With a GPU the per-pixel sums stay there. Only the pixels that can be masked come back - those the + // tests below could pass (see GetCandidates) - into the host arrays, which are zero everywhere else, + // so the tests below run on them unchanged. + if (gpu) { + const auto c = gpu->GetCandidates(frames, min_valid, spacing_deg > 0.0, k_chance); + for (size_t j = 0; j < c.index.size(); j++) { + const size_t i = c.index[j]; + n_lit[i] = c.n_lit[j]; + n_error[i] = c.n_error[j]; + n_error_ring_ok[i] = c.n_error_ring_ok[j]; + sum_value[i] = c.sum_value[j]; + error_level_sum[i] = c.error_level_sum[j]; + } + for (size_t k = 0; k < key_frames.size(); k++) { + key_frames[k] = static_cast(c.key_frames[k]); + key_level_sum[k] = c.key_level_sum[k]; + } + } +#endif std::vector persistent(width * height, 0); ParallelChunks(static_cast(height), nthreads, [&](int y0, int y1) { for (size_t y = y0; y < static_cast(y1); y++) diff --git a/rugnux/HotPixels.h b/rugnux/HotPixels.h index 63e6bd7a8..b837ba32e 100644 --- a/rugnux/HotPixels.h +++ b/rugnux/HotPixels.h @@ -94,6 +94,10 @@ public: // `scratch` above. The per-pixel sums are then kept on the device and replace the host's when the // mask is read, so a finder is fed one way or the other, not both. Thread safe. void AddDeviceImage(const int32_t *device_image, HotPixelFinderGPU::Frame &frame); + + // Build the device half now rather than with the first device frame, so the workers do not wait on + // it one behind the other. + void PrepareDevice(); #endif // The mask, from the frames added so far. oscillation_deg is the rotation per image and @@ -132,11 +136,11 @@ private: std::unique_ptr n_error_ring_ok; std::unique_ptr error_level_sum; #ifdef JFJOCH_USE_CUDA - std::unique_ptr gpu; // built by the first device frame + std::unique_ptr gpu; // built by PrepareDevice or the first device frame #endif // Each ring-sector's level and lit threshold, from the frame's order statistics, and the frame's - // share of the per-key sums. The host and the device path both come through here. + // share of the per-key sums. The device does the same in HotPixelsGPU.cu, with the same roundings. void AddLevels(const std::vector §or_level, const std::vector &ring_level, const std::vector &ring_spread, const std::vector &ring_ok, std::vector &level, std::vector &threshold); diff --git a/rugnux/HotPixelsGPU.cu b/rugnux/HotPixelsGPU.cu index 2691c221f..dc6346dc6 100644 --- a/rugnux/HotPixelsGPU.cu +++ b/rugnux/HotPixelsGPU.cu @@ -3,6 +3,8 @@ #include "HotPixelsGPU.h" +#include + #include "../common/JFJochException.h" namespace { @@ -145,6 +147,84 @@ __global__ void accumulate_kernel(const int32_t *__restrict__ image, const int32 } } +// Each key's level and lit threshold from the frame's order statistics, and the frame's share of the +// per-key sums - HotPixelFinder::AddImage and AddLevels, step for step. The threshold is +// fma(nsigma, noise, level) + offset, the two roundings the host takes (see AddLevels). +__global__ void levels_kernel(size_t nkeys, int sectors, HotPixelLevelRules rules, const uint32_t *__restrict__ count, + const int32_t *__restrict__ sector_median, const int32_t *__restrict__ ring_median, + const int32_t *__restrict__ ring_mad, int32_t *__restrict__ level, + float *__restrict__ threshold, char *__restrict__ ring_ok, + uint32_t *__restrict__ key_frames, int64_t *__restrict__ key_level_sum) { + const size_t k = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (k >= nkeys) return; + const size_t r = k / sectors; + uint32_t n = 0; + for (int s = 0; s < sectors; s++) + n += count[r * sectors + s]; + const bool ok = n >= static_cast(rules.min_ring_pixels); + const int32_t ring_level = ok ? ring_median[r] : 0; + const float ring_spread = ok ? 1.4826f * static_cast(ring_mad[r]) : 0.0f; + const int32_t sector_level = count[k] >= static_cast(rules.min_sector_pixels) ? sector_median[k] : 0; + const int32_t lv = max(ring_level, sector_level); + const float root = sqrtf(static_cast(max(lv, 0))); + const float noise = root < ring_spread ? ring_spread : root; + level[k] = lv; + threshold[k] = __fadd_rn(__fmaf_rn(rules.lit_nsigma, noise, static_cast(lv)), rules.lit_offset); + if (k % sectors == 0) + ring_ok[r] = ok; + if (ok) { + atomicAdd(&key_frames[k], 1u); + atomicAdd(reinterpret_cast(&key_level_sum[k]), + static_cast(static_cast(lv))); + } +} + +__device__ int valid_frames(const int32_t *key, const uint32_t *key_frames, const uint16_t *n_error_ring_ok, size_t i) { + return static_cast(key_frames[key[i]]) - static_cast(n_error_ring_ok[i]); +} + +__global__ void chance_kernel(size_t npixels, int sectors, const int32_t *__restrict__ key, + const uint32_t *__restrict__ key_frames, const uint16_t *__restrict__ n_lit, + const uint16_t *__restrict__ n_error_ring_ok, unsigned long long *__restrict__ lit, + unsigned long long *__restrict__ seen) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= npixels || key[i] < 0) return; + const int nv = valid_frames(key, key_frames, n_error_ring_ok, i); + if (nv > 0 && 2 * static_cast(n_lit[i]) <= nv) { + atomicAdd(&lit[key[i] / sectors], static_cast(n_lit[i])); + atomicAdd(&seen[key[i] / sectors], static_cast(nv)); + } +} + +// Writes the candidates' indices from `out` on when `out` is given, and counts them either way. +__global__ void candidate_kernel(size_t npixels, int sectors, uint32_t frames, int min_valid, bool spacing_ok, + const int32_t *__restrict__ key, const uint32_t *__restrict__ key_frames, + const uint16_t *__restrict__ n_lit, const uint16_t *__restrict__ n_error, + const uint16_t *__restrict__ n_error_ring_ok, const int *__restrict__ k_chance, + uint32_t *__restrict__ count, uint32_t *__restrict__ out) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= npixels || key[i] < 0) return; + const int nv = valid_frames(key, key_frames, n_error_ring_ok, i); + const bool error = 2 * static_cast(n_error[i]) > frames; + const int lit = n_lit[i]; + const bool persistent = spacing_ok && nv >= min_valid && lit > 0 + && lit >= min(max(2, k_chance[key[i] / sectors]), nv); + if (!error && !persistent) return; + const uint32_t slot = atomicAdd(count, 1u); + if (out) out[slot] = static_cast(i); +} + +__global__ void iota_kernel(size_t n, uint32_t *__restrict__ out) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i < n) out[i] = static_cast(i); +} + +template +__global__ void gather_kernel(size_t n, const uint32_t *__restrict__ index, const T *__restrict__ in, T *__restrict__ out) { + const size_t j = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (j < n) out[j] = in[index[j]]; +} + // The shared tables and sums are filled on a stream of their own, added to on the workers' streams // and downloaded on the NULL stream, so they are allocated synchronously rather than from the pool: // a pooled buffer is freed on the thread's allocation stream, which none of those is ordered before. @@ -156,22 +236,18 @@ constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; } // namespace HotPixelFinderGPU::HotPixelFinderGPU(const int32_t *host_key, size_t npixels, - const std::vector &host_key_begin, int nrings, int sectors) + const std::vector &host_key_begin, int nrings, int sectors, + const HotPixelLevelRules &rules) : npixels(npixels), nkeys(static_cast(nrings) * sectors), nrings(nrings), - sectors(sectors), - key(npixels, ALLOC), pixels_by_key(host_key_begin.back(), ALLOC), key_begin(host_key_begin.size(), ALLOC), + sectors(sectors), rules(rules), + key(npixels, ALLOC), pixels_by_key(std::max(host_key_begin.back(), 1), ALLOC), + key_begin(host_key_begin.size(), ALLOC), n_lit(npixels, ALLOC), n_error(npixels, ALLOC), n_error_ring_ok(npixels, ALLOC), - sum_value(npixels, ALLOC), error_level_sum(npixels, ALLOC) { - std::vector pixels(host_key_begin.back()); - std::vector filled(host_key_begin.begin(), host_key_begin.end() - 1); - for (size_t i = 0; i < npixels; i++) - if (host_key[i] >= 0) - pixels[filled[host_key[i]]++] = static_cast(i); - + sum_value(npixels, ALLOC), error_level_sum(npixels, ALLOC), + key_frames(std::max(nkeys, 1), ALLOC), key_level_sum(std::max(nkeys, 1), ALLOC) { + cuda_err(cudaEventCreateWithFlags(&last_accumulate, cudaEventDisableTiming)); CudaStream stream; cuda_err(cudaMemcpyAsync(key, host_key, npixels * sizeof(int32_t), cudaMemcpyHostToDevice, stream)); - cuda_err(cudaMemcpyAsync(pixels_by_key, pixels.data(), pixels.size() * sizeof(uint32_t), - cudaMemcpyHostToDevice, stream)); cuda_err(cudaMemcpyAsync(key_begin, host_key_begin.data(), host_key_begin.size() * sizeof(uint32_t), cudaMemcpyHostToDevice, stream)); cuda_err(cudaMemsetAsync(n_lit, 0, npixels * sizeof(uint16_t), stream)); @@ -179,16 +255,34 @@ HotPixelFinderGPU::HotPixelFinderGPU(const int32_t *host_key, size_t npixels, cuda_err(cudaMemsetAsync(n_error_ring_ok, 0, npixels * sizeof(uint16_t), stream)); cuda_err(cudaMemsetAsync(sum_value, 0, npixels * sizeof(int64_t), stream)); cuda_err(cudaMemsetAsync(error_level_sum, 0, npixels * sizeof(int64_t), stream)); - cuda_err(cudaStreamSynchronize(stream)); + cuda_err(cudaMemsetAsync(key_frames, 0, std::max(nkeys, 1) * sizeof(uint32_t), stream)); + cuda_err(cudaMemsetAsync(key_level_sum, 0, std::max(nkeys, 1) * sizeof(int64_t), stream)); + + // The unmasked pixels grouped by key: a stable sort of the pixel indices by key, so each key's + // pixels stay in pixel order. A masked pixel's key, -1, is the largest as unsigned and sorts last. + { + CudaDevicePtr index(npixels, ALLOC), sorted_key(npixels, ALLOC), sorted_index(npixels, ALLOC); + iota_kernel<<((npixels + THREADS - 1) / THREADS), THREADS, 0, stream>>>(npixels, index); + cuda_err(cudaGetLastError()); + const auto *keys_in = reinterpret_cast(key.get()); + size_t bytes = 0; + cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes, keys_in, sorted_key.get(), index.get(), + sorted_index.get(), npixels, 0, 32, stream)); + CudaDevicePtr scratch(bytes, ALLOC); + cuda_err(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, keys_in, sorted_key.get(), index.get(), + sorted_index.get(), npixels, 0, 32, stream)); + if (host_key_begin.back() > 0) + cuda_err(cudaMemcpyAsync(pixels_by_key, sorted_index, host_key_begin.back() * sizeof(uint32_t), + cudaMemcpyDeviceToDevice, stream)); + cuda_err(cudaStreamSynchronize(stream)); + } } -void HotPixelFinderGPU::Statistics(const int32_t *device_image, Frame &frame, std::vector &count, - std::vector §or_median, std::vector &ring_median, - std::vector &ring_mad) { - count.resize(nkeys); - sector_median.resize(nkeys); - ring_median.resize(nrings); - ring_mad.resize(nrings); +HotPixelFinderGPU::~HotPixelFinderGPU() { + if (last_accumulate) cudaEventDestroy(last_accumulate); +} + +void HotPixelFinderGPU::Add(const int32_t *device_image, Frame &frame) { if (nrings == 0) return; if (!frame.count.get()) { @@ -201,44 +295,101 @@ void HotPixelFinderGPU::Statistics(const int32_t *device_image, Frame &frame, st frame.ring_ok = CudaDevicePtr(nrings); } const cudaStream_t stream = *frame.stream; - sector_kernel<<(nkeys), THREADS, 0, stream>>>(device_image, pixels_by_key, key_begin, frame.count, - frame.sector_median); + sector_kernel<<(nkeys), THREADS, 0, stream>>>(device_image, pixels_by_key, key_begin, + frame.count, frame.sector_median); cuda_err(cudaGetLastError()); ring_kernel<<>>(device_image, pixels_by_key, key_begin, frame.count, sectors, frame.ring_median, frame.ring_mad); cuda_err(cudaGetLastError()); - - cuda_err(cudaMemcpyAsync(count.data(), frame.count, nkeys * sizeof(uint32_t), cudaMemcpyDeviceToHost, stream)); - cuda_err(cudaMemcpyAsync(sector_median.data(), frame.sector_median, nkeys * sizeof(int32_t), - cudaMemcpyDeviceToHost, stream)); - cuda_err(cudaMemcpyAsync(ring_median.data(), frame.ring_median, nrings * sizeof(int32_t), - cudaMemcpyDeviceToHost, stream)); - cuda_err(cudaMemcpyAsync(ring_mad.data(), frame.ring_mad, nrings * sizeof(int32_t), - cudaMemcpyDeviceToHost, stream)); - cuda_err(cudaStreamSynchronize(stream)); -} - -void HotPixelFinderGPU::Accumulate(const int32_t *device_image, Frame &frame, const std::vector &level, - const std::vector &threshold, const std::vector &ring_ok) { - const cudaStream_t stream = *frame.stream; - cuda_err(cudaMemcpyAsync(frame.level, level.data(), nkeys * sizeof(int32_t), cudaMemcpyHostToDevice, stream)); - cuda_err(cudaMemcpyAsync(frame.threshold, threshold.data(), nkeys * sizeof(float), cudaMemcpyHostToDevice, - stream)); - cuda_err(cudaMemcpyAsync(frame.ring_ok, ring_ok.data(), nrings * sizeof(char), cudaMemcpyHostToDevice, stream)); + levels_kernel<<((nkeys + THREADS - 1) / THREADS), THREADS, 0, stream>>>( + nkeys, sectors, rules, frame.count, frame.sector_median, frame.ring_median, frame.ring_mad, + frame.level, frame.threshold, frame.ring_ok, key_frames, key_level_sum); + cuda_err(cudaGetLastError()); std::lock_guard lock(accumulate_mutex); + cuda_err(cudaStreamWaitEvent(stream, last_accumulate, 0)); accumulate_kernel<<((npixels + THREADS - 1) / THREADS), THREADS, 0, stream>>>( device_image, key, npixels, sectors, frame.level, frame.threshold, frame.ring_ok, n_lit, n_error, sum_value, n_error_ring_ok, error_level_sum); cuda_err(cudaGetLastError()); + cuda_err(cudaEventRecord(last_accumulate, stream)); +} + +void HotPixelFinderGPU::ChanceCounts(std::vector &lit, std::vector &seen) { + CudaStream stream; + cuda_err(cudaStreamWaitEvent(stream, last_accumulate, 0)); + const size_t rings = std::max(nrings, 1); + CudaDevicePtr d_lit(rings, ALLOC), d_seen(rings, ALLOC); + cuda_err(cudaMemsetAsync(d_lit, 0, rings * sizeof(unsigned long long), stream)); + cuda_err(cudaMemsetAsync(d_seen, 0, rings * sizeof(unsigned long long), stream)); + chance_kernel<<((npixels + THREADS - 1) / THREADS), THREADS, 0, stream>>>( + npixels, sectors, key, key_frames, n_lit, n_error_ring_ok, d_lit, d_seen); + cuda_err(cudaGetLastError()); + lit.resize(nrings); + seen.resize(nrings); + cuda_err(cudaMemcpyAsync(lit.data(), d_lit, nrings * sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(seen.data(), d_seen, nrings * sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); cuda_err(cudaStreamSynchronize(stream)); } -void HotPixelFinderGPU::Download(uint16_t *host_n_lit, uint16_t *host_n_error, int64_t *host_sum_value, - uint16_t *host_n_error_ring_ok, int64_t *host_error_level_sum) { - cuda_err(cudaMemcpy(host_n_lit, n_lit, npixels * sizeof(uint16_t), cudaMemcpyDeviceToHost)); - cuda_err(cudaMemcpy(host_n_error, n_error, npixels * sizeof(uint16_t), cudaMemcpyDeviceToHost)); - cuda_err(cudaMemcpy(host_sum_value, sum_value, npixels * sizeof(int64_t), cudaMemcpyDeviceToHost)); - cuda_err(cudaMemcpy(host_n_error_ring_ok, n_error_ring_ok, npixels * sizeof(uint16_t), cudaMemcpyDeviceToHost)); - cuda_err(cudaMemcpy(host_error_level_sum, error_level_sum, npixels * sizeof(int64_t), cudaMemcpyDeviceToHost)); +HotPixelFinderGPU::Candidates HotPixelFinderGPU::GetCandidates(uint32_t frames, int min_valid, bool spacing_ok, + const std::vector &k_chance) { + CudaStream stream; + cuda_err(cudaStreamWaitEvent(stream, last_accumulate, 0)); + const unsigned blocks = static_cast((npixels + THREADS - 1) / THREADS); + CudaDevicePtr d_k_chance(std::max(k_chance.size(), 1), ALLOC); + CudaDevicePtr d_count(1, ALLOC); + if (!k_chance.empty()) + cuda_err(cudaMemcpyAsync(d_k_chance, k_chance.data(), k_chance.size() * sizeof(int), cudaMemcpyHostToDevice, + stream)); + cuda_err(cudaMemsetAsync(d_count, 0, sizeof(uint32_t), stream)); + candidate_kernel<<>>(npixels, sectors, frames, min_valid, spacing_ok, key, key_frames, + n_lit, n_error, n_error_ring_ok, d_k_chance, d_count, nullptr); + cuda_err(cudaGetLastError()); + uint32_t n = 0; + cuda_err(cudaMemcpyAsync(&n, d_count, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); + + Candidates c; + c.key_frames.resize(nkeys); + c.key_level_sum.resize(nkeys); + if (nkeys > 0) { + cuda_err(cudaMemcpyAsync(c.key_frames.data(), key_frames, nkeys * sizeof(uint32_t), cudaMemcpyDeviceToHost, + stream)); + cuda_err(cudaMemcpyAsync(c.key_level_sum.data(), key_level_sum, nkeys * sizeof(int64_t), + cudaMemcpyDeviceToHost, stream)); + } + if (n > 0) { + CudaDevicePtr index(n, ALLOC); + CudaDevicePtr g_lit(n, ALLOC), g_error(n, ALLOC), g_error_ring_ok(n, ALLOC); + CudaDevicePtr g_sum(n, ALLOC), g_error_level(n, ALLOC); + cuda_err(cudaMemsetAsync(d_count, 0, sizeof(uint32_t), stream)); + candidate_kernel<<>>(npixels, sectors, frames, min_valid, spacing_ok, key, + key_frames, n_lit, n_error, n_error_ring_ok, d_k_chance, + d_count, index); + cuda_err(cudaGetLastError()); + const unsigned gb = (n + THREADS - 1) / THREADS; + gather_kernel<<>>(n, index, n_lit.get(), g_lit.get()); + gather_kernel<<>>(n, index, n_error.get(), g_error.get()); + gather_kernel<<>>(n, index, n_error_ring_ok.get(), g_error_ring_ok.get()); + gather_kernel<<>>(n, index, sum_value.get(), g_sum.get()); + gather_kernel<<>>(n, index, error_level_sum.get(), g_error_level.get()); + cuda_err(cudaGetLastError()); + c.index.resize(n); + c.n_lit.resize(n); + c.n_error.resize(n); + c.n_error_ring_ok.resize(n); + c.sum_value.resize(n); + c.error_level_sum.resize(n); + cuda_err(cudaMemcpyAsync(c.index.data(), index, n * sizeof(uint32_t), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(c.n_lit.data(), g_lit, n * sizeof(uint16_t), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(c.n_error.data(), g_error, n * sizeof(uint16_t), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(c.n_error_ring_ok.data(), g_error_ring_ok, n * sizeof(uint16_t), + cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(c.sum_value.data(), g_sum, n * sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(c.error_level_sum.data(), g_error_level, n * sizeof(int64_t), + cudaMemcpyDeviceToHost, stream)); + } + cuda_err(cudaStreamSynchronize(stream)); + return c; } diff --git a/rugnux/HotPixelsGPU.h b/rugnux/HotPixelsGPU.h index f062e64cf..e68b6e4f0 100644 --- a/rugnux/HotPixelsGPU.h +++ b/rugnux/HotPixelsGPU.h @@ -10,33 +10,53 @@ #include "../image_analysis/indexing/CUDAMemHelpers.h" -// The device half of HotPixelFinder, for frames already preprocessed on the GPU: the per-frame order -// statistics (each ring-sector's median, each ring's median and median absolute deviation) and the -// per-pixel sums run where the image already is, and only the per-key statistics - tens of thousands -// of numbers - come to the host, which turns them into levels and thresholds with the very code the -// host path uses. Every statistic is an exact order statistic of integers and every sum an integer, -// so the sums, and with them the mask, are identical to what HotPixelFinder::AddImage produces. +// What a frame's levels and lit thresholds are made of (HotPixelFinder's constants), handed to the +// device so the two halves cannot drift apart. +struct HotPixelLevelRules { + int min_sector_pixels; + int min_ring_pixels; + float lit_nsigma; + float lit_offset; +}; + +// The device half of HotPixelFinder, for frames already preprocessed on the GPU. Everything a frame adds +// stays on the device: the per-frame order statistics (each ring-sector's median, each ring's median +// and median absolute deviation), the levels and thresholds made from them, and the per-pixel and +// per-key sums. Every statistic is an exact order statistic of integers, the threshold is computed +// with the same rounding steps as the host's (see HotPixelFinder::AddLevels) and every sum is an +// integer, so the sums are identical to what HotPixelFinder::AddImage produces. When the mask is read +// only what can decide it comes back: per-ring counts for the chance rate, then the few pixels that can +// be persistent or carry the error value. class HotPixelFinderGPU { const size_t npixels; const size_t nkeys; const int nrings; const int sectors; + const HotPixelLevelRules rules; CudaDevicePtr key; // ring * sectors + sector of each pixel, -1 masked - CudaDevicePtr pixels_by_key; // the unmasked pixels, grouped by key + CudaDevicePtr pixels_by_key; // the unmasked pixels, grouped by key, in pixel order CudaDevicePtr key_begin; // where each key's pixels start in pixels_by_key - // The per-pixel sums, exactly those of HotPixelFinder. + // The per-pixel and per-key sums, exactly those of HotPixelFinder. CudaDevicePtr n_lit, n_error, n_error_ring_ok; CudaDevicePtr sum_value, error_level_sum; - // Frames are selected on their workers' streams in parallel, but each pixel's sums are plain - // read-modify-writes, so one frame at a time adds to them. + CudaDevicePtr key_frames; + CudaDevicePtr key_level_sum; + + // Each pixel's sums are plain read-modify-writes, so the frames add to them one after another: + // every accumulation waits on the device for the one before it (an event, not the host). std::mutex accumulate_mutex; + cudaEvent_t last_accumulate = nullptr; public: // One worker's buffers, on the stream its frames are preprocessed on. struct Frame { explicit Frame(std::shared_ptr stream) : stream(std::move(stream)) {} + // Its frames are queued, not waited for (Add), so the buffers below must not be freed before + // the stream has used them. + ~Frame() { if (stream) cudaStreamSynchronize(*stream); } + Frame(Frame &&) = default; std::shared_ptr stream; CudaDevicePtr count; // valid pixels per key CudaDevicePtr sector_median; // per key @@ -46,21 +66,32 @@ public: CudaDevicePtr ring_ok; }; + // The keys of HotPixelFinder: one per pixel, and where each key's pixels start among the unmasked + // pixels sorted by key. HotPixelFinderGPU(const int32_t *key, size_t npixels, const std::vector &key_begin, int nrings, - int sectors); + int sectors, const HotPixelLevelRules &rules); + ~HotPixelFinderGPU(); + HotPixelFinderGPU(const HotPixelFinderGPU &) = delete; + HotPixelFinderGPU &operator=(const HotPixelFinderGPU &) = delete; - // The lower median of the valid values of each key (count[k] of them, 0 where there are none), and - // of each ring the median and the lower median of the absolute deviations from it. - void Statistics(const int32_t *device_image, Frame &frame, std::vector &count, - std::vector §or_median, std::vector &ring_median, - std::vector &ring_mad); + // Add one frame, preprocessed on frame.stream, as HotPixelFinder::AddImage does. Queued on that + // stream; nothing is waited for on the host. + void Add(const int32_t *device_image, Frame &frame); - // Add the frame to the per-pixel sums, with each key's level and lit threshold and each ring's - // verdict on whether it has a level at all - as HotPixelFinder::AddImage does. - void Accumulate(const int32_t *device_image, Frame &frame, const std::vector &level, - const std::vector &threshold, const std::vector &ring_ok); + // Per ring, over the pixels lit on no more than half of their valid frames: the lit frames and the + // valid frames, summed (HotPixelFinder::GetMask's chance rate). Waits for every frame added. + void ChanceCounts(std::vector &lit, std::vector &seen); - // The per-pixel sums, npixels each. - void Download(uint16_t *n_lit, uint16_t *n_error, int64_t *sum_value, uint16_t *n_error_ring_ok, - int64_t *error_level_sum); + // The pixels that can be masked: those holding the error value on more than half of `frames`, and, + // where spacing_ok, those lit on at least min(max(2, k_chance[ring]), valid frames) of at least + // min_valid valid frames - a persistent pixel is lit on at least that many, because one reflection + // explains at least two. For each, its index and per-pixel sums; and every key's sums. + struct Candidates { + std::vector index; + std::vector n_lit, n_error, n_error_ring_ok; + std::vector sum_value, error_level_sum; + std::vector key_frames; + std::vector key_level_sum; + }; + Candidates GetCandidates(uint32_t frames, int min_valid, bool spacing_ok, const std::vector &k_chance); }; diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index d99d8d8e1..10605a01a 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -1669,6 +1669,11 @@ void Rugnux::MaskDefectivePixels(int start_image, const std::vector &sample const size_t nthreads = static_cast(std::max(config_.nthreads, 1)); HotPixelFinder finder(about, pixel_mask_, nthreads); +#ifdef JFJOCH_USE_CUDA + // Built here, once, so the workers' first frames do not queue behind it. + if (get_gpu_count() > 0) + finder.PrepareDevice(); +#endif std::atomic next{0}; std::vector> futures; const size_t nworkers = std::min(nthreads, std::min(PRESCAN_MAX_WORKERS, sample.size()));