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

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
#include <numeric>
#include <string>

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

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

namespace {

// The grid in bricks of 2 x 2 x 2 points, brick (i/2, j/2, l/2) holding point (i%2, j%2, l%2) at
// (i%2) + 2 (j%2) + 4 (l%2); the bricks in the grid's order. The kernel's index is 32-bit.
size_t BrickedPoints(const gemmi::FPhiGrid<float> &g) {
    return static_cast<size_t>((g.nu + 1) / 2) * ((g.nv + 1) / 2) * ((g.nw + 1) / 2) * 8;
}

constexpr size_t THREADS = 256;
// Orientations per command buffer: half a second of the GPU at most on the largest transforms.
constexpr size_t BATCH = 256;

struct Params {
    int32_t nu, nv, bu, bv;
    uint32_t nops, bins;
    float box, blur;
};

struct Kernel {
    id<MTLDevice> device = nil;   // nil: no Metal device
    id<MTLComputePipelineState> pipeline = 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];
        options.languageVersion = MTLLanguageVersion3_0;
        NSError *error = nil;
        id<MTLLibrary> library = [k.device newLibraryWithSource:@(mr_metal_source) options:options error:&error];
        if (library == nil)
            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
                                  std::string("rotation function kernel: ") + error.localizedDescription.UTF8String);
        id<MTLFunction> function = [library newFunctionWithName:@"rotation_function"];
        k.pipeline = [k.device newComputePipelineStateWithFunction:function error:&error];
        if (k.pipeline == nil)
            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
                                  std::string("rotation_function: ") + error.localizedDescription.UTF8String);
        return k;
    }();
    return kernel;
}

id<MTLBuffer> Upload(id<MTLDevice> device, const void *data, size_t bytes) {
    id<MTLBuffer> b = [device newBufferWithBytes:data length:std::max<size_t>(bytes, 4)
                                         options:MTLResourceStorageModeShared];
    if (b == nil)
        throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory");
    return b;
}

}  // namespace

struct MolecularReplacementMetal::State {
    id<MTLCommandQueue> queue = nil;
    id<MTLBuffer> grid = nil;
    Params params{};
};

bool MolecularReplacementMetal::Available(const MolecularTransform &ft) {
    if (gpus_are_hidden())
        return false;
    const Kernel &k = GetKernel();
    const size_t points = BrickedPoints(ft.Grid());
    return k.device != nil && points <= std::numeric_limits<uint32_t>::max()
           && points * sizeof(std::complex<float>) <= k.device.maxBufferLength;
}

MolecularReplacementMetal::MolecularReplacementMetal(const MolecularTransform &ft, size_t nthreads)
    : state_(std::make_unique<State>()) {
    const Kernel &k = GetKernel();
    const gemmi::FPhiGrid<float> &g = ft.Grid();
    const int bu = (g.nu + 1) / 2, bv = (g.nv + 1) / 2, bw = (g.nw + 1) / 2;
    state_->params = {g.nu, g.nv, bu, bv, 0, 0, static_cast<float>(ft.Box()), static_cast<float>(ft.Blur())};
    state_->queue = [k.device newCommandQueue];
    // Points past the grid's edge, in its last bricks, are never read.
    state_->grid = [k.device newBufferWithLength:BrickedPoints(g) * sizeof(std::complex<float>)
                                         options:MTLResourceStorageModeShared];
    if (state_->queue == nil || state_->grid == nil)
        throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory");
    auto *out = static_cast<std::complex<float> *>(state_->grid.contents);
    ParallelFor(bw, nthreads, [&](int bl) {
        for (int l = 2 * bl; l < std::min(2 * bl + 2, g.nw); l++)
            for (int j = 0; j < g.nv; j++)
                for (int i = 0; i < g.nu; i++) {
                    const size_t brick = (i >> 1) + static_cast<size_t>(bu) * ((j >> 1) + static_cast<size_t>(bv) * (l >> 1));
                    out[brick * 8 + (i & 1) + 2 * (j & 1) + 4 * (l & 1)] =
                            g.data[i + static_cast<size_t>(g.nu) * (j + static_cast<size_t>(g.nv) * l)];
                }
    });
}

MolecularReplacementMetal::~MolecularReplacementMetal() = default;

std::vector<double> MolecularReplacementMetal::RotationFunction(const SearchReflections &refl,
                                                                const gemmi::UnitCell &cell,
                                                                const gemmi::SpaceGroup &sg,
                                                                const std::vector<gemmi::Mat33> &rotations) const {
    const Kernel &k = GetKernel();
    const size_t n = refl.hkl.size();
    const std::vector<gemmi::Vec3> s = CopyVectors(refl, cell, sg);
    const size_t nops = s.size() / n;
    double w_mean = 0;
    for (double e2 : refl.e2)
        w_mean += e2;
    w_mean /= static_cast<double>(n);
    double w_ss = 0;
    for (double e2 : refl.e2)
        w_ss += (e2 - w_mean) * (e2 - w_mean);

    // The reflections in the order of their bins, each bin's from bin_first[b] to bin_first[b + 1].
    std::vector<size_t> order(n);
    std::iota(order.begin(), order.end(), 0);
    std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { return refl.bin[a] < refl.bin[b]; });
    std::vector<uint32_t> bin_first(refl.bins + 1, 0);
    for (size_t i = 0; i < n; i++)
        bin_first[refl.bin[i] + 1]++;
    for (int b = 0; b < refl.bins; b++)
        bin_first[b + 1] += bin_first[b];
    std::vector<float> sf(3 * n * nops), w(n);
    for (size_t i = 0; i < n; i++) {
        for (size_t j = 0; j < nops; j++) {
            const gemmi::Vec3 &v = s[order[i] * nops + j];
            sf[3 * (i * nops + j)] = static_cast<float>(v.x);
            sf[3 * (i * nops + j) + 1] = static_cast<float>(v.y);
            sf[3 * (i * nops + j) + 2] = static_cast<float>(v.z);
        }
        w[i] = static_cast<float>(refl.e2[order[i]] - w_mean);
    }
    // R^T as Metal's float3x3: three columns of four floats.
    std::vector<float> rot(12 * rotations.size(), 0.0f);
    for (size_t r = 0; r < rotations.size(); r++)
        for (int col = 0; col < 3; col++)
            for (int row = 0; row < 3; row++)
                rot[12 * r + 4 * col + row] = static_cast<float>(rotations[r].a[col][row]);

    Params params = state_->params;
    params.nops = static_cast<uint32_t>(nops);
    params.bins = static_cast<uint32_t>(refl.bins);
    id<MTLBuffer> d_s = Upload(k.device, sf.data(), sf.size() * sizeof(float));
    id<MTLBuffer> d_w = Upload(k.device, w.data(), w.size() * sizeof(float));
    id<MTLBuffer> d_bin_first = Upload(k.device, bin_first.data(), bin_first.size() * sizeof(uint32_t));
    id<MTLBuffer> d_rot = Upload(k.device, rot.data(), rot.size() * sizeof(float));
    id<MTLBuffer> d_out = [k.device newBufferWithLength:std::max<size_t>(rotations.size() * refl.bins * 3 * sizeof(float), 4)
                                                options:MTLResourceStorageModeShared];
    if (d_out == nil)
        throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory");
    std::vector<id<MTLCommandBuffer>> running;
    for (size_t first = 0; first < rotations.size(); first += BATCH) {
        const size_t count = std::min(BATCH, rotations.size() - first);
        id<MTLCommandBuffer> cb = [state_->queue commandBuffer];
        id<MTLComputeCommandEncoder> e = [cb computeCommandEncoder];
        [e setComputePipelineState:k.pipeline];
        [e setBuffer:state_->grid offset:0 atIndex:0];
        [e setBytes:&params length:sizeof(params) atIndex:1];
        [e setBuffer:d_rot offset:first * 12 * sizeof(float) atIndex:2];
        [e setBuffer:d_s offset:0 atIndex:3];
        [e setBuffer:d_w offset:0 atIndex:4];
        [e setBuffer:d_bin_first offset:0 atIndex:5];
        [e setBuffer:d_out offset:first * refl.bins * 3 * sizeof(float) atIndex:6];
        [e dispatchThreadgroups:MTLSizeMake(count, 1, 1) threadsPerThreadgroup:MTLSizeMake(THREADS, 1, 1)];
        [e endEncoding];
        [cb commit];
        running.push_back(cb);
    }
    for (id<MTLCommandBuffer> cb : running) {
        [cb waitUntilCompleted];
        if (cb.status != MTLCommandBufferStatusCompleted)
            throw JFJochException(JFJochExceptionCategory::GPUMetalError,
                                  cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed");
    }

    // The correlation of E^2 - 1 with ic over its bin's mean, from the bins' sums: x = ic / mean_b.
    const auto *sums = static_cast<const float *>(d_out.contents);
    std::vector<double> out(rotations.size());
    for (size_t r = 0; r < rotations.size(); r++) {
        double x_sum = 0, x_ss = 0, xy = 0;
        for (int b = 0; b < refl.bins; b++) {
            const float *o = sums + (r * refl.bins + b) * 3;
            const double mean = o[0] / static_cast<double>(bin_first[b + 1] - bin_first[b]);
            x_sum += o[0] / mean;
            x_ss += o[1] / (mean * mean);
            xy += o[2] / mean;
        }
        const double sxx = x_ss - x_sum * x_sum / static_cast<double>(n);
        out[r] = sxx > 0 && w_ss > 0 ? xy / std::sqrt(sxx * w_ss) : 0.0;
    }
    return out;
}