Jungfraujoch/tests/AdaptiveSpotFinderGPUTest.cpp

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"

#ifdef JFJOCH_USE_CUDA

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstring>

#include "../common/AzimuthalIntegrationMapping.h"
#include "../common/AzimuthalIntegrationProfile.h"
#include "../image_analysis/azint/AzIntEngineCPU.h"
#include "../image_analysis/azint/AzIntEngineGPU.h"
#include "../image_analysis/spot_finding/AdaptiveSpotFinderCPU.h"
#include "../image_analysis/spot_finding/AdaptiveSpotFinderGPU.h"
#include "../image_analysis/spot_finding/ImageSpotFinderGPU.h"
#include "../image_analysis/spot_finding/AdaptiveThreshold.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"

namespace {

// Build a realistic full-detector azimuthal-integration mapping (JF4M, ~4.5 MP) whose q-range spans
// most of the detector, so the timing runs over a representative pixel count.
DiffractionExperiment MakeExperiment() {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(80).BeamX_pxl(1030).BeamY_pxl(1080);
    x.QSpacingForAzimInt_recipA(0.05).QRangeForAzimInt_recipA(0.05, 5.0);
    return x;
}

// Deterministic image: a low, slightly rippled background (well below any adaptive threshold) plus a
// grid of bright multi-pixel blobs that both finders must recover identically.
void FillTestImage(ImagePreprocessorBuffer &buffer, const DiffractionExperiment &x) {
    const size_t w = x.GetXPixelsNum();
    const size_t h = x.GetYPixelsNum();
    for (size_t i = 0; i < w * h; i++)
        buffer[i] = 8 + static_cast<int32_t>(i % 5);   // background 8..12 (mean 10)

    // Bright 3x3 blobs on a coarse grid, kept clear of the edges and the beam centre.
    for (size_t row = 300; row < h - 300; row += 450) {
        for (size_t col = 300; col < w - 300; col += 450) {
            for (int dr = -1; dr <= 1; dr++)
                for (int dc = -1; dc <= 1; dc++)
                    buffer[(row + dr) * w + (col + dc)] = 200;
        }
    }
}

SpotFindingSettings AdaptiveSettings() {
    SpotFindingSettings s{};
    s.adaptive_threshold = true;
    s.false_pixels_per_frame = 100.0f;
    s.min_pix_per_spot = 1;
    s.max_pix_per_spot = 50;
    s.high_resolution_limit = 0.0f;      // no resolution gate for the parity test
    s.low_resolution_limit = 1.0e6f;
    s.high_res_gap_Q_recipA = std::nullopt;
    return s;
}

std::vector<std::pair<int, int>> SortedCoords(const std::vector<DiffractionSpot> &spots) {
    std::vector<std::pair<int, int>> out;
    out.reserve(spots.size());
    for (const auto &s : spots)
        out.emplace_back(static_cast<int>(std::lround(s.RawCoord().y)),
                         static_cast<int>(std::lround(s.RawCoord().x)));
    std::sort(out.begin(), out.end());
    return out;
}

} // namespace

// Spot-finding functionality: the fused GPU engine must reproduce the reference CPU adaptive finder's
// spot list. The two share AdaptiveThreshold.h and the host connected-component extractor. They do
// NOT accumulate identically: the CPU sums each ring serially in double, while the GPU stages a
// block's contribution in float before reducing across blocks in double (see the comment on the
// kernel). So ring sigma can differ in the last bits, and since detection compares integer pixel
// values against the threshold, a threshold that crosses an integer flips every pixel of that value
// in the ring at once. That is the difference this test is bounding.
TEST_CASE("AdaptiveSpotFinderGPU_SpotFindingParity", "[AdaptiveSpotFinderGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_SpotFindingParity");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    FillTestImage(buffer, x);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    // The engines run on non-blocking streams, which do not wait for this NULL-stream copy: a pageable
    // H2D cudaMemcpy returns once the source is staged, with the DMA still in flight.
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    std::vector<bool> res_mask(x.GetPixelsNum(), false);
    const SpotFindingSettings settings = AdaptiveSettings();

    AdaptiveSpotFinderCPU cpu(mapping);
    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderGPU gpu(mapping, stream);

    cpu.SetResolutionMask(res_mask);
    gpu.SetResolutionMask(res_mask);
    const auto cpu_spots = cpu.Run(buffer, settings);
    const auto gpu_spots = gpu.Run(buffer, settings);

    INFO("cpu spots=" << cpu_spots.size() << " gpu spots=" << gpu_spots.size());
    REQUIRE(cpu_spots.size() > 0);
    REQUIRE(cpu_spots.size() == gpu_spots.size());
    CHECK(SortedCoords(cpu_spots) == SortedCoords(gpu_spots));
}

// Azimuthal-integration functionality: the profile the fused engine computes as a byproduct of the
// same pass must match a standalone GPU azimuthal integrator over the same image.
TEST_CASE("AdaptiveSpotFinderGPU_AzimuthalIntegration", "[AdaptiveSpotFinderGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_AzimuthalIntegration");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    FillTestImage(buffer, x);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    // The engines run on non-blocking streams, which do not wait for this NULL-stream copy: a pageable
    // H2D cudaMemcpy returns once the source is staged, with the DMA still in flight.
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    std::vector<bool> res_mask(x.GetPixelsNum(), false);
    const SpotFindingSettings settings = AdaptiveSettings();

    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderGPU gpu(mapping, stream);
    gpu.Run(buffer, settings);

    AzIntEngineGPU azint(mapping, stream);
    AzimuthalIntegrationProfile ref_profile(mapping);
    azint.Run(buffer, ref_profile);

    const auto ref = ref_profile.GetResult();
    const auto got = gpu.GetProfile().GetResult();
    const auto ref_count = ref_profile.GetPixelCount();
    const auto got_count = gpu.GetProfile().GetPixelCount();
    REQUIRE(ref.size() == got.size());
    REQUIRE(ref_count == got_count);   // identical per-ring pixel counts (same valid-pixel binning)
    for (size_t b = 0; b < ref.size(); b++) {
        if (std::isnan(ref[b])) {
            CHECK(std::isnan(got[b]));
        } else {
            CHECK(got[b] == Catch::Approx(ref[b]).epsilon(0.01).margin(0.02));
        }
    }
}

// The GPU azimuthal integration against the CPU one, on a pixel count that is not a multiple of four (the
// kernel reads four pixels at a time and does the rest one by one) and with masked and saturated pixels
// in it. The per-ring pixel counts are integers and must agree exactly; the float sums only to rounding.
TEST_CASE("AzIntEngineGPU_MatchesCPU", "[AdaptiveSpotFinderGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AzIntEngineGPU_MatchesCPU");
        return;
    }

    DiffractionExperiment x(DetDECTRIS(1031, 1063, "Test", {}));
    x.DetectorDistance_mm(80).BeamX_pxl(515).BeamY_pxl(530);
    x.QSpacingForAzimInt_recipA(0.05).QRangeForAzimInt_recipA(0.05, 5.0);
    REQUIRE(x.GetPixelsNum() % 4 != 0);
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    for (size_t i = 0; i < x.GetPixelsNum(); i++)
        buffer[i] = (i % 997 == 0) ? INT32_MIN : (i % 1009 == 0) ? INT32_MAX
                                                                 : 8 + static_cast<int32_t>((i * 7919) % 23);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    AzimuthalIntegrationProfile cpu_profile(mapping), gpu_profile(mapping);
    AzIntEngineCPU(mapping).Run(buffer, cpu_profile);
    AzIntEngineGPU gpu_engine(mapping, std::make_shared<CudaStream>());
    gpu_engine.Run(buffer, gpu_profile);

    // The GPU sums are fixed point, so a re-run on the same image repeats them exactly.
    for (int repeat = 0; repeat < 20; repeat++) {
        AzimuthalIntegrationProfile again(mapping);
        gpu_engine.Run(buffer, again);
        const auto a = again.GetResult(), g = gpu_profile.GetResult();
        REQUIRE(a.size() == g.size());
        REQUIRE(std::memcmp(a.data(), g.data(), a.size() * sizeof(float)) == 0);
    }

    REQUIRE(gpu_profile.GetPixelCount() == cpu_profile.GetPixelCount());
    const auto ref = cpu_profile.GetResult();
    const auto got = gpu_profile.GetResult();
    REQUIRE(ref.size() == got.size());
    for (size_t b = 0; b < ref.size(); b++) {
        if (std::isnan(ref[b]))
            CHECK(std::isnan(got[b]));
        else
            CHECK(got[b] == Catch::Approx(ref[b]).epsilon(1e-5));
    }
}

// The ring sums are built by atomics, which arrive in an arbitrary order, so the same frame has to be
// re-run to show the engine agrees with itself: detection is a hard "value >= threshold" on integer
// counts, and a threshold that wobbles between runs flips pixels on the boundary and with them the size
// of a connected component. Two runs, same spot list.
TEST_CASE("AdaptiveSpotFinderGPU_RunToRunReproducible", "[AdaptiveSpotFinderGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_RunToRunReproducible");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    FillTestImage(buffer, x);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    // The engines run on non-blocking streams, which do not wait for this NULL-stream copy: a pageable
    // H2D cudaMemcpy returns once the source is staged, with the DMA still in flight.
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    std::vector<bool> res_mask(x.GetPixelsNum(), false);
    const SpotFindingSettings settings = AdaptiveSettings();

    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderGPU gpu(mapping, stream);

    // The ring accumulators are exact integers, so the threshold does not depend on the order the
    // block atomics arrive in and the spot list has to be bit-identical every time - not merely
    // close. Repeat enough times to give a scheduling-dependent threshold a chance to show itself:
    // the effect it used to have was ~1 changed observation in a million, so a handful of repeats on
    // a quiet background would not have caught it.
    // The azimuthal profile computed in the same pass has to repeat bit for bit as well: its
    // corrected sums are fixed point for that reason. Compared as bytes, since empty rings are NaN.
    const auto first = gpu.Run(buffer, settings);
    REQUIRE(first.size() > 0);
    const auto first_profile = gpu.GetProfile().GetResult();
    const auto first_std = gpu.GetProfile().GetStd();
    for (int repeat = 0; repeat < 50; repeat++) {
        const auto again = gpu.Run(buffer, settings);
        REQUIRE(again.size() == first.size());
        REQUIRE(SortedCoords(again) == SortedCoords(first));
        const auto profile = gpu.GetProfile().GetResult();
        const auto std_dev = gpu.GetProfile().GetStd();
        REQUIRE(std::memcmp(profile.data(), first_profile.data(), profile.size() * sizeof(float)) == 0);
        REQUIRE(std::memcmp(std_dev.data(), first_std.data(), std_dev.size() * sizeof(float)) == 0);
    }
}

// The threshold is computed from sums of int32 pixel values, so the two engines can agree EXACTLY
// rather than approximately - and that is the property worth locking, because it is what makes the
// GPU path's spot list independent of how the reduction happened to be scheduled.
TEST_CASE("AdaptiveSpotFinderGPU_RingStatsMatchCPUExactly", "[AdaptiveSpotFinderGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_RingStatsMatchCPUExactly");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    FillTestImage(buffer, x);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    const SpotFindingSettings settings = AdaptiveSettings();

    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderGPU gpu(mapping, stream);
    AdaptiveSpotFinderCPU cpu(mapping);

    const auto gpu_spots = gpu.Run(buffer, settings);
    const auto cpu_spots = cpu.Run(buffer, settings);
    REQUIRE(gpu_spots.size() == cpu_spots.size());
    REQUIRE(SortedCoords(gpu_spots) == SortedCoords(cpu_spots));
}

// The per-ring background the GPU engine reported before its sigma clips were taken from a value
// histogram: a plain pass and two 3-sigma clip passes, each over every pixel of the image, with the
// clip bounds as the fused multiply-adds the image-pass kernel computes them with.
std::vector<float> ImagePassRingBackground(const ImagePreprocessorBuffer &image,
                                           const AzimuthalIntegrationMapping &mapping) {
    const auto &pixel_to_bin = mapping.GetPixelToBin();
    const size_t nbins = mapping.GetBinNumber();
    std::vector<float> mean(nbins, 0.0f), sigma(nbins, 0.0f);
    std::vector<unsigned long long> sum(nbins), sum2(nbins);
    std::vector<uint32_t> count(nbins);
    for (int pass = 0; pass < 3; pass++) {
        std::fill(sum.begin(), sum.end(), 0);
        std::fill(sum2.begin(), sum2.end(), 0);
        std::fill(count.begin(), count.end(), 0);
        for (size_t i = 0; i < pixel_to_bin.size(); i++) {
            const int32_t v = image[i];
            const uint16_t b = pixel_to_bin[i];
            if (v == INT32_MIN || v == INT32_MAX || b >= nbins) continue;
            const float fv = static_cast<float>(v);
            if (pass > 0 && (fv < std::fma(-3.0f, sigma[b], mean[b]) || fv > std::fma(3.0f, sigma[b], mean[b])))
                continue;
            sum[b] += static_cast<unsigned long long>(static_cast<long long>(v));
            sum2[b] += static_cast<unsigned long long>(static_cast<long long>(v) * v);
            count[b] += 1;
        }
        for (size_t b = 0; b < nbins; b++) {
            if (count[b] == 0) continue;
            const double m = static_cast<double>(static_cast<long long>(sum[b])) / count[b];
            const double var = std::max(0.0, static_cast<double>(sum2[b]) / count[b] - m * m);
            mean[b] = static_cast<float>(m);
            sigma[b] = static_cast<float>(std::sqrt(var));
        }
    }
    std::vector<float> bkg(nbins);
    for (size_t b = 0; b < nbins; b++)
        bkg[b] = (count[b] < adaptive_threshold::MIN_RING_PIXELS) ? NAN : mean[b];
    return bkg;
}

// The sigma clips are taken from the plain pass's per-ring value histogram, with the values outside it
// on an overflow list - and, when that list does not fit, from the image after all. Every route has to
// give the ring statistics the image passes give, to the bit: they are integer sums of the same values.
TEST_CASE("AdaptiveSpotFinderGPU_HistogramClipMatchesImagePasses", "[AdaptiveSpotFinderGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_HistogramClipMatchesImagePasses");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);
    const auto &pixel_to_bin = mapping.GetPixelToBin();
    const size_t npix = x.GetPixelsNum();

    const SpotFindingSettings settings = AdaptiveSettings();
    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderGPU gpu(mapping, stream);

    const auto check = [&](ImagePreprocessorBufferGPU &buffer) {
        REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                           npix * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
        REQUIRE(cudaDeviceSynchronize() == cudaSuccess);
        gpu.Run(buffer, settings);
        const auto expected = ImagePassRingBackground(buffer, mapping);
        const auto &got = gpu.GetRingBackground();
        REQUIRE(got.size() == expected.size());
        size_t differing = 0, rings = 0;
        for (size_t b = 0; b < got.size(); b++) {
            if (std::isnan(expected[b])) {
                if (!std::isnan(got[b])) differing++;
            } else {
                rings++;
                if (got[b] != expected[b]) differing++;
            }
        }
        CHECK(rings > 10);
        CHECK(differing == 0);
    };
    const auto count_outside = [&](const ImagePreprocessorBuffer &buffer) {
        size_t n = 0;
        for (size_t i = 0; i < npix; i++)
            if (pixel_to_bin[i] < mapping.GetBinNumber() && buffer[i] != INT32_MIN && buffer[i] != INT32_MAX
                && (buffer[i] < 0 || buffer[i] >= AdaptiveSpotFinderGPU::HIST_VALUES))
                n++;
        return n;
    };

    ImagePreprocessorBufferGPU buffer(npix);

    SECTION("Values beyond the histogram, on the overflow list") {
        // Two rings far above the histogram's range and one straddling zero, next to rings inside it;
        // with the bright blobs, a few percent of the pixels at most.
        FillTestImage(buffer, x);
        for (size_t i = 0; i < npix; i++) {
            const uint16_t b = pixel_to_bin[i];
            if (b == 20 || b == 21)
                buffer[i] = 1020 + static_cast<int32_t>(i % 13);   // across the top edge
            else if (b == 30)
                buffer[i] = -4 + static_cast<int32_t>(i % 9);      // across zero
        }
        const size_t outside = count_outside(buffer);
        REQUIRE(outside > 0);
        REQUIRE(outside <= npix / 32);
        check(buffer);
    }

    SECTION("Overflow list too long, clipped from the image") {
        for (size_t i = 0; i < npix; i++)
            buffer[i] = 1500 + static_cast<int32_t>((i * 7) % 61) + ((i % 997 == 0) ? 3000 : 0);
        REQUIRE(count_outside(buffer) > npix / 32);
        check(buffer);
    }
}

TEST_CASE("AdaptiveSpotFinderGPU_Speed", "[AdaptiveSpotFinderGPU][.benchmark]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_Speed");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    FillTestImage(buffer, x);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    // The engines run on non-blocking streams, which do not wait for this NULL-stream copy: a pageable
    // H2D cudaMemcpy returns once the source is staged, with the DMA still in flight.
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    std::vector<bool> res_mask(x.GetPixelsNum(), false);
    const SpotFindingSettings settings = AdaptiveSettings();

    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderCPU cpu(mapping);
    AdaptiveSpotFinderGPU gpu_fused(mapping, stream);
    ImageSpotFinderGPU gpu_classic(x.GetXPixelsNum(), x.GetYPixelsNum(), stream);
    AzIntEngineGPU azint(mapping, stream);
    AzimuthalIntegrationProfile profile(mapping);

    const int warmup = 5, iters = 40;
    auto bench = [&](const char *name, auto &&fn) {
        for (int i = 0; i < warmup; i++) fn();
        const auto t0 = std::chrono::steady_clock::now();
        for (int i = 0; i < iters; i++) fn();
        const auto t1 = std::chrono::steady_clock::now();
        const double ms = std::chrono::duration<double, std::milli>(t1 - t0).count() / iters;
        WARN(name << ": " << ms << " ms/frame");
        return ms;
    };

    const double t_azint = bench("GPU azint (standalone)", [&] { azint.Run(buffer, profile); });
    const double t_cpu = bench("CPU adaptive spot finding", [&] { cpu.Run(buffer, settings); });
    const double t_classic = bench("GPU classic spot finding (local-box)", [&] { gpu_classic.Run(buffer, settings); });
    const double t_fused = bench("GPU adaptive FUSED (azint + spot finding)", [&] { gpu_fused.Run(buffer, settings); });

    WARN("standard adaptive path (GPU azint + CPU adaptive) = " << (t_azint + t_cpu)
         << " ms/frame   vs   fused GPU = " << t_fused << " ms/frame   (speedup "
         << (t_azint + t_cpu) / t_fused << "x)");
    WARN("fused GPU vs GPU classic finder alone (no azint): " << t_fused << " vs " << t_classic << " ms/frame");
}

#endif