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)