spot_finding: find connected components on the GPU

The spot finder flagged strong pixels on the device and then labelled them on the host, so every frame sent the packed bitmask back - 2.26 MB on a large detector - and the host walked all of it to recover a few hundred pixels. Do the labelling on the device instead: compact the bitmask into a flat-index-sorted list, find each pixel's backward neighbours by binary search, union them lock-free with path halving, then label, accumulate and filter in one kernel. Only the spot list comes back, and only one stream synchronisation per frame. The gain in the ordinary case is modest - about a quarter off per-image spot finding - because the host algorithm is genuinely fast on a normal frame. What justifies it is the frame that is not ordinary. The host labels a sorted sparse list through a window spanning two detector lines, so its cost is quadratic in how many strong pixels share a line. A lit band of detector rows - a hot module, a panel edge - costs 33 ms at two rows and 377 ms at fifteen, all of it under the pixel cap that was supposed to bound this, and none of it maskable when the cause is a diffraction ring rather than a defect: a ring runs tangent to a row at its top and bottom, which is exactly the shape that hurts. The device version is flat at 0.05 to 0.64 ms across every geometry tried, so an online run no longer stalls a quarter of a second on an ice ring. Rejecting an over-cap frame is now free too, since the count is known before any pixel is written. Also label once and filter three times. The per-image minimum-pixel search runs the extraction at three settings, but that setting only decides which components are kept - it does not change the components - so the search itself need not be repeated. This helps the host path as much as the device one. The resolution mask moves to the device as a bit mask, uploaded when the limits change rather than per frame, since the compaction needs it there. Parity is asserted permanently rather than argued: five cases covering realistic frames, occupancy from a hundred pixels to past the cap, the pathological geometries including rings, the resolution mask, and a hundred-repeat determinism check - requiring the same partition, the same spot order, and identical counts. The centroid is a float sum and therefore order-dependent, so the device walks each component from its root in ascending order and fuses its multiply-add the way the host's does; note that whether the host fuses at all depends on the architecture flags, so exact centroid equality is asserted where the compiler fuses and a two-ulp bound otherwise. Making those accumulators integer would remove that dependence entirely and is worth doing separately. Regression set: all 37 crystals identical to the last printed digit. Unit suite passes with the new cases. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-02 13:34:43 +02:00
parent e4d70f0e55
commit 4bdb229fb8
21 changed files with 1030 additions and 87 deletions
@@ -0,0 +1,370 @@
+// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+// The GPU spot extractor (SpotExtractorGPU) has to reproduce the host connected-component search
+// (StrongPixelSet::sparseccl, driven through ImageSpotFinder::ExtractComponentsHost) EXACTLY - the
+// same partition of the strong pixels into components, the same number of spots, the same spot
+// ORDER, and the same pixel count, photon sum, maximum and centroid sums for each of them. If the
+// two ever drift apart, every downstream result (indexing, integration, merging) moves with them
+// and nothing else in the suite would notice.
+//
+// Both extractors are driven from the SAME strong-pixel bit buffer here, so this compares only the
+// extraction; the detection stages have their own parity cases
+// (ImageSpotFinder_CPU_GPU_Parity, AdaptiveSpotFinderGPU_SpotFindingParity).
+//
+// The frames are generated, not recorded: dataset-derived fixtures are not allowed in the tree. The
+// blob-size and density distributions were chosen to match what real 18 MP frames produce
+// (~1e3 - 1.5e4 strong pixels, spots of a few to a few tens of pixels).
+
+#include <catch2/catch_all.hpp>
+#include "../common/CUDAWrapper.h"
+
+#ifdef JFJOCH_USE_CUDA
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <random>
+#include <vector>
+
+#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"
+#include "../image_analysis/spot_finding/ImageSpotFinder.h"
+#include "../image_analysis/spot_finding/SpotExtractorGPU.h"
+
+namespace {
+
+// A detector large enough for the geometries below to be representative, small enough that a case
+// costs milliseconds. The pathological shapes are what set the size, not the spot count.
+constexpr int32_t WIDTH = 2048;
+constexpr int32_t HEIGHT = 2048;
+
+SpotFindingSettings ExtractionSettings(int64_t max_pix = 50) {
+    SpotFindingSettings s{};
+    s.min_pix_per_spot = 1;
+    s.max_pix_per_spot = max_pix;
+    return s;
+}
+
+// A finder whose detection stage is "use the bit buffer I was handed". It exists so both extractors
+// can be given byte-identical input.
+class PlantedSpotFinder : public ImageSpotFinder {
+public:
+    PlantedSpotFinder() : ImageSpotFinder(WIDTH, HEIGHT) {}
+    void Detect(const ImagePreprocessorBuffer &, const SpotFindingSettings &) override {}
+    void Plant(const std::vector<uint32_t> &strong_index) {
+        std::fill(output_buffer.begin(), output_buffer.end(), 0);
+        for (uint32_t flat : strong_index)
+            output_buffer[flat / 32] |= 1u << (flat % 32);
+    }
+    const std::vector<uint32_t> &BitBuffer() const { return output_buffer; }
+    size_t Words() const { return OutputSize(); }
+};
+
+// ---------------------------------------------------------------------------------------------
+// Frame geometries
+// ---------------------------------------------------------------------------------------------
+struct Frame {
+    std::string name;
+    std::vector<uint32_t> index;       // flat pixel index, ascending
+    std::vector<int32_t> value;
+};
+
+Frame Finish(std::string name, std::vector<uint32_t> index, uint64_t seed) {
+    std::sort(index.begin(), index.end());
+    index.erase(std::unique(index.begin(), index.end()), index.end());
+    Frame f{std::move(name), std::move(index), {}};
+    // Values spanning the range a preprocessed image really carries, including negatives (a
+    // pedestal-subtracted pixel can go below zero) - the centroid sums must match on those too.
+    std::mt19937_64 rng(seed);
+    f.value.reserve(f.index.size());
+    for (size_t i = 0; i < f.index.size(); i++)
+        f.value.push_back(static_cast<int32_t>(rng() % 40000) - 200);
+    return f;
+}
+
+// Compact blobs scattered over the detector: what a diffraction pattern looks like to the extractor.
+Frame BlobFrame(const std::string &name, int nblob, uint64_t seed) {
+    std::mt19937_64 rng(seed);
+    std::vector<uint32_t> index;
+    for (int b = 0; b < nblob; b++) {
+        const int col = 20 + static_cast<int>(rng() % (WIDTH - 40));
+        const int line = 20 + static_cast<int>(rng() % (HEIGHT - 40));
+        const int radius = 1 + static_cast<int>(rng() % 3);     // spots of ~4 to ~30 pixels
+        for (int dl = -radius; dl <= radius; dl++)
+            for (int dc = -radius; dc <= radius; dc++)
+                if (dl * dl + dc * dc <= radius * radius)
+                    index.push_back(static_cast<uint32_t>((line + dl) * WIDTH + (col + dc)));
+    }
+    return Finish(name, std::move(index), seed ^ 0x9e37u);
+}
+
+// A lit band of whole detector rows - a hot module, a panel edge, and the limit a diffraction ring
+// approaches where it runs tangent to a row. This is the shape the host's two-row sliding window
+// turns quadratic on.
+Frame BandFrame(const std::string &name, int nrows) {
+    std::vector<uint32_t> index;
+    for (int line = 500; line < 500 + nrows; line++)
+        for (int col = 0; col < WIDTH; col++)
+            index.push_back(static_cast<uint32_t>(line * WIDTH + col));
+    return Finish(name, std::move(index), 0x51ed3u + nrows);
+}
+
+// Concentric diffraction rings. Unlike a hot column, a ring is SIGNAL and cannot be masked away, and
+// at its top and bottom it produces exactly the long horizontal run that the band case models.
+Frame RingFrame(const std::string &name, int radius, int half_width, int nring) {
+    const double cx = WIDTH * 0.5, cy = HEIGHT * 0.5;
+    std::vector<uint32_t> index;
+    for (int line = 0; line < HEIGHT; line++) {
+        for (int col = 0; col < WIDTH; col++) {
+            const double d = std::sqrt((col - cx) * (col - cx) + (line - cy) * (line - cy));
+            for (int k = 0; k < nring; k++) {
+                if (std::fabs(d - (radius + k * 130)) <= half_width) {
+                    index.push_back(static_cast<uint32_t>(line * WIDTH + col));
+                    break;
+                }
+            }
+        }
+    }
+    return Finish(name, std::move(index), 0x21f9u + radius * 31 + nring);
+}
+
+Frame SolidBlockFrame(const std::string &name, int side) {
+    std::vector<uint32_t> index;
+    for (int line = 300; line < 300 + side; line++)
+        for (int col = 300; col < 300 + side; col++)
+            index.push_back(static_cast<uint32_t>(line * WIDTH + col));
+    return Finish(name, std::move(index), 0x7c1u + side);
+}
+
+// Isolated single pixels at the density a very noisy frame reaches: many components, all tiny.
+Frame SparseNoiseFrame(const std::string &name, int npixel, uint64_t seed) {
+    std::mt19937_64 rng(seed);
+    std::vector<uint32_t> index;
+    for (int i = 0; i < npixel; i++)
+        index.push_back(static_cast<uint32_t>(rng() % (static_cast<uint64_t>(WIDTH) * HEIGHT)));
+    return Finish(name, std::move(index), seed ^ 0x1234u);
+}
+
+// ---------------------------------------------------------------------------------------------
+// Harness: run both extractors on one frame
+// ---------------------------------------------------------------------------------------------
+class ParityRunner {
+    PlantedSpotFinder finder;
+    ImagePreprocessorBufferGPU image;
+    std::shared_ptr<CudaStream> stream;
+    SpotExtractorGPU extractor;
+    CudaDevicePtr<uint32_t> gpu_strong;
+
+public:
+    ParityRunner()
+        : image(static_cast<size_t>(WIDTH) * HEIGHT),
+          stream(std::make_shared<CudaStream>()),
+          extractor(WIDTH, HEIGHT, stream),
+          gpu_strong(finder.Words()) {}
+
+    void SetResolutionMask(const std::vector<bool> &mask) {
+        finder.SetResolutionMask(mask);
+        std::vector<uint32_t> packed((static_cast<size_t>(WIDTH) * HEIGHT + 31) / 32, 0);
+        for (size_t i = 0; i < mask.size(); i++)
+            if (mask[i]) packed[i / 32] |= 1u << (i % 32);
+        extractor.SetResolutionMask(packed);
+    }
+
+    void Load(const Frame &f) {
+        auto &buffer = image.getBuffer();
+        std::fill(buffer.begin(), buffer.end(), 0);
+        for (size_t i = 0; i < f.index.size(); i++)
+            buffer[f.index[i]] = f.value[i];
+        REQUIRE(cudaMemcpy(image.getGPUBuffer(), buffer.data(), buffer.size() * sizeof(int32_t),
+                           cudaMemcpyHostToDevice) == cudaSuccess);
+        finder.Plant(f.index);
+        REQUIRE(cudaMemcpy(gpu_strong.get(), finder.BitBuffer().data(),
+                           finder.Words() * sizeof(uint32_t), cudaMemcpyHostToDevice) == cudaSuccess);
+        REQUIRE(cudaDeviceSynchronize() == cudaSuccess);
+    }
+
+    std::vector<DiffractionSpot> Cpu(const SpotFindingSettings &settings) {
+        return finder.ExtractComponents(image, settings);
+    }
+    std::vector<DiffractionSpot> Gpu(const SpotFindingSettings &settings) {
+        std::vector<DiffractionSpot> spots;
+        extractor.Extract(gpu_strong.get(), image.getGPUBuffer(), settings, spots);
+        return spots;
+    }
+};
+
+// How many representable floats apart two values are.
+int64_t UlpDistance(float a, float b) {
+    int32_t ia, ib;
+    memcpy(&ia, &a, sizeof(ia));
+    memcpy(&ib, &b, sizeof(ib));
+    if (ia < 0) ia = INT32_MIN - ia;      // map to a monotone ordering across the sign
+    if (ib < 0) ib = INT32_MIN - ib;
+    return std::abs(static_cast<int64_t>(ia) - static_cast<int64_t>(ib));
+}
+
+// Everything that decides which spots exist and what they weigh is compared EXACTLY: the number of
+// spots, their order, and each one's pixel count, photon sum and maximum. Those are integers, and a
+// difference in any of them is a difference in the partition.
+//
+// The centroid is a float sum, and its last bit is a property of the BUILD rather than of either
+// implementation: gcc contracts DiffractionSpot::AddPixel into an FMA whenever the flags allow it
+// (the CI -march=x86-64-v3 does), while a baseline -march, or MSVC with its default /fp:precise,
+// cannot. The extractor uses __fmaf_rn, so it is bit-exact against a host that fuses and one ulp
+// off one that does not. An ulp bound catches a real divergence - which moves a centroid by pixels,
+// not by 1e-4 of one - while staying true whichever way the host was built.
+void RequireIdentical(const std::string &what,
+                      const std::vector<DiffractionSpot> &cpu,
+                      const std::vector<DiffractionSpot> &gpu) {
+    INFO(what);
+    REQUIRE(gpu.size() == cpu.size());
+    for (size_t i = 0; i < cpu.size(); i++) {
+        INFO("spot " << i << " of " << cpu.size());
+        REQUIRE(gpu[i].PixelCount() == cpu[i].PixelCount());
+        REQUIRE(gpu[i].Count() == cpu[i].Count());
+        REQUIRE(gpu[i].MaxCount() == cpu[i].MaxCount());
+        // RawCoord divides the sums by the photon count, so comparing it compares the sums; a spot
+        // whose photons sum to zero reports (0,0) on both sides by the same branch.
+        const Coord c = cpu[i].RawCoord(), g = gpu[i].RawCoord();
+#ifdef __FMA__
+        REQUIRE(memcmp(&g.x, &c.x, sizeof(float)) == 0);
+        REQUIRE(memcmp(&g.y, &c.y, sizeof(float)) == 0);
+#else
+        REQUIRE(UlpDistance(g.x, c.x) <= 2);
+        REQUIRE(UlpDistance(g.y, c.y) <= 2);
+#endif
+    }
+}
+
+std::vector<Frame> RealisticFrames() {
+    std::vector<Frame> frames;
+    // Spot counts a real 18 MP frame produces, scaled to this detector: a weak still, a good
+    // rotation frame, a heavily-diffracting one.
+    frames.push_back(BlobFrame("blobs_sparse", 60, 11));
+    frames.push_back(BlobFrame("blobs_typical", 600, 22));
+    frames.push_back(BlobFrame("blobs_dense", 2500, 33));
+    return frames;
+}
+
+std::vector<Frame> PathologicalFrames() {
+    std::vector<Frame> frames;
+    frames.push_back(BandFrame("band_2_rows", 2));
+    frames.push_back(BandFrame("band_8_rows", 8));
+    frames.push_back(BandFrame("band_31_rows", 31));       // just under the 65534-strong-pixel cap
+    frames.push_back(SolidBlockFrame("solid_block_120", 120));
+    frames.push_back(SolidBlockFrame("solid_block_240", 240));
+    frames.push_back(RingFrame("ring_r700_w3", 700, 1, 1));
+    frames.push_back(RingFrame("ring_r700_w9", 700, 4, 1));
+    frames.push_back(RingFrame("ring_r300_w3_x4", 300, 1, 4));
+    frames.push_back(RingFrame("ring_r950_w5", 950, 2, 1));
+    frames.push_back(SparseNoiseFrame("noise_20k", 20000, 44));
+    return frames;
+}
+
+} // namespace
+
+TEST_CASE("SpotExtractorGPU_Parity_RealisticFrames", "[SpotExtractorGPU]") {
+    if (get_gpu_count() == 0) {
+        WARN("No CUDA GPU present. Skipping SpotExtractorGPU_Parity_RealisticFrames");
+        return;
+    }
+    ParityRunner runner;
+    const auto settings = ExtractionSettings();
+    for (const auto &f : RealisticFrames()) {
+        runner.Load(f);
+        RequireIdentical(f.name + " (" + std::to_string(f.index.size()) + " strong pixels)",
+                         runner.Cpu(settings), runner.Gpu(settings));
+    }
+}
+
+// The occupancy range the extractor has to hold up over, from a nearly empty frame to one pixel
+// short of the give-up threshold, plus the first frame past it (where both sides must return
+// nothing rather than disagree).
+TEST_CASE("SpotExtractorGPU_Parity_OccupancySweep", "[SpotExtractorGPU]") {
+    if (get_gpu_count() == 0) {
+        WARN("No CUDA GPU present. Skipping SpotExtractorGPU_Parity_OccupancySweep");
+        return;
+    }
+    ParityRunner runner;
+    const auto settings = ExtractionSettings();
+    for (int target : {100, 300, 1000, 3000, 10000, 30000, 65000, 70000}) {
+        // Isolated pixels give the most components per strong pixel, so this is also the worst case
+        // for the labelling and the compaction.
+        Frame f = SparseNoiseFrame("occupancy_" + std::to_string(target), target, 900u + target);
+        runner.Load(f);
+        RequireIdentical(f.name + " (" + std::to_string(f.index.size()) + " strong pixels)",
+                         runner.Cpu(settings), runner.Gpu(settings));
+    }
+}
+
+TEST_CASE("SpotExtractorGPU_Parity_Pathological", "[SpotExtractorGPU]") {
+    if (get_gpu_count() == 0) {
+        WARN("No CUDA GPU present. Skipping SpotExtractorGPU_Parity_Pathological");
+        return;
+    }
+    ParityRunner runner;
+    for (const auto &f : PathologicalFrames()) {
+        runner.Load(f);
+        // Both the production max-pix and an effectively unlimited one: the second forces the big
+        // components to be summed as well, which the first deliberately skips.
+        for (int64_t max_pix : {int64_t(50), int64_t(1000000)}) {
+            const auto settings = ExtractionSettings(max_pix);
+            RequireIdentical(f.name + " (" + std::to_string(f.index.size()) + " strong pixels, max_pix "
+                                     + std::to_string(max_pix) + ")",
+                             runner.Cpu(settings), runner.Gpu(settings));
+        }
+    }
+}
+
+TEST_CASE("SpotExtractorGPU_Parity_ResolutionMask", "[SpotExtractorGPU]") {
+    if (get_gpu_count() == 0) {
+        WARN("No CUDA GPU present. Skipping SpotExtractorGPU_Parity_ResolutionMask");
+        return;
+    }
+    ParityRunner runner;
+    const auto settings = ExtractionSettings();
+    // A resolution band, the shape UpdateMaskResolution really produces: excluding it has to split
+    // and drop components identically on both sides.
+    std::vector<bool> mask(static_cast<size_t>(WIDTH) * HEIGHT, false);
+    const double cx = WIDTH * 0.5, cy = HEIGHT * 0.5;
+    for (int line = 0; line < HEIGHT; line++)
+        for (int col = 0; col < WIDTH; col++) {
+            const double d = std::sqrt((col - cx) * (col - cx) + (line - cy) * (line - cy));
+            if (d < 250.0 || d > 900.0)
+                mask[static_cast<size_t>(line) * WIDTH + col] = true;
+        }
+    runner.SetResolutionMask(mask);
+    for (const auto &f : RealisticFrames()) {
+        runner.Load(f);
+        RequireIdentical(f.name + " with resolution mask", runner.Cpu(settings), runner.Gpu(settings));
+    }
+    runner.Load(RingFrame("ring_r700_w3_masked", 700, 1, 1));
+    RequireIdentical("ring with resolution mask", runner.Cpu(settings), runner.Gpu(settings));
+}
+
+// The extraction is full of atomics and a lock-free union-find, so the same frame has to come back
+// byte-identical however the blocks happen to interleave. Anything less and the spot list - and
+// every result downstream of it - would wander between runs of the same data.
+TEST_CASE("SpotExtractorGPU_Determinism", "[SpotExtractorGPU]") {
+    if (get_gpu_count() == 0) {
+        WARN("No CUDA GPU present. Skipping SpotExtractorGPU_Determinism");
+        return;
+    }
+    ParityRunner runner;
+    const auto settings = ExtractionSettings();
+    for (const auto &f : {BlobFrame("blobs_typical", 600, 22), RingFrame("ring_r700_w3", 700, 1, 1),
+                          BandFrame("band_8_rows", 8)}) {
+        runner.Load(f);
+        const auto reference = runner.Gpu(settings);
+        for (int repeat = 0; repeat < 100; repeat++) {
+            const auto again = runner.Gpu(settings);
+            INFO(f.name << ", repeat " << repeat);
+            REQUIRE(again.size() == reference.size());
+            REQUIRE(memcmp(again.data(), reference.data(), again.size() * sizeof(DiffractionSpot)) == 0);
+        }
+    }
+}
+
+#endif