From ba3f7547520365622bddc91d2308594a7fde886f Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 14 Oct 2026 19:42:31 +0200 Subject: [PATCH 1/4] RESULT-CHANGING (owner approved, ulp level): CPU profile fit in the GPU engine's arithmetic Pass B of BraggIntegrationEngineCPU now does what BraggIntegrationEngineGPU's fit does: profiles in float, the Gaussian from expf with the widths' reciprocals taken once and normalised by a reciprocal, and the reweighted fit in float with one reciprocal of the variance for its three sums (the GPU's __frcp_rn). Per-reflection geometry and the final sigma stay in double. The logic was already the same on both; the CPU now differs from the GPU only in summation order. Measured (CPU-only, Apple M6): Pass B loops 1.2x faster on P cores; cytochrome C 1.9%, myoglobin 1.5% faster end to end, 4nwv integration steps 3%. Intensities move by ~1e-6 relative; merged statistics in the fourth digit (error model b 1.3259e-3 -> 1.3258e-3, R_work 0.2440 -> 0.2439). Co-Authored-By: Claude Opus 5.5 --- .../BraggIntegrationEngineCPU.cpp | 79 +++++++++++-------- 1 file changed, 44 insertions(+), 35 deletions(-) diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp index b62ab2aa9..4b7f785a9 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp @@ -537,22 +537,22 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, }; // Normalised empirical profile (sum = 1), the average grid over the strong spots of a shell. // ProfileGaussian does not use it - it rebuilds a per-reflection Gaussian in Pass B. + // Stored in float, as the fit below reads it (and as the GPU engine keeps it). auto build_profile = [&](const std::vector &grid) { - std::vector P(GG, 0.0); + std::vector P(GG, 0.0f); double sum = 0.0; - for (int k = 0; k < GG; ++k) { - P[k] = std::max(0.0, grid[k]); - sum += P[k]; - } + for (int k = 0; k < GG; ++k) + sum += std::max(0.0, grid[k]); if (sum > 0.0) - for (double &p : P) p /= sum; + for (int k = 0; k < GG; ++k) + P[k] = static_cast(std::max(0.0, grid[k]) / sum); return P; }; - const std::vector global_P = empirical && global_n > 0 ? build_profile(global_grid) - : std::vector(GG, 0.0); + const std::vector global_P = empirical && global_n > 0 ? build_profile(global_grid) + : std::vector(GG, 0.0f); const Sigma2 global_sigma2 = widths(global_mom); - std::vector> shell_P(N_SHELL, global_P); + std::vector> shell_P(N_SHELL, global_P); std::vector shell_sigma2(N_SHELL, global_sigma2); for (int s = 0; s < N_SHELL; ++s) { if (shell_n[s] < MIN_STRONG_PER_SHELL) continue; @@ -562,11 +562,14 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, // --- Pass B: profile-fit each reflection (Kabsch, de-biased variance v = B + I*P; iterate). The // reweighting is the Kabsch/Otwinowski iteration: Kabsch, Acta Cryst D66, 133-144 (2010); - // Otwinowski & Minor, Methods Enzymol 276, 307-326 (1997). --- - std::vector Pbuf; + // Otwinowski & Minor, Methods Enzymol 276, 307-326 (1997). + // In single precision, with the GPU engine's arithmetic (BraggIntegrationEngineGPU.cu, fit): + // the profile from expf with the widths' reciprocals taken once, normalised by a reciprocal, + // and one reciprocal of the variance for the three weights. --- + std::vector Pbuf; // The pixels the fit reads, gathered once while p_valid walks the grid: the profile value and the // background-subtracted count, in the grid's order. - struct FitPixel { double Pp, signal; }; + struct FitPixel { float Pp, signal; }; std::vector fit_pixels; for (size_t i = 0; i < npredicted; ++i) { const auto &rh = rough[i]; @@ -584,7 +587,7 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, const int sh = rh.shell < 0 ? 0 : rh.shell; int Rf = R; - const std::vector *Pvec = &shell_P[sh]; + const std::vector *Pvec = &shell_P[sh]; const BraggStencil st = MakeBraggStencil(predicted[i].predicted_x, predicted[i].predicted_y, stencil); if (!empirical) { const double rx = predicted[i].predicted_x - beam_x, ry = predicted[i].predicted_y - beam_y; @@ -625,17 +628,21 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, prefetch_row(img.address(static_cast(y) * W + gx0), img.address(static_cast(y) * W + gx1)); const int Gf = 2 * Rf + 1; - Pbuf.assign(static_cast(Gf) * Gf, 0.0); - double gs = 0.0; + Pbuf.assign(static_cast(Gf) * Gf, 0.0f); + const float fxf = static_cast(fx), fyf = static_cast(fy); + const float uxf = static_cast(ux), uyf = static_cast(uy); + const float inv_2s2r = 0.5f / static_cast(s2r), inv_2s2t = 0.5f / static_cast(s2t); + float gs = 0.0f; for (int dy = -Rf; dy <= Rf; ++dy) for (int dx = -Rf; dx <= Rf; ++dx) { - const double ex = dx - fx, ey = dy - fy; - const double rad = ex * ux + ey * uy, tn = -ex * uy + ey * ux; - const double g = std::exp(-rad * rad / (2.0 * s2r) - tn * tn / (2.0 * s2t)); + const float ex = dx - fxf, ey = dy - fyf; + const float rad = ex * uxf + ey * uyf, tn = -ex * uyf + ey * uxf; + const float g = std::exp(-rad * rad * inv_2s2r - tn * tn * inv_2s2t); Pbuf[(dy + Rf) * Gf + (dx + Rf)] = g; gs += g; } - for (double &p : Pbuf) p /= gs; + const float inv_gs = 1.0f / gs; + for (float &p : Pbuf) p *= inv_gs; Pvec = &Pbuf; } @@ -648,13 +655,13 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, // pixels, what Reject cuts on. m_read / m_all is the fraction of the r1 disk the // summation seed the runaway guard compares against actually summed; with nothing missing // and nothing excluded it is 1 and the guard is untouched. --- - double p_grid = 0.0, p_valid = 0.0, p_own = 0.0, m_all = 0.0, m_read = 0.0; - double p_peak = 0.0, p_lost_peak = 0.0; + float p_grid = 0.0f, p_valid = 0.0f, p_own = 0.0f, m_all = 0.0f, m_read = 0.0f; + float p_peak = 0.0f, p_lost_peak = 0.0f; fit_pixels.clear(); for (int dy = -Rf; dy <= Rf; ++dy) for (int dx = -Rf; dx <= Rf; ++dx) { - const double Pp = (*Pvec)[(dy + Rf) * Gf + (dx + Rf)]; - if (Pp <= 0.0) continue; + const float Pp = (*Pvec)[(dy + Rf) * Gf + (dx + Rf)]; + if (Pp <= 0.0f) continue; const int x = rh.cx + dx, y = rh.cy + dy; if (x < 0 || y < 0 || x >= W || y >= H) continue; // The region the summation seed was taken over (pass A), so the guard below @@ -677,7 +684,7 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, if (own) p_own += Pp; if (in_disk && (own || !exclude)) m_read += Pp; if (own || !exclude) - fit_pixels.push_back({Pp, static_cast(img[y * W + x]) - rh.bkg}); + fit_pixels.push_back({Pp, static_cast(img[y * W + x]) - static_cast(rh.bkg)}); } if (p_valid < overlap_min_peak * p_grid) continue; // A hole in the profile's PEAK is a different defect from a hole in its wings, and the mass @@ -687,21 +694,23 @@ std::vector BraggIntegrationEngineCPU::RunImpl(const Sampler &img, if (p_lost_peak > MINPK_MAX_MISSING_PEAK * p_peak) continue; if (overlap == OverlapMode::Reject && p_own < overlap_min_peak) continue; - const double B = std::max(rh.bkg, PIXEL_VARIANCE_FLOOR); - double I = rh.I, den = 0.0, wsum = 0.0; + const float B = static_cast(std::max(rh.bkg, PIXEL_VARIANCE_FLOOR)); + const float v_min = static_cast(WEIGHT_VARIANCE_MIN_FRACTION) * B; + float If = static_cast(rh.I), den = 0.0f, wsum = 0.0f; for (int iter = 0; iter < 4; ++iter) { - double num = 0.0; - den = 0.0; - wsum = 0.0; + float num = 0.0f; + den = 0.0f; + wsum = 0.0f; for (const auto &[Pp, signal] : fit_pixels) { - const double v = std::max(B + I * Pp, WEIGHT_VARIANCE_MIN_FRACTION * B); - num += Pp * signal / v; - den += Pp * Pp / v; - wsum += Pp / v; + const float iv = 1.0f / std::max(B + If * Pp, v_min); + num += Pp * signal * iv; + den += Pp * Pp * iv; + wsum += Pp * iv; } - if (den > 0.0) I = num / den; else break; + if (den > 0.0f) If = num / den; else break; } - if (!(den > 0.0)) continue; + if (!(den > 0.0f)) continue; + double I = If; // Guard against profile-fit runaways: on a weak / near-zero reflection the reweighted Kabsch // iteration has no real peak to lock onto and can manufacture intensity the box sum never sees. From f05e67b2b4f8087a17702b5122abdf0d29a7b328 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 14 Oct 2026 23:18:42 +0200 Subject: [PATCH 2/4] 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 From 4e177bdb2190743a0ceb40bdf2c6e1412a4cbee0 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 07:50:16 +0200 Subject: [PATCH 3/4] RotationScaleMerge: CPU partial scaling loop reads its working copy in order, output unchanged The loop's corr now lives in the group CSR's order next to the references (the partials in no group after them), so the group means read both sequentially instead of gathering corr through the permutation, and the corr update walks the same arrays. FitPerFrameG's filters, which do not change over the loop, are applied once into a per-frame ScalingFit list with the reference-independent factors precomputed (FitScalingFrames). ReferenceWeight is inline, and the group means are no longer pre-filled with NaN serially. Same values, same operations in the same order: p.hkl and the merged and unmerged MTZ data are byte-identical on lysozyme, cytochrome C, myoglobin, a small-molecule set and 4nwv. On 4nwv (Apple M6, 12 threads) the partial scaling iterations take 3-4x less time, ~18 s -> ~6 s over 168 iterations. Co-Authored-By: Claude Opus 5.5 --- .../scale_merge/RotationScaleMerge.cpp | 121 ++++++++++++------ .../scale_merge/RotationScaleMerge.h | 40 ++++-- 2 files changed, 109 insertions(+), 52 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 6840ec09e..08bbd1b51 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -275,6 +275,12 @@ namespace { // the loop shrank every scale by 10-30% per round for thirty rounds under the robust fit and // settled in fourteen under the plain one. With the weights the reference uses, the plain // alternating fit is exact coordinate descent on one objective and cannot climb. + double ScaleFromSums(double num, double den, double &info) { + const double G = den > 0.0 ? num / den : NAN; + info = den; + return std::isfinite(G) ? std::max(0.0, G) : 1.0; + } + double SolveScale(const std::vector &obs, double &info) { double num = 0.0, den = 0.0; for (const auto &o : obs) { @@ -282,9 +288,7 @@ namespace { num += w2 * o.coeff * o.Iobs; den += w2 * o.coeff * o.coeff; } - const double G = den > 0.0 ? num / den : NAN; - info = den; - return std::isfinite(G) ? std::max(0.0, G) : 1.0; + return ScaleFromSums(num, den, info); } double median_of(std::vector &v) { @@ -1512,18 +1516,6 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { return n_groups; } -double RotationScaleMerge::ReferenceWeight(int32_t group, float corr, float partiality, float I, float sigma, - double min_partiality, float &I_corr) { - if (group < 0) return 0.0; - if (!(corr > 0.0f) || !std::isfinite(corr)) return 0.0; - if (partiality < min_partiality) return 0.0; - const float Ic = I * corr; - const float sigma_corr = sigma * corr; - if (!std::isfinite(Ic) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0; - I_corr = Ic; - return 1.0 / (static_cast(sigma_corr) * sigma_corr); -} - void RotationScaleMerge::ReduceFullGroupMeans(int n_groups, std::vector &out_mean) const { // Inverse-variance per-group mean of I*corr = the merge reference (no cell mask: the scaling // reference is built without a reference cell - only the final merge applies it). A group per task @@ -1565,7 +1557,7 @@ void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vector &refs, const std::vector &corr, std::vector &out_mean) const { - out_mean.assign(n_groups, NAN); + out_mean.resize(n_groups); ParallelChunks(n_groups, ThreadsForWork(corr.size(), nthreads), [&](int glo, int ghi) { for (int g = glo; g < ghi; ++g) { double sw = 0.0, swI = 0.0; @@ -1573,13 +1565,12 @@ void RotationScaleMerge::ReduceScalingGroupMeans(int n_groups, const std::vector for (int p = lo; p < hi; ++p) { const ScalingRef &r = refs[p]; float I_corr = 0.0f; - const double w = ReferenceWeight(g, corr[partial_group_perm[p]], r.partiality, r.I, r.sigma, - min_partiality, I_corr); + const double w = ReferenceWeight(g, corr[p], r.partiality, r.I, r.sigma, min_partiality, I_corr); if (w <= 0.0) continue; sw += w; swI += w * I_corr; } - if (sw > 0.0) out_mean[g] = swI / sw; + out_mean[g] = sw > 0.0 ? swI / sw : NAN; } }); } @@ -1709,6 +1700,31 @@ void RotationScaleMerge::FitPerFrameG(const std::vector &obs, const std::vect frame_info_scratch = std::move(info); } +void RotationScaleMerge::FitScalingFrames(const std::vector &fit, const std::vector &fit_start, + const std::vector &group_mean_in, std::vector &g) { + const int n_fr = static_cast(fit_start.size()) - 1; + std::vector scaled(n_fr, 0); + std::vector info(n_fr, 0.0); + ParallelFor(n_fr, nthreads, [&](int f) { + double num = 0.0, den = 0.0; + size_t n = 0; + for (int i = fit_start[f]; i < fit_start[f + 1]; ++i) { + const ScalingFit &o = fit[i]; + const double mean = group_mean_in[o.group]; + if (!std::isfinite(mean)) continue; + const double coeff = o.coeff_per_mean * mean; + num += o.w2 * coeff * o.Iobs; + den += o.w2 * coeff * coeff; + ++n; + } + if (n < MIN_REFLECTIONS) return; + g[f] = ScaleFromSums(num, den, info[f]); + scaled[f] = 1; + }); + frame_scaled_scratch = std::move(scaled); + frame_info_scratch = std::move(info); +} + void RotationScaleMerge::ShrinkToRestrained(const char *what, const std::vector &g_restrained, const std::vector &fitted, const std::vector &frame_obs_count, @@ -6405,28 +6421,49 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool logger->Info("Per-frame scaling (partials): the same grouping and settings as an earlier merge, " "so its {} iterations stand", partial_loop.iterations); } else { - // The loop works on its own compact copy of the partials (ScalingObs in the header) and hands - // corr back to them when it is done. + // The loop works on its own compact copy of the partials (ScalingRef and ScalingFit in the + // header) and hands corr back to them when it is done. scaling_index[q] is the partial at + // position q of the copy: the CSR's group order, then the partials in no group. const int n_partials = static_cast(partials.size()); const size_t partial_threads = ThreadsForWork(partials.size(), nthreads); - std::vector scaling_obs(n_partials); + const int n_grouped = partial_group_start.empty() + ? 0 : partial_group_start.back() + partial_group_count.back(); + std::vector scaling_index(partial_group_perm.get(), partial_group_perm.get() + n_grouped); + for (int i = 0; i < n_partials; ++i) + if (partials[i].group < 0) scaling_index.push_back(i); + std::vector scaling_refs(n_partials); std::vector scaling_corr(n_partials); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - const Obs &o = partials[i]; - scaling_obs[i] = {o.frame, o.group, o.I, o.sigma, o.partiality, o.prescaling_corr, o.zeta, - o.on_ice}; - scaling_corr[i] = o.corr; + for (int q = lo; q < hi; ++q) { + const Obs &o = partials[scaling_index[q]]; + scaling_refs[q] = {o.frame, o.I, o.sigma, o.partiality, o.prescaling_corr}; + scaling_corr[q] = o.corr; } }); - std::vector scaling_refs(partial_group_start.empty() - ? 0 : partial_group_start.back() + partial_group_count.back()); - ParallelChunks(n_groups, partial_threads, [&](int glo, int ghi) { - for (int g = glo; g < ghi; ++g) - for (int p = partial_group_start[g]; p < partial_group_start[g] + partial_group_count[g]; ++p) { - const ScalingObs &o = scaling_obs[partial_group_perm[p]]; - scaling_refs[p] = {o.I, o.sigma, o.partiality}; - } + + // FitPerFrameG's filters, all but the reference mean's, frame by frame. + auto fitted_obs = [&](const Obs &o) { + return o.group >= 0 && !o.on_ice && o.partiality >= min_partiality + && std::isfinite(o.zeta) && o.zeta > 0.0f; + }; + const int n_fit_frames = static_cast(frame_start.size()); + std::vector fit_start(n_fit_frames + 1, 0); + ParallelFor(n_fit_frames, nthreads, [&](int f) { + fit_start[f + 1] = static_cast(std::count_if( + partials.begin() + frame_start[f], partials.begin() + frame_start[f] + frame_count[f], + fitted_obs)); + }); + for (int f = 0; f < n_fit_frames; ++f) fit_start[f + 1] += fit_start[f]; + std::vector scaling_fit(fit_start.back()); + ParallelFor(n_fit_frames, nthreads, [&](int f) { + int32_t k = fit_start[f]; + for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) { + const Obs &o = partials[i]; + if (!fitted_obs(o)) continue; + const double weight = SafeInv(o.sigma, 1.0); + scaling_fit[k++] = {o.group, o.partiality * SafeInv(o.prescaling_corr, 1.0), + static_cast(o.I), weight * weight}; + } }); // RunScalingLoop rescales after every iteration before anything reads corr, so the iteration's @@ -6437,14 +6474,14 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool partial_loop = RunScalingLoop("partials", g_partial, frame_count, smooth_window, [&] { ReduceScalingGroupMeans(n_groups, scaling_refs, scaling_corr, partial_mean); - FitPerFrameG(scaling_obs, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial); + FitScalingFrames(scaling_fit, fit_start, partial_mean, g_partial); g_fitted = g_partial; frame_fitted = frame_scaled_scratch; }, [&](const std::vector &apply, const std::vector &ratio) { ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - const ScalingObs &o = scaling_obs[i]; - float &corr = scaling_corr[i]; + for (int q = lo; q < hi; ++q) { + const ScalingRef &o = scaling_refs[q]; + float &corr = scaling_corr[q]; if (frame_fitted[o.frame]) { // UpdateCorr const double denom = static_cast(o.partiality) * g_fitted[o.frame]; if (std::isfinite(o.prescaling_corr) && std::isfinite(denom) && denom > 0.0) @@ -6459,7 +6496,11 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool }); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) partials[i].corr = scaling_corr[i]; + for (int q = lo; q < hi; ++q) partials[scaling_index[q]].corr = scaling_corr[q]; + }); + // The memo keeps corr in the partials' own order. + ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { + for (int i = lo; i < hi; ++i) scaling_corr[i] = partials[i].corr; }); PartialLoopMemo m{key, partial_loop, frame_scaled_scratch, g_partial, std::move(scaling_corr)}; if (partial_loop_memos.size() == MAX_PARTIAL_LOOP_MEMOS) diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 34baebcd9..cb77e956b 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -616,22 +616,35 @@ private: // The CPU partial scaling loop's working copy of the partials. It passes over them several times // per iteration and reads a handful of fields each time, so it works on these rather than on the - // eighty-byte Obs: ScalingObs holds, in the partials' own order, what FitPerFrameG and the corr - // update read; ScalingRef holds, in the CSR's group order, what the group means read besides corr, - // so that only corr is gathered. The same values, so the same arithmetic in the same order. - struct ScalingObs { - int32_t frame, group; - float I, sigma, partiality, prescaling_corr, zeta; - uint8_t on_ice; - }; + // eighty-byte Obs. ScalingRef holds what the group means and the corr update read, in the CSR's + // group order with the partials in no group after them, and corr is kept beside it in the same + // order - so neither pass gathers. ScalingFit holds, frame by frame, the observations + // FitScalingFrames fits: FitPerFrameG's filters, which do not change over the loop, applied once, + // and the factors of its sums that do not depend on the reference worked out once. The same + // values, so the same arithmetic in the same order. struct ScalingRef { - float I, sigma, partiality; + int32_t frame; + float I, sigma, partiality, prescaling_corr; + }; + struct ScalingFit { + int32_t group; + double coeff_per_mean; // partiality / prescaling_corr: FitPerFrameG's coeff for a mean of 1 + double Iobs, w2; }; // The observation's inverse-variance weight in the scaling reference, and its scaled intensity; // 0 when the observation is not in the reference at all. static double ReferenceWeight(int32_t group, float corr, float partiality, float I, float sigma, - double min_partiality, float &I_corr); + double min_partiality, float &I_corr) { + if (group < 0) return 0.0; + if (!(corr > 0.0f) || !std::isfinite(corr)) return 0.0; + if (partiality < min_partiality) return 0.0; + const float Ic = I * corr; + const float sigma_corr = sigma * corr; + if (!std::isfinite(Ic) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0; + I_corr = Ic; + return 1.0 / (static_cast(sigma_corr) * sigma_corr); + } static double ReferenceWeight(const Obs &o, double min_partiality, float &I_corr) { return ReferenceWeight(o.group, o.corr, o.partiality, o.I, o.sigma, min_partiality, I_corr); } @@ -640,17 +653,20 @@ private: void ReduceFullGroupMeans(int n_groups, std::vector &out_mean) const; // The same over the partials, using the CSR above - the same sums in the same order. void ReducePartialGroupMeans(int n_groups, std::vector &out_mean) const; - // The same again over the scaling loop's copy: refs in group order, corr in the partials' order. + // The same again over the scaling loop's copy, refs and corr both in group order. void ReduceScalingGroupMeans(int n_groups, const std::vector &refs, const std::vector &corr, std::vector &out_mean) const; // Robust per-frame G fit (IRLS, Cauchy k=3), unity=false uses the rotation partiality, unity=true the // scale-fulls (partiality already folded in). Reads out_mean[group] as the reference intensity. - // O is Obs or ScalingObs. template void FitPerFrameG(const std::vector &obs, const std::vector &fstart, const std::vector &fcount, const std::vector &group_mean_in, bool unity, std::vector &g); + // FitPerFrameG(unity = false) over the scaling loop's ScalingFit, frame f's at + // [fit_start[f], fit_start[f + 1]). + void FitScalingFrames(const std::vector &fit, const std::vector &fit_start, + const std::vector &group_mean_in, std::vector &g); // The alternating per-frame scale fit run to convergence: `iterate` fits every frame once against // the current reference (leaving G in `g` and the frames it fitted in frame_scaled_scratch), From 229a432a66bffb6ff4e3627fa5d9872304b4612c Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 08:18:43 +0200 Subject: [PATCH 4/4] RESULT-CHANGING (owner approved, ulp level): Metal profile fit, CPU Gaussian from a shared exp Pass B of BraggIntegrationEngineCPU is split where a device can take over: per reflection, what the fit needs is worked out in double (BraggFitSetup), the profile, its readable mass and the reweighted fit run in single precision (FitProfile -> BraggFitOut), and the error, the summation guard and the result are finished in double. BraggIntegrationEngineMetal runs the middle as a Metal kernel, one thread per reflection, summing in the CPU's order (FitOnDevice); the CUDA engine is untouched. The Gaussian's exp is now BraggExp (BraggExp.h), on the CPU and on the GPU: no GPU exp reproduces the host's expf, and BraggExp uses only +, -, * and exact power-of-two scaling, so both compile it to the same bits. Checked on the Apple M6 over every float in [-110, 89] (2.2e9 arguments), with subnormal products and sums matching too. It is closer to correctly rounded than macOS expf (max 0.53 ulp; 0.2% of arguments below -1 not correctly rounded vs 0.62%). This is the one result change: per observation |dI|/sigma p99.9 <= 1.7e-6, max 1e-4, 10-30x below what ba3f75475 moved. Merged statistics (resolution, completeness, CC1/2, R_meas, I/sigma, ISa, space group) are identical on lysozyme, cytochrome C, myoglobin, a small-molecule set and 4nwv; only tNCS fields that flip on any last-digit change move (cytochrome C's pseudo-translation flag, whose modulation ratio sits at the 2.5 threshold, turns off). The Metal engine stays bit-identical to the CPU engine: BraggIntegrationEngineMetalTest passes unchanged, and catches planted kernel bugs (another exp, 3 iterations, a dropped lost-peak test, <= for the disk, inverted exclude logic). Measured wall time, Apple M6 (CPU pass B -> GPU pass B): lysozyme 24.7 -> 18.6 s, myoglobin 39.8 -> 35.7 s, cytochrome C 59.3 -> 49.0 s, 4nwv 150.2 -> 131.3 s; the small-molecule set, sparse frames, 92.5 -> 95.6 s (noise of +-3 s). Metal's fast::exp in the kernel instead measured the same on every set (18.6, 35.8, 49.1, 131.2 s), so it would buy nothing for the bitwise check it would cost. On the CPU alone BraggExp costs nothing measurable (cytochrome C --no-gpu 72.8 -> 73.1 s). Co-Authored-By: Claude Opus 5.5 --- docs/ACKNOWLEDGEMENT.md | 7 + docs/GPU_DATA_ANALYSIS.md | 11 +- image_analysis/bragg_integration/BraggExp.h | 63 +++++ .../BraggIntegrationEngineCPU.cpp | 256 +++++++++++------- .../BraggIntegrationEngineCPU.h | 18 ++ .../BraggIntegrationEngineMetal.h | 23 +- .../BraggIntegrationEngineMetal.metal | 86 ++++++ .../BraggIntegrationEngineMetal.mm | 64 ++++- image_analysis/bragg_integration/BraggPassB.h | 55 ++++ .../bragg_integration/CMakeLists.txt | 7 +- 10 files changed, 470 insertions(+), 120 deletions(-) create mode 100644 image_analysis/bragg_integration/BraggExp.h create mode 100644 image_analysis/bragg_integration/BraggPassB.h diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index be74d7ca6..9a4afd7d4 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -523,6 +523,13 @@ corrected once by the exact remainder, which returns the correctly rounded quoti P. W. Markstein, "Computation of elementary functions on the IBM RISC System/6000 processor" (1990), IBM J. Res. Dev. 34, 111-119. +The single-precision exp the Bragg integration's profile fit shares between the CPU and an Apple GPU +(`BraggExp.h`) is built from textbook error-free floating-point arithmetic: the argument reduction of +W. J. Cody and W. Waite, "Software Manual for the Elementary Functions" (1980), Prentice-Hall; the +exact product (with Veltkamp's split) and Fast2Sum of T. J. Dekker, "A floating-point technique for +extending the available precision" (1971), Numer. Math. 18, 224-242; and the 2Sum of D. E. Knuth, +"The Art of Computer Programming, Vol. 2: Seminumerical Algorithms", section 4.2.2. + This software uses the Viridis, Magma and Inferno colormaps from Matplotlib under its BSD-compatible license. J. D. Hunter, "Matplotlib: A 2D graphics environment" (2007), Comput. Sci. Eng. 9, 90-95 [doi:10.1109/MCSE.2007.55](https://doi.org/10.1109/MCSE.2007.55). diff --git a/docs/GPU_DATA_ANALYSIS.md b/docs/GPU_DATA_ANALYSIS.md index 8fab936d7..4573c4a69 100644 --- a/docs/GPU_DATA_ANALYSIS.md +++ b/docs/GPU_DATA_ANALYSIS.md @@ -186,9 +186,10 @@ structure factors above; the likelihood scoring runs on the host. The search is 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. + pixel work moved to the GPU: the neighbour mask, the ownership of shared pixels, the box + summation's walk over every reflection, the background clip and the single-precision profile fit. + Everything computed in double and the profile learning stay on the CPU, as Apple GPUs have no + double precision. The fit's Gaussian uses an exp of our own (`BraggExp.h`) on both, as no GPU exp + matches the CPU's. 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/bragg_integration/BraggExp.h b/image_analysis/bragg_integration/BraggExp.h new file mode 100644 index 000000000..818afd4dc --- /dev/null +++ b/image_analysis/bragg_integration/BraggExp.h @@ -0,0 +1,63 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// exp in single precision, the same bits wherever it is compiled: the host compiler (no FMA contraction, +// see CMakeLists.txt) and the Metal compiler (fast math off) both see only IEEE +, -, * rounded to +// nearest and exact scaling by a power of two, so the profile fit's Gaussian is the same on the CPU and +// on an Apple GPU, whose exp no host function reproduces. It is closer to correctly rounded than macOS's +// expf: at most 0.53 ulp for a normal result, and away from a correctly rounded float for 0.2% of +// arguments below -1 and 0.004% in [-1, 0]. +// +// exp(x) = 2^k exp(r), k = rint(x / ln2), r = x - k ln2 with ln2 split three ways, the first two parts +// of 16 bits so that k times them is exact (Cody & Waite). r is carried as rr + re, rr^2 is formed exactly +// as zh + zl (Dekker's product with Veltkamp's split) and the sums into 1 + rr + rr^2/2 are error-free +// (2Sum, Fast2Sum), so the only rounding that matters is the last one. The rest of the series is a +// degree-4 minimax fit, relative error 1.3e-10. + +#ifdef __METAL_VERSION__ +inline float BraggBitsToFloat(int i) { return as_type(i); } +#else +#include +#include +inline float BraggBitsToFloat(int32_t i) { return std::bit_cast(i); } +#endif + +// p * 2^k, rounded once where the result is subnormal. +inline float BraggScale2(float p, int k) { + if (k < -125) + return (p * BraggBitsToFloat((k + 100 + 127) << 23)) * BraggBitsToFloat((-100 + 127) << 23); + return p * BraggBitsToFloat((k + 127) << 23); +} + +// Following Cody & Waite (1980) Software Manual for the Elementary Functions; Dekker (1971) Numer. Math. +// 18, 224-242; Knuth, The Art of Computer Programming Vol. 2, 4.2.2 (2Sum). +inline float BraggExp(float x) { + if (x < -104.0f) return 0.0f; // below half the smallest subnormal + if (x > 88.72283f) return BraggBitsToFloat(0x7f800000); // overflows + const float shifter = 12582912.0f; // 1.5 * 2^23: adding it rounds to an integer + const float kf = (x * 1.44269504f + shifter) - shifter; + const int k = static_cast(kf); + const float r1 = x - kf * 0.693145751953125f; // exact + const float t2 = kf * 1.4286197256296873e-06f; // exact + const float rr = r1 - t2; // r1 - t2 = rr + e2 (2Sum) + const float bv0 = rr - r1; + const float e2 = (r1 - (rr - bv0)) + (-t2 - bv0); + const float re = e2 - kf * -1.2905320041778356e-11f; + const float sp = rr * 4097.0f; // rr^2 = zh + zl + const float rh = sp - (sp - rr); + const float rl = rr - rh; + const float zh = rr * rr; + const float zl = ((rh * rh - zh) + (rh * rl + rh * rl)) + rl * rl; + const float h = 0.1666666567325592f + rr * (0.04166644066572189f + rr * (0.008333569392561913f + + rr * (0.0013937620678916574f + rr * 0.0001974606711883098f))); + const float c = 0.5f * zl + ((re + re * rr) + (zh * rr) * h); + const float hz = 0.5f * zh; // rr + zh/2 = s1 + e1 (2Sum) + const float s1 = rr + hz; + const float bv1 = s1 - rr; + const float e1 = (rr - (s1 - bv1)) + (hz - bv1); + const float t1 = 1.0f + s1; // 1 + s1 = t1 + e3 (Fast2Sum, |s1| < 1) + const float e3 = (1.0f - t1) + s1; + return BraggScale2(t1 + (e3 + (e1 + c)), k); +} diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp index b488aff79..aa574ae17 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp @@ -429,6 +429,107 @@ void BraggIntegrationEngineCPU::FinishPassA(PassA &a, size_t i, const Reflection a.rough[i] = out; } +bool BraggIntegrationEngineCPU::FitKept(const BraggFitOut &o) const { + if (o.p_valid < overlap_min_peak * o.p_grid) return false; + // A hole in the profile's PEAK is a different defect from a hole in its wings, and the mass + // fraction above cannot tell them apart - the peak of a broad spot is a few percent of the + // mass, so MINPK passes a reflection that has lost the one part of the profile its amplitude + // is determined by. See the header. + if (o.p_lost_peak > MINPK_MAX_MISSING_PEAK * o.p_peak) return false; + if (overlap == OverlapMode::Reject && o.p_own < overlap_min_peak) return false; + return true; +} + +template +BraggFitOut BraggIntegrationEngineCPU::FitProfile(const Sampler &img, const OwnerMap &owner_map, + const BraggFitSetup &s, size_t i, + const std::vector &shell_profiles, + std::vector &cells, + std::vector &fit_pixels) const { + const int W = static_cast(xpixel), H = static_cast(ypixel); + const bool exclude = overlap == OverlapMode::Exclude; + const bool use_owner = overlap != OverlapMode::Off; + const int Rf = s.Rf, Gf = 2 * Rf + 1; + BraggFitOut o{}; + + const float *P; + if (s.shell >= 0) + P = shell_profiles.data() + static_cast(s.shell) * GG; + else { + // The grid's image rows, asked for now (as in pass A): building the profile below takes long + // enough for them to arrive before the fit reads them. + const int gx0 = std::max(0, s.cx - Rf), gx1 = std::min(W - 1, s.cx + Rf); + for (int y = std::max(0, s.cy - Rf); y <= std::min(H - 1, s.cy + Rf); ++y) + prefetch_row(img.address(static_cast(y) * W + gx0), + img.address(static_cast(y) * W + gx1)); + cells.resize(static_cast(Gf) * Gf); + float gs = 0.0f; + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float g = BraggGaussian(dx, dy, s.fx, s.fy, s.ux, s.uy, s.inv_2s2r, s.inv_2s2t); + cells[(dy + Rf) * Gf + (dx + Rf)] = g; + gs += g; + } + const float inv_gs = 1.0f / gs; + for (float &p : cells) p *= inv_gs; + P = cells.data(); + } + + // --- How much of the expected profile the fit can actually see. p_valid is the readable + // fraction of it - XDS's MINPK, dials' valid_foreground_threshold - measured against the + // mass that falls on the detector at all, so a reflection is judged on the pixels that + // exist and not on where the sensor ends. p_own is the same quantity over neighbour-owned + // pixels, what Reject cuts on. m_read / m_all is the fraction of the r1 disk the + // summation seed the runaway guard compares against actually summed; with nothing missing + // and nothing excluded it is 1 and the guard is untouched. --- + const BraggStencil st = MakeBraggStencil(s.px, s.py, stencil); + fit_pixels.clear(); + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float Pp = P[(dy + Rf) * Gf + (dx + Rf)]; + if (Pp <= 0.0f) continue; + const int x = s.cx + dx, y = s.cy + dy; + if (x < 0 || y < 0 || x >= W || y >= H) continue; + // The region the summation seed was taken over (pass A), so the guard compares the two + // over the same pixels. + const bool in_disk = s.footprint + ? BraggStencilInSignal(st, x - s.px, y - s.py, + (x - s.px) * (x - s.px) + (y - s.py) * (y - s.py), r1_sq) + : dx * dx + dy * dy < r1_sq; + o.p_grid += Pp; + o.p_peak = std::max(o.p_peak, Pp); + if (in_disk) o.m_all += Pp; + const int32_t v = img[y * W + x]; + if (!valid(v)) { + o.p_lost_peak = std::max(o.p_lost_peak, Pp); + continue; + } + o.p_valid += Pp; + const bool own = !use_owner || BraggOwnedBy(owner_map.Get(x, y), static_cast(i)); + if (own) o.p_own += Pp; + if (in_disk && (own || !exclude)) o.m_read += Pp; + if (own || !exclude) + fit_pixels.push_back({Pp, static_cast(v) - s.bkg}); + } + if (!FitKept(o)) return o; + + float If = s.I0; + for (int iter = 0; iter < 4; ++iter) { + float num = 0.0f; + o.den = 0.0f; + o.wsum = 0.0f; + for (const auto &[Pp, signal] : fit_pixels) { + const float iv = 1.0f / std::max(s.B + If * Pp, s.v_min); + num += Pp * signal * iv; + o.den += Pp * Pp * iv; + o.wsum += Pp * iv; + } + if (o.den > 0.0f) If = num / o.den; else break; + } + o.I = If; + return o; +} + template std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler &img, const OwnerMap &owner_map, const std::vector &predicted, @@ -578,16 +679,21 @@ std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler & // --- Pass B: profile-fit each reflection (Kabsch, de-biased variance v = B + I*P; iterate). The // reweighting is the Kabsch/Otwinowski iteration: Kabsch, Acta Cryst D66, 133-144 (2010); // Otwinowski & Minor, Methods Enzymol 276, 307-326 (1997). - // In single precision, with the GPU engine's arithmetic (BraggIntegrationEngineGPU.cu, fit): - // the profile from expf with the widths' reciprocals taken once, normalised by a reciprocal, - // and one reciprocal of the variance for the three weights. --- - std::vector Pbuf; - // The pixels the fit reads, gathered once while p_valid walks the grid: the profile value and the - // background-subtracted count, in the grid's order. - struct FitPixel { float Pp, signal; }; - std::vector fit_pixels; + // Per reflection in three steps: what the fit needs, worked out in double (BraggFitSetup); the + // fit in single precision (FitProfile, or the Metal engine's kernel), with the GPU engine's + // arithmetic (BraggIntegrationEngineGPU.cu, fit) - the profile from exp with the widths' + // reciprocals taken once, normalised by a reciprocal, and one reciprocal of the variance for the + // three weights; and its error, the summation guard and the result, in double. --- + std::vector shell_profiles; // the learned profiles, shell by shell (ProfileEmpirical) + if (empirical) + for (int s = 0; s < N_SHELL; ++s) + shell_profiles.insert(shell_profiles.end(), shell_P[s].begin(), shell_P[s].end()); + std::vector setups(npredicted); + uint32_t cells = 0; for (size_t i = 0; i < npredicted; ++i) { const auto &rh = rough[i]; + BraggFitSetup &fs = setups[i]; + fs = {}; if (!rh.ok) continue; // An overloaded reflection is not measured by any fit - its brightest pixels are what is // missing. It is passed on with its box sum and the flag, for the merge to drop its rocking @@ -600,10 +706,15 @@ std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler & continue; } const int sh = rh.shell < 0 ? 0 : rh.shell; - - int Rf = R; - const std::vector *Pvec = &shell_P[sh]; const BraggStencil st = MakeBraggStencil(predicted[i].predicted_x, predicted[i].predicted_y, stencil); + fs.fit = 1; + fs.px = predicted[i].predicted_x; + fs.py = predicted[i].predicted_y; + fs.cx = rh.cx; + fs.cy = rh.cy; + fs.Rf = R; + fs.shell = sh; + fs.footprint = st.fp_s2t > 0.0f; if (!empirical) { const double rx = predicted[i].predicted_x - beam_x, ry = predicted[i].predicted_y - beam_y; const double Rpx = std::hypot(rx, ry); @@ -632,100 +743,41 @@ std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler & s2r = std::max(s2r, st.fp_s2r); elong = true; } - // Build the Gaussian per reflection, centred on the sub-pixel predicted position and (when + // The Gaussian is built per reflection, centred on the sub-pixel predicted position and (when // needed) elongated, on a grid grown to hold the wider of its two axes. - const double fx = predicted[i].predicted_x - rh.cx, fy = predicted[i].predicted_y - rh.cy; - Rf = elong ? std::min(3 * R, static_cast(std::ceil(r2 + 2.0 * std::sqrt(std::max(s2r, s2t))))) : R; - // The grid's image rows, asked for now (as in pass A): building the profile below takes - // long enough for them to arrive before the fit reads them. - const int gx0 = std::max(0, rh.cx - Rf), gx1 = std::min(W - 1, rh.cx + Rf); - for (int y = std::max(0, rh.cy - Rf); y <= std::min(H - 1, rh.cy + Rf); ++y) - prefetch_row(img.address(static_cast(y) * W + gx0), - img.address(static_cast(y) * W + gx1)); - const int Gf = 2 * Rf + 1; - Pbuf.assign(static_cast(Gf) * Gf, 0.0f); - const float fxf = static_cast(fx), fyf = static_cast(fy); - const float uxf = static_cast(ux), uyf = static_cast(uy); - const float inv_2s2r = 0.5f / static_cast(s2r), inv_2s2t = 0.5f / static_cast(s2t); - float gs = 0.0f; - for (int dy = -Rf; dy <= Rf; ++dy) - for (int dx = -Rf; dx <= Rf; ++dx) { - const float ex = dx - fxf, ey = dy - fyf; - const float rad = ex * uxf + ey * uyf, tn = -ex * uyf + ey * uxf; - const float g = std::exp(-rad * rad * inv_2s2r - tn * tn * inv_2s2t); - Pbuf[(dy + Rf) * Gf + (dx + Rf)] = g; - gs += g; - } - const float inv_gs = 1.0f / gs; - for (float &p : Pbuf) p *= inv_gs; - Pvec = &Pbuf; + fs.Rf = elong ? std::min(3 * R, static_cast(std::ceil(r2 + 2.0 * std::sqrt(std::max(s2r, s2t))))) : R; + fs.shell = -1; + fs.cells = cells; + cells += (2 * fs.Rf + 1) * (2 * fs.Rf + 1); + fs.fx = static_cast(predicted[i].predicted_x - rh.cx); + fs.fy = static_cast(predicted[i].predicted_y - rh.cy); + fs.ux = static_cast(ux); + fs.uy = static_cast(uy); + fs.inv_2s2r = 0.5f / static_cast(s2r); + fs.inv_2s2t = 0.5f / static_cast(s2t); } + fs.bkg = static_cast(rh.bkg); + fs.B = static_cast(std::max(rh.bkg, PIXEL_VARIANCE_FLOOR)); + fs.v_min = static_cast(WEIGHT_VARIANCE_MIN_FRACTION) * fs.B; + fs.I0 = static_cast(rh.I); + } - const int Gf = 2 * Rf + 1; + std::vector fits(npredicted); + if (!FitOnDevice(setups, shell_profiles, fits)) { + std::vector cells_buf; + std::vector fit_pixels; + for (size_t i = 0; i < npredicted; ++i) + if (setups[i].fit) + fits[i] = FitProfile(img, owner_map, setups[i], i, shell_profiles, cells_buf, fit_pixels); + } - // --- How much of the expected profile the fit can actually see. p_valid is the readable - // fraction of it - XDS's MINPK, dials' valid_foreground_threshold - measured against the - // mass that falls on the detector at all, so a reflection is judged on the pixels that - // exist and not on where the sensor ends. p_own is the same quantity over neighbour-owned - // pixels, what Reject cuts on. m_read / m_all is the fraction of the r1 disk the - // summation seed the runaway guard compares against actually summed; with nothing missing - // and nothing excluded it is 1 and the guard is untouched. --- - float p_grid = 0.0f, p_valid = 0.0f, p_own = 0.0f, m_all = 0.0f, m_read = 0.0f; - float p_peak = 0.0f, p_lost_peak = 0.0f; - fit_pixels.clear(); - for (int dy = -Rf; dy <= Rf; ++dy) - for (int dx = -Rf; dx <= Rf; ++dx) { - const float Pp = (*Pvec)[(dy + Rf) * Gf + (dx + Rf)]; - if (Pp <= 0.0f) continue; - const int x = rh.cx + dx, y = rh.cy + dy; - if (x < 0 || y < 0 || x >= W || y >= H) continue; - // The region the summation seed was taken over (pass A), so the guard below - // compares the two over the same pixels. - const bool in_disk = st.fp_s2t > 0.0f - ? BraggStencilInSignal(st, x - predicted[i].predicted_x, y - predicted[i].predicted_y, - (x - predicted[i].predicted_x) * (x - predicted[i].predicted_x) - + (y - predicted[i].predicted_y) * (y - predicted[i].predicted_y), - r1_sq) - : dx * dx + dy * dy < r1_sq; - p_grid += Pp; - p_peak = std::max(p_peak, Pp); - if (in_disk) m_all += Pp; - if (!valid(img[y * W + x])) { - p_lost_peak = std::max(p_lost_peak, Pp); - continue; - } - p_valid += Pp; - const bool own = clean(x, y, i); - if (own) p_own += Pp; - if (in_disk && (own || !exclude)) m_read += Pp; - if (own || !exclude) - fit_pixels.push_back({Pp, static_cast(img[y * W + x]) - static_cast(rh.bkg)}); - } - if (p_valid < overlap_min_peak * p_grid) continue; - // A hole in the profile's PEAK is a different defect from a hole in its wings, and the mass - // fraction above cannot tell them apart - the peak of a broad spot is a few percent of the - // mass, so MINPK passes a reflection that has lost the one part of the profile its amplitude - // is determined by. See the header. - if (p_lost_peak > MINPK_MAX_MISSING_PEAK * p_peak) continue; - if (overlap == OverlapMode::Reject && p_own < overlap_min_peak) continue; - - const float B = static_cast(std::max(rh.bkg, PIXEL_VARIANCE_FLOOR)); - const float v_min = static_cast(WEIGHT_VARIANCE_MIN_FRACTION) * B; - float If = static_cast(rh.I), den = 0.0f, wsum = 0.0f; - for (int iter = 0; iter < 4; ++iter) { - float num = 0.0f; - den = 0.0f; - wsum = 0.0f; - for (const auto &[Pp, signal] : fit_pixels) { - const float iv = 1.0f / std::max(B + If * Pp, v_min); - num += Pp * signal * iv; - den += Pp * Pp * iv; - wsum += Pp * iv; - } - if (den > 0.0f) If = num / den; else break; - } - if (!(den > 0.0f)) continue; - double I = If; + for (size_t i = 0; i < npredicted; ++i) { + if (!setups[i].fit) continue; + const auto &rh = rough[i]; + const BraggFitOut &fo = fits[i]; + if (!FitKept(fo) || !(fo.den > 0.0f)) continue; + const float den = fo.den, wsum = fo.wsum; + double I = fo.I; // Guard against profile-fit runaways: on a weak / near-zero reflection the reweighted Kabsch // iteration has no real peak to lock onto and can manufacture intensity the box sum never sees. @@ -746,7 +798,7 @@ std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler & // The seed is a sum over the disk the box sum actually read, so when pixels are missing from // both - excluded to a neighbour, or unreadable - the fit's full-profile intensity has to be // scaled down to that same disk before the two are comparable. Nothing dropped gives 1. - const double guard_scale = m_all > 0.0 ? m_read / m_all : 1.0; + const double guard_scale = fo.m_all > 0.0f ? fo.m_read / fo.m_all : 1.0; if (std::abs(I * guard_scale - rh.I) > PROFILE_SUMMATION_MAX_NSIGMA * rh.sigma) { I = rh.I; sigma = rh.sigma; diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h index b9ee60116..e72497831 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h @@ -5,6 +5,7 @@ #include "BraggIntegrationEngine.h" #include "BraggPassA.h" +#include "BraggPassB.h" #include "../../common/PixelMask.h" class CompressedImage; @@ -62,6 +63,17 @@ class BraggIntegrationEngineCPU : public BraggIntegrationEngine { std::vector RunImpl(const Sampler &img, const std::vector &predicted, size_t npredicted, int64_t image_number); + // Pass B's single-precision part for reflection i, as the Metal engine's kernel does it: the profile + // (the shell's learned one, or the Gaussian, built in `cells`), how much of it is readable, and the + // fit. fit_pixels is scratch: the pixels the fit reads, gathered once. + struct FitPixel { float Pp, signal; }; + template + BraggFitOut FitProfile(const Sampler &img, const OwnerMap &owner_map, const BraggFitSetup &s, size_t i, + const std::vector &shell_profiles, std::vector &cells, + std::vector &fit_pixels) const; + // Whether enough of a reflection's profile is readable, and its own, for the fit to be kept. + bool FitKept(const BraggFitOut &o) const; + 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). @@ -118,6 +130,12 @@ protected: // 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. + // Pass B's single-precision part for every reflection with setups[i].fit, on a device - the + // learned profiles side by side in shell_profiles, empty for the Gaussian. False where there is no + // device, and the CPU fits them instead. + virtual bool FitOnDevice(const std::vector &setups, const std::vector &shell_profiles, + std::vector &fits) { return false; } + template std::vector ProfilePasses(const Sampler &img, const OwnerMap &owner_map, const std::vector &predicted, size_t npredicted, diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h index 6f4c19d74..ffacb973f 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h @@ -8,23 +8,26 @@ #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. +// Apple GPU engine (macOS only): BraggIntegrationEngineCPU with its pixel work - the neighbour mask, the +// owner map, pass A's walk over every box, the background clip and pass B's single-precision profile +// fit - run as Metal kernels, one thread per reflection. Everything computed in double and the profile +// learning stay on the CPU, in the CPU engine's own code: Apple GPUs have no double precision. // // 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. +// visits pixels in, the fit's float sums are taken in the CPU's order, one reflection per thread, and +// the arithmetic deciding which pixel is which is the same source (BraggStencil.h, BraggPassB.h, and +// BraggExp.h for the Gaussian, as no GPU exp matches the host's), 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; + bool FitOnDevice(const std::vector &setups, const std::vector &shell_profiles, + std::vector &fits) override; + 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, diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal index 7ebf36fd3..4d2632fa3 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal @@ -187,3 +187,89 @@ kernel void clip_sums(device const float2 *xy [[buffer(0)]], } out[i] = c; } + +// BraggIntegrationEngineCPU::FitProfile. The CPU gathers the pixels the fit reads once; here every +// iteration reads them again, the same pixels in the same order. +kernel void profile_fit(device const BraggFitSetup *setups [[buffer(0)]], + device const int32_t *img [[buffer(1)]], + device const uint32_t *owner [[buffer(2)]], + device const float *shell_profiles [[buffer(3)]], + device float *cells [[buffer(4)]], + constant char *stencil_raw [[buffer(5)]], + constant BraggPassBFrame &f [[buffer(6)]], + device BraggFitOut *out [[buffer(7)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n || !setups[i].fit) + return; + const BraggFitSetup s = setups[i]; + const int Rf = s.Rf, Gf = 2 * Rf + 1; + BraggFitOut o = {}; + + device const float *P; + if (s.shell >= 0) + P = shell_profiles + s.shell * (2 * f.R + 1) * (2 * f.R + 1); + else { + device float *g = cells + s.cells; + float gs = 0.0f; + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float v = BraggGaussian(dx, dy, s.fx, s.fy, s.ux, s.uy, s.inv_2s2r, s.inv_2s2t); + g[(dy + Rf) * Gf + (dx + Rf)] = v; + gs += v; + } + const float inv_gs = 1.0f / gs; + for (int k = 0; k < Gf * Gf; ++k) g[k] *= inv_gs; + P = g; + } + + const BraggStencil st = MakeBraggStencil(s.px, s.py, stencil_params(stencil_raw)); + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float Pp = P[(dy + Rf) * Gf + (dx + Rf)]; + if (Pp <= 0.0f) continue; + const int x = s.cx + dx, y = s.cy + dy; + if (x < 0 || y < 0 || x >= f.width || y >= f.height) continue; + const bool in_disk = s.footprint + ? BraggStencilInSignal(st, x - s.px, y - s.py, + (x - s.px) * (x - s.px) + (y - s.py) * (y - s.py), f.r1_sq) + : dx * dx + dy * dy < f.r1_sq; + o.p_grid += Pp; + o.p_peak = o.p_peak < Pp ? Pp : o.p_peak; // std::max + if (in_disk) o.m_all += Pp; + if (!valid_px(img[y * f.width + x])) { + o.p_lost_peak = o.p_lost_peak < Pp ? Pp : o.p_lost_peak; + continue; + } + o.p_valid += Pp; + const bool own = !f.use_owner || BraggOwnedBy(owner[y * f.width + x], (int) i); + if (own) o.p_own += Pp; + if (in_disk && (own || !f.exclude)) o.m_read += Pp; + } + + // FitKept is double arithmetic: the CPU decides it, and the fit runs for every reflection. + float If = s.I0; + for (int iter = 0; iter < 4; ++iter) { + float num = 0.0f; + o.den = 0.0f; + o.wsum = 0.0f; + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float Pp = P[(dy + Rf) * Gf + (dx + Rf)]; + if (Pp <= 0.0f) continue; + const int x = s.cx + dx, y = s.cy + dy; + if (x < 0 || y < 0 || x >= f.width || y >= f.height) continue; + const int32_t v = img[y * f.width + x]; + if (!valid_px(v)) continue; + if (f.exclude && f.use_owner && !BraggOwnedBy(owner[y * f.width + x], (int) i)) continue; + const float signal = (float) v - s.bkg; + const float d = s.B + If * Pp; + const float iv = 1.0f / (d < s.v_min ? s.v_min : d); // std::max + num += Pp * signal * iv; + o.den += Pp * Pp * iv; + o.wsum += Pp * iv; + } + if (o.den > 0.0f) If = num / o.den; else break; + } + o.I = If; + out[i] = o; +} diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm index 9c67d7687..85ffafbe7 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm @@ -18,6 +18,8 @@ using namespace bragg_engine; // 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_exp; +extern const char *const bragg_metal_source_pass_b; extern const char *const bragg_metal_source_kernels; namespace { @@ -26,7 +28,7 @@ static_assert(MASK_TAIL == 1 && MASK_FLUX == 2, "BraggIntegrationEngineMetal.met struct Kernels { id device = nil; // nil: no Metal device - id mark = nil, claim = nil, box = nil, clip = nil; + id mark = nil, claim = nil, box = nil, clip = nil, fit = nil; }; // Compiled once per process: a worker makes an engine per pass, and compiling takes a fraction of a @@ -48,6 +50,7 @@ const Kernels &GetKernels() { "#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_exp + bragg_metal_source_pass_b + bragg_metal_source_kernels; MTLCompileOptions *options = [MTLCompileOptions new]; // Fast math off: sqrt and division then round correctly and nothing is reassociated, as on the @@ -72,6 +75,7 @@ const Kernels &GetKernels() { k.claim = pipeline("claim_owner"); k.box = pipeline("box_sums"); k.clip = pipeline("clip_sums"); + k.fit = pipeline("profile_fit"); return k; }(); return kernels; @@ -106,6 +110,12 @@ struct BraggIntegrationEngineMetal::Device { // Per reflection, grown with the frames. size_t capacity = 0; id xy = nil, level = nil, sums = nil, clips = nil, thr = nil, clip_out = nil; + id fit_setups = nil, fit_out = nil; + // Pass B's Gaussians, grown with the frames, and the learned profiles. + size_t cells_capacity = 0, profiles_capacity = 0; + id cells = nil, profiles = nil; + // The frame's image while Run hands pass B to FitOnDevice; nil otherwise (the CPU engine's own Run). + id img = nil; }; bool BraggIntegrationEngineMetal::Available(const DiffractionExperiment &experiment) { @@ -140,6 +150,7 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor if (image.size() != npixel || npredicted == 0) return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number); + dev->img = nil; @autoreleasepool { const Kernels &k = GetKernels(); // The GPU reads the host's own memory, so the image is read where it lies - which Metal allows @@ -182,6 +193,8 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor dev->clips = buffer(n); dev->thr = buffer(n * sizeof(int32_t)); dev->clip_out = buffer(n * sizeof(BraggClipSums)); + dev->fit_setups = buffer(n * sizeof(BraggFitSetup)); + dev->fit_out = buffer(n * sizeof(BraggFitOut)); dev->capacity = n; } @@ -313,7 +326,9 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor clips[i] ? clip_out[i] : BraggClipSums{}, no_trim); } const FlatOwnerMap owner_map{use_owner ? static_cast(dev->owner.contents) : nullptr, xpixel}; + dev->img = img; auto ret = ProfilePasses(BufferSampler{image.data()}, owner_map, predicted, npredicted, a, image_number); + dev->img = nil; // --- Unmark this frame's regions and give back its claims, while the CPU goes on --- dev->unmark = [dev->queue commandBuffer]; @@ -341,3 +356,50 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor return ret; } } + +bool BraggIntegrationEngineMetal::FitOnDevice(const std::vector &setups, + const std::vector &shell_profiles, + std::vector &fits) { + if (dev->img == nil) + return false; + @autoreleasepool { + const Kernels &k = GetKernels(); + const size_t n = setups.size(); + size_t n_cells = 0; + for (const BraggFitSetup &s : setups) + if (s.fit && s.shell < 0) + n_cells = std::max(n_cells, s.cells + static_cast(2 * s.Rf + 1) * (2 * s.Rf + 1)); + auto grow = [&](__strong id &b, size_t &capacity, size_t bytes) { + if (capacity < bytes) { + capacity = bytes + bytes / 2; + b = [k.device newBufferWithLength:std::max(capacity, 4) options:MTLResourceStorageModeShared]; + } + }; + grow(dev->cells, dev->cells_capacity, n_cells * sizeof(float)); + grow(dev->profiles, dev->profiles_capacity, shell_profiles.size() * sizeof(float)); + std::copy(setups.begin(), setups.end(), static_cast(dev->fit_setups.contents)); + std::copy(shell_profiles.begin(), shell_profiles.end(), static_cast(dev->profiles.contents)); + + const bool use_owner = overlap != OverlapMode::Off; + const BraggPassBFrame f{static_cast(xpixel), static_cast(ypixel), r1_sq, R, use_owner, + overlap == OverlapMode::Exclude, static_cast(n)}; + + id cb = [dev->queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:k.fit]; + [enc setBuffer:dev->fit_setups offset:0 atIndex:0]; + [enc setBuffer:dev->img offset:0 atIndex:1]; + [enc setBuffer:(use_owner ? dev->owner : dev->mask) offset:0 atIndex:2]; // not read without an owner map + [enc setBuffer:dev->profiles offset:0 atIndex:3]; + [enc setBuffer:dev->cells offset:0 atIndex:4]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:5]; + [enc setBytes:&f length:sizeof(f) atIndex:6]; + [enc setBuffer:dev->fit_out offset:0 atIndex:7]; + [enc dispatchThreads:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; + [enc endEncoding]; + Wait(cb); + const auto *out = static_cast(dev->fit_out.contents); + std::copy(out, out + n, fits.begin()); + } + return true; +} diff --git a/image_analysis/bragg_integration/BraggPassB.h b/image_analysis/bragg_integration/BraggPassB.h new file mode 100644 index 000000000..e03981a5e --- /dev/null +++ b/image_analysis/bragg_integration/BraggPassB.h @@ -0,0 +1,55 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// Pass B of BraggIntegrationEngineCPU, the profile fit, as the data it hands on. The CPU works out in +// double what one reflection's fit needs (BraggFitSetup), the fit itself is single precision +// (BraggFitOut), and the CPU finishes it in double. The Metal engine runs the middle on the GPU, so this +// is compiled by the host and by the Metal compiler (no default member initialisers, as in BraggPassA.h). + +#ifndef __METAL_VERSION__ +#include + +#include "BraggExp.h" // the Metal engine puts it in front +#endif + +struct BraggFitSetup { + float px, py; // predicted position + int32_t cx, cy; // profile grid centre + int32_t Rf; // grid half-size + int32_t shell; // the learned profile of this shell, or -1: the Gaussian below + uint32_t cells; // where the Metal engine keeps the Gaussian, Gf^2 floats + int32_t fit; // 0: not fitted (no pass A result, or overloaded) + int32_t footprint; // the stencil has a measured footprint: it decides the signal disk + float fx, fy; // sub-pixel offset of the prediction from the grid centre + float ux, uy; // radial direction + float inv_2s2r, inv_2s2t; // 1 / (2 sigma^2), radially and tangentially + float bkg; // background per pixel + float B, v_min; // the variance model's floor, and the least variance a weight takes + float I0; // the box sum, where the fit starts +}; + +struct BraggFitOut { + float p_grid, p_valid, p_own; // profile mass on the detector, readable, and this reflection's own + float p_peak, p_lost_peak; // the largest profile value, and the largest one unreadable + float m_all, m_read; // mass in the signal disk, and the part the box sum read + float I, den, wsum; // the fit; den = 0 when it did not run or failed +}; + +// What the Metal engine's fit kernel is told about the frame, beside the stencil. +struct BraggPassBFrame { + int32_t width, height; + float r1_sq; + int32_t R; // the learned profiles' grid half-size + int32_t use_owner, exclude; // as in BraggPassAFrame + uint32_t n; // reflections +}; + +// The Gaussian profile's (unnormalised) value at grid offset (dx, dy). +inline float BraggGaussian(int dx, int dy, float fx, float fy, float ux, float uy, float inv_2s2r, + float inv_2s2t) { + const float ex = dx - fx, ey = dy - fy; + const float rad = ex * ux + ey * uy, tn = -ex * uy + ey * ux; + return BraggExp(-rad * rad * inv_2s2r - tn * tn * inv_2s2t); +} diff --git a/image_analysis/bragg_integration/CMakeLists.txt b/image_analysis/bragg_integration/CMakeLists.txt index f344c142e..2a74ecd82 100644 --- a/image_analysis/bragg_integration/CMakeLists.txt +++ b/image_analysis/bragg_integration/CMakeLists.txt @@ -4,6 +4,8 @@ ADD_LIBRARY(JFJochBraggIntegration STATIC BraggIntegrationEngineCPU.cpp BraggIntegrationEngineCPU.h BraggPassA.h + BraggPassB.h + BraggExp.h SpotFootprint.cpp SpotFootprint.h Regression.h @@ -26,9 +28,10 @@ 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 + # source is the stencil, exp and pass headers the CPU engine compiles 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_PARTS stencil BraggStencil.h pass_a BraggPassA.h exp BraggExp.h pass_b BraggPassB.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)