Jungfraujoch/image_analysis/structure_refinement/MolecularReplacementMetal.h

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

#pragma once

// The rotation function of MolecularReplacement.h on an Apple GPU (macOS only): the same trilinear interpolation
// of the same molecular transform, the same bins and the same correlation, in float on the device. It agrees with
// the CPU to rounding, not bit for bit. The device sums every bin's |F|^2, |F|^4 and |F|^2 (E^2 - 1) for each
// orientation in a fixed order, and the host makes the correlation of them in double.
//
// The device reads its own copy of the transform, in 2 x 2 x 2 bricks: the eight corners of an interpolation then
// lie on fewer cache lines than in the grid's rows, which the GPU, waiting on memory as the CPU does, reads 1.3x
// faster. The copy takes as much memory as the transform and is held as long as this object.

#include <memory>
#include <vector>

#include "MolecularReplacement.h"

class MolecularReplacementMetal {
public:
    // Whether the rotation function of `ft` can run on the GPU: a Metal device the process may use (rugnux
    // --no-gpu hides it, see hide_gpus) whose buffers hold the transform.
    static bool Available(const MolecularTransform &ft);

    // Copies the transform into the device's layout. Throws JFJochException on a Metal failure, as every call
    // below does.
    MolecularReplacementMetal(const MolecularTransform &ft, size_t nthreads);
    ~MolecularReplacementMetal();

    // RotationFunction() of MolecularReplacement.h; may be called from several threads at once.
    std::vector<double> RotationFunction(const SearchReflections &refl, const gemmi::UnitCell &cell,
                                         const gemmi::SpaceGroup &sg, const std::vector<gemmi::Mat33> &rotations) const;

private:
    struct State;
    std::unique_ptr<State> state_;
};