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 2330dbb93..50971f3ad 100644 --- a/image_analysis/geom_refinement/CMakeLists.txt +++ b/image_analysis/geom_refinement/CMakeLists.txt @@ -29,6 +29,7 @@ ADD_LIBRARY(JFJochGeomRefinement STATIC LMSolver.cpp LMSolver.h Dual.h + BackgroundBand.h PostRefine.cpp PostRefine.h GeometryRefiner.cpp @@ -42,7 +43,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) @@ -51,3 +53,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/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 30a449fae..94a9325af 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -1702,6 +1702,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())); 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