Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngineMetal.mm

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#import <Metal/Metal.h>

#include "BraggIntegrationEngineMetal.h"

#include <algorithm>
#include <string>
#include <unistd.h>

#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"

using namespace bragg_engine;

// The kernels' source: BraggStencil.h, BraggPassA.h and BraggIntegrationEngineMetal.metal, copied in
// at configure time (CMakeLists.txt) so the binary needs no file beside it.
extern const char *const bragg_metal_source_stencil;
extern const char *const bragg_metal_source_pass_a;
extern const char *const bragg_metal_source_exp;
extern const char *const bragg_metal_source_pass_b;
extern const char *const bragg_metal_source_kernels;

namespace {

static_assert(MASK_TAIL == 1 && MASK_FLUX == 2, "BraggIntegrationEngineMetal.metal spells the two marks out");

struct Kernels {
    id<MTLDevice> device = nil;   // nil: no Metal device
    id<MTLComputePipelineState> 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
// second. The Metal compiler ships with macOS, so nothing has to be installed to build or run this.
const Kernels &GetKernels() {
    static const Kernels kernels = [] {
        Kernels k;
        k.device = MTLCreateSystemDefaultDevice();
        if (k.device == nil)
            return k;
        // No contraction of a*b+c into a fused multiply-add, as everywhere on the host (CMakeLists.txt).
        // Metal would otherwise contract wherever its backend saw fit. BraggStencil.h calls C's math
        // functions; the defines are Metal's spellings of them.
        const std::string source = std::string("#include <metal_stdlib>\n"
                                               "using namespace metal;\n"
                                               "#pragma clang fp contract(off)\n"
                                               "#define sqrtf(x) sqrt(x)\n"
                                               "#define fmaxf(a, b) fmax(a, b)\n"
                                               "#define floorf(x) floor(x)\n"
                                               "#define lroundf(x) ((long) round(x))\n")
                                   + bragg_metal_source_stencil + bragg_metal_source_pass_a
                                   + bragg_metal_source_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
        // CPU. (MTLCompileOptions.mathMode says the same from macOS 15; the build targets 13.) And the
        // language of macOS 13 on every macOS, so the kernels mean the same wherever they compile.
        options.fastMathEnabled = NO;
        options.languageVersion = MTLLanguageVersion3_0;
        NSError *error = nil;
        id<MTLLibrary> library = [k.device newLibraryWithSource:@(source.c_str()) options:options error:&error];
        if (library == nil)
            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
                                  std::string("Bragg integration kernels: ") + error.localizedDescription.UTF8String);
        auto pipeline = [&](const char *name) {
            id<MTLFunction> function = [library newFunctionWithName:@(name)];
            id<MTLComputePipelineState> p = [k.device newComputePipelineStateWithFunction:function error:&error];
            if (p == nil)
                throw JFJochException(JFJochExceptionCategory::GPUMetalError,
                                      std::string(name) + ": " + error.localizedDescription.UTF8String);
            return p;
        };
        k.mark = pipeline("mark_neighbours");
        k.claim = pipeline("claim_owner");
        k.box = pipeline("box_sums");
        k.clip = pipeline("clip_sums");
        k.fit = pipeline("profile_fit");
        return k;
    }();
    return kernels;
}

void WaitCommitted(id<MTLCommandBuffer> cb) {
    [cb waitUntilCompleted];
    if (cb.status != MTLCommandBufferStatusCompleted)
        throw JFJochException(JFJochExceptionCategory::GPUMetalError,
                              cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed");
}

void Wait(id<MTLCommandBuffer> cb) {
    [cb commit];
    WaitCommitted(cb);
}

}   // namespace

struct BraggIntegrationEngineMetal::Device {
    id<MTLCommandQueue> queue = nil;
    id<MTLBuffer> static_mask = nil;
    // Frame-sized, made on the first image: a worker's analysis is built for every pass, and most of
    // them never integrate. A frame unmarks what it marked once the CPU has read them (`unmark`, which
    // the next frame waits for): clearing all of both took 1.5 ms of a 16M-pixel frame, longer than
    // integrating a sparse one. `clear` is true once an unmark completed; the first frame, and one
    // after a frame that threw part-way through, clear the whole of both instead.
    id<MTLBuffer> mask = nil, owner = nil;
    id<MTLCommandBuffer> unmark = nil;
    bool clear = false;
    id<MTLBuffer> rad_sum = nil, rad_cnt = nil;   // the radial background curve, as the clip kernel adds it up
    // Per reflection, grown with the frames.
    size_t capacity = 0;
    id<MTLBuffer> xy = nil, level = nil, sums = nil, clips = nil, thr = nil, clip_out = nil;
    id<MTLBuffer> 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<MTLBuffer> cells = nil, profiles = nil;
    // The frame's image while Run hands pass B to FitOnDevice; nil otherwise (the CPU engine's own Run).
    id<MTLBuffer> img = nil;
};

bool BraggIntegrationEngineMetal::Available(const DiffractionExperiment &experiment) {
    const auto settings = experiment.GetBraggIntegrationSettings();
    const bool trimmed_mean = !(settings.GetBackgroundClipNSigma() > 0.0f)
                              && settings.GetBackgroundTrimFraction() > 0.0f;
    return !gpus_are_hidden() && !trimmed_mean && GetKernels().device != nil;
}

std::string BraggIntegrationEngineMetal::DeviceName() {
    if (gpus_are_hidden() || GetKernels().device == nil)
        return {};
    return GetKernels().device.name.UTF8String;
}

BraggIntegrationEngineMetal::BraggIntegrationEngineMetal(const DiffractionExperiment &experiment,
                                                         const PixelMask &mask)
    : BraggIntegrationEngineCPU(experiment, mask), dev(std::make_unique<Device>()) {
    const Kernels &k = GetKernels();
    dev->queue = [k.device newCommandQueue];
    dev->static_mask = [k.device newBufferWithBytes:static_mask.data()
                                             length:std::max<size_t>(static_mask.size() * sizeof(uint32_t), 4)
                                            options:MTLResourceStorageModeShared];
}

BraggIntegrationEngineMetal::~BraggIntegrationEngineMetal() = default;

std::vector<Reflection> BraggIntegrationEngineMetal::Run(const ImagePreprocessorBuffer &image,
                                                         const std::vector<Reflection> &predicted,
                                                         size_t npredicted, int64_t image_number) {
    // An image of another size, and a frame with nothing predicted, the CPU engine answers as it would.
    if (image.size() != npixel || npredicted == 0)
        return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number);

    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
        // for memory starting on a page, as macOS's allocator hands out any buffer of a frame's size.
        // Should one not be, the CPU engine integrates it: the same result.
        const size_t page = getpagesize(), bytes = npixel * sizeof(int32_t);
        id<MTLBuffer> img = nil;
        if (reinterpret_cast<uintptr_t>(image.data()) % page == 0)
            img = [k.device newBufferWithBytesNoCopy:const_cast<int32_t *>(image.data())
                                              length:(bytes + page - 1) / page * page
                                             options:MTLResourceStorageModeShared
                                         deallocator:nil];
        if (img == nil)
            return BraggIntegrationEngineCPU::Run(image, predicted, npredicted, image_number);

        counts.predicted += npredicted;
        const bool use_owner = overlap != OverlapMode::Off;
        if (dev->unmark != nil) {   // it reads xy, which is overwritten below
            id<MTLCommandBuffer> unmark = dev->unmark;
            dev->unmark = nil;
            dev->clear = false;   // unless it completed
            WaitCommitted(unmark);
            dev->clear = true;
        }
        auto buffer = [&](size_t length) {
            return [k.device newBufferWithLength:std::max<size_t>(length, 4) options:MTLResourceStorageModeShared];
        };
        if (dev->mask == nil) {
            dev->mask = buffer(npixel);
            if (use_owner)
                dev->owner = buffer(npixel * sizeof(uint32_t));
        }
        const bool was_clear = dev->clear;
        dev->clear = false;
        if (dev->capacity < npredicted) {
            const size_t n = npredicted + npredicted / 2;
            dev->xy = buffer(n * 2 * sizeof(float));
            dev->level = buffer(n);
            dev->sums = buffer(n * sizeof(BraggBoxSums));
            dev->clips = buffer(n);
            dev->thr = buffer(n * sizeof(int32_t));
            dev->clip_out = buffer(n * sizeof(BraggClipSums));
            dev->fit_setups = buffer(n * sizeof(BraggFitSetup));
            dev->fit_out = buffer(n * sizeof(BraggFitOut));
            dev->capacity = n;
        }

        auto *xy = static_cast<float *>(dev->xy.contents);
        auto *level = static_cast<uint8_t *>(dev->level.contents);
        for (size_t i = 0; i < npredicted; ++i) {
            xy[2 * i] = predicted[i].predicted_x;
            xy[2 * i + 1] = predicted[i].predicted_y;
            level[i] = predicted[i].partiality >= NEIGHBOUR_MASK_MIN_PARTIALITY ? MASK_FLUX : MASK_TAIL;
        }

        PassA a = StartPassA(npredicted);
        BraggPassAFrame f{};
        f.width = static_cast<int32_t>(xpixel);
        f.height = static_cast<int32_t>(ypixel);
        f.r1_sq = r1_sq;
        f.r2_sq = r2_sq;
        f.r3_sq = r3_sq;
        f.claim = claim;
        f.inv_claim = inv_claim;
        f.use_owner = use_owner;
        f.exclude = overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum;
        f.n_rad = bkg_radial ? a.n_rad : 0;
        f.static_mask_words = static_cast<uint32_t>(static_mask.size());
        f.n = static_cast<uint32_t>(npredicted);

        const MTLSize grid = MTLSizeMake(npredicted, 1, 1), group = MTLSizeMake(64, 1, 1);

        // --- The neighbour mask, the owner map and pass A's box walk ---
        id<MTLCommandBuffer> cb = [dev->queue commandBuffer];
        id<MTLBlitCommandEncoder> blit;
        if (!was_clear) {
            blit = [cb blitCommandEncoder];
            [blit fillBuffer:dev->mask range:NSMakeRange(0, npixel) value:0];
            if (use_owner)   // every byte 0xff: BRAGG_OWNER_NONE
                [blit fillBuffer:dev->owner range:NSMakeRange(0, npixel * sizeof(uint32_t)) value:0xff];
            [blit endEncoding];
        }
        const uint32_t claim_not_release = 0;
        id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoder];
        for (const uint8_t want : {MASK_TAIL, MASK_FLUX}) {
            [enc setComputePipelineState:k.mark];
            [enc setBuffer:dev->xy offset:0 atIndex:0];
            [enc setBuffer:dev->level offset:0 atIndex:1];
            [enc setBuffer:dev->mask offset:0 atIndex:2];
            [enc setBytes:&stencil length:sizeof(stencil) atIndex:3];
            [enc setBytes:&f length:sizeof(f) atIndex:4];
            [enc setBytes:&want length:sizeof(want) atIndex:5];
            [enc dispatchThreads:grid threadsPerThreadgroup:group];
        }
        if (use_owner) {
            [enc setComputePipelineState:k.claim];
            [enc setBuffer:dev->xy offset:0 atIndex:0];
            [enc setBuffer:dev->owner offset:0 atIndex:1];
            [enc setBytes:&f length:sizeof(f) atIndex:2];
            [enc setBytes:&claim_not_release length:sizeof(claim_not_release) atIndex:3];
            [enc dispatchThreads:grid threadsPerThreadgroup:group];
        }
        [enc setComputePipelineState:k.box];
        [enc setBuffer:dev->xy offset:0 atIndex:0];
        [enc setBuffer:img offset:0 atIndex:1];
        [enc setBuffer:dev->mask offset:0 atIndex:2];
        [enc setBuffer:(use_owner ? dev->owner : dev->mask) offset:0 atIndex:3];   // not read without an owner map
        [enc setBuffer:dev->static_mask offset:0 atIndex:4];
        [enc setBytes:&stencil length:sizeof(stencil) atIndex:5];
        [enc setBytes:&f length:sizeof(f) atIndex:6];
        [enc setBuffer:dev->sums offset:0 atIndex:7];
        [enc dispatchThreads:grid threadsPerThreadgroup:group];
        [enc endEncoding];
        Wait(cb);

        // --- The clip thresholds, which are double arithmetic, on the CPU; the clip on the GPU ---
        const auto *sums = static_cast<const BraggBoxSums *>(dev->sums.contents);
        auto *clips = static_cast<uint8_t *>(dev->clips.contents);
        auto *thr = static_cast<int32_t *>(dev->thr.contents);
        bool any_clip = false;
        for (size_t i = 0; i < npredicted; ++i) {
            clips[i] = ClipsRing(sums[i]);
            if (clips[i]) {
                thr[i] = ClipThreshold(sums[i]);
                any_clip = true;
            }
        }
        if (any_clip) {
            if (bkg_radial && dev->rad_sum == nil) {
                dev->rad_sum = buffer(a.n_rad * 2 * sizeof(uint32_t));
                dev->rad_cnt = buffer(a.n_rad * sizeof(int32_t));
            }
            cb = [dev->queue commandBuffer];
            if (bkg_radial) {
                blit = [cb blitCommandEncoder];
                [blit fillBuffer:dev->rad_sum range:NSMakeRange(0, a.n_rad * 2 * sizeof(uint32_t)) value:0];
                [blit fillBuffer:dev->rad_cnt range:NSMakeRange(0, a.n_rad * sizeof(int32_t)) value:0];
                [blit endEncoding];
            }
            enc = [cb computeCommandEncoder];
            [enc setComputePipelineState:k.clip];
            [enc setBuffer:dev->xy offset:0 atIndex:0];
            [enc setBuffer:img offset:0 atIndex:1];
            [enc setBuffer:dev->mask offset:0 atIndex:2];
            [enc setBytes:&stencil length:sizeof(stencil) atIndex:3];
            [enc setBytes:&f length:sizeof(f) atIndex:4];
            [enc setBuffer:dev->clips offset:0 atIndex:5];
            [enc setBuffer:dev->thr offset:0 atIndex:6];
            [enc setBuffer:dev->clip_out offset:0 atIndex:7];
            [enc setBuffer:(bkg_radial ? dev->rad_sum : dev->clips) offset:0 atIndex:8];   // not written without the curve
            [enc setBuffer:(bkg_radial ? dev->rad_cnt : dev->clips) offset:0 atIndex:9];
            [enc dispatchThreads:grid threadsPerThreadgroup:group];
            [enc endEncoding];
            Wait(cb);
            if (bkg_radial) {
                const auto *w = static_cast<const uint32_t *>(dev->rad_sum.contents);
                const auto *c = static_cast<const int32_t *>(dev->rad_cnt.contents);
                for (int b = 0; b < a.n_rad; ++b) {
                    a.rad_sum[b] = static_cast<int64_t>(w[2 * b] | static_cast<uint64_t>(w[2 * b + 1]) << 32);
                    a.rad_cnt[b] = c[b];
                }
            }
        }

        ++frames_on_gpu;

        // --- The rest is the CPU engine's ---
        const auto *clip_out = static_cast<const BraggClipSums *>(dev->clip_out.contents);
        std::vector<int32_t> no_trim;   // the trimmed mean is the CPU engine's alone (Available)
        for (size_t i = 0; i < npredicted; ++i) {
            const Reflection &r = predicted[i];
            FinishPassA(a, i, r, MakeBraggStencil(r.predicted_x, r.predicted_y, stencil), sums[i],
                        clips[i] ? clip_out[i] : BraggClipSums{}, no_trim);
        }
        const FlatOwnerMap owner_map{use_owner ? static_cast<const uint32_t *>(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];
        enc = [dev->unmark computeCommandEncoder];
        const uint8_t unmark = 0;
        [enc setComputePipelineState:k.mark];
        [enc setBuffer:dev->xy offset:0 atIndex:0];
        [enc setBuffer:dev->level offset:0 atIndex:1];
        [enc setBuffer:dev->mask offset:0 atIndex:2];
        [enc setBytes:&stencil length:sizeof(stencil) atIndex:3];
        [enc setBytes:&f length:sizeof(f) atIndex:4];
        [enc setBytes:&unmark length:sizeof(unmark) atIndex:5];
        [enc dispatchThreads:grid threadsPerThreadgroup:group];
        if (use_owner) {
            const uint32_t release = 1;
            [enc setComputePipelineState:k.claim];
            [enc setBuffer:dev->xy offset:0 atIndex:0];
            [enc setBuffer:dev->owner offset:0 atIndex:1];
            [enc setBytes:&f length:sizeof(f) atIndex:2];
            [enc setBytes:&release length:sizeof(release) atIndex:3];
            [enc dispatchThreads:grid threadsPerThreadgroup:group];
        }
        [enc endEncoding];
        [dev->unmark commit];
        return ret;
    }
}

bool BraggIntegrationEngineMetal::FitOnDevice(const std::vector<BraggFitSetup> &setups,
                                              const std::vector<float> &shell_profiles,
                                              std::vector<BraggFitOut> &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<size_t>(n_cells, s.cells + static_cast<size_t>(2 * s.Rf + 1) * (2 * s.Rf + 1));
        auto grow = [&](__strong id<MTLBuffer> &b, size_t &capacity, size_t bytes) {
            if (capacity < bytes) {
                capacity = bytes + bytes / 2;
                b = [k.device newBufferWithLength:std::max<size_t>(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<BraggFitSetup *>(dev->fit_setups.contents));
        std::copy(shell_profiles.begin(), shell_profiles.end(), static_cast<float *>(dev->profiles.contents));

        const bool use_owner = overlap != OverlapMode::Off;
        const BraggPassBFrame f{static_cast<int32_t>(xpixel), static_cast<int32_t>(ypixel), r1_sq, R, use_owner,
            overlap == OverlapMode::Exclude, static_cast<uint32_t>(n)};

        id<MTLCommandBuffer> cb = [dev->queue commandBuffer];
        id<MTLComputeCommandEncoder> 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<const BraggFitOut *>(dev->fit_out.contents);
        std::copy(out, out + n, fits.begin());
    }
    return true;
}