Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngineMetal.metal

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// 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)]],
                            device const BraggPanel *panels [[buffer(6)]],
                            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;
    // Only on the reflection's own panel, as the CPU engine marks.
    const BraggPanel pn = panels[i];
    const BraggStencil st = MakeBraggStencil(pn, sp);
    const int x0 = max(pn.x0, (int) floor(px - st.ex_in - 1.0f));
    const int x1 = min(pn.x1, (int) ceil(px + st.ex_in + 1.0f));
    const int y0 = max(pn.y0, (int) floor(py - st.ey_in - 1.0f));
    const int y1 = min(pn.y1, (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)]],
                        device const BraggPanel *panels [[buffer(4)]],
                        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 BraggPanel pn = panels[i];
    const int x0 = max(pn.x0, (int) floor(px - f.claim));
    const int x1 = min(pn.x1, (int) ceil(px + f.claim));
    const int y0 = max(pn.y0, (int) floor(py - f.claim));
    const int y1 = min(pn.y1, (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)]],
                     device const BraggPanel *panels [[buffer(8)]],
                     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 BraggPanel pn = panels[i];
    const BraggStencil st = MakeBraggStencil(pn, 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;
            // Off the reflection's panel is a gap: unreadable, and not lost to the flux.
            const bool on = BraggPanelHolds(pn, x, y);
            const int32_t v = on ? img[idx] : PX_MASKED;
            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 && on && !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)]],
                      device const BraggPanel *panels [[buffer(10)]],
                      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 BraggPanel pn = panels[i];
    const BraggStencil st = MakeBraggStencil(pn, 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 = BraggPanelHolds(pn, x, y) ? img[idx] : PX_MASKED;
            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;
}

// 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)]],
                        device const BraggPanel *panels [[buffer(8)]],
                        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 = {};

    // A learned profile is in the plane's axes (BraggProfileCell); a Gaussian is built in the panel's.
    const BraggPanel pn = panels[i];
    const bool learned = s.shell >= 0;
    device const float *P;
    if (learned)
        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(pn, stencil_params(stencil_raw));
    for (int dy = -Rf; dy <= Rf; ++dy)
        for (int dx = -Rf; dx <= Rf; ++dx) {
            const float Pp = P[learned ? BraggProfileCell(pn, dx, dy, Rf) : (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 (!BraggPanelHolds(pn, x, y) || !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[learned ? BraggProfileCell(pn, dx, dy, Rf) : (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 || !BraggPanelHolds(pn, x, y)) 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;
}