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