From 229a432a66bffb6ff4e3627fa5d9872304b4612c Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 08:18:43 +0200 Subject: [PATCH] RESULT-CHANGING (owner approved, ulp level): Metal profile fit, CPU Gaussian from a shared exp Pass B of BraggIntegrationEngineCPU is split where a device can take over: per reflection, what the fit needs is worked out in double (BraggFitSetup), the profile, its readable mass and the reweighted fit run in single precision (FitProfile -> BraggFitOut), and the error, the summation guard and the result are finished in double. BraggIntegrationEngineMetal runs the middle as a Metal kernel, one thread per reflection, summing in the CPU's order (FitOnDevice); the CUDA engine is untouched. The Gaussian's exp is now BraggExp (BraggExp.h), on the CPU and on the GPU: no GPU exp reproduces the host's expf, and BraggExp uses only +, -, * and exact power-of-two scaling, so both compile it to the same bits. Checked on the Apple M6 over every float in [-110, 89] (2.2e9 arguments), with subnormal products and sums matching too. It is closer to correctly rounded than macOS expf (max 0.53 ulp; 0.2% of arguments below -1 not correctly rounded vs 0.62%). This is the one result change: per observation |dI|/sigma p99.9 <= 1.7e-6, max 1e-4, 10-30x below what ba3f75475 moved. Merged statistics (resolution, completeness, CC1/2, R_meas, I/sigma, ISa, space group) are identical on lysozyme, cytochrome C, myoglobin, a small-molecule set and 4nwv; only tNCS fields that flip on any last-digit change move (cytochrome C's pseudo-translation flag, whose modulation ratio sits at the 2.5 threshold, turns off). The Metal engine stays bit-identical to the CPU engine: BraggIntegrationEngineMetalTest passes unchanged, and catches planted kernel bugs (another exp, 3 iterations, a dropped lost-peak test, <= for the disk, inverted exclude logic). Measured wall time, Apple M6 (CPU pass B -> GPU pass B): lysozyme 24.7 -> 18.6 s, myoglobin 39.8 -> 35.7 s, cytochrome C 59.3 -> 49.0 s, 4nwv 150.2 -> 131.3 s; the small-molecule set, sparse frames, 92.5 -> 95.6 s (noise of +-3 s). Metal's fast::exp in the kernel instead measured the same on every set (18.6, 35.8, 49.1, 131.2 s), so it would buy nothing for the bitwise check it would cost. On the CPU alone BraggExp costs nothing measurable (cytochrome C --no-gpu 72.8 -> 73.1 s). Co-Authored-By: Claude Opus 5.5 --- docs/ACKNOWLEDGEMENT.md | 7 + docs/GPU_DATA_ANALYSIS.md | 11 +- image_analysis/bragg_integration/BraggExp.h | 63 +++++ .../BraggIntegrationEngineCPU.cpp | 256 +++++++++++------- .../BraggIntegrationEngineCPU.h | 18 ++ .../BraggIntegrationEngineMetal.h | 23 +- .../BraggIntegrationEngineMetal.metal | 86 ++++++ .../BraggIntegrationEngineMetal.mm | 64 ++++- image_analysis/bragg_integration/BraggPassB.h | 55 ++++ .../bragg_integration/CMakeLists.txt | 7 +- 10 files changed, 470 insertions(+), 120 deletions(-) create mode 100644 image_analysis/bragg_integration/BraggExp.h create mode 100644 image_analysis/bragg_integration/BraggPassB.h diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index be74d7ca6..9a4afd7d4 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -523,6 +523,13 @@ corrected once by the exact remainder, which returns the correctly rounded quoti P. W. Markstein, "Computation of elementary functions on the IBM RISC System/6000 processor" (1990), IBM J. Res. Dev. 34, 111-119. +The single-precision exp the Bragg integration's profile fit shares between the CPU and an Apple GPU +(`BraggExp.h`) is built from textbook error-free floating-point arithmetic: the argument reduction of +W. J. Cody and W. Waite, "Software Manual for the Elementary Functions" (1980), Prentice-Hall; the +exact product (with Veltkamp's split) and Fast2Sum of T. J. Dekker, "A floating-point technique for +extending the available precision" (1971), Numer. Math. 18, 224-242; and the 2Sum of D. E. Knuth, +"The Art of Computer Programming, Vol. 2: Seminumerical Algorithms", section 4.2.2. + This software uses the Viridis, Magma and Inferno colormaps from Matplotlib under its BSD-compatible license. J. D. Hunter, "Matplotlib: A 2D graphics environment" (2007), Comput. Sci. Eng. 9, 90-95 [doi:10.1109/MCSE.2007.55](https://doi.org/10.1109/MCSE.2007.55). diff --git a/docs/GPU_DATA_ANALYSIS.md b/docs/GPU_DATA_ANALYSIS.md index 8fab936d7..4573c4a69 100644 --- a/docs/GPU_DATA_ANALYSIS.md +++ b/docs/GPU_DATA_ANALYSIS.md @@ -186,9 +186,10 @@ structure factors above; the likelihood scoring runs on the host. The search is A macOS build has no CUDA, and one engine runs on the Apple GPU instead: - **Bragg integration** (`BraggIntegrationEngineMetal.mm`, `.metal`) - the CPU engine with its - integer pixel work moved to the GPU: the neighbour mask, the ownership of shared pixels, the box - summation's walk over every reflection and the background clip. Everything computed in double, the - profile learning and the profile fit stay on the CPU, as Apple GPUs have no double precision. The - same result as the CPU engine **to the bit**, which a macOS CI step checks on every push; - `rugnux --no-gpu` keeps it off the GPU. + pixel work moved to the GPU: the neighbour mask, the ownership of shared pixels, the box + summation's walk over every reflection, the background clip and the single-precision profile fit. + Everything computed in double and the profile learning stay on the CPU, as Apple GPUs have no + double precision. The fit's Gaussian uses an exp of our own (`BraggExp.h`) on both, as no GPU exp + matches the CPU's. The same result as the CPU engine **to the bit**, which a macOS CI step checks on + every push; `rugnux --no-gpu` keeps it off the GPU. [ยง9.2](CPU_DATA_ANALYSIS_INTEGRATION.md#92-box-summation-seed-and-fallback) diff --git a/image_analysis/bragg_integration/BraggExp.h b/image_analysis/bragg_integration/BraggExp.h new file mode 100644 index 000000000..818afd4dc --- /dev/null +++ b/image_analysis/bragg_integration/BraggExp.h @@ -0,0 +1,63 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// exp in single precision, the same bits wherever it is compiled: the host compiler (no FMA contraction, +// see CMakeLists.txt) and the Metal compiler (fast math off) both see only IEEE +, -, * rounded to +// nearest and exact scaling by a power of two, so the profile fit's Gaussian is the same on the CPU and +// on an Apple GPU, whose exp no host function reproduces. It is closer to correctly rounded than macOS's +// expf: at most 0.53 ulp for a normal result, and away from a correctly rounded float for 0.2% of +// arguments below -1 and 0.004% in [-1, 0]. +// +// exp(x) = 2^k exp(r), k = rint(x / ln2), r = x - k ln2 with ln2 split three ways, the first two parts +// of 16 bits so that k times them is exact (Cody & Waite). r is carried as rr + re, rr^2 is formed exactly +// as zh + zl (Dekker's product with Veltkamp's split) and the sums into 1 + rr + rr^2/2 are error-free +// (2Sum, Fast2Sum), so the only rounding that matters is the last one. The rest of the series is a +// degree-4 minimax fit, relative error 1.3e-10. + +#ifdef __METAL_VERSION__ +inline float BraggBitsToFloat(int i) { return as_type(i); } +#else +#include +#include +inline float BraggBitsToFloat(int32_t i) { return std::bit_cast(i); } +#endif + +// p * 2^k, rounded once where the result is subnormal. +inline float BraggScale2(float p, int k) { + if (k < -125) + return (p * BraggBitsToFloat((k + 100 + 127) << 23)) * BraggBitsToFloat((-100 + 127) << 23); + return p * BraggBitsToFloat((k + 127) << 23); +} + +// Following Cody & Waite (1980) Software Manual for the Elementary Functions; Dekker (1971) Numer. Math. +// 18, 224-242; Knuth, The Art of Computer Programming Vol. 2, 4.2.2 (2Sum). +inline float BraggExp(float x) { + if (x < -104.0f) return 0.0f; // below half the smallest subnormal + if (x > 88.72283f) return BraggBitsToFloat(0x7f800000); // overflows + const float shifter = 12582912.0f; // 1.5 * 2^23: adding it rounds to an integer + const float kf = (x * 1.44269504f + shifter) - shifter; + const int k = static_cast(kf); + const float r1 = x - kf * 0.693145751953125f; // exact + const float t2 = kf * 1.4286197256296873e-06f; // exact + const float rr = r1 - t2; // r1 - t2 = rr + e2 (2Sum) + const float bv0 = rr - r1; + const float e2 = (r1 - (rr - bv0)) + (-t2 - bv0); + const float re = e2 - kf * -1.2905320041778356e-11f; + const float sp = rr * 4097.0f; // rr^2 = zh + zl + const float rh = sp - (sp - rr); + const float rl = rr - rh; + const float zh = rr * rr; + const float zl = ((rh * rh - zh) + (rh * rl + rh * rl)) + rl * rl; + const float h = 0.1666666567325592f + rr * (0.04166644066572189f + rr * (0.008333569392561913f + + rr * (0.0013937620678916574f + rr * 0.0001974606711883098f))); + const float c = 0.5f * zl + ((re + re * rr) + (zh * rr) * h); + const float hz = 0.5f * zh; // rr + zh/2 = s1 + e1 (2Sum) + const float s1 = rr + hz; + const float bv1 = s1 - rr; + const float e1 = (rr - (s1 - bv1)) + (hz - bv1); + const float t1 = 1.0f + s1; // 1 + s1 = t1 + e3 (Fast2Sum, |s1| < 1) + const float e3 = (1.0f - t1) + s1; + return BraggScale2(t1 + (e3 + (e1 + c)), k); +} diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp index b488aff79..aa574ae17 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.cpp @@ -429,6 +429,107 @@ void BraggIntegrationEngineCPU::FinishPassA(PassA &a, size_t i, const Reflection a.rough[i] = out; } +bool BraggIntegrationEngineCPU::FitKept(const BraggFitOut &o) const { + if (o.p_valid < overlap_min_peak * o.p_grid) return false; + // A hole in the profile's PEAK is a different defect from a hole in its wings, and the mass + // fraction above cannot tell them apart - the peak of a broad spot is a few percent of the + // mass, so MINPK passes a reflection that has lost the one part of the profile its amplitude + // is determined by. See the header. + if (o.p_lost_peak > MINPK_MAX_MISSING_PEAK * o.p_peak) return false; + if (overlap == OverlapMode::Reject && o.p_own < overlap_min_peak) return false; + return true; +} + +template +BraggFitOut BraggIntegrationEngineCPU::FitProfile(const Sampler &img, const OwnerMap &owner_map, + const BraggFitSetup &s, size_t i, + const std::vector &shell_profiles, + std::vector &cells, + std::vector &fit_pixels) const { + const int W = static_cast(xpixel), H = static_cast(ypixel); + const bool exclude = overlap == OverlapMode::Exclude; + const bool use_owner = overlap != OverlapMode::Off; + const int Rf = s.Rf, Gf = 2 * Rf + 1; + BraggFitOut o{}; + + const float *P; + if (s.shell >= 0) + P = shell_profiles.data() + static_cast(s.shell) * GG; + else { + // The grid's image rows, asked for now (as in pass A): building the profile below takes long + // enough for them to arrive before the fit reads them. + const int gx0 = std::max(0, s.cx - Rf), gx1 = std::min(W - 1, s.cx + Rf); + for (int y = std::max(0, s.cy - Rf); y <= std::min(H - 1, s.cy + Rf); ++y) + prefetch_row(img.address(static_cast(y) * W + gx0), + img.address(static_cast(y) * W + gx1)); + cells.resize(static_cast(Gf) * Gf); + float gs = 0.0f; + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float g = BraggGaussian(dx, dy, s.fx, s.fy, s.ux, s.uy, s.inv_2s2r, s.inv_2s2t); + cells[(dy + Rf) * Gf + (dx + Rf)] = g; + gs += g; + } + const float inv_gs = 1.0f / gs; + for (float &p : cells) p *= inv_gs; + P = cells.data(); + } + + // --- How much of the expected profile the fit can actually see. p_valid is the readable + // fraction of it - XDS's MINPK, dials' valid_foreground_threshold - measured against the + // mass that falls on the detector at all, so a reflection is judged on the pixels that + // exist and not on where the sensor ends. p_own is the same quantity over neighbour-owned + // pixels, what Reject cuts on. m_read / m_all is the fraction of the r1 disk the + // summation seed the runaway guard compares against actually summed; with nothing missing + // and nothing excluded it is 1 and the guard is untouched. --- + const BraggStencil st = MakeBraggStencil(s.px, s.py, stencil); + fit_pixels.clear(); + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float Pp = P[(dy + Rf) * Gf + (dx + Rf)]; + if (Pp <= 0.0f) continue; + const int x = s.cx + dx, y = s.cy + dy; + if (x < 0 || y < 0 || x >= W || y >= H) continue; + // The region the summation seed was taken over (pass A), so the guard compares the two + // over the same pixels. + const bool in_disk = s.footprint + ? BraggStencilInSignal(st, x - s.px, y - s.py, + (x - s.px) * (x - s.px) + (y - s.py) * (y - s.py), r1_sq) + : dx * dx + dy * dy < r1_sq; + o.p_grid += Pp; + o.p_peak = std::max(o.p_peak, Pp); + if (in_disk) o.m_all += Pp; + const int32_t v = img[y * W + x]; + if (!valid(v)) { + o.p_lost_peak = std::max(o.p_lost_peak, Pp); + continue; + } + o.p_valid += Pp; + const bool own = !use_owner || BraggOwnedBy(owner_map.Get(x, y), static_cast(i)); + if (own) o.p_own += Pp; + if (in_disk && (own || !exclude)) o.m_read += Pp; + if (own || !exclude) + fit_pixels.push_back({Pp, static_cast(v) - s.bkg}); + } + if (!FitKept(o)) return o; + + float If = s.I0; + for (int iter = 0; iter < 4; ++iter) { + float num = 0.0f; + o.den = 0.0f; + o.wsum = 0.0f; + for (const auto &[Pp, signal] : fit_pixels) { + const float iv = 1.0f / std::max(s.B + If * Pp, s.v_min); + num += Pp * signal * iv; + o.den += Pp * Pp * iv; + o.wsum += Pp * iv; + } + if (o.den > 0.0f) If = num / o.den; else break; + } + o.I = If; + return o; +} + template std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler &img, const OwnerMap &owner_map, const std::vector &predicted, @@ -578,16 +679,21 @@ std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler & // --- Pass B: profile-fit each reflection (Kabsch, de-biased variance v = B + I*P; iterate). The // reweighting is the Kabsch/Otwinowski iteration: Kabsch, Acta Cryst D66, 133-144 (2010); // Otwinowski & Minor, Methods Enzymol 276, 307-326 (1997). - // In single precision, with the GPU engine's arithmetic (BraggIntegrationEngineGPU.cu, fit): - // the profile from expf with the widths' reciprocals taken once, normalised by a reciprocal, - // and one reciprocal of the variance for the three weights. --- - std::vector Pbuf; - // The pixels the fit reads, gathered once while p_valid walks the grid: the profile value and the - // background-subtracted count, in the grid's order. - struct FitPixel { float Pp, signal; }; - std::vector fit_pixels; + // Per reflection in three steps: what the fit needs, worked out in double (BraggFitSetup); the + // fit in single precision (FitProfile, or the Metal engine's kernel), with the GPU engine's + // arithmetic (BraggIntegrationEngineGPU.cu, fit) - the profile from exp with the widths' + // reciprocals taken once, normalised by a reciprocal, and one reciprocal of the variance for the + // three weights; and its error, the summation guard and the result, in double. --- + std::vector shell_profiles; // the learned profiles, shell by shell (ProfileEmpirical) + if (empirical) + for (int s = 0; s < N_SHELL; ++s) + shell_profiles.insert(shell_profiles.end(), shell_P[s].begin(), shell_P[s].end()); + std::vector setups(npredicted); + uint32_t cells = 0; for (size_t i = 0; i < npredicted; ++i) { const auto &rh = rough[i]; + BraggFitSetup &fs = setups[i]; + fs = {}; if (!rh.ok) continue; // An overloaded reflection is not measured by any fit - its brightest pixels are what is // missing. It is passed on with its box sum and the flag, for the merge to drop its rocking @@ -600,10 +706,15 @@ std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler & continue; } const int sh = rh.shell < 0 ? 0 : rh.shell; - - int Rf = R; - const std::vector *Pvec = &shell_P[sh]; const BraggStencil st = MakeBraggStencil(predicted[i].predicted_x, predicted[i].predicted_y, stencil); + fs.fit = 1; + fs.px = predicted[i].predicted_x; + fs.py = predicted[i].predicted_y; + fs.cx = rh.cx; + fs.cy = rh.cy; + fs.Rf = R; + fs.shell = sh; + fs.footprint = st.fp_s2t > 0.0f; if (!empirical) { const double rx = predicted[i].predicted_x - beam_x, ry = predicted[i].predicted_y - beam_y; const double Rpx = std::hypot(rx, ry); @@ -632,100 +743,41 @@ std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler & s2r = std::max(s2r, st.fp_s2r); elong = true; } - // Build the Gaussian per reflection, centred on the sub-pixel predicted position and (when + // The Gaussian is built per reflection, centred on the sub-pixel predicted position and (when // needed) elongated, on a grid grown to hold the wider of its two axes. - const double fx = predicted[i].predicted_x - rh.cx, fy = predicted[i].predicted_y - rh.cy; - Rf = elong ? std::min(3 * R, static_cast(std::ceil(r2 + 2.0 * std::sqrt(std::max(s2r, s2t))))) : R; - // The grid's image rows, asked for now (as in pass A): building the profile below takes - // long enough for them to arrive before the fit reads them. - const int gx0 = std::max(0, rh.cx - Rf), gx1 = std::min(W - 1, rh.cx + Rf); - for (int y = std::max(0, rh.cy - Rf); y <= std::min(H - 1, rh.cy + Rf); ++y) - prefetch_row(img.address(static_cast(y) * W + gx0), - img.address(static_cast(y) * W + gx1)); - const int Gf = 2 * Rf + 1; - Pbuf.assign(static_cast(Gf) * Gf, 0.0f); - const float fxf = static_cast(fx), fyf = static_cast(fy); - const float uxf = static_cast(ux), uyf = static_cast(uy); - const float inv_2s2r = 0.5f / static_cast(s2r), inv_2s2t = 0.5f / static_cast(s2t); - float gs = 0.0f; - for (int dy = -Rf; dy <= Rf; ++dy) - for (int dx = -Rf; dx <= Rf; ++dx) { - const float ex = dx - fxf, ey = dy - fyf; - const float rad = ex * uxf + ey * uyf, tn = -ex * uyf + ey * uxf; - const float g = std::exp(-rad * rad * inv_2s2r - tn * tn * inv_2s2t); - Pbuf[(dy + Rf) * Gf + (dx + Rf)] = g; - gs += g; - } - const float inv_gs = 1.0f / gs; - for (float &p : Pbuf) p *= inv_gs; - Pvec = &Pbuf; + fs.Rf = elong ? std::min(3 * R, static_cast(std::ceil(r2 + 2.0 * std::sqrt(std::max(s2r, s2t))))) : R; + fs.shell = -1; + fs.cells = cells; + cells += (2 * fs.Rf + 1) * (2 * fs.Rf + 1); + fs.fx = static_cast(predicted[i].predicted_x - rh.cx); + fs.fy = static_cast(predicted[i].predicted_y - rh.cy); + fs.ux = static_cast(ux); + fs.uy = static_cast(uy); + fs.inv_2s2r = 0.5f / static_cast(s2r); + fs.inv_2s2t = 0.5f / static_cast(s2t); } + fs.bkg = static_cast(rh.bkg); + fs.B = static_cast(std::max(rh.bkg, PIXEL_VARIANCE_FLOOR)); + fs.v_min = static_cast(WEIGHT_VARIANCE_MIN_FRACTION) * fs.B; + fs.I0 = static_cast(rh.I); + } - const int Gf = 2 * Rf + 1; + std::vector fits(npredicted); + if (!FitOnDevice(setups, shell_profiles, fits)) { + std::vector cells_buf; + std::vector fit_pixels; + for (size_t i = 0; i < npredicted; ++i) + if (setups[i].fit) + fits[i] = FitProfile(img, owner_map, setups[i], i, shell_profiles, cells_buf, fit_pixels); + } - // --- How much of the expected profile the fit can actually see. p_valid is the readable - // fraction of it - XDS's MINPK, dials' valid_foreground_threshold - measured against the - // mass that falls on the detector at all, so a reflection is judged on the pixels that - // exist and not on where the sensor ends. p_own is the same quantity over neighbour-owned - // pixels, what Reject cuts on. m_read / m_all is the fraction of the r1 disk the - // summation seed the runaway guard compares against actually summed; with nothing missing - // and nothing excluded it is 1 and the guard is untouched. --- - float p_grid = 0.0f, p_valid = 0.0f, p_own = 0.0f, m_all = 0.0f, m_read = 0.0f; - float p_peak = 0.0f, p_lost_peak = 0.0f; - fit_pixels.clear(); - for (int dy = -Rf; dy <= Rf; ++dy) - for (int dx = -Rf; dx <= Rf; ++dx) { - const float Pp = (*Pvec)[(dy + Rf) * Gf + (dx + Rf)]; - if (Pp <= 0.0f) continue; - const int x = rh.cx + dx, y = rh.cy + dy; - if (x < 0 || y < 0 || x >= W || y >= H) continue; - // The region the summation seed was taken over (pass A), so the guard below - // compares the two over the same pixels. - const bool in_disk = st.fp_s2t > 0.0f - ? BraggStencilInSignal(st, x - predicted[i].predicted_x, y - predicted[i].predicted_y, - (x - predicted[i].predicted_x) * (x - predicted[i].predicted_x) - + (y - predicted[i].predicted_y) * (y - predicted[i].predicted_y), - r1_sq) - : dx * dx + dy * dy < r1_sq; - p_grid += Pp; - p_peak = std::max(p_peak, Pp); - if (in_disk) m_all += Pp; - if (!valid(img[y * W + x])) { - p_lost_peak = std::max(p_lost_peak, Pp); - continue; - } - p_valid += Pp; - const bool own = clean(x, y, i); - if (own) p_own += Pp; - if (in_disk && (own || !exclude)) m_read += Pp; - if (own || !exclude) - fit_pixels.push_back({Pp, static_cast(img[y * W + x]) - static_cast(rh.bkg)}); - } - if (p_valid < overlap_min_peak * p_grid) continue; - // A hole in the profile's PEAK is a different defect from a hole in its wings, and the mass - // fraction above cannot tell them apart - the peak of a broad spot is a few percent of the - // mass, so MINPK passes a reflection that has lost the one part of the profile its amplitude - // is determined by. See the header. - if (p_lost_peak > MINPK_MAX_MISSING_PEAK * p_peak) continue; - if (overlap == OverlapMode::Reject && p_own < overlap_min_peak) continue; - - const float B = static_cast(std::max(rh.bkg, PIXEL_VARIANCE_FLOOR)); - const float v_min = static_cast(WEIGHT_VARIANCE_MIN_FRACTION) * B; - float If = static_cast(rh.I), den = 0.0f, wsum = 0.0f; - for (int iter = 0; iter < 4; ++iter) { - float num = 0.0f; - den = 0.0f; - wsum = 0.0f; - for (const auto &[Pp, signal] : fit_pixels) { - const float iv = 1.0f / std::max(B + If * Pp, v_min); - num += Pp * signal * iv; - den += Pp * Pp * iv; - wsum += Pp * iv; - } - if (den > 0.0f) If = num / den; else break; - } - if (!(den > 0.0f)) continue; - double I = If; + for (size_t i = 0; i < npredicted; ++i) { + if (!setups[i].fit) continue; + const auto &rh = rough[i]; + const BraggFitOut &fo = fits[i]; + if (!FitKept(fo) || !(fo.den > 0.0f)) continue; + const float den = fo.den, wsum = fo.wsum; + double I = fo.I; // Guard against profile-fit runaways: on a weak / near-zero reflection the reweighted Kabsch // iteration has no real peak to lock onto and can manufacture intensity the box sum never sees. @@ -746,7 +798,7 @@ std::vector BraggIntegrationEngineCPU::ProfilePasses(const Sampler & // The seed is a sum over the disk the box sum actually read, so when pixels are missing from // both - excluded to a neighbour, or unreadable - the fit's full-profile intensity has to be // scaled down to that same disk before the two are comparable. Nothing dropped gives 1. - const double guard_scale = m_all > 0.0 ? m_read / m_all : 1.0; + const double guard_scale = fo.m_all > 0.0f ? fo.m_read / fo.m_all : 1.0; if (std::abs(I * guard_scale - rh.I) > PROFILE_SUMMATION_MAX_NSIGMA * rh.sigma) { I = rh.I; sigma = rh.sigma; diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h index b9ee60116..e72497831 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h +++ b/image_analysis/bragg_integration/BraggIntegrationEngineCPU.h @@ -5,6 +5,7 @@ #include "BraggIntegrationEngine.h" #include "BraggPassA.h" +#include "BraggPassB.h" #include "../../common/PixelMask.h" class CompressedImage; @@ -62,6 +63,17 @@ class BraggIntegrationEngineCPU : public BraggIntegrationEngine { std::vector RunImpl(const Sampler &img, const std::vector &predicted, size_t npredicted, int64_t image_number); + // Pass B's single-precision part for reflection i, as the Metal engine's kernel does it: the profile + // (the shell's learned one, or the Gaussian, built in `cells`), how much of it is readable, and the + // fit. fit_pixels is scratch: the pixels the fit reads, gathered once. + struct FitPixel { float Pp, signal; }; + template + BraggFitOut FitProfile(const Sampler &img, const OwnerMap &owner_map, const BraggFitSetup &s, size_t i, + const std::vector &shell_profiles, std::vector &cells, + std::vector &fit_pixels) const; + // Whether enough of a reflection's profile is readable, and its own, for the fit to be kept. + bool FitKept(const BraggFitOut &o) const; + protected: // The run's pixel mask, packed 32 pixels to a word (PixelMask::GetPackedMask). An unreadable pixel // it does not explain was unreadable on this frame only - an overload (see Reflection::overloaded). @@ -118,6 +130,12 @@ protected: // Everything after pass A: the radial background correction, then either the box-sum results or // the profile learning and pass B. Defined for BufferSampler with either owner map, and for the // CompressedImage samplers with the tiled one. + // Pass B's single-precision part for every reflection with setups[i].fit, on a device - the + // learned profiles side by side in shell_profiles, empty for the Gaussian. False where there is no + // device, and the CPU fits them instead. + virtual bool FitOnDevice(const std::vector &setups, const std::vector &shell_profiles, + std::vector &fits) { return false; } + template std::vector ProfilePasses(const Sampler &img, const OwnerMap &owner_map, const std::vector &predicted, size_t npredicted, diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h index 6f4c19d74..ffacb973f 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h @@ -8,23 +8,26 @@ #include "BraggIntegrationEngineCPU.h" -// Apple GPU engine (macOS only): BraggIntegrationEngineCPU with the pixel work that is all integer - the -// neighbour mask, the owner map, pass A's walk over every box and the background clip - run as Metal -// kernels, one thread per reflection. Everything computed in double, the profile learning and pass B -// stay on the CPU, in the CPU engine's own code. Apple GPUs have no double precision, and pass B's -// Gaussian needs exp, which no GPU function matches the CPU's expf in bit for bit. +// Apple GPU engine (macOS only): BraggIntegrationEngineCPU with its pixel work - the neighbour mask, the +// owner map, pass A's walk over every box, the background clip and pass B's single-precision profile +// fit - run as Metal kernels, one thread per reflection. Everything computed in double and the profile +// learning stay on the CPU, in the CPU engine's own code: Apple GPUs have no double precision. // // The result is the CPU engine's, bit for bit: the integer sums do not depend on the order the GPU -// visits pixels in, and the stencil arithmetic deciding which pixel is which is the same source -// (BraggStencil.h), compiled for the GPU with fast math off - sqrt and division then round correctly - -// and, as on the host, with no a*b+c contracted into a fused multiply-add. The CPU and the GPU read the -// same memory, so the preprocessed image is not copied. tests/BraggIntegrationEngineMetalTest.cpp holds -// the engine to all of this. +// visits pixels in, the fit's float sums are taken in the CPU's order, one reflection per thread, and +// the arithmetic deciding which pixel is which is the same source (BraggStencil.h, BraggPassB.h, and +// BraggExp.h for the Gaussian, as no GPU exp matches the host's), compiled for the GPU with fast math +// off - sqrt and division then round correctly - and, as on the host, with no a*b+c contracted into a +// fused multiply-add. The CPU and the GPU read the same memory, so the preprocessed image is not +// copied. tests/BraggIntegrationEngineMetalTest.cpp holds the engine to all of this. class BraggIntegrationEngineMetal : public BraggIntegrationEngineCPU { struct Device; std::unique_ptr dev; uint64_t frames_on_gpu = 0; + bool FitOnDevice(const std::vector &setups, const std::vector &shell_profiles, + std::vector &fits) override; + public: // Whether this engine can integrate this experiment: a Metal device the process may use (rugnux // --no-gpu hides it, see hide_gpus) and settings the kernels cover - not the trimmed-mean background, diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal index 7ebf36fd3..4d2632fa3 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal @@ -187,3 +187,89 @@ kernel void clip_sums(device const float2 *xy [[buffer(0)]], } out[i] = c; } + +// BraggIntegrationEngineCPU::FitProfile. The CPU gathers the pixels the fit reads once; here every +// iteration reads them again, the same pixels in the same order. +kernel void profile_fit(device const BraggFitSetup *setups [[buffer(0)]], + device const int32_t *img [[buffer(1)]], + device const uint32_t *owner [[buffer(2)]], + device const float *shell_profiles [[buffer(3)]], + device float *cells [[buffer(4)]], + constant char *stencil_raw [[buffer(5)]], + constant BraggPassBFrame &f [[buffer(6)]], + device BraggFitOut *out [[buffer(7)]], + uint i [[thread_position_in_grid]]) { + if (i >= f.n || !setups[i].fit) + return; + const BraggFitSetup s = setups[i]; + const int Rf = s.Rf, Gf = 2 * Rf + 1; + BraggFitOut o = {}; + + device const float *P; + if (s.shell >= 0) + P = shell_profiles + s.shell * (2 * f.R + 1) * (2 * f.R + 1); + else { + device float *g = cells + s.cells; + float gs = 0.0f; + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float v = BraggGaussian(dx, dy, s.fx, s.fy, s.ux, s.uy, s.inv_2s2r, s.inv_2s2t); + g[(dy + Rf) * Gf + (dx + Rf)] = v; + gs += v; + } + const float inv_gs = 1.0f / gs; + for (int k = 0; k < Gf * Gf; ++k) g[k] *= inv_gs; + P = g; + } + + const BraggStencil st = MakeBraggStencil(s.px, s.py, stencil_params(stencil_raw)); + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float Pp = P[(dy + Rf) * Gf + (dx + Rf)]; + if (Pp <= 0.0f) continue; + const int x = s.cx + dx, y = s.cy + dy; + if (x < 0 || y < 0 || x >= f.width || y >= f.height) continue; + const bool in_disk = s.footprint + ? BraggStencilInSignal(st, x - s.px, y - s.py, + (x - s.px) * (x - s.px) + (y - s.py) * (y - s.py), f.r1_sq) + : dx * dx + dy * dy < f.r1_sq; + o.p_grid += Pp; + o.p_peak = o.p_peak < Pp ? Pp : o.p_peak; // std::max + if (in_disk) o.m_all += Pp; + if (!valid_px(img[y * f.width + x])) { + o.p_lost_peak = o.p_lost_peak < Pp ? Pp : o.p_lost_peak; + continue; + } + o.p_valid += Pp; + const bool own = !f.use_owner || BraggOwnedBy(owner[y * f.width + x], (int) i); + if (own) o.p_own += Pp; + if (in_disk && (own || !f.exclude)) o.m_read += Pp; + } + + // FitKept is double arithmetic: the CPU decides it, and the fit runs for every reflection. + float If = s.I0; + for (int iter = 0; iter < 4; ++iter) { + float num = 0.0f; + o.den = 0.0f; + o.wsum = 0.0f; + for (int dy = -Rf; dy <= Rf; ++dy) + for (int dx = -Rf; dx <= Rf; ++dx) { + const float Pp = P[(dy + Rf) * Gf + (dx + Rf)]; + if (Pp <= 0.0f) continue; + const int x = s.cx + dx, y = s.cy + dy; + if (x < 0 || y < 0 || x >= f.width || y >= f.height) continue; + const int32_t v = img[y * f.width + x]; + if (!valid_px(v)) continue; + if (f.exclude && f.use_owner && !BraggOwnedBy(owner[y * f.width + x], (int) i)) continue; + const float signal = (float) v - s.bkg; + const float d = s.B + If * Pp; + const float iv = 1.0f / (d < s.v_min ? s.v_min : d); // std::max + num += Pp * signal * iv; + o.den += Pp * Pp * iv; + o.wsum += Pp * iv; + } + if (o.den > 0.0f) If = num / o.den; else break; + } + o.I = If; + out[i] = o; +} diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm index 9c67d7687..85ffafbe7 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm +++ b/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm @@ -18,6 +18,8 @@ using namespace bragg_engine; // at configure time (CMakeLists.txt) so the binary needs no file beside it. extern const char *const bragg_metal_source_stencil; extern const char *const bragg_metal_source_pass_a; +extern const char *const bragg_metal_source_exp; +extern const char *const bragg_metal_source_pass_b; extern const char *const bragg_metal_source_kernels; namespace { @@ -26,7 +28,7 @@ static_assert(MASK_TAIL == 1 && MASK_FLUX == 2, "BraggIntegrationEngineMetal.met struct Kernels { id device = nil; // nil: no Metal device - id mark = nil, claim = nil, box = nil, clip = nil; + id mark = nil, claim = nil, box = nil, clip = nil, fit = nil; }; // Compiled once per process: a worker makes an engine per pass, and compiling takes a fraction of a @@ -48,6 +50,7 @@ const Kernels &GetKernels() { "#define floorf(x) floor(x)\n" "#define lroundf(x) ((long) round(x))\n") + bragg_metal_source_stencil + bragg_metal_source_pass_a + + bragg_metal_source_exp + bragg_metal_source_pass_b + bragg_metal_source_kernels; MTLCompileOptions *options = [MTLCompileOptions new]; // Fast math off: sqrt and division then round correctly and nothing is reassociated, as on the @@ -72,6 +75,7 @@ const Kernels &GetKernels() { k.claim = pipeline("claim_owner"); k.box = pipeline("box_sums"); k.clip = pipeline("clip_sums"); + k.fit = pipeline("profile_fit"); return k; }(); return kernels; @@ -106,6 +110,12 @@ struct BraggIntegrationEngineMetal::Device { // Per reflection, grown with the frames. size_t capacity = 0; id xy = nil, level = nil, sums = nil, clips = nil, thr = nil, clip_out = nil; + id fit_setups = nil, fit_out = nil; + // Pass B's Gaussians, grown with the frames, and the learned profiles. + size_t cells_capacity = 0, profiles_capacity = 0; + id cells = nil, profiles = nil; + // The frame's image while Run hands pass B to FitOnDevice; nil otherwise (the CPU engine's own Run). + id img = nil; }; bool BraggIntegrationEngineMetal::Available(const DiffractionExperiment &experiment) { @@ -140,6 +150,7 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor if (image.size() != npixel || npredicted == 0) return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number); + dev->img = nil; @autoreleasepool { const Kernels &k = GetKernels(); // The GPU reads the host's own memory, so the image is read where it lies - which Metal allows @@ -182,6 +193,8 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor dev->clips = buffer(n); dev->thr = buffer(n * sizeof(int32_t)); dev->clip_out = buffer(n * sizeof(BraggClipSums)); + dev->fit_setups = buffer(n * sizeof(BraggFitSetup)); + dev->fit_out = buffer(n * sizeof(BraggFitOut)); dev->capacity = n; } @@ -313,7 +326,9 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor clips[i] ? clip_out[i] : BraggClipSums{}, no_trim); } const FlatOwnerMap owner_map{use_owner ? static_cast(dev->owner.contents) : nullptr, xpixel}; + dev->img = img; auto ret = ProfilePasses(BufferSampler{image.data()}, owner_map, predicted, npredicted, a, image_number); + dev->img = nil; // --- Unmark this frame's regions and give back its claims, while the CPU goes on --- dev->unmark = [dev->queue commandBuffer]; @@ -341,3 +356,50 @@ std::vector BraggIntegrationEngineMetal::Run(const ImagePreprocessor return ret; } } + +bool BraggIntegrationEngineMetal::FitOnDevice(const std::vector &setups, + const std::vector &shell_profiles, + std::vector &fits) { + if (dev->img == nil) + return false; + @autoreleasepool { + const Kernels &k = GetKernels(); + const size_t n = setups.size(); + size_t n_cells = 0; + for (const BraggFitSetup &s : setups) + if (s.fit && s.shell < 0) + n_cells = std::max(n_cells, s.cells + static_cast(2 * s.Rf + 1) * (2 * s.Rf + 1)); + auto grow = [&](__strong id &b, size_t &capacity, size_t bytes) { + if (capacity < bytes) { + capacity = bytes + bytes / 2; + b = [k.device newBufferWithLength:std::max(capacity, 4) options:MTLResourceStorageModeShared]; + } + }; + grow(dev->cells, dev->cells_capacity, n_cells * sizeof(float)); + grow(dev->profiles, dev->profiles_capacity, shell_profiles.size() * sizeof(float)); + std::copy(setups.begin(), setups.end(), static_cast(dev->fit_setups.contents)); + std::copy(shell_profiles.begin(), shell_profiles.end(), static_cast(dev->profiles.contents)); + + const bool use_owner = overlap != OverlapMode::Off; + const BraggPassBFrame f{static_cast(xpixel), static_cast(ypixel), r1_sq, R, use_owner, + overlap == OverlapMode::Exclude, static_cast(n)}; + + id cb = [dev->queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:k.fit]; + [enc setBuffer:dev->fit_setups offset:0 atIndex:0]; + [enc setBuffer:dev->img offset:0 atIndex:1]; + [enc setBuffer:(use_owner ? dev->owner : dev->mask) offset:0 atIndex:2]; // not read without an owner map + [enc setBuffer:dev->profiles offset:0 atIndex:3]; + [enc setBuffer:dev->cells offset:0 atIndex:4]; + [enc setBytes:&stencil length:sizeof(stencil) atIndex:5]; + [enc setBytes:&f length:sizeof(f) atIndex:6]; + [enc setBuffer:dev->fit_out offset:0 atIndex:7]; + [enc dispatchThreads:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(64, 1, 1)]; + [enc endEncoding]; + Wait(cb); + const auto *out = static_cast(dev->fit_out.contents); + std::copy(out, out + n, fits.begin()); + } + return true; +} diff --git a/image_analysis/bragg_integration/BraggPassB.h b/image_analysis/bragg_integration/BraggPassB.h new file mode 100644 index 000000000..e03981a5e --- /dev/null +++ b/image_analysis/bragg_integration/BraggPassB.h @@ -0,0 +1,55 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// Pass B of BraggIntegrationEngineCPU, the profile fit, as the data it hands on. The CPU works out in +// double what one reflection's fit needs (BraggFitSetup), the fit itself is single precision +// (BraggFitOut), and the CPU finishes it in double. The Metal engine runs the middle on the GPU, so this +// is compiled by the host and by the Metal compiler (no default member initialisers, as in BraggPassA.h). + +#ifndef __METAL_VERSION__ +#include + +#include "BraggExp.h" // the Metal engine puts it in front +#endif + +struct BraggFitSetup { + float px, py; // predicted position + int32_t cx, cy; // profile grid centre + int32_t Rf; // grid half-size + int32_t shell; // the learned profile of this shell, or -1: the Gaussian below + uint32_t cells; // where the Metal engine keeps the Gaussian, Gf^2 floats + int32_t fit; // 0: not fitted (no pass A result, or overloaded) + int32_t footprint; // the stencil has a measured footprint: it decides the signal disk + float fx, fy; // sub-pixel offset of the prediction from the grid centre + float ux, uy; // radial direction + float inv_2s2r, inv_2s2t; // 1 / (2 sigma^2), radially and tangentially + float bkg; // background per pixel + float B, v_min; // the variance model's floor, and the least variance a weight takes + float I0; // the box sum, where the fit starts +}; + +struct BraggFitOut { + float p_grid, p_valid, p_own; // profile mass on the detector, readable, and this reflection's own + float p_peak, p_lost_peak; // the largest profile value, and the largest one unreadable + float m_all, m_read; // mass in the signal disk, and the part the box sum read + float I, den, wsum; // the fit; den = 0 when it did not run or failed +}; + +// What the Metal engine's fit kernel is told about the frame, beside the stencil. +struct BraggPassBFrame { + int32_t width, height; + float r1_sq; + int32_t R; // the learned profiles' grid half-size + int32_t use_owner, exclude; // as in BraggPassAFrame + uint32_t n; // reflections +}; + +// The Gaussian profile's (unnormalised) value at grid offset (dx, dy). +inline float BraggGaussian(int dx, int dy, float fx, float fy, float ux, float uy, float inv_2s2r, + float inv_2s2t) { + const float ex = dx - fx, ey = dy - fy; + const float rad = ex * ux + ey * uy, tn = -ex * uy + ey * ux; + return BraggExp(-rad * rad * inv_2s2r - tn * tn * inv_2s2t); +} diff --git a/image_analysis/bragg_integration/CMakeLists.txt b/image_analysis/bragg_integration/CMakeLists.txt index f344c142e..2a74ecd82 100644 --- a/image_analysis/bragg_integration/CMakeLists.txt +++ b/image_analysis/bragg_integration/CMakeLists.txt @@ -4,6 +4,8 @@ ADD_LIBRARY(JFJochBraggIntegration STATIC BraggIntegrationEngineCPU.cpp BraggIntegrationEngineCPU.h BraggPassA.h + BraggPassB.h + BraggExp.h SpotFootprint.cpp SpotFootprint.h Regression.h @@ -26,9 +28,10 @@ ENDIF() IF (APPLE) # The Metal engine compiles its kernels when the program starts, from source carried in the binary: # macOS has the Metal compiler built in, while the offline one is an Xcode download of its own. The - # source is the stencil header the CPU engine compiles, the pass-A header and the kernels, copied + # source is the stencil, exp and pass headers the CPU engine compiles and the kernels, copied # into a generated .cpp here; editing one of them re-runs the configure. - SET(METAL_PARTS stencil BraggStencil.h pass_a BraggPassA.h kernels BraggIntegrationEngineMetal.metal) + SET(METAL_PARTS stencil BraggStencil.h pass_a BraggPassA.h exp BraggExp.h pass_b BraggPassB.h + kernels BraggIntegrationEngineMetal.metal) SET(METAL_SOURCE "// Generated from image_analysis/bragg_integration by its CMakeLists.txt - do not edit.\n") WHILE (METAL_PARTS) LIST(POP_FRONT METAL_PARTS name file)