Jungfraujoch/image_analysis/structure_refinement/ModelDensityMetal.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 "ModelDensityMetal.h"

#include <algorithm>
#include <atomic>
#include <cstdint>
#include <cstring>
#include <limits>
#include <string>
#include <unistd.h>

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

// The kernel's source, ModelDensityMetal.metal, copied in at configure time (CMakeLists.txt) so the binary needs
// no file beside it.
extern const char *const model_density_metal_source;

namespace {

constexpr int BRICK = 4;   // as in the kernel

std::atomic<size_t> grids_on_gpu{0};

static_assert(sizeof(ModelDensityMetalAtom) == 47 * 4, "the kernel's Atom is 47 words, none of them padding");

struct Kernel {
    id<MTLDevice> device = nil;   // nil: no Metal device
    id<MTLComputePipelineState> pipeline = nil;
    id<MTLCommandQueue> queue = nil;
};

// Compiled once per process, the first time it is asked for.
const Kernel &GetKernel() {
    static const Kernel kernel = [] {
        Kernel k;
        k.device = MTLCreateSystemDefaultDevice();
        if (k.device == nil)
            return k;
        MTLCompileOptions *options = [MTLCompileOptions new];
        // Fast math off: nothing reassociated or approximated, as on the CPU (MTLCompileOptions.mathMode says
        // the same from macOS 15; the build targets 13). The language of macOS 13 on every macOS.
        options.fastMathEnabled = NO;
        options.languageVersion = MTLLanguageVersion3_0;
        NSError *error = nil;
        id<MTLLibrary> library = [k.device newLibraryWithSource:@(model_density_metal_source) options:options
                                                          error:&error];
        if (library == nil)
            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
                                  std::string("model density kernel: ") + error.localizedDescription.UTF8String);
        id<MTLFunction> function = [library newFunctionWithName:@"model_density"];
        k.pipeline = [k.device newComputePipelineStateWithFunction:function error:&error];
        if (k.pipeline == nil)
            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
                                  std::string("model_density: ") + error.localizedDescription.UTF8String);
        k.queue = [k.device newCommandQueue];
        return k;
    }();
    return kernel;
}

id<MTLBuffer> NewBuffer(id<MTLDevice> device, size_t bytes) {
    id<MTLBuffer> b = [device newBufferWithLength:std::max<size_t>(bytes, 4) options:MTLResourceStorageModeShared];
    if (b == nil)
        throw JFJochException(JFJochExceptionCategory::GPUMetalError, "model density: out of GPU memory");
    return b;
}

// The bricks along an axis of n points that a box's points first .. first + span (wrapped, span < n, first in
// [0, n)) fall in, each once, into out; how many.
int AxisBricks(int first, int span, int n, int *out) {
    int count = 0;
    for (int p = first; p <= first + span;) {
        const int q = p < n ? p : p - n;
        const int b = q / BRICK;
        if (std::find(out, out + count, b) == out + count)
            out[count++] = b;
        p += std::min(BRICK - q % BRICK, n - q);   // to the next brick, or past the axis' last point
    }
    return count;
}

}  // namespace

bool ModelDensityMetal(const ModelDensityMetalGrid &grid, const std::vector<ModelDensityMetalAtom> &atoms,
                       float *out, size_t nthreads) {
    // In a pool of its own: the command buffer and its encoder are autoreleased, and the threads this runs on have
    // none, so they - and every buffer they hold - would otherwise live as long as the thread.
    @autoreleasepool {
        const size_t points = static_cast<size_t>(grid.nu) * grid.nv * grid.nw;
        if (gpus_are_hidden() || points > std::numeric_limits<uint32_t>::max())
            return false;
        const int n3[3] = {grid.nu, grid.nv, grid.nw};
        for (const ModelDensityMetalAtom &a : atoms)
            for (int k = 0; k < 3; k++)
                if (2 * a.d[k] + 1 > n3[k])
                    return false;   // a box wider than the cell: gemmi's walk adds both images of the atom
        const Kernel &k = GetKernel();
        if (k.device == nil)
            return false;

        // Where each atom's box starts along each axis, wrapped into the grid, and its width less one.
        id<MTLBuffer> d_boxes = NewBuffer(k.device, atoms.size() * 6 * sizeof(int32_t));
        auto *boxes = static_cast<int32_t *>(d_boxes.contents);
        for (size_t i = 0; i < atoms.size(); i++)
            for (int a = 0; a < 3; a++) {
                const int first_point = (atoms[i].c[a] - atoms[i].d[a]) % n3[a];
                boxes[6 * i + a] = first_point < 0 ? first_point + n3[a] : first_point;
                boxes[6 * i + 3 + a] = 2 * atoms[i].d[a];
            }
        // Each brick's atoms in model order: the (brick, atom) pairs counted, then filled, each worker taking the
        // bricks of a slab of w - a contiguous range of them - and the atoms in order.
        const int nb3[3] = {(grid.nu + BRICK - 1) / BRICK, (grid.nv + BRICK - 1) / BRICK,
                            (grid.nw + BRICK - 1) / BRICK};
        const size_t bricks = static_cast<size_t>(nb3[0]) * nb3[1] * nb3[2];
        std::vector<uint32_t> first(bricks + 1, 0);
        const auto pairs_of_slab = [&](int z0, int z1, const auto &pair) {
            std::vector<int> bu(nb3[0]), bv(nb3[1]), bw(nb3[2]);
            for (size_t i = 0; i < atoms.size(); i++) {
                const int32_t *box = boxes + 6 * i;
                const int nw = AxisBricks(box[2], box[5], grid.nw, bw.data());
                int nv = -1, nu = 0;
                for (int kw = 0; kw < nw; kw++) {
                    if (bw[kw] < z0 || bw[kw] >= z1)
                        continue;
                    if (nv < 0) {
                        nv = AxisBricks(box[1], box[4], grid.nv, bv.data());
                        nu = AxisBricks(box[0], box[3], grid.nu, bu.data());
                    }
                    for (int kv = 0; kv < nv; kv++)
                        for (int ku = 0; ku < nu; ku++)
                            pair(bu[ku] + static_cast<size_t>(nb3[0]) * (bv[kv] + static_cast<size_t>(nb3[1]) * bw[kw]),
                                 i);
                }
            }
        };
        ParallelChunks(nb3[2], nthreads, [&](int z0, int z1) {
            pairs_of_slab(z0, z1, [&](size_t brick, size_t) { first[brick + 1]++; });
        });
        for (size_t b = 0; b < bricks; b++)
            first[b + 1] += first[b];
        id<MTLBuffer> d_atoms = NewBuffer(k.device, atoms.size() * sizeof(ModelDensityMetalAtom));
        id<MTLBuffer> d_pairs = NewBuffer(k.device, first[bricks] * sizeof(uint32_t));
        id<MTLBuffer> d_first = NewBuffer(k.device, first.size() * sizeof(uint32_t));
        std::memcpy(d_atoms.contents, atoms.data(), atoms.size() * sizeof(ModelDensityMetalAtom));
        std::memcpy(d_first.contents, first.data(), first.size() * sizeof(uint32_t));
        auto *pairs = static_cast<uint32_t *>(d_pairs.contents);
        ParallelChunks(nb3[2], nthreads, [&](int z0, int z1) {
            pairs_of_slab(z0, z1, [&](size_t brick, size_t i) { pairs[first[brick]++] = static_cast<uint32_t>(i); });
        });
        // The fill moved every brick's start to the next one's: d_first keeps the starts.

        // The kernel writes into `out` itself where its pages can be handed to the GPU as they are; into a buffer
        // of its own, copied out, otherwise.
        const size_t page = static_cast<size_t>(getpagesize());
        const size_t bytes = points * sizeof(float);
        id<MTLBuffer> d_out = nil;
        if (reinterpret_cast<uintptr_t>(out) % page == 0)
            d_out = [k.device newBufferWithBytesNoCopy:out length:(bytes + page - 1) / page * page
                                               options:MTLResourceStorageModeShared deallocator:nil];
        const bool in_place = d_out != nil;
        if (!in_place)
            d_out = NewBuffer(k.device, bytes);

        id<MTLCommandBuffer> cb = [k.queue commandBuffer];
        id<MTLComputeCommandEncoder> e = [cb computeCommandEncoder];
        [e setComputePipelineState:k.pipeline];
        [e setBuffer:d_atoms offset:0 atIndex:0];
        [e setBytes:&grid length:sizeof(grid) atIndex:1];
        [e setBuffer:d_pairs offset:0 atIndex:2];
        [e setBuffer:d_first offset:0 atIndex:3];
        [e setBuffer:d_out offset:0 atIndex:4];
        [e setBuffer:d_boxes offset:0 atIndex:5];
        [e dispatchThreadgroups:MTLSizeMake(bricks, 1, 1) threadsPerThreadgroup:MTLSizeMake(BRICK, BRICK, BRICK)];
        [e endEncoding];
        [cb commit];
        [cb waitUntilCompleted];
        if (cb.status != MTLCommandBufferStatusCompleted)
            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
                                  cb.error ? cb.error.localizedDescription.UTF8String
                                           : "model density: command buffer failed");
        if (!in_place)
            std::memcpy(out, d_out.contents, bytes);
        grids_on_gpu++;
        return true;
    }
}

size_t ModelDensityMetalGrids() {
    return grids_on_gpu;
}