Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngineMetal.h

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

#pragma once

#include <memory>
#include <string>

#include "BraggIntegrationEngineCPU.h"

// 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, 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<Device> dev;
    uint64_t frames_on_gpu = 0;

    bool FitOnDevice(const std::vector<BraggFitSetup> &setups, const std::vector<float> &shell_profiles,
                     std::vector<BraggFitOut> &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,
    // which only the CPU engine sorts the ring for.
    static bool Available(const DiffractionExperiment &experiment);
    // The Metal device's name, for the log; empty where there is none or the process may not use it.
    static std::string DeviceName();

    BraggIntegrationEngineMetal(const DiffractionExperiment &experiment, const PixelMask &mask);
    ~BraggIntegrationEngineMetal() override;

    using BraggIntegrationEngineCPU::Run;   // the CompressedImage overload stays the CPU engine's

    std::vector<Reflection> Run(const ImagePreprocessorBuffer &image,
                                const std::vector<Reflection> &predicted, size_t npredicted,
                                int64_t image_number) override;

    // The frames whose pixels the GPU walked, rather than the CPU engine this one falls back to.
    [[nodiscard]] uint64_t FramesOnGPU() const { return frames_on_gpu; }
};