From f05e67b2b4f8087a17702b5122abdf0d29a7b328 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 14 Oct 2026 23:18:42 +0200 Subject: [PATCH] Metal Bragg integration engine on macOS, output unchanged (bit-identical to the CPU engine) BraggIntegrationEngineMetal is BraggIntegrationEngineCPU with its integer pixel work run as Metal kernels, one thread per reflection: the neighbour mask, the owner map, pass A's walk over every box and the background clip. Everything in double, the profile learning and pass B stay on the CPU, in the CPU engine's own code - Apple GPUs have no double, and no GPU exp matches macOS expf bit for bit. The sums are integers, so the GPU's order does not matter; the stencil arithmetic is the same BraggStencil.h, compiled at run time (MSL 3.0, no offline Metal toolchain needed) with fast math off (sqrt and division correctly rounded, measured) and fp contraction off, as the host already builds (-ffp-contract=off, now extended to Objective-C++). The image is read in place (unified memory). The CPU engine is split so both engines share the finishing code: pass A counts into BraggBoxSums (integer; the double ring sums were exact integers already), the clip threshold is its integer floor, FinishPassA and ProfilePasses are shared. Verified byte-identical p.hkl and MTZ data before/after the split and with the Metal engine on lysozyme, cytochrome C, myoglobin, 4nwv and a small-molecule set. Measured on an Apple M6 (wall, CPU engine -> Metal): 4nwv 180.5 -> 151.2 s, cytochrome C 72.7 -> 59.8 s, lysozyme 29.2 -> 22.5 s, myoglobin 43.3 -> 41.2 s; the sparse small-molecule set +1.5 s (few reflections per frame leave the GPU idle). CI: a macOS step runs tests/BraggIntegrationEngineMetalTest.cpp, which requires every float of every reflection and every count to match the CPU engine over 19 setting combinations on random crowded frames (negative pixels, exact-pixel ties, hot, masked and overloaded pixels), three frames per engine; it fails without a Metal device and if the GPU path fell back to the CPU. Ten planted kernel bugs were each caught. rugnux --no-gpu hides the Metal device too (gpus_are_hidden). Co-Authored-By: Claude Opus 5.5 --- .gitea/workflows/build_and_test.yml | 10 +- CMakeLists.txt | 13 +- common/CUDAWrapper.cpp | 13 +- common/CUDAWrapper.cu | 4 + common/CUDAWrapper.h | 3 + common/JFJochException.h | 5 +- docs/GPU_DATA_ANALYSIS.md | 12 + image_analysis/MXAnalysisWithoutFPGA.cpp | 11 +- .../BraggIntegrationEngineCPU.cpp | 409 +++++++++--------- .../BraggIntegrationEngineCPU.h | 67 ++- .../BraggIntegrationEngineMetal.h | 47 ++ .../BraggIntegrationEngineMetal.metal | 189 ++++++++ .../BraggIntegrationEngineMetal.mm | 343 +++++++++++++++ image_analysis/bragg_integration/BraggPassA.h | 45 ++ .../bragg_integration/BraggStencil.h | 51 ++- .../bragg_integration/CMakeLists.txt | 27 ++ rugnux/rugnux_cli.cpp | 12 + tests/BraggIntegrationEngineMetalTest.cpp | 265 ++++++++++++ tests/CMakeLists.txt | 1 + 19 files changed, 1304 insertions(+), 223 deletions(-) create mode 100644 image_analysis/bragg_integration/BraggIntegrationEngineMetal.h create mode 100644 image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal create mode 100644 image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm create mode 100644 image_analysis/bragg_integration/BraggPassA.h create mode 100644 tests/BraggIntegrationEngineMetalTest.cpp diff --git a/.gitea/workflows/build_and_test.yml b/.gitea/workflows/build_and_test.yml index b69ad25e5..516f54555 100644 --- a/.gitea/workflows/build_and_test.yml +++ b/.gitea/workflows/build_and_test.yml @@ -269,7 +269,15 @@ jobs: shell: bash run: | cmake --build build-macos --target jfjoch_portable_test -j "$(sysctl -n hw.ncpu)" - cd build-macos/tests && ./jfjoch_portable_test "[portable]" + cd build-macos/tests && ./jfjoch_portable_test "[portable]~[Metal]" + # The Metal Bragg integration engine has to give the CPU engine's result bit for bit + # (BraggIntegrationEngineMetal.h): every reflection's every float, on random crowded frames + # through every integrator, overlap and background mode. Its own step, so a mismatch is named as + # such; it fails rather than skips on a machine without a Metal device, as this is the only job + # that can check it. + - name: Metal integration engine is bit-identical to the CPU engine + shell: bash + run: cd build-macos/tests && ./jfjoch_portable_test "[Metal]" # One build, two artifacts: the viewer .dmg and rugnux as .tar.gz (CPACK_PROJECT_CONFIG_FILE # in CMakeLists.txt gives each generator its components). The XDS plugin is Linux-only. - name: Build viewer DMG and rugnux tgz diff --git a/CMakeLists.txt b/CMakeLists.txt index 1875640bc..54e09879b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -21,6 +21,14 @@ INCLUDE(GNUInstallDirs) SET(CMAKE_CXX_STANDARD 20) SET(CMAKE_CXX_STANDARD_REQUIRED True) +# The Metal Bragg integration engine (image_analysis/bragg_integration) is Objective-C++. Enabled here, +# above every target that links it, as enable_language asks. +IF (APPLE) + ENABLE_LANGUAGE(OBJCXX) + SET(CMAKE_OBJCXX_STANDARD 20) + SET(CMAKE_OBJCXX_STANDARD_REQUIRED True) +ENDIF() + # MSVC gets C++ exception handling only from /EHsc, and CMake puts that flag in the DEFAULT value of # CMAKE_CXX_FLAGS ("/DWIN32 /D_WINDOWS /GR /EHsc"), which is used only when the configure does not # set CMAKE_CXX_FLAGS itself. CI does (-DCMAKE_CXX_FLAGS="/arch:AVX"), and that drops /EHsc: MSVC then @@ -47,9 +55,10 @@ ENDIF() # came out differently from the same source and data. Off, every x86-64 level computes the same bits # (measured: identical p.mtz on nine battery sets, no measurable time cost), and aarch64, where FMA is # always there and GCC contracts by default, follows the same arithmetic. MSVC does not contract under -# its default /fp:precise. Host code of .cu sources gets it through nvcc. +# its default /fp:precise. Host code of .cu sources gets it through nvcc, and the Metal kernels say so in +# their source (BraggIntegrationEngineMetal.mm). IF (NOT MSVC) - ADD_COMPILE_OPTIONS($<$:-ffp-contract=off>) + ADD_COMPILE_OPTIONS($<$:-ffp-contract=off>) ADD_COMPILE_OPTIONS($<$:-Xcompiler=-ffp-contract=off>) ENDIF() # The same for device code. nvcc contracts by default and does it separately for every architecture diff --git a/common/CUDAWrapper.cpp b/common/CUDAWrapper.cpp index 0da306fcf..2c6e404ae 100644 --- a/common/CUDAWrapper.cpp +++ b/common/CUDAWrapper.cpp @@ -3,6 +3,7 @@ #include "CUDAWrapper.h" +#include #include #include @@ -93,7 +94,17 @@ int32_t get_gpu_count() { return 0; } -void hide_gpus() {} +namespace { + std::atomic gpus_hidden{false}; +} + +void hide_gpus() { + gpus_hidden = true; +} + +bool gpus_are_hidden() { + return gpus_hidden; +} std::vector get_gpu_names() { return {}; diff --git a/common/CUDAWrapper.cu b/common/CUDAWrapper.cu index 4df5a9daa..04923e107 100644 --- a/common/CUDAWrapper.cu +++ b/common/CUDAWrapper.cu @@ -26,6 +26,10 @@ void hide_gpus() { gpus_hidden = true; } +bool gpus_are_hidden() { + return gpus_hidden; +} + int32_t get_gpu_count() { if (gpus_hidden) return 0; diff --git a/common/CUDAWrapper.h b/common/CUDAWrapper.h index 76f5cf23a..5ee6b1d5a 100644 --- a/common/CUDAWrapper.h +++ b/common/CUDAWrapper.h @@ -16,6 +16,9 @@ int32_t get_gpu_count(); // CUDA_VISIBLE_DEVICES= does, asked for by the program (rugnux --no-gpu). Call it before anything uses // CUDA; there is no way back. void hide_gpus(); +// Whether hide_gpus() was called - for a GPU path that is not CUDA's (the Metal integration engine), so +// --no-gpu keeps the process off every GPU. +bool gpus_are_hidden(); // Names of the visible GPUs, in device order and one entry per device, so repeated cards repeat. // Empty without CUDA and on a machine with no device, which is also what get_gpu_count() == 0 says. diff --git a/common/JFJochException.h b/common/JFJochException.h index fa3c64809..50882528c 100644 --- a/common/JFJochException.h +++ b/common/JFJochException.h @@ -45,7 +45,8 @@ UDPError, CommunicationError, SimplonError, CalibrationError, -IndexerError +IndexerError, +GPUMetalError }; class JFJochException : public std::exception { @@ -128,6 +129,8 @@ protected: return "Detector calibration error"; case JFJochExceptionCategory::IndexerError: return "Indexer error"; + case JFJochExceptionCategory::GPUMetalError: + return "Metal (GPU) error"; default: return ""; } diff --git a/docs/GPU_DATA_ANALYSIS.md b/docs/GPU_DATA_ANALYSIS.md index b48b2d260..8fab936d7 100644 --- a/docs/GPU_DATA_ANALYSIS.md +++ b/docs/GPU_DATA_ANALYSIS.md @@ -180,3 +180,15 @@ The placements are then refined, scored and checked against their null with the structure factors above; the likelihood scoring runs on the host. The search is described in [§14.10](CPU_DATA_ANALYSIS_DECISIONS.md#1410-molecular-replacement-placing-a-model-the-rigid-body-cannot), its keys and decision rules in [Molecular replacement](RUGNUX_REPORT.md#molecular-replacement). + +## Apple GPUs (Metal) + +A macOS build has no CUDA, and one engine runs on the Apple GPU instead: + +- **Bragg integration** (`BraggIntegrationEngineMetal.mm`, `.metal`) - the CPU engine with its + integer pixel work moved to the GPU: the neighbour mask, the ownership of shared pixels, the box + summation's walk over every reflection and the background clip. Everything computed in double, the + profile learning and the profile fit stay on the CPU, as Apple GPUs have no double precision. The + same result as the CPU engine **to the bit**, which a macOS CI step checks on every push; + `rugnux --no-gpu` keeps it off the GPU. + [§9.2](CPU_DATA_ANALYSIS_INTEGRATION.md#92-box-summation-seed-and-fallback) diff --git a/image_analysis/MXAnalysisWithoutFPGA.cpp b/image_analysis/MXAnalysisWithoutFPGA.cpp index 67a29a0ea..61b9884f2 100644 --- a/image_analysis/MXAnalysisWithoutFPGA.cpp +++ b/image_analysis/MXAnalysisWithoutFPGA.cpp @@ -19,6 +19,9 @@ #include "spot_finding/ImageSpotFinderCPU.h" #include "spot_finding/AdaptiveSpotFinderCPU.h" #include "bragg_integration/BraggIntegrationEngineCPU.h" +#ifdef __APPLE__ +#include "bragg_integration/BraggIntegrationEngineMetal.h" +#endif #ifdef JFJOCH_USE_CUDA #include "azint/AzIntEngineGPU.h" #include "roi/ROIIntegrationGPU.h" @@ -53,7 +56,13 @@ MXAnalysisWithoutFPGA::MXAnalysisWithoutFPGA(const DiffractionExperiment &in_exp auto cpu_preprocessor = std::make_unique(in_experiment, in_mask); preprocessor_cpu = cpu_preprocessor.get(); preprocessor = std::move(cpu_preprocessor); - bragg_engine = std::make_unique(in_experiment, in_mask); +#ifdef __APPLE__ + // The CPU engine with its integer pixel work on the Apple GPU: the same result, bit for bit. + if (BraggIntegrationEngineMetal::Available(in_experiment)) + bragg_engine = std::make_unique(in_experiment, in_mask); + else +#endif + bragg_engine = std::make_unique(in_experiment, in_mask); if (experiment.ROI().size() >= 1) roi = std::make_unique(experiment); #ifdef JFJOCH_USE_CUDA diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp index 4b7f785a9..b488aff79 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp @@ -36,13 +36,6 @@ inline void prefetch_row(const void *first, const void *last) { prefetch(last); } -// Identity sampler over the preprocessed int32 buffer (already in that convention). -struct BufferSampler { - const int32_t *p; - int32_t operator[](size_t i) const { return p[i]; } - const void *address(size_t i) const { return p + i; } -}; - // Sampler over a raw detector image of pixel type T: masked pixels carry the type minimum, saturated // the type maximum (the FPGA image has no lossy-codec +/-1 band). Only the pixels actually read - the // reflection disks - are converted, so there is no whole-image pass. @@ -65,9 +58,9 @@ struct ImageSampler { BraggIntegrationEngineCPU::BraggIntegrationEngineCPU(const DiffractionExperiment &experiment, const PixelMask &mask) : BraggIntegrationEngine(experiment), - static_mask(mask.GetPackedMask()), refl_mask(static_cast(xpixel), static_cast(ypixel), 0), - owner(static_cast(xpixel), static_cast(ypixel), BRAGG_OWNER_NONE) {} + owner(static_cast(xpixel), static_cast(ypixel), BRAGG_OWNER_NONE), + static_mask(mask.GetPackedMask()) {} void BraggIntegrationEngineCPU::MarkNeighbours(const std::vector &predicted, size_t npredicted) { const int W = static_cast(xpixel), H = static_cast(ypixel); @@ -133,23 +126,13 @@ template std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, const std::vector &predicted, size_t npredicted, int64_t image_number) { - std::vector results(npredicted); if (npredicted == 0) - return Finalize(predicted, npredicted, results, image_number); + return Finalize(predicted, npredicted, std::vector(npredicted), image_number); counts.predicted += npredicted; const int W = static_cast(xpixel), H = static_cast(ypixel); const bool do_clip = bkg_clip_nsigma > 0.0f && mode != IntegratorMode::BoxSum; - // Symmetric trimmed-mean background fraction (BraggIntegrationSettings, rugnux --background-trim): - // replaces the r2..r3 ring MEAN with an f-trimmed mean, robust to the high-side contamination - // (neighbour wings, tails) that biases the plain mean up and makes it over-subtract weak high-angle - // reflections. The base ctor has already forced it to 0 whenever the clip below is in force, which - // is the default - the two are alternatives. - const double bkg_trim_frac = bkg_trim; - - auto grid_idx = [this](int dx, int dy) { return (dy + R) * G + (dx + R); }; - MarkNeighbours(predicted, npredicted); owner.Clear(); @@ -184,20 +167,6 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, }; // --- Pass A: box-sum every reflection (rough I, background, centroid, strong flag). --- - struct Rough { - double I = 0.0, sigma = NAN, bkg = 0.0, obs_x = 0.0, obs_y = 0.0; - double bkg_var = 0.0; // variance of the background ESTIMATE itself, bkg / n_bkg - double var_bkg = 0.0; // total non-signal variance carried to the merge - int64_t I_sum = 0; // kept so I can be rebuilt after the radial background correction - int n_inner = 0; - int n_disk = 0, n_own = 0; // signal-disk pixels, and how many of them are this reflection's - int r_bin = 0; // rounded distance from the beam centre, indexes the radial curve - int k_bin = 0; // which radial-background kernel this reflection's stencil needs - int cx = 0, cy = 0, shell = -1; - bool ok = false, strong = false, has_obs = false; - bool full = false; // every pixel of the signal disk was readable - bool overloaded = false; // a pixel of the signal disk was saturated - }; // A saturated pixel, or one unreadable on this frame alone - the overload marker of a detector // that writes its error value for a pixel it could not count (EIGER), which the preprocessor turns // into a masked pixel like a gap's. The run's mask holds the gaps, so whatever is unreadable beyond @@ -205,26 +174,13 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, auto static_masked = [&](size_t idx) { return idx / 32 < static_mask.size() && ((static_mask[idx / 32] >> (idx % 32)) & 1U); }; - std::vector rough(npredicted); - double inv_d2_min = std::numeric_limits::max(), inv_d2_max = 0.0; + PassA a = StartPassA(npredicted); std::vector bkg_vals; // reused per reflection for the trimmed-mean background (idea 1) // The ring pixels the background clip below reads, with their radial offsets, gathered while pass - // A walks the box anyway - in the order the clip sums them, so it sums the same values. + // A walks the box anyway. struct ClipPixel { int32_t px; float rad; }; std::vector clip_ring; - // Radial background curve, accumulated from the annulus pixels this pass already reads. A pixel's - // radius is the reflection's radius plus the pixel's projection on the beam->reflection direction, - // so no per-pixel sqrt is needed. - const int n_rad = static_cast(std::ceil(std::hypot(std::max(beam_x, W - beam_x), - std::max(beam_y, H - beam_y)))) + 2; - std::vector rad_sum; - std::vector rad_cnt; - if (bkg_radial) { - rad_sum.assign(n_rad, 0.0); - rad_cnt.assign(n_rad, 0); - } - // The box pass A reads: the OUTER ellipse, tightly - taking the largest semi-axis in both // directions instead would read up to 60% more pixels for nothing. struct Box { int x0, x1, y0, y1; }; @@ -247,28 +203,18 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, img.address(static_cast(y) * W + b.x1)); } const auto &r = predicted[i]; - Rough out; // This reflection's stencil: the r1 signal disk, and the r2..r3 ring elongated radially by // the analytic smear. const BraggStencil st = MakeBraggStencil(r.predicted_x, r.predicted_y, stencil); const Box box = box_of(r, st); const int x0 = box.x0, x1 = box.x1, y0 = box.y0, y1 = box.y1; - // Both from the stencil's own radius, so nothing downstream is derived from a second one. - out.r_bin = std::clamp(static_cast(std::lround(st.r0)), 0, n_rad - 1); - out.k_bin = BraggStencilKernelIndex(st, n_kern); - - int64_t I_sum = 0, I_sum_x = 0, I_sum_y = 0, n_inner = 0, n_inner_valid = 0; - int64_t x_sum = 0, y_sum = 0; // positions of the pixels behind I_sum, for the centroid - int n_disk = 0, n_own = 0; // pixels in the signal disk, and how many are this reflection's - double bkg_sum = 0.0; - int n_bkg = 0; - // Readable ring pixels, and those of them no predicted neighbour's signal region covers. - // refl_mask marks d.inner < r2_sq and the ring is d.inner >= r2_sq, so the two are - // complementary and a masked ring pixel is always someone else's - never this reflection's + // n_ring / n_clean: readable ring pixels, and those of them no predicted neighbour's signal + // region covers. refl_mask marks d.inner < r2_sq and the ring is d.inner >= r2_sq, so the two + // are complementary and a masked ring pixel is always someone else's - never this reflection's // own core. That makes the difference an exact count of what the pattern's density took, // separable from what the detector took. - int n_ring = 0, n_clean = 0; + BraggBoxSums s{}; bkg_vals.clear(); clip_ring.clear(); for (int y = y0; y <= y1; ++y) @@ -282,154 +228,223 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, // mode it leaves the disk entirely - the sum, the pixel count the background is // subtracted with, and the all-or-nothing validity gate alike. The box sum only // counts how many were lost, which is all it can act on. - ++n_disk; + ++s.n_disk; if (overlap != OverlapMode::Off) { - if (clean(x, y, i)) ++n_own; + if (clean(x, y, i)) ++s.n_own; else if (exclude) continue; } - ++n_inner; + ++s.n_inner; if (px == INT32_MAX || (px == INT32_MIN && !static_masked(static_cast(y) * W + x))) - out.overloaded = true; + s.overloaded = 1; if (!valid(px)) continue; - I_sum += px; - I_sum_x += static_cast(x) * px; - I_sum_y += static_cast(y) * px; + s.I_sum += px; + s.I_sum_x += static_cast(x) * px; + s.I_sum_y += static_cast(y) * px; // Position sums over the very same pixels, so the background can be taken out of // the centroid below. It is not known yet - the ring is read in this same loop. - x_sum += x; - y_sum += y; - ++n_inner_valid; + s.x_sum += x; + s.y_sum += y; + ++s.n_inner_valid; } else if (d.inner >= r2_sq && d.outer < r3_sq) { if (!valid(px)) continue; - ++n_ring; + ++s.n_ring; const uint8_t m = refl_mask.Get(x, y); - if (m == 0) ++n_clean; + if (m == 0) ++s.n_clean; if (m == MASK_FLUX) continue; - bkg_sum += static_cast(px); - if (bkg_trim_frac > 0.0) bkg_vals.push_back(px); - ++n_bkg; + s.bkg_sum += px; + if (bkg_trim > 0.0f) bkg_vals.push_back(px); + ++s.n_bkg; } } - if (use_owner && n_own < n_disk) { - ++counts.overlapped; - if (n_own < overlap_min_peak * n_disk) ++counts.overlapped_major; - } - - int n_bkg_used = n_bkg; // pixels behind the FINAL background value (trim/clip shrink it) - // A masked, untrusted, gapped or overloaded pixel inside the signal disk used to discard the - // reflection outright. A profile fit does not need it to: the fit is the amplitude of a - // NORMALISED profile, so leaving pixels out renormalises the estimator by construction and - // widens sigma by the information they carried - the same argument Exclude makes for a - // neighbour's pixels. Pass B keeps the reflection only while enough of the expected profile - // survived to constrain the amplitude (XDS's MINPK, dials' valid_foreground_threshold). A box - // sum has no profile to renormalise with, so there it stays all or nothing. - const bool full = n_inner_valid == n_inner; - // A ring left with five or fewer readable pixels cannot estimate a background, so the - // reflection is dropped whole - the one thing the stencil geometry does to the DATA rather - // than to a measurement. Counted here because it is the only direct evidence of a radius that has - // outgrown the pattern it is integrating (BraggIntegrationCounts). - const bool keep_partial = full || mode != IntegratorMode::BoxSum || out.overloaded; - if (keep_partial && n_bkg <= 5) ++counts.bkg_starved; - // Would every predicted neighbour, the tails included, leave the ring starved that the - // detector alone would not? That is the pattern's density, what the guard on a widened - // radius reads (BraggIntegrationCounts). - if (keep_partial && n_clean <= 5 && n_ring > 5) ++counts.bkg_starved_by_neighbour; - if (keep_partial && n_bkg > 5) { - out.bkg = bkg_sum / n_bkg; - if (bkg_trim_frac > 0.0 && bkg_vals.size() > 5 - && bkg_vals.size() <= static_cast(bragg_engine::BKG_TRIM_MAX)) { - // Symmetric trimmed mean over the background ring (idea 1): drop the lowest and highest - // bkg_trim_frac of the pixels, average the rest. Robust to the high-side contamination - // that biases the plain ring mean and makes it over-subtract at high resolution. - std::sort(bkg_vals.begin(), bkg_vals.end()); - const size_t lo = static_cast(bkg_vals.size() * bkg_trim_frac); - const size_t hi = bkg_vals.size() - lo; - if (hi > lo) { - double s = 0.0; - for (size_t t = lo; t < hi; ++t) s += bkg_vals[t]; - out.bkg = s / static_cast(hi - lo); - n_bkg_used = static_cast(hi - lo); - } - } else if (do_clip) { - // One high-outlier sigma-clip pass on the background ring: reject pixels above - // mean + n*sqrt(gain*mean) to strip a neighbour core or a zinger that biases the mean. - // The gain is 1 for counted photons, so this is then exactly the Poisson clip. - const double thr = out.bkg + bkg_clip_nsigma * std::sqrt(static_cast(bkg_pixel_gain) - * std::max(out.bkg, 1.0)); - double s = 0.0; - int n = 0; - for (const auto &[px, rad] : clip_ring) - if (static_cast(px) <= thr) { - s += px; - ++n; - if (bkg_radial) { - // The radial curve is binned on the TRUE detector radius, so the - // offset here is the unshrunk radial projection. - const int b = std::clamp(static_cast(std::lround(st.r0 + rad)), 0, n_rad - 1); - rad_sum[b] += static_cast(px); - ++rad_cnt[b]; - } + // One high-outlier sigma-clip pass on the background ring (see ClipThreshold), and the radial + // background curve from the pixels it keeps. + BraggClipSums clip{}; + if (ClipsRing(s)) { + const int32_t thr = ClipThreshold(s); + for (const auto &[px, rad] : clip_ring) + if (px <= thr) { + clip.sum += px; + ++clip.n; + if (bkg_radial) { + // The radial curve is binned on the TRUE detector radius, so the + // offset here is the unshrunk radial projection. + const int b = std::clamp(static_cast(std::lround(st.r0 + rad)), 0, a.n_rad - 1); + a.rad_sum[b] += px; + ++a.rad_cnt[b]; } - if (n > 5) { out.bkg = s / n; n_bkg_used = n; } - } - // The sum is over the pixels actually READ, so that is the count the background is - // subtracted with; with nothing missing it is the whole disk, exactly as before. - out.I = static_cast(I_sum) - static_cast(n_inner_valid) * out.bkg; - // I = I_sum - n_inner*bkg, and bkg is itself estimated from n_bkg_used pixels, so its - // error enters n_inner times over: var(I) = I_sum + n_inner^2 * bkg/n_bkg_used. Leaving - // the second term out understates sigma by sqrt(1 + n_inner/n_bkg) - 1.109x at the - // default r1=4/r2=6/r3=10 stencil, on every reflection of every dataset. - out.bkg_var = out.bkg / n_bkg_used; - out.var_bkg = static_cast(n_inner_valid) * out.bkg - + static_cast(n_inner_valid) * n_inner_valid * out.bkg_var; - out.I_sum = I_sum; - out.n_inner = static_cast(n_inner_valid); - out.n_disk = n_disk; - out.n_own = n_own; - const double var_bkg_term = static_cast(n_inner_valid) * n_inner_valid * out.bkg_var; - out.sigma = 1.0; - if (I_sum > 0) { - out.sigma = std::max(out.sigma, std::sqrt(static_cast(I_sum) + var_bkg_term)); - // The centroid of the SIGNAL, not of the disk. Weighting by the raw counts weights by - // signal plus background, and the background is flat over a disk centred on the - // PREDICTION - so its own centroid is the prediction exactly, and it pulls the answer - // there. The measured displacement comes out shrunk by I/(I + n*bkg), which is worst - // where the background dominates, i.e. at high resolution. Post-refinement fits the - // beam and the distance on these centroids, so it under-corrects by that factor and - // does so resolution-dependently. Subtracting a flat pedestal from a first moment is - // exact: sum(x*(px-bkg)) = I_sum_x - bkg*x_sum over the same pixels. - const double net = static_cast(I_sum) - n_inner_valid * out.bkg; - if (net > 0.0) { - out.obs_x = (static_cast(I_sum_x) - out.bkg * x_sum) / net; - out.obs_y = (static_cast(I_sum_y) - out.bkg * y_sum) / net; - } else { - // Nothing above background to take a centroid of; the raw one is all there is. - out.obs_x = static_cast(I_sum_x) / static_cast(I_sum); - out.obs_y = static_cast(I_sum_y) / static_cast(I_sum); } - // A disk with a hole in it gives a centroid pulled away from the hole, and the hole - // sits at a fixed place on the detector - post-refinement would read that as geometry. - out.has_obs = full; - } - out.cx = static_cast(std::lround(r.predicted_x)); - out.cy = static_cast(std::lround(r.predicted_y)); - out.ok = true; - out.full = full; - // The profile, its resolution shells and their widths are learned from COMPLETE - // reflections: a partial one has no intensity to normalise its grid by and a hole to - // stack into it. So the learning sees exactly what it saw before this rescue existed. - out.strong = full && out.sigma > 0.0 && out.I / out.sigma >= STRONG_I_OVER_SIGMA; - if (full && r.d > 0.0f) { - const double inv_d2 = 1.0 / (static_cast(r.d) * r.d); - inv_d2_min = std::min(inv_d2_min, inv_d2); - inv_d2_max = std::max(inv_d2_max, inv_d2); - } } - rough[i] = out; + FinishPassA(a, i, r, st, s, clip, bkg_vals); } + return ProfilePasses(img, owner, predicted, npredicted, a, image_number); +} + +BraggIntegrationEngineCPU::PassA BraggIntegrationEngineCPU::StartPassA(size_t npredicted) const { + const int W = static_cast(xpixel), H = static_cast(ypixel); + PassA a; + a.rough.resize(npredicted); + // A pixel's radius is the reflection's radius plus the pixel's projection on the beam->reflection + // direction, so no per-pixel sqrt is needed. + a.n_rad = static_cast(std::ceil(std::hypot(std::max(beam_x, W - beam_x), + std::max(beam_y, H - beam_y)))) + 2; + if (bkg_radial) { + a.rad_sum.assign(a.n_rad, 0); + a.rad_cnt.assign(a.n_rad, 0); + } + a.inv_d2_min = std::numeric_limits::max(); + a.inv_d2_max = 0.0; + return a; +} + +bool BraggIntegrationEngineCPU::ClipsRing(const BraggBoxSums &s) const { + // The clip is the alternative to the trimmed mean, so where it is on the trim is off, and the + // profile modes keep every reflection with a background (FinishPassA). + return bkg_clip_nsigma > 0.0f && mode != IntegratorMode::BoxSum && s.n_bkg > 5; +} + +int32_t BraggIntegrationEngineCPU::ClipThreshold(const BraggBoxSums &s) const { + // Reject pixels above mean + n*sqrt(gain*mean) to strip a neighbour core or a zinger that biases + // the mean. The gain is 1 for counted photons, so this is then exactly the Poisson clip. A pixel + // value is an integer, so it is at or below the threshold exactly when it is at or below its floor. + const double bkg = static_cast(s.bkg_sum) / s.n_bkg; + const double thr = bkg + bkg_clip_nsigma * std::sqrt(static_cast(bkg_pixel_gain) * std::max(bkg, 1.0)); + return static_cast(std::clamp(std::floor(thr), static_cast(INT32_MIN), + static_cast(INT32_MAX))); +} + +void BraggIntegrationEngineCPU::FinishPassA(PassA &a, size_t i, const Reflection &r, const BraggStencil &st, + const BraggBoxSums &s, const BraggClipSums &clip, + std::vector &bkg_vals) { + Rough out; + // Both from the stencil's own radius, so nothing downstream is derived from a second one. + out.r_bin = std::clamp(static_cast(std::lround(st.r0)), 0, a.n_rad - 1); + out.k_bin = BraggStencilKernelIndex(st, n_kern); + out.overloaded = s.overloaded != 0; + + if (overlap != OverlapMode::Off && s.n_own < s.n_disk) { + ++counts.overlapped; + if (s.n_own < overlap_min_peak * s.n_disk) ++counts.overlapped_major; + } + + int n_bkg_used = s.n_bkg; // pixels behind the FINAL background value (trim/clip shrink it) + // A masked, untrusted, gapped or overloaded pixel inside the signal disk used to discard the + // reflection outright. A profile fit does not need it to: the fit is the amplitude of a + // NORMALISED profile, so leaving pixels out renormalises the estimator by construction and + // widens sigma by the information they carried - the same argument Exclude makes for a + // neighbour's pixels. Pass B keeps the reflection only while enough of the expected profile + // survived to constrain the amplitude (XDS's MINPK, dials' valid_foreground_threshold). A box + // sum has no profile to renormalise with, so there it stays all or nothing. + const bool full = s.n_inner_valid == s.n_inner; + // A ring left with five or fewer readable pixels cannot estimate a background, so the + // reflection is dropped whole - the one thing the stencil geometry does to the DATA rather + // than to a measurement. Counted here because it is the only direct evidence of a radius that has + // outgrown the pattern it is integrating (BraggIntegrationCounts). + const bool keep_partial = full || mode != IntegratorMode::BoxSum || out.overloaded; + if (keep_partial && s.n_bkg <= 5) ++counts.bkg_starved; + // Would every predicted neighbour, the tails included, leave the ring starved that the + // detector alone would not? That is the pattern's density, what the guard on a widened + // radius reads (BraggIntegrationCounts). + if (keep_partial && s.n_clean <= 5 && s.n_ring > 5) ++counts.bkg_starved_by_neighbour; + if (keep_partial && s.n_bkg > 5) { + out.bkg = static_cast(s.bkg_sum) / s.n_bkg; + // Symmetric trimmed-mean background fraction (BraggIntegrationSettings, rugnux --background-trim): + // replaces the r2..r3 ring MEAN with an f-trimmed mean, robust to the high-side contamination + // (neighbour wings, tails) that biases the plain mean up and makes it over-subtract weak + // high-angle reflections. The base ctor has already forced it to 0 whenever the clip is in + // force, which is the default - the two are alternatives. + const double bkg_trim_frac = bkg_trim; + if (bkg_trim_frac > 0.0 && bkg_vals.size() > 5 + && bkg_vals.size() <= static_cast(bragg_engine::BKG_TRIM_MAX)) { + // Symmetric trimmed mean over the background ring (idea 1): drop the lowest and highest + // bkg_trim_frac of the pixels, average the rest. Robust to the high-side contamination + // that biases the plain ring mean and makes it over-subtract at high resolution. + std::sort(bkg_vals.begin(), bkg_vals.end()); + const size_t lo = static_cast(bkg_vals.size() * bkg_trim_frac); + const size_t hi = bkg_vals.size() - lo; + if (hi > lo) { + double sum = 0.0; + for (size_t t = lo; t < hi; ++t) sum += bkg_vals[t]; + out.bkg = sum / static_cast(hi - lo); + n_bkg_used = static_cast(hi - lo); + } + } else if (ClipsRing(s)) { + if (clip.n > 5) { out.bkg = static_cast(clip.sum) / clip.n; n_bkg_used = clip.n; } + } + // The sum is over the pixels actually READ, so that is the count the background is + // subtracted with; with nothing missing it is the whole disk, exactly as before. + out.I = static_cast(s.I_sum) - static_cast(s.n_inner_valid) * out.bkg; + // I = I_sum - n_inner*bkg, and bkg is itself estimated from n_bkg_used pixels, so its + // error enters n_inner times over: var(I) = I_sum + n_inner^2 * bkg/n_bkg_used. Leaving + // the second term out understates sigma by sqrt(1 + n_inner/n_bkg) - 1.109x at the + // default r1=4/r2=6/r3=10 stencil, on every reflection of every dataset. + out.bkg_var = out.bkg / n_bkg_used; + out.var_bkg = static_cast(s.n_inner_valid) * out.bkg + + static_cast(s.n_inner_valid) * s.n_inner_valid * out.bkg_var; + out.I_sum = s.I_sum; + out.n_inner = static_cast(s.n_inner_valid); + out.n_disk = s.n_disk; + out.n_own = s.n_own; + const double var_bkg_term = static_cast(s.n_inner_valid) * s.n_inner_valid * out.bkg_var; + out.sigma = 1.0; + if (s.I_sum > 0) { + out.sigma = std::max(out.sigma, std::sqrt(static_cast(s.I_sum) + var_bkg_term)); + // The centroid of the SIGNAL, not of the disk. Weighting by the raw counts weights by + // signal plus background, and the background is flat over a disk centred on the + // PREDICTION - so its own centroid is the prediction exactly, and it pulls the answer + // there. The measured displacement comes out shrunk by I/(I + n*bkg), which is worst + // where the background dominates, i.e. at high resolution. Post-refinement fits the + // beam and the distance on these centroids, so it under-corrects by that factor and + // does so resolution-dependently. Subtracting a flat pedestal from a first moment is + // exact: sum(x*(px-bkg)) = I_sum_x - bkg*x_sum over the same pixels. + const double net = static_cast(s.I_sum) - s.n_inner_valid * out.bkg; + if (net > 0.0) { + out.obs_x = (static_cast(s.I_sum_x) - out.bkg * s.x_sum) / net; + out.obs_y = (static_cast(s.I_sum_y) - out.bkg * s.y_sum) / net; + } else { + // Nothing above background to take a centroid of; the raw one is all there is. + out.obs_x = static_cast(s.I_sum_x) / static_cast(s.I_sum); + out.obs_y = static_cast(s.I_sum_y) / static_cast(s.I_sum); + } + // A disk with a hole in it gives a centroid pulled away from the hole, and the hole + // sits at a fixed place on the detector - post-refinement would read that as geometry. + out.has_obs = full; + } + out.cx = static_cast(std::lround(r.predicted_x)); + out.cy = static_cast(std::lround(r.predicted_y)); + out.ok = true; + out.full = full; + // The profile, its resolution shells and their widths are learned from COMPLETE + // reflections: a partial one has no intensity to normalise its grid by and a hole to + // stack into it. So the learning sees exactly what it saw before this rescue existed. + out.strong = full && out.sigma > 0.0 && out.I / out.sigma >= STRONG_I_OVER_SIGMA; + if (full && r.d > 0.0f) { + const double inv_d2 = 1.0 / (static_cast(r.d) * r.d); + a.inv_d2_min = std::min(a.inv_d2_min, inv_d2); + a.inv_d2_max = std::max(a.inv_d2_max, inv_d2); + } + } + a.rough[i] = out; +} + +template +std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler &img, const OwnerMap &owner_map, + const std::vector &predicted, + size_t npredicted, PassA &a, + int64_t image_number) { + const int W = static_cast(xpixel), H = static_cast(ypixel); + std::vector results(npredicted); + std::vector &rough = a.rough; + const bool exclude = overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum; + const bool use_owner = overlap != OverlapMode::Off; + auto clean = [&](int x, int y, size_t i) { + return !use_owner + || BraggOwnedBy(owner_map.Get(x, y), static_cast(i)); + }; + auto grid_idx = [this](int dx, int dy) { return (dy + R) * G + (dx + R); }; + // --- Radial background curvature correction. The annulus mean is blind to the curvature of the // radial background (a linear background cancels between the concentric disk and annulus), so // correct it by mean_annulus(B) - mean_disk(B) taken from the curve just accumulated. Reads no @@ -443,8 +458,8 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, const float *kern = k_diff.data() + static_cast(rh.k_bin) * k_len; double corr = 0.0; for (int k = 0; k < k_len; ++k) { - const int b = std::clamp(rh.r_bin + k - k_off, 0, n_rad - 1); - const double v = rad_cnt[b] > 0 ? rad_sum[b] / rad_cnt[b] : rh.bkg; + const int b = std::clamp(rh.r_bin + k - k_off, 0, a.n_rad - 1); + const double v = a.rad_cnt[b] > 0 ? static_cast(a.rad_sum[b]) / a.rad_cnt[b] : rh.bkg; corr += static_cast(kern[k]) * v; } rh.bkg -= corr; // annulus mean -> mean over the signal disk @@ -473,8 +488,8 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, } auto shell_of = [&](float d) { - if (!(d > 0.0f) || inv_d2_max <= inv_d2_min) return 0; - const double t = (1.0 / (static_cast(d) * d) - inv_d2_min) / (inv_d2_max - inv_d2_min); + if (!(d > 0.0f) || a.inv_d2_max <= a.inv_d2_min) return 0; + const double t = (1.0 / (static_cast(d) * d) - a.inv_d2_min) / (a.inv_d2_max - a.inv_d2_min); return std::clamp(static_cast(t * N_SHELL), 0, N_SHELL - 1); }; for (size_t i = 0; i < npredicted; ++i) @@ -749,6 +764,10 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, return Finalize(predicted, npredicted, results, image_number); } +template std::vector BraggIntegrationEngineCPU::ProfilePasses( + const BraggIntegrationEngineCPU::BufferSampler &, const BraggIntegrationEngineCPU::FlatOwnerMap &, + const std::vector &, size_t, BraggIntegrationEngineCPU::PassA &, int64_t); + std::vector BraggIntegrationEngineCPU::Run(const ImagePreprocessorBuffer &image, const std::vector &predicted, size_t npredicted, int64_t image_number) { diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h index 95b0a943a..b9ee60116 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h @@ -4,13 +4,14 @@ #pragma once #include "BraggIntegrationEngine.h" +#include "BraggPassA.h" #include "../../common/PixelMask.h" class CompressedImage; // Plain-C++ reference/fallback engine: a faithful serial re-expression of BraggIntegrate2D (box // sum) and ProfileIntegrate2D (Kabsch profile fit) reading the preprocessed int32 image. Also the -// numeric oracle the CUDA engine is checked against. +// numeric oracle the CUDA engine is checked against, and the engine the Metal engine extends. class BraggIntegrationEngineCPU : public BraggIntegrationEngine { // Core integrator, templated on a pixel sampler so it reads either the preprocessed int32 buffer // or a raw CompressedImage of any pixel type - both presented per-pixel in the INT32_MIN(masked)/ @@ -53,9 +54,6 @@ class BraggIntegrationEngineCPU : public BraggIntegrationEngine { }; TiledFrame refl_mask; TiledFrame owner; - // The run's pixel mask, packed 32 pixels to a word (PixelMask::GetPackedMask). An unreadable pixel - // it does not explain was unreadable on this frame only - an overload (see Reflection::overloaded). - std::vector static_mask; // Marks every predicted reflection's r2 region in refl_mask (MASK_FLUX / MASK_TAIL). void MarkNeighbours(const std::vector &predicted, size_t npredicted); @@ -64,6 +62,67 @@ class BraggIntegrationEngineCPU : public BraggIntegrationEngine { std::vector RunImpl(const Sampler &img, const std::vector &predicted, size_t npredicted, int64_t image_number); +protected: + // The run's pixel mask, packed 32 pixels to a word (PixelMask::GetPackedMask). An unreadable pixel + // it does not explain was unreadable on this frame only - an overload (see Reflection::overloaded). + std::vector static_mask; + + // Identity sampler over the preprocessed int32 buffer (already in the engine's convention). + struct BufferSampler { + const int32_t *p; + int32_t operator[](size_t i) const { return p[i]; } + const void *address(size_t i) const { return p + i; } + }; + // The owner map as one frame-sized array, as the Metal engine builds it (TiledFrame's Get otherwise). + struct FlatOwnerMap { + const uint32_t *p; + size_t width; + uint32_t Get(int x, int y) const { return p[static_cast(y) * width + x]; } + }; + + // Pass A's result for one reflection: the box sum, its background, centroid and strong flag. + struct Rough { + double I = 0.0, sigma = NAN, bkg = 0.0, obs_x = 0.0, obs_y = 0.0; + double bkg_var = 0.0; // variance of the background ESTIMATE itself, bkg / n_bkg + double var_bkg = 0.0; // total non-signal variance carried to the merge + int64_t I_sum = 0; // kept so I can be rebuilt after the radial background correction + int n_inner = 0; + int n_disk = 0, n_own = 0; // signal-disk pixels, and how many of them are this reflection's + int r_bin = 0; // rounded distance from the beam centre, indexes the radial curve + int k_bin = 0; // which radial-background kernel this reflection's stencil needs + int cx = 0, cy = 0, shell = -1; + bool ok = false, strong = false, has_obs = false; + bool full = false; // every pixel of the signal disk was readable + bool overloaded = false; // a pixel of the signal disk was saturated + }; + // Everything pass A hands the profile passes. + struct PassA { + std::vector rough; + // Radial background curve, accumulated from the ring pixels the background clip keeps. Integer + // sums, so the order the pixels arrive in does not matter. + int n_rad = 0; + std::vector rad_sum; + std::vector rad_cnt; + double inv_d2_min = 0.0, inv_d2_max = 0.0; // resolution range of the complete reflections + }; + PassA StartPassA(size_t npredicted) const; + + // Whether pass A sigma-clips this reflection's background ring, and the clip threshold: the ring + // pixels at or below it are kept. A pixel value is an integer, so the threshold is one too. + bool ClipsRing(const BraggBoxSums &s) const; + int32_t ClipThreshold(const BraggBoxSums &s) const; + // Pass A's result for reflection i from what its box walk counted, and from its clip sums where + // ClipsRing. bkg_vals holds the ring pixels for the trimmed-mean background, and is sorted here. + void FinishPassA(PassA &a, size_t i, const Reflection &r, const BraggStencil &st, const BraggBoxSums &s, + const BraggClipSums &clip, std::vector &bkg_vals); + // Everything after pass A: the radial background correction, then either the box-sum results or + // the profile learning and pass B. Defined for BufferSampler with either owner map, and for the + // CompressedImage samplers with the tiled one. + template + std::vector ProfilePasses(const Sampler &img, const OwnerMap &owner_map, + const std::vector &predicted, size_t npredicted, + PassA &a, int64_t image_number); + public: BraggIntegrationEngineCPU(const DiffractionExperiment &experiment, const PixelMask &mask); diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h new file mode 100644 index 000000000..6f4c19d74 --- /dev/null +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h @@ -0,0 +1,47 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include + +#include "BraggIntegrationEngineCPU.h" + +// Apple GPU engine (macOS only): BraggIntegrationEngineCPU with the pixel work that is all integer - the +// neighbour mask, the owner map, pass A's walk over every box and the background clip - run as Metal +// kernels, one thread per reflection. Everything computed in double, the profile learning and pass B +// stay on the CPU, in the CPU engine's own code. Apple GPUs have no double precision, and pass B's +// Gaussian needs exp, which no GPU function matches the CPU's expf in bit for bit. +// +// The result is the CPU engine's, bit for bit: the integer sums do not depend on the order the GPU +// visits pixels in, and the stencil arithmetic deciding which pixel is which is the same source +// (BraggStencil.h), compiled for the GPU with fast math off - sqrt and division then round correctly - +// and, as on the host, with no a*b+c contracted into a fused multiply-add. The CPU and the GPU read the +// same memory, so the preprocessed image is not copied. tests/BraggIntegrationEngineMetalTest.cpp holds +// the engine to all of this. +class BraggIntegrationEngineMetal : public BraggIntegrationEngineCPU { + struct Device; + std::unique_ptr dev; + uint64_t frames_on_gpu = 0; + +public: + // Whether this engine can integrate this experiment: a Metal device the process may use (rugnux + // --no-gpu hides it, see hide_gpus) and settings the kernels cover - not the trimmed-mean background, + // which only the CPU engine sorts the ring for. + static bool Available(const DiffractionExperiment &experiment); + // The Metal device's name, for the log; empty where there is none or the process may not use it. + static std::string DeviceName(); + + BraggIntegrationEngineMetal(const DiffractionExperiment &experiment, const PixelMask &mask); + ~BraggIntegrationEngineMetal() override; + + using BraggIntegrationEngineCPU::Run; // the CompressedImage overload stays the CPU engine's + + std::vector Run(const ImagePreprocessorBuffer &image, + const std::vector &predicted, size_t npredicted, + int64_t image_number) override; + + // The frames whose pixels the GPU walked, rather than the CPU engine this one falls back to. + [[nodiscard]] uint64_t FramesOnGPU() const { return frames_on_gpu; } +}; diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal new file mode 100644 index 000000000..7ebf36fd3 --- /dev/null +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal @@ -0,0 +1,189 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +// Kernels of BraggIntegrationEngineMetal. Compiled at run time with fast math off, after BraggStencil.h +// and BraggPassA.h (BraggIntegrationEngineMetal.mm puts the source together). +// +// Each kernel is one of BraggIntegrationEngineCPU's loops over a reflection, one thread per reflection, +// written as the CPU engine writes it. The stencil arithmetic comes from the same header and rounds as +// the CPU does, every sum is an integer, and the two masks are a maximum and a minimum - so neither the +// order a thread visits its pixels in nor the order the threads run in changes a result. + +constant int32_t PX_MASKED = -2147483647 - 1; // INT32_MIN: the preprocessor's masked pixel +constant int32_t PX_SATURATED = 2147483647; // INT32_MAX: and its saturated one +constant uint8_t MASK_TAIL = 1; // bragg_engine::MASK_TAIL / MASK_FLUX +constant uint8_t MASK_FLUX = 2; + +inline bool valid_px(int32_t v) { return v != PX_MASKED && v != PX_SATURATED; } + +// The stencil parameters arrive as bytes: Metal takes no struct with default member initialisers as a +// buffer type, and BraggStencilParams has them. +inline BraggStencilParams stencil_params(constant char *raw) { + return *(constant const BraggStencilParams *) raw; +} + +// MarkNeighbours: every reflection marks its r2 region. One dispatch per level, TAIL first, so a pixel +// two regions share ends up with the larger mark, as the CPU's max gives. want = 0 has every reflection +// unmark its region, which leaves the mask clear for the next frame. +kernel void mark_neighbours(device const float2 *xy [[buffer(0)]], + device const uint8_t *level [[buffer(1)]], + device uint8_t *mask [[buffer(2)]], + constant char *stencil_raw [[buffer(3)]], + constant BraggPassAFrame &f [[buffer(4)]], + constant uint8_t &want [[buffer(5)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n || (want != 0 && level[i] != want)) + return; + const BraggStencilParams sp = stencil_params(stencil_raw); + const float px = xy[i].x, py = xy[i].y; + const BraggStencil st = MakeBraggStencil(px, py, sp); + const int x0 = max(0, (int) floor(px - st.ex_in - 1.0f)); + const int x1 = min(f.width - 1, (int) ceil(px + st.ex_in + 1.0f)); + const int y0 = max(0, (int) floor(py - st.ey_in - 1.0f)); + const int y1 = min(f.height - 1, (int) ceil(py + st.ey_in + 1.0f)); + for (int y = y0; y <= y1; ++y) + for (int x = x0; x <= x1; ++x) { + const BraggStencilDist d = BraggStencilDistances(st, x - px, y - py); + if (d.inner < f.r2_sq) + mask[y * f.width + x] = want; + } +} + +// The owner map: every reflection claims the pixels within `claim` of its centre with an atomic minimum +// of its (distance, index) key, so the nearest centre wins whatever order the threads run in. release +// gives the same pixels back to nobody, which leaves the map clear for the next frame. +kernel void claim_owner(device const float2 *xy [[buffer(0)]], + device atomic_uint *owner [[buffer(1)]], + constant BraggPassAFrame &f [[buffer(2)]], + constant uint &release [[buffer(3)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n) + return; + const float px = xy[i].x, py = xy[i].y; + const float claim_sq = f.claim * f.claim; + const int x0 = max(0, (int) floor(px - f.claim)); + const int x1 = min(f.width - 1, (int) ceil(px + f.claim)); + const int y0 = max(0, (int) floor(py - f.claim)); + const int y1 = min(f.height - 1, (int) ceil(py + f.claim)); + for (int y = y0; y <= y1; ++y) + for (int x = x0; x <= x1; ++x) { + const float dx = x - px, dy = y - py; + const float d2 = dx * dx + dy * dy; + if (d2 >= claim_sq) + continue; + if (release) + atomic_store_explicit(&owner[y * f.width + x], BRAGG_OWNER_NONE, memory_order_relaxed); + else + atomic_fetch_min_explicit(&owner[y * f.width + x], BraggOwnerKey(sqrt(d2), f.inv_claim, (int) i), + memory_order_relaxed); + } +} + +// Pass A's walk over the box of the outer ellipse: the signal-disk and background-ring counts. +kernel void box_sums(device const float2 *xy [[buffer(0)]], + device const int32_t *img [[buffer(1)]], + device const uint8_t *mask [[buffer(2)]], + device const uint32_t *owner [[buffer(3)]], + device const uint32_t *static_mask [[buffer(4)]], + constant char *stencil_raw [[buffer(5)]], + constant BraggPassAFrame &f [[buffer(6)]], + device BraggBoxSums *sums [[buffer(7)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n) + return; + const BraggStencilParams sp = stencil_params(stencil_raw); + const float px = xy[i].x, py = xy[i].y; + const BraggStencil st = MakeBraggStencil(px, py, sp); + const int x0 = max(0, (int) floor(px - st.ex_out - 1.0f)); + const int x1 = min(f.width - 1, (int) ceil(px + st.ex_out + 1.0f)); + const int y0 = max(0, (int) floor(py - st.ey_out - 1.0f)); + const int y1 = min(f.height - 1, (int) ceil(py + st.ey_out + 1.0f)); + BraggBoxSums s = {}; + for (int y = y0; y <= y1; ++y) + for (int x = x0; x <= x1; ++x) { + const BraggStencilDist d = BraggStencilDistances(st, x - px, y - py); + const int idx = y * f.width + x; + const int32_t v = img[idx]; + if (BraggStencilInSignal(st, x - px, y - py, d.signal, f.r1_sq)) { + ++s.n_disk; + if (f.use_owner) { + if (BraggOwnedBy(owner[idx], (int) i)) ++s.n_own; + else if (f.exclude) continue; + } + ++s.n_inner; + // Unreadable here but not in the run's mask: lost to the flux on this frame. + const bool masked = (uint) idx / 32 < f.static_mask_words + && ((static_mask[(uint) idx / 32] >> ((uint) idx % 32)) & 1u); + if (v == PX_SATURATED || (v == PX_MASKED && !masked)) + s.overloaded = 1; + if (!valid_px(v)) continue; + s.I_sum += v; + s.I_sum_x += (int64_t) x * v; + s.I_sum_y += (int64_t) y * v; + s.x_sum += x; + s.y_sum += y; + ++s.n_inner_valid; + } else if (d.inner >= f.r2_sq && d.outer < f.r3_sq) { + if (!valid_px(v)) continue; + ++s.n_ring; + const uint8_t m = mask[idx]; + if (m == 0) ++s.n_clean; + if (m == MASK_FLUX) continue; + s.bkg_sum += v; + ++s.n_bkg; + } + } + sums[i] = s; +} + +// A 64-bit integer sum kept as two 32-bit words, low first: Metal has no 64-bit atomic add. The carry +// out of the low word is known to the thread that caused it, so the pair adds up exactly. +inline void atomic_add_i64(device atomic_uint *w, int64_t v) { + const uint64_t u = (uint64_t) v; + const uint lo = (uint) u, hi = (uint) (u >> 32); + const uint old = atomic_fetch_add_explicit(&w[0], lo, memory_order_relaxed); + atomic_fetch_add_explicit(&w[1], hi + (old + lo < old ? 1u : 0u), memory_order_relaxed); +} + +// The background clip: the ring pixels at or below the threshold the CPU computed from box_sums, and +// the radial background curve from the same pixels. +kernel void clip_sums(device const float2 *xy [[buffer(0)]], + device const int32_t *img [[buffer(1)]], + device const uint8_t *mask [[buffer(2)]], + constant char *stencil_raw [[buffer(3)]], + constant BraggPassAFrame &f [[buffer(4)]], + device const uint8_t *clips [[buffer(5)]], + device const int32_t *thr [[buffer(6)]], + device BraggClipSums *out [[buffer(7)]], + device atomic_uint *rad_sum [[buffer(8)]], + device atomic_int *rad_cnt [[buffer(9)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n || !clips[i]) + return; + const BraggStencilParams sp = stencil_params(stencil_raw); + const float px = xy[i].x, py = xy[i].y; + const BraggStencil st = MakeBraggStencil(px, py, sp); + const int x0 = max(0, (int) floor(px - st.ex_out - 1.0f)); + const int x1 = min(f.width - 1, (int) ceil(px + st.ex_out + 1.0f)); + const int y0 = max(0, (int) floor(py - st.ey_out - 1.0f)); + const int y1 = min(f.height - 1, (int) ceil(py + st.ey_out + 1.0f)); + const int32_t t = thr[i]; + BraggClipSums c = {}; + for (int y = y0; y <= y1; ++y) + for (int x = x0; x <= x1; ++x) { + const BraggStencilDist d = BraggStencilDistances(st, x - px, y - py); + const int idx = y * f.width + x; + const int32_t v = img[idx]; + if (!(d.inner >= f.r2_sq && d.outer < f.r3_sq && valid_px(v) && mask[idx] != MASK_FLUX && v <= t)) + continue; + c.sum += v; + ++c.n; + if (f.n_rad > 0) { + // Binned on the true detector radius: the unshrunk radial projection. + const int b = clamp((int) round(st.r0 + d.rad), 0, f.n_rad - 1); + atomic_add_i64(&rad_sum[2 * b], v); + atomic_fetch_add_explicit(&rad_cnt[b], 1, memory_order_relaxed); + } + } + out[i] = c; +} diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm new file mode 100644 index 000000000..9c67d7687 --- /dev/null +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm @@ -0,0 +1,343 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#import + +#include "BraggIntegrationEngineMetal.h" + +#include +#include +#include + +#include "../../common/CUDAWrapper.h" +#include "../../common/JFJochException.h" + +using namespace bragg_engine; + +// The kernels' source: BraggStencil.h, BraggPassA.h and BraggIntegrationEngineMetal.metal, copied in +// at configure time (CMakeLists.txt) so the binary needs no file beside it. +extern const char *const bragg_metal_source_stencil; +extern const char *const bragg_metal_source_pass_a; +extern const char *const bragg_metal_source_kernels; + +namespace { + +static_assert(MASK_TAIL == 1 && MASK_FLUX == 2, "BraggIntegrationEngineMetal.metal spells the two marks out"); + +struct Kernels { + id device = nil; // nil: no Metal device + id mark = nil, claim = nil, box = nil, clip = nil; +}; + +// Compiled once per process: a worker makes an engine per pass, and compiling takes a fraction of a +// second. The Metal compiler ships with macOS, so nothing has to be installed to build or run this. +const Kernels &GetKernels() { + static const Kernels kernels = [] { + Kernels k; + k.device = MTLCreateSystemDefaultDevice(); + if (k.device == nil) + return k; + // No contraction of a*b+c into a fused multiply-add, as everywhere on the host (CMakeLists.txt). + // Metal would otherwise contract wherever its backend saw fit. BraggStencil.h calls C's math + // functions; the defines are Metal's spellings of them. + const std::string source = std::string("#include \n" + "using namespace metal;\n" + "#pragma clang fp contract(off)\n" + "#define sqrtf(x) sqrt(x)\n" + "#define fmaxf(a, b) fmax(a, b)\n" + "#define floorf(x) floor(x)\n" + "#define lroundf(x) ((long) round(x))\n") + + bragg_metal_source_stencil + bragg_metal_source_pass_a + + bragg_metal_source_kernels; + MTLCompileOptions *options = [MTLCompileOptions new]; + // Fast math off: sqrt and division then round correctly and nothing is reassociated, as on the + // CPU. (MTLCompileOptions.mathMode says the same from macOS 15; the build targets 13.) And the + // language of macOS 13 on every macOS, so the kernels mean the same wherever they compile. + options.fastMathEnabled = NO; + options.languageVersion = MTLLanguageVersion3_0; + NSError *error = nil; + id library = [k.device newLibraryWithSource:@(source.c_str()) options:options error:&error]; + if (library == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + std::string("Bragg integration kernels: ") + error.localizedDescription.UTF8String); + auto pipeline = [&](const char *name) { + id function = [library newFunctionWithName:@(name)]; + id p = [k.device newComputePipelineStateWithFunction:function error:&error]; + if (p == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + std::string(name) + ": " + error.localizedDescription.UTF8String); + return p; + }; + k.mark = pipeline("mark_neighbours"); + k.claim = pipeline("claim_owner"); + k.box = pipeline("box_sums"); + k.clip = pipeline("clip_sums"); + return k; + }(); + return kernels; +} + +void WaitCommitted(id cb) { + [cb waitUntilCompleted]; + if (cb.status != MTLCommandBufferStatusCompleted) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed"); +} + +void Wait(id cb) { + [cb commit]; + WaitCommitted(cb); +} + +} // namespace + +struct BraggIntegrationEngineMetal::Device { + id queue = nil; + id static_mask = nil; + // Frame-sized, made on the first image: a worker's analysis is built for every pass, and most of + // them never integrate. A frame unmarks what it marked once the CPU has read them (`unmark`, which + // the next frame waits for): clearing all of both took 1.5 ms of a 16M-pixel frame, longer than + // integrating a sparse one. `clear` is true once an unmark completed; the first frame, and one + // after a frame that threw part-way through, clear the whole of both instead. + id mask = nil, owner = nil; + id unmark = nil; + bool clear = false; + id rad_sum = nil, rad_cnt = nil; // the radial background curve, as the clip kernel adds it up + // Per reflection, grown with the frames. + size_t capacity = 0; + id xy = nil, level = nil, sums = nil, clips = nil, thr = nil, clip_out = nil; +}; + +bool BraggIntegrationEngineMetal::Available(const DiffractionExperiment &experiment) { + const auto settings = experiment.GetBraggIntegrationSettings(); + const bool trimmed_mean = !(settings.GetBackgroundClipNSigma() > 0.0f) + && settings.GetBackgroundTrimFraction() > 0.0f; + return !gpus_are_hidden() && !trimmed_mean && GetKernels().device != nil; +} + +std::string BraggIntegrationEngineMetal::DeviceName() { + if (gpus_are_hidden() || GetKernels().device == nil) + return {}; + return GetKernels().device.name.UTF8String; +} + +BraggIntegrationEngineMetal::BraggIntegrationEngineMetal(const DiffractionExperiment &experiment, + const PixelMask &mask) + : BraggIntegrationEngineCPU(experiment, mask), dev(std::make_unique()) { + const Kernels &k = GetKernels(); + dev->queue = [k.device newCommandQueue]; + dev->static_mask = [k.device newBufferWithBytes:static_mask.data() + length:std::max(static_mask.size() * sizeof(uint32_t), 4) + options:MTLResourceStorageModeShared]; +} + +BraggIntegrationEngineMetal::~BraggIntegrationEngineMetal() = default; + +std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessorBuffer &image, + const std::vector &predicted, + size_t npredicted, int64_t image_number) { + // An image of another size, and a frame with nothing predicted, the CPU engine answers as it would. + if (image.size() != npixel || npredicted == 0) + return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number); + + @autoreleasepool { + const Kernels &k = GetKernels(); + // The GPU reads the host's own memory, so the image is read where it lies - which Metal allows + // for memory starting on a page, as macOS's allocator hands out any buffer of a frame's size. + // Should one not be, the CPU engine integrates it: the same result. + const size_t page = getpagesize(), bytes = npixel * sizeof(int32_t); + id img = nil; + if (reinterpret_cast(image.data()) % page == 0) + img = [k.device newBufferWithBytesNoCopy:const_cast(image.data()) + length:(bytes + page - 1) / page * page + options:MTLResourceStorageModeShared + deallocator:nil]; + if (img == nil) + return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number); + + counts.predicted += npredicted; + const bool use_owner = overlap != OverlapMode::Off; + if (dev->unmark != nil) { // it reads xy, which is overwritten below + id unmark = dev->unmark; + dev->unmark = nil; + dev->clear = false; // unless it completed + WaitCommitted(unmark); + dev->clear = true; + } + auto buffer = [&](size_t length) { + return [k.device newBufferWithLength:std::max(length, 4) options:MTLResourceStorageModeShared]; + }; + if (dev->mask == nil) { + dev->mask = buffer(npixel); + if (use_owner) + dev->owner = buffer(npixel * sizeof(uint32_t)); + } + const bool was_clear = dev->clear; + dev->clear = false; + if (dev->capacity < npredicted) { + const size_t n = npredicted + npredicted / 2; + dev->xy = buffer(n * 2 * sizeof(float)); + dev->level = buffer(n); + dev->sums = buffer(n * sizeof(BraggBoxSums)); + dev->clips = buffer(n); + dev->thr = buffer(n * sizeof(int32_t)); + dev->clip_out = buffer(n * sizeof(BraggClipSums)); + dev->capacity = n; + } + + auto *xy = static_cast(dev->xy.contents); + auto *level = static_cast(dev->level.contents); + for (size_t i = 0; i < npredicted; ++i) { + xy[2 * i] = predicted[i].predicted_x; + xy[2 * i + 1] = predicted[i].predicted_y; + level[i] = predicted[i].partiality >= NEIGHBOUR_MASK_MIN_PARTIALITY ? MASK_FLUX : MASK_TAIL; + } + + PassA a = StartPassA(npredicted); + BraggPassAFrame f{}; + f.width = static_cast(xpixel); + f.height = static_cast(ypixel); + f.r1_sq = r1_sq; + f.r2_sq = r2_sq; + f.r3_sq = r3_sq; + f.claim = claim; + f.inv_claim = inv_claim; + f.use_owner = use_owner; + f.exclude = overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum; + f.n_rad = bkg_radial ? a.n_rad : 0; + f.static_mask_words = static_cast(static_mask.size()); + f.n = static_cast(npredicted); + + const MTLSize grid = MTLSizeMake(npredicted, 1, 1), group = MTLSizeMake(64, 1, 1); + + // --- The neighbour mask, the owner map and pass A's box walk --- + id cb = [dev->queue commandBuffer]; + id blit; + if (!was_clear) { + blit = [cb blitCommandEncoder]; + [blit fillBuffer:dev->mask range:NSMakeRange(0, npixel) value:0]; + if (use_owner) // every byte 0xff: BRAGG_OWNER_NONE + [blit fillBuffer:dev->owner range:NSMakeRange(0, npixel * sizeof(uint32_t)) value:0xff]; + [blit endEncoding]; + } + const uint32_t claim_not_release = 0; + id enc = [cb computeCommandEncoder]; + for (const uint8_t want : {MASK_TAIL, MASK_FLUX}) { + [enc setComputePipelineState:k.mark]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:dev->level offset:0 atIndex:1]; + [enc setBuffer:dev->mask offset:0 atIndex:2]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:3]; + [enc setBytes:&f length:sizeof(f) atIndex:4]; + [enc setBytes:&want length:sizeof(want) atIndex:5]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + } + if (use_owner) { + [enc setComputePipelineState:k.claim]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:dev->owner offset:0 atIndex:1]; + [enc setBytes:&f length:sizeof(f) atIndex:2]; + [enc setBytes:&claim_not_release length:sizeof(claim_not_release) atIndex:3]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + } + [enc setComputePipelineState:k.box]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:img offset:0 atIndex:1]; + [enc setBuffer:dev->mask offset:0 atIndex:2]; + [enc setBuffer:(use_owner ? dev->owner : dev->mask) offset:0 atIndex:3]; // not read without an owner map + [enc setBuffer:dev->static_mask offset:0 atIndex:4]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:5]; + [enc setBytes:&f length:sizeof(f) atIndex:6]; + [enc setBuffer:dev->sums offset:0 atIndex:7]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + [enc endEncoding]; + Wait(cb); + + // --- The clip thresholds, which are double arithmetic, on the CPU; the clip on the GPU --- + const auto *sums = static_cast(dev->sums.contents); + auto *clips = static_cast(dev->clips.contents); + auto *thr = static_cast(dev->thr.contents); + bool any_clip = false; + for (size_t i = 0; i < npredicted; ++i) { + clips[i] = ClipsRing(sums[i]); + if (clips[i]) { + thr[i] = ClipThreshold(sums[i]); + any_clip = true; + } + } + if (any_clip) { + if (bkg_radial && dev->rad_sum == nil) { + dev->rad_sum = buffer(a.n_rad * 2 * sizeof(uint32_t)); + dev->rad_cnt = buffer(a.n_rad * sizeof(int32_t)); + } + cb = [dev->queue commandBuffer]; + if (bkg_radial) { + blit = [cb blitCommandEncoder]; + [blit fillBuffer:dev->rad_sum range:NSMakeRange(0, a.n_rad * 2 * sizeof(uint32_t)) value:0]; + [blit fillBuffer:dev->rad_cnt range:NSMakeRange(0, a.n_rad * sizeof(int32_t)) value:0]; + [blit endEncoding]; + } + enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:k.clip]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:img offset:0 atIndex:1]; + [enc setBuffer:dev->mask offset:0 atIndex:2]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:3]; + [enc setBytes:&f length:sizeof(f) atIndex:4]; + [enc setBuffer:dev->clips offset:0 atIndex:5]; + [enc setBuffer:dev->thr offset:0 atIndex:6]; + [enc setBuffer:dev->clip_out offset:0 atIndex:7]; + [enc setBuffer:(bkg_radial ? dev->rad_sum : dev->clips) offset:0 atIndex:8]; // not written without the curve + [enc setBuffer:(bkg_radial ? dev->rad_cnt : dev->clips) offset:0 atIndex:9]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + [enc endEncoding]; + Wait(cb); + if (bkg_radial) { + const auto *w = static_cast(dev->rad_sum.contents); + const auto *c = static_cast(dev->rad_cnt.contents); + for (int b = 0; b < a.n_rad; ++b) { + a.rad_sum[b] = static_cast(w[2 * b] | static_cast(w[2 * b + 1]) << 32); + a.rad_cnt[b] = c[b]; + } + } + } + + ++frames_on_gpu; + + // --- The rest is the CPU engine's --- + const auto *clip_out = static_cast(dev->clip_out.contents); + std::vector no_trim; // the trimmed mean is the CPU engine's alone (Available) + for (size_t i = 0; i < npredicted; ++i) { + const Reflection &r = predicted[i]; + FinishPassA(a, i, r, MakeBraggStencil(r.predicted_x, r.predicted_y, stencil), sums[i], + clips[i] ? clip_out[i] : BraggClipSums{}, no_trim); + } + const FlatOwnerMap owner_map{use_owner ? static_cast(dev->owner.contents) : nullptr, xpixel}; + auto ret = ProfilePasses(BufferSampler{image.data()}, owner_map, predicted, npredicted, a, image_number); + + // --- Unmark this frame's regions and give back its claims, while the CPU goes on --- + dev->unmark = [dev->queue commandBuffer]; + enc = [dev->unmark computeCommandEncoder]; + const uint8_t unmark = 0; + [enc setComputePipelineState:k.mark]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:dev->level offset:0 atIndex:1]; + [enc setBuffer:dev->mask offset:0 atIndex:2]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:3]; + [enc setBytes:&f length:sizeof(f) atIndex:4]; + [enc setBytes:&unmark length:sizeof(unmark) atIndex:5]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + if (use_owner) { + const uint32_t release = 1; + [enc setComputePipelineState:k.claim]; + [enc setBuffer:dev->xy offset:0 atIndex:0]; + [enc setBuffer:dev->owner offset:0 atIndex:1]; + [enc setBytes:&f length:sizeof(f) atIndex:2]; + [enc setBytes:&release length:sizeof(release) atIndex:3]; + [enc dispatchThreads:grid threadsPerThreadgroup:group]; + } + [enc endEncoding]; + [dev->unmark commit]; + return ret; + } +} diff --git a/image_analysis/bragg_integration/BraggPassA.h b/image_analysis/bragg_integration/BraggPassA.h new file mode 100644 index 000000000..8b1b10979 --- /dev/null +++ b/image_analysis/bragg_integration/BraggPassA.h @@ -0,0 +1,45 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// Pass A of BraggIntegrationEngineCPU, as the data it hands on: what the walk over one reflection's box +// counts, and the clip sums of its background ring. Integers only, so the counts do not depend on the +// order the pixels are visited in: the Metal engine walks the boxes on the GPU and gets the CPU's +// counts exactly. Compiled by the host and by the Metal compiler, so the two agree on the layout - +// which is also why there are no default member initialisers (Metal does not take a struct with them +// as a buffer type); zero one with {}. + +#ifndef __METAL_VERSION__ +#include +#endif + +struct BraggBoxSums { + int64_t I_sum, I_sum_x, I_sum_y; // readable signal-disk pixels, and their first moments + int64_t x_sum, y_sum; // the positions behind I_sum, to take the background out of the centroid + int64_t bkg_sum; // ring pixels kept for the background + int64_t n_inner, n_inner_valid; // signal-disk pixels kept, and how many of them were readable + int32_t n_disk, n_own; // signal-disk pixels, and how many of them are this reflection's + int32_t n_bkg; // ring pixels behind bkg_sum + int32_t n_ring, n_clean; // readable ring pixels, and those no predicted neighbour covers + int32_t overloaded; // a signal-disk pixel was saturated, or unreadable on this frame only +}; + +// The ring pixels at or below a reflection's clip threshold (BraggIntegrationEngineCPU::ClipThreshold). +struct BraggClipSums { + int64_t sum; + int32_t n; + int32_t unused; +}; + +// What the Metal kernels are told about the frame, beside the stencil (BraggStencilParams). +struct BraggPassAFrame { + int32_t width, height; + float r1_sq, r2_sq, r3_sq; + float claim, inv_claim; // how far a reflection claims pixels in the owner map + int32_t use_owner; // an overlap treatment is on, so the owner map is built and read + int32_t exclude; // ... and a pixel a neighbour owns leaves the signal disk + int32_t n_rad; // radial background bins; 0 when the curve is not accumulated + uint32_t static_mask_words; // length of the run's packed pixel mask + uint32_t n; // reflections +}; diff --git a/image_analysis/bragg_integration/BraggStencil.h b/image_analysis/bragg_integration/BraggStencil.h index c92122fb0..00030799d 100644 --- a/image_analysis/bragg_integration/BraggStencil.h +++ b/image_analysis/bragg_integration/BraggStencil.h @@ -49,15 +49,19 @@ // gives q = 0 and hence d2_inner == d2_outer == d2 EXACTLY, in floating point, not merely to within // rounding. k_sigma = 0 therefore reproduces the old circular stencil bit for bit. // -// This header is compiled by both the host compiler and nvcc, so the geometry has exactly one -// definition: the two engines classify by the same rule, built from the same inputs. They are not -// bit-identical to each other - nvcc contracts a*b+c into an FMA and the host baseline has no FMA -// instruction to contract into - so a pixel within a rounding of a boundary can still fall either -// way, which is what the parity test's tolerances are for. +// This header is compiled by the host compiler, by nvcc and by the Metal compiler, so the geometry has +// exactly one definition: the engines classify by the same rule, built from the same inputs. The CUDA +// engine is not bit-identical to the CPU one - nvcc contracts a*b+c into an FMA and the host baseline +// has no FMA instruction to contract into - so a pixel within a rounding of a boundary can still fall +// either way, which is what its parity test's tolerances are for. The Metal engine is: neither the host +// nor the Metal compiler contracts anything, and with fast math off Metal rounds sqrt and division +// correctly, as the CPU does (BraggIntegrationEngineMetal.mm). // ============================================================================= +#ifndef __METAL_VERSION__ #include #include +#endif #ifdef __CUDACC__ #define BRAGG_STENCIL_HD __host__ __device__ inline @@ -65,14 +69,24 @@ #define BRAGG_STENCIL_HD inline #endif +// Metal keeps a program-scope constant in its constant address space, and a reference names the space +// it points into (here always the thread's own), so both say so. +#ifdef __METAL_VERSION__ +#define BRAGG_STENCIL_CONSTANT constant +#define BRAGG_STENCIL_THREAD thread +#else +#define BRAGG_STENCIL_CONSTANT +#define BRAGG_STENCIL_THREAD +#endif + // The measured spot footprint (SpotFootprint.h): the radial and tangential standard deviations of a // spot [px], tabulated against distance from the beam centre in bins of fp_bin_px. Held here, by value, // because both engines build every stencil from these parameters - the GPU gets the table with the // kernel arguments and needs no buffer of its own. -constexpr int BRAGG_FOOTPRINT_MAX_BINS = 16; +BRAGG_STENCIL_CONSTANT constexpr int BRAGG_FOOTPRINT_MAX_BINS = 16; // A spot whose footprint reaches this many of its standard deviations past r1 has outgrown the disk, and // gets the footprint stencil below. -constexpr float BRAGG_FOOTPRINT_NSIGMA = 3.0f; +BRAGG_STENCIL_CONSTANT constexpr float BRAGG_FOOTPRINT_NSIGMA = 3.0f; // How far that stencil reaches, in the footprint's standard deviations: the summation ellipse ends and the // background ring starts there, and the profile is fitted out to it. Further than the 3 sigma that decides // whether a spot outgrew the disk, because a wide spot is not a Gaussian: mosaic streaks and diffuse halos @@ -80,7 +94,7 @@ constexpr float BRAGG_FOOTPRINT_NSIGMA = 3.0f; // corrections applied to both integrations, so only the integration differs, 4 sigma closed about a third // of the gap to XDS's 3D summation on the haloed small-molecule set and was neutral on the streaked one; // 5 sigma was no better than 4. -constexpr float BRAGG_FOOTPRINT_REACH = 4.0f; +BRAGG_STENCIL_CONSTANT constexpr float BRAGG_FOOTPRINT_REACH = 4.0f; // Fixed per-experiment inputs to the stencil law (mirrors BraggIntegrationEngine's members). struct BraggStencilParams { @@ -116,7 +130,7 @@ struct BraggStencil { // Radial growth of the ring at Rpx, in pixels: k_sigma times the bandwidth streak, capped. Shared // so the stencil and the radial-background kernel table agree on what "grown by this much" means. -BRAGG_STENCIL_HD float BraggStencilGrow_px(float Rpx, const BraggStencilParams &p) { +BRAGG_STENCIL_HD float BraggStencilGrow_px(float Rpx, const BRAGG_STENCIL_THREAD BraggStencilParams &p) { if (!(p.k_sigma > 0.0f)) return 0.0f; // The cap bounds what a mis-declared bandwidth can do to the bounding box. A caller that // forgets to set it gets the uncapped growth rather than silently no growth at all. @@ -125,7 +139,8 @@ BRAGG_STENCIL_HD float BraggStencilGrow_px(float Rpx, const BraggStencilParams & } // The footprint at Rpx: linear between bin centres, flat beyond the first and the last. -BRAGG_STENCIL_HD void BraggFootprintAt(float Rpx, const BraggStencilParams &p, float &s_rad, float &s_tan) { +BRAGG_STENCIL_HD void BraggFootprintAt(float Rpx, const BRAGG_STENCIL_THREAD BraggStencilParams &p, + BRAGG_STENCIL_THREAD float &s_rad, BRAGG_STENCIL_THREAD float &s_tan) { const float t = Rpx / p.fp_bin_px - 0.5f; int i = (int) floorf(t); float f = t - (float) i; @@ -139,7 +154,7 @@ BRAGG_STENCIL_HD void BraggFootprintAt(float Rpx, const BraggStencilParams &p, f // Growth of the ring beyond r2 that the footprint asks for along one axis: none while the spot fits // the r1 disk - there the disk resolves the spot and the stencil stays what it always was - and // otherwise out to BRAGG_FOOTPRINT_REACH of it, capped as the bandwidth growth is. -BRAGG_STENCIL_HD float BraggFootprintGrow_px(float sigma, const BraggStencilParams &p) { +BRAGG_STENCIL_HD float BraggFootprintGrow_px(float sigma, const BRAGG_STENCIL_THREAD BraggStencilParams &p) { const float g = BRAGG_FOOTPRINT_REACH * sigma - p.r2; if (!(g > 0.0f)) return 0.0f; return (p.max_grow > 0.0f && g > p.max_grow) ? p.max_grow : g; @@ -147,7 +162,7 @@ BRAGG_STENCIL_HD float BraggFootprintGrow_px(float sigma, const BraggStencilPara // The largest growth any reflection on a detector reaching r_max can get, radially - what the // buffers that hold a ring are sized by. -BRAGG_STENCIL_HD float BraggStencilMaxGrow_px(float r_max, const BraggStencilParams &p) { +BRAGG_STENCIL_HD float BraggStencilMaxGrow_px(float r_max, const BRAGG_STENCIL_THREAD BraggStencilParams &p) { float g = BraggStencilGrow_px(r_max, p); for (int i = 0; i < p.fp_n; ++i) { const float s = fmaxf(p.fp_sigma_rad[i], p.fp_sigma_tan[i]); @@ -156,7 +171,7 @@ BRAGG_STENCIL_HD float BraggStencilMaxGrow_px(float r_max, const BraggStencilPar return g; } -BRAGG_STENCIL_HD BraggStencil MakeBraggStencil(float px_x, float px_y, const BraggStencilParams &p) { +BRAGG_STENCIL_HD BraggStencil MakeBraggStencil(float px_x, float px_y, const BRAGG_STENCIL_THREAD BraggStencilParams &p) { BraggStencil s; const float rx = px_x - p.beam_x, ry = px_y - p.beam_y; const float r0 = sqrtf(rx * rx + ry * ry); @@ -212,7 +227,7 @@ BRAGG_STENCIL_HD BraggStencil MakeBraggStencil(float px_x, float px_y, const Bra // Taken from the stencil, so the kernel and the ring it corrects are built from the same radius - // deriving it from a separately computed r0 let the two engines round differently and pick // different rows. -BRAGG_STENCIL_HD int BraggStencilKernelIndex(const BraggStencil &s, int n_kern) { +BRAGG_STENCIL_HD int BraggStencilKernelIndex(const BRAGG_STENCIL_THREAD BraggStencil &s, int n_kern) { if (n_kern <= 1) return 0; int idx = (int) lroundf(s.grow); if (idx < 0) idx = 0; @@ -234,8 +249,8 @@ struct BraggStencilDist { // minimum, so the nearest centre wins whatever order the writes arrive in and the lowest index breaks // a tie. Eight bits of distance over the claim radius is 1/40 px at the default r2 = 6 px, far finer // than the prediction itself. A pixel nobody claimed is nobody's neighbour, hence the owner's. -constexpr uint32_t BRAGG_OWNER_NONE = 0xffffffffu; -constexpr uint32_t BRAGG_OWNER_IDX_MASK = 0x00ffffffu; +BRAGG_STENCIL_CONSTANT constexpr uint32_t BRAGG_OWNER_NONE = 0xffffffffu; +BRAGG_STENCIL_CONSTANT constexpr uint32_t BRAGG_OWNER_IDX_MASK = 0x00ffffffu; BRAGG_STENCIL_HD int BraggOwnerQuant(float d, float inv_claim) { int q = (int) (d * inv_claim * 255.0f); @@ -258,7 +273,7 @@ BRAGG_STENCIL_HD bool BraggOwnedBy(uint32_t owner, int index) { // intensity at 1, 2 and 3 px, and the merge fell apart with it. What the margin removes is the // reflection's own profile, not the neighbour's tail, so the plain nearest-centre split is the rule. -BRAGG_STENCIL_HD BraggStencilDist BraggStencilDistances(const BraggStencil &s, float ddx, float ddy) { +BRAGG_STENCIL_HD BraggStencilDist BraggStencilDistances(const BRAGG_STENCIL_THREAD BraggStencil &s, float ddx, float ddy) { BraggStencilDist d; d.rad = ddx * s.ux + ddy * s.uy; d.signal = ddx * ddx + ddy * ddy; @@ -276,7 +291,7 @@ BRAGG_STENCIL_HD BraggStencilDist BraggStencilDistances(const BraggStencil &s, f // The signal region: the r1 disk, and for a spot that outgrows it the whole footprint ellipse too, so // the summation - the seed of the profile fit, and what the fit falls back to - holds the spot rather // than its core. Without a footprint this is the plain disk test, bit for bit. -BRAGG_STENCIL_HD bool BraggStencilInSignal(const BraggStencil &s, float ddx, float ddy, float d2, float r1_sq) { +BRAGG_STENCIL_HD bool BraggStencilInSignal(const BRAGG_STENCIL_THREAD BraggStencil &s, float ddx, float ddy, float d2, float r1_sq) { if (d2 < r1_sq) return true; if (!(s.fp_s2t > 0.0f)) return false; const float rad = ddx * s.ux + ddy * s.uy, tn = -ddx * s.uy + ddy * s.ux; diff --git a/image_analysis/bragg_integration/CMakeLists.txt b/image_analysis/bragg_integration/CMakeLists.txt index ac34468cf..f344c142e 100644 --- a/image_analysis/bragg_integration/CMakeLists.txt +++ b/image_analysis/bragg_integration/CMakeLists.txt @@ -3,6 +3,7 @@ ADD_LIBRARY(JFJochBraggIntegration STATIC BraggIntegrationEngine.h BraggIntegrationEngineCPU.cpp BraggIntegrationEngineCPU.h + BraggPassA.h SpotFootprint.cpp SpotFootprint.h Regression.h @@ -21,3 +22,29 @@ IF (JFJOCH_CUDA_AVAILABLE) TARGET_SOURCES(JFJochBraggIntegration PRIVATE ../indexing/CUDAMemHelpers.h BraggIntegrationEngineGPU.cu BraggIntegrationEngineGPU.h) ENDIF() + +IF (APPLE) + # The Metal engine compiles its kernels when the program starts, from source carried in the binary: + # macOS has the Metal compiler built in, while the offline one is an Xcode download of its own. The + # source is the stencil header the CPU engine compiles, the pass-A header and the kernels, copied + # into a generated .cpp here; editing one of them re-runs the configure. + SET(METAL_PARTS stencil BraggStencil.h pass_a BraggPassA.h kernels BraggIntegrationEngineMetal.metal) + SET(METAL_SOURCE "// Generated from image_analysis/bragg_integration by its CMakeLists.txt - do not edit.\n") + WHILE (METAL_PARTS) + LIST(POP_FRONT METAL_PARTS name file) + FILE(READ ${CMAKE_CURRENT_SOURCE_DIR}/${file} text) + STRING(APPEND METAL_SOURCE "extern const char *const bragg_metal_source_${name};\n" + "const char *const bragg_metal_source_${name} = R\"BRAGG_METAL(\n// ${file}\n${text})BRAGG_METAL\";\n") + SET_PROPERTY(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/${file}) + ENDWHILE() + # Written only when it changed, so a configure does not rebuild the library. + FILE(WRITE ${CMAKE_CURRENT_BINARY_DIR}/BraggIntegrationEngineMetalSource.cpp.tmp "${METAL_SOURCE}") + CONFIGURE_FILE(${CMAKE_CURRENT_BINARY_DIR}/BraggIntegrationEngineMetalSource.cpp.tmp + ${CMAKE_CURRENT_BINARY_DIR}/BraggIntegrationEngineMetalSource.cpp COPYONLY) + + TARGET_SOURCES(JFJochBraggIntegration PRIVATE BraggIntegrationEngineMetal.mm BraggIntegrationEngineMetal.h + BraggIntegrationEngineMetal.metal ${CMAKE_CURRENT_BINARY_DIR}/BraggIntegrationEngineMetalSource.cpp) + SET_SOURCE_FILES_PROPERTIES(BraggIntegrationEngineMetal.mm PROPERTIES COMPILE_OPTIONS -fobjc-arc) + SET_SOURCE_FILES_PROPERTIES(BraggIntegrationEngineMetal.metal PROPERTIES HEADER_FILE_ONLY ON) + TARGET_LINK_LIBRARIES(JFJochBraggIntegration "-framework Metal" "-framework Foundation") +ENDIF() diff --git a/rugnux/rugnux_cli.cpp b/rugnux/rugnux_cli.cpp index 9295795a2..0390cfedd 100644 --- a/rugnux/rugnux_cli.cpp +++ b/rugnux/rugnux_cli.cpp @@ -72,6 +72,9 @@ #include "../writer/WriteModel.h" #include "ResultReport.h" #include "RugnuxConsole.h" +#ifdef __APPLE__ +#include "../image_analysis/bragg_integration/BraggIntegrationEngineMetal.h" +#endif // Spots kept per image (the strongest ones) and handed to indexing. Offline reprocessing is not // bound by the online spot budget, so this is rugnux's own default rather than the 250 the @@ -866,6 +869,12 @@ static int RunRugnux(int argc, char **argv, bool no_gpu) { RunProvenance provenance; provenance.command_line = JoinCommandLine(argc, argv); provenance.gpu_count = get_gpu_count(); +#ifdef __APPLE__ + // Without CUDA, the Bragg integration's integer pixel work runs on the Apple GPU (MXAnalysisWithoutFPGA). + const std::string metal_device = BraggIntegrationEngineMetal::DeviceName(); +#else + const std::string metal_device; +#endif provenance.gpu_description = get_gpu_description(); // The terminal shows the run as steps (RugnuxConsole) and every Logger writes to the log file @@ -907,6 +916,7 @@ static int RunRugnux(int argc, char **argv, bool no_gpu) { logger.Info("Command line: {}", provenance.command_line); logger.Info("GPU: {}", no_gpu ? std::string("not used (--no-gpu) - running on the CPU") : provenance.gpu_count > 0 ? fmt::format("{} ({} visible)", provenance.gpu_description, provenance.gpu_count) + : !metal_device.empty() ? fmt::format("{} through Metal, for the Bragg integration only - the rest runs on the CPU", metal_device) : std::string("none visible - running on the CPU")); RegisterHDF5Filter(); @@ -1662,6 +1672,8 @@ static int RunRugnux(int argc, char **argv, bool no_gpu) { : provenance.gpu_count > 0 ? fmt::format("{}{}, {} threads", provenance.gpu_count > 1 ? fmt::format("{} x ", provenance.gpu_count) : "", provenance.gpu_description, nthreads) + : !metal_device.empty() + ? fmt::format("{} CPU threads, Bragg integration on {} (Metal)", nthreads, metal_device) : fmt::format("no GPU visible - {} CPU threads, which is much slower", nthreads)); console->Rule('-'); diff --git a/tests/BraggIntegrationEngineMetalTest.cpp b/tests/BraggIntegrationEngineMetalTest.cpp new file mode 100644 index 000000000..703dcfdd1 --- /dev/null +++ b/tests/BraggIntegrationEngineMetalTest.cpp @@ -0,0 +1,265 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#ifdef __APPLE__ + +#include +#include +#include +#include +#include + +#include "../common/BraggIntegrationSettings.h" +#include "../common/DetectorSetup.h" +#include "../common/DiffractionExperiment.h" +#include "../common/PixelMask.h" +#include "../common/Reflection.h" +#include "../image_analysis/bragg_integration/BraggIntegrationEngineCPU.h" +#include "../image_analysis/bragg_integration/BraggIntegrationEngineMetal.h" + +// The Metal engine's whole claim is that it IS the CPU engine, bit for bit (BraggIntegrationEngineMetal.h). +// So nothing here has a tolerance: every float of every reflection is compared as a bit pattern, and +// every count as an integer. A Mac without a Metal device fails rather than skips - the macOS CI job is +// the only place this can be checked, and a skip there would check nothing. + +namespace { + +struct Frame { + std::vector image; + std::vector predicted; +}; + +// A crowded random frame: predictions at random sub-pixel positions (so the pixel classification meets +// every kind of rounding near the stencil boundaries), a quarter of them on whole and half pixels (so it +// meets exact ties with the radii too), with random partialities (rocking-curve tails and flux, which +// mask the rings differently), about a third of them with a spot, overlapping where they fall close. On +// a background that rises with the radius (the radial correction has a curve to fit) with noise that +// takes it below zero near the beam, as a pedestal-subtracted detector does (a negative pixel is what +// carries into the high word of the GPU's 64-bit radial sums), hot pixels in the rings (so the +// background clip rejects something), and unreadable pixels: some in the run's mask, some saturated, +// some unreadable on this frame alone (both read as an overload), inside spots as well as between them. +Frame RandomFrame(const DiffractionExperiment &x, uint32_t seed, const std::vector &masked) { + const int W = static_cast(x.GetXPixelsNum()), H = static_cast(x.GetYPixelsNum()); + const float beam_x = x.GetBeamX_pxl(), beam_y = x.GetBeamY_pxl(); + std::mt19937 rng(seed); + std::uniform_real_distribution U(0.0f, 1.0f); + Frame f; + f.image.resize(static_cast(W) * H); + for (int y = 0; y < H; ++y) + for (int x_ = 0; x_ < W; ++x_) { + const float r = std::hypot(x_ - beam_x, y - beam_y); + const float mean = 0.5f + 0.01f * r; + f.image[static_cast(y) * W + x_] = static_cast(std::lround(mean + (U(rng) - 0.5f) * 8.0f)); + } + const int nrefl = W * H / 400; + int hkl = 1; + for (int n = 0; n < nrefl; ++n) { + Reflection r{}; + r.h = hkl; r.k = -hkl; r.l = hkl % 7; ++hkl; + r.predicted_x = 2.0f + U(rng) * (W - 4); + r.predicted_y = 2.0f + U(rng) * (H - 4); + if (n % 4 == 1) { + r.predicted_x = std::round(r.predicted_x * 2.0f) / 2.0f; + r.predicted_y = std::round(r.predicted_y); + } + const float u = U(rng); + r.partiality = u < 0.3f ? 0.01f * u : u; // a third of them rocking-curve tails + r.d = 1.2f + 3.0f * U(rng); + r.prescaling_corr = 1.0f; + f.predicted.push_back(r); + if (U(rng) < 0.35f) { + const float amp = 20.0f + 3000.0f * U(rng) * U(rng); + const float s = 0.8f + 1.4f * U(rng); + for (int dy = -8; dy <= 8; ++dy) + for (int dx = -8; dx <= 8; ++dx) { + const int px = static_cast(std::lround(r.predicted_x)) + dx; + const int py = static_cast(std::lround(r.predicted_y)) + dy; + if (px < 0 || py < 0 || px >= W || py >= H) continue; + const float ex = px - r.predicted_x, ey = py - r.predicted_y; + f.image[static_cast(py) * W + px] += static_cast( + std::lround(amp * std::exp(-(ex * ex + ey * ey) / (2.0f * s * s)))); + } + } + } + std::uniform_int_distribution P(0, f.image.size() - 1); + for (int k = 0; k < W * H / 500; ++k) + f.image[P(rng)] += 200 + static_cast(5000 * U(rng)); // hot pixels + for (int k = 0; k < W * H / 3000; ++k) + f.image[P(rng)] = U(rng) < 0.5f ? INT32_MIN : INT32_MAX; // unreadable on this frame + for (size_t i = 0; i < f.predicted.size(); i += 13) { // ... inside spots too + const auto &r = f.predicted[i]; + const size_t idx = static_cast(std::lround(r.predicted_y)) * W + std::lround(r.predicted_x) + i % 3; + if (idx < f.image.size()) f.image[idx] = i % 2 ? INT32_MIN : INT32_MAX; + } + for (size_t idx : masked) + f.image[idx] = INT32_MIN; // what the run's mask holds + return f; +} + +struct Case { + IntegratorMode mode = IntegratorMode::ProfileGaussian; + OverlapMode overlap = OverlapMode::Exclude; + bool radial = false; + std::optional bandwidth; + float stencil_k = 0.0f; + bool footprint = false; + float clip_nsigma = 4.0f; +}; + +DiffractionExperiment MakeExperiment(const DetectorSetup &det, const Case &c) { + DiffractionExperiment x(det); + x.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV) + .BeamX_pxl(0.4f * static_cast(x.GetXPixelsNum())) + .BeamY_pxl(0.55f * static_cast(x.GetYPixelsNum())); + x.BandwidthFWHM(c.bandwidth); + BraggIntegrationSettings s; + s.Integrator(c.mode).Overlap(c.overlap).BackgroundRadialCorrection(c.radial).StencilKSigma(c.stencil_k) + .BackgroundClipNSigma(c.clip_nsigma); + if (c.footprint) { + // Wider than the r1 disk can hold from the middle out, so both stencils occur on one frame. + SpotFootprint fp; + fp.bin_px = 100.0f; + for (int b = 0; b < 8; ++b) { + fp.sigma_rad.push_back(0.9f + 0.35f * b); + fp.sigma_tan.push_back(0.8f + 0.3f * b); + } + s.Footprint(fp); + } + x.ImportBraggIntegrationSettings(s); + return x; +} + +uint32_t Bits(float f) { return std::bit_cast(f); } + +void RequireSame(const std::vector &cpu, const std::vector &metal) { + REQUIRE(metal.size() == cpu.size()); + for (size_t i = 0; i < cpu.size(); ++i) { + const Reflection &a = cpu[i], &b = metal[i]; + INFO("reflection " << i << " hkl " << a.h << " " << a.k << " " << a.l + << " I " << a.I << " / " << b.I << " bkg " << a.bkg << " / " << b.bkg); + REQUIRE(a.h == b.h); + REQUIRE(a.k == b.k); + REQUIRE(a.l == b.l); + REQUIRE(Bits(a.image_number) == Bits(b.image_number)); + REQUIRE(Bits(a.I) == Bits(b.I)); + REQUIRE(Bits(a.sigma) == Bits(b.sigma)); + REQUIRE(Bits(a.bkg) == Bits(b.bkg)); + REQUIRE(Bits(a.var_bkg) == Bits(b.var_bkg)); + REQUIRE(Bits(a.observed_x) == Bits(b.observed_x)); + REQUIRE(Bits(a.observed_y) == Bits(b.observed_y)); + REQUIRE(Bits(a.image_scale_corr) == Bits(b.image_scale_corr)); + REQUIRE(Bits(a.prescaling_corr) == Bits(b.prescaling_corr)); + REQUIRE(a.observed == b.observed); + REQUIRE(a.overloaded == b.overloaded); + } +} + +void RequireSame(const BraggIntegrationCounts &a, const BraggIntegrationCounts &b) { + CHECK(a.predicted == b.predicted); + CHECK(a.bkg_starved == b.bkg_starved); + CHECK(a.bkg_starved_by_neighbour == b.bkg_starved_by_neighbour); + CHECK(a.profile_fallback == b.profile_fallback); + CHECK(a.overlapped == b.overlapped); + CHECK(a.overlapped_major == b.overlapped_major); +} + +// Three frames through one pair of engines - so whatever the Metal engine keeps between frames (its +// masks, its buffers) has to be left as the CPU engine leaves its own - each compared exactly. +// Returns how many reflections came out, so a case can show it was not vacuous. +size_t CompareFrames(const DetectorSetup &det, const Case &c) { + const DiffractionExperiment x = MakeExperiment(det, c); + REQUIRE(BraggIntegrationEngineMetal::Available(x)); + const size_t W = x.GetXPixelsNum(); + + // A few pixels of the run's mask, under spots and between them: unreadable there, but not overloads. + PixelMask mask(x); + std::vector user(x.GetPixelsNum(), 0); + std::vector masked; + for (size_t k = 0; k < x.GetPixelsNum(); k += 7919) { + user[k] = 1; + masked.push_back(k); + } + mask.LoadUserMask(x, user); + + BraggIntegrationEngineCPU cpu(x, mask); + BraggIntegrationEngineMetal metal(x, mask); + ImagePreprocessorBuffer image(x.GetPixelsNum()); + + size_t n_out = 0; + for (uint32_t frame = 0; frame < 3; ++frame) { + INFO("frame " << frame); + const Frame f = RandomFrame(x, 1000 * frame + static_cast(W), masked); + std::copy(f.image.begin(), f.image.end(), image.data()); + // The auto radial mode switches per image; here it is on for the radial cases, every image. + cpu.BackgroundRadial(c.radial); + metal.BackgroundRadial(c.radial); + const auto out_cpu = cpu.Run(image, f.predicted, f.predicted.size(), frame); + const auto out_metal = metal.Run(image, f.predicted, f.predicted.size(), frame); + RequireSame(out_cpu, out_metal); + n_out += out_cpu.size(); + } + RequireSame(cpu.Counts(), metal.Counts()); + // Not the CPU engine compared with itself: every frame went through the kernels. + REQUIRE(metal.FramesOnGPU() == 3); + return n_out; +} + +} // namespace + +TEST_CASE("BraggIntegrationEngineMetal_BitIdenticalToCPU", "[portable][Integration][Metal]") { + // A JUNGFRAU 1M, and a small odd-sized detector, where most boxes meet an edge. + const DetectorSetup large = DetJF(2); + const DetectorSetup small = DetDECTRIS(211, 197, "small", {}); + + std::vector> cases; + for (IntegratorMode mode : {IntegratorMode::BoxSum, IntegratorMode::ProfileGaussian, IntegratorMode::ProfileEmpirical}) + for (OverlapMode overlap : {OverlapMode::Off, OverlapMode::Reject, OverlapMode::Exclude}) + for (bool radial : {false, true}) { + Case c; + c.mode = mode; + c.overlap = overlap; + c.radial = radial; + cases.emplace_back("mode " + std::to_string(static_cast(mode)) + " overlap " + + std::to_string(static_cast(overlap)) + (radial ? " radial" : ""), c); + } + // Rings elongated by the bandwidth, and the measured footprint: per-reflection stencils. + Case broadband; + broadband.bandwidth = 0.01f; + broadband.stencil_k = 3.0f; + broadband.radial = true; + cases.emplace_back("broadband stencil radial", broadband); + Case footprint; + footprint.footprint = true; + footprint.radial = true; + cases.emplace_back("footprint radial", footprint); + Case footprint_box = footprint; + footprint_box.mode = IntegratorMode::BoxSum; + cases.emplace_back("footprint box sum", footprint_box); + // No clip: the plain ring mean, and no clip kernel at all. + Case no_clip; + no_clip.clip_nsigma = 0.0f; + cases.emplace_back("no clip", no_clip); + + for (const auto &[name, c] : cases) { + DYNAMIC_SECTION(name) { + CHECK(CompareFrames(large, c) > 1000); + CHECK(CompareFrames(small, c) > 50); + } + } +} + +// The trimmed-mean background sorts each ring, which only the CPU engine does: such a run keeps the CPU +// engine (MXAnalysisWithoutFPGA asks Available). +TEST_CASE("BraggIntegrationEngineMetal_TrimmedMeanStaysOnCPU", "[portable][Integration][Metal]") { + Case trim; + trim.clip_nsigma = 0.0f; + DiffractionExperiment x = MakeExperiment(DetJF(2), trim); + BraggIntegrationSettings s = x.GetBraggIntegrationSettings(); + s.BackgroundTrimFraction(0.1f); + x.ImportBraggIntegrationSettings(s); + CHECK(!BraggIntegrationEngineMetal::Available(x)); +} + +#endif diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 6f4632394..2ce3fe5d2 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -20,6 +20,7 @@ ADD_EXECUTABLE(jfjoch_portable_test EXCLUDE_FROM_ALL SpotFootprintTest.cpp BraggIntegrationEngineCompressedImageTest.cpp BraggIntegrationEngineCPUTest.cpp + BraggIntegrationEngineMetalTest.cpp WriteReflectionsTest.cpp LoadReferenceMtzFreeFlagsTest.cpp LoadReferenceSfCifTest.cpp