MolecularReplacement: the rotation function on the Apple GPU (Metal), equal to rounding

On macOS, where a Metal device may be used (not under --no-gpu), each merge's rotation stage - the sphere
and the peaks' refinement - runs on the GPU: the same trilinear interpolation of the same transform in
float, one threadgroup per orientation, every bin's sums of |F|^2, |F|^4 and |F|^2 (E^2 - 1) in a fixed
order, and the correlation made of them in double on the host. The GPU reads its own copy of the
transform in 2 x 2 x 2 bricks (1.3x faster than the grid's rows, as the eight corners then lie on fewer
cache lines), made for the rotation stage and given back before the translation function; the seeds'
orientations later on stay on the CPU. Where the GPU fails, the stage goes on on the CPU with a warning.
The kernel is MolecularReplacementMetal.metal, carried in the binary as the Bragg integration's are.
Elsewhere nothing changes.

Measured on an Apple M6 (16 GB), alternating with rc176 + the CPU prefetch, two runs each:
- 4nwv (the timing benchmark's moved model): rotation 78.9 -> 30.7 s, wall 344.5 -> 298.8 s; peak
  footprint 16.1-16.8 -> 18.5 GB.
- 9s02: rotation 2.99 -> 0.28 s, wall 79.2 -> 76.4 s.
p.hkl, the placed model and the maps are identical to the CPU's on both, and so is every line of the
molecular-replacement log (RF Z, TF Z, R values). The RF itself agrees with the CPU's to 1e-4, with the
same best orientation (MolecularReplacementMetal_MatchesCPU, [portable]).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018My7tqt8a4LyKKjP3FhScT

image_analysis/structure_refinement/MolecularReplacement.cpp

@@ -38,6 +38,9 @@
 #include "ModelValidation.h"
 #include "RigidBodyRefine.h"
 #include "MolecularReplacementGPU.h"
+#ifdef __APPLE__
+#include "MolecularReplacementMetal.h"
+#endif
 #ifdef JFJOCH_USE_CUDA
 #include "RigidBodyGPU.h"
 #include "RigidBodyGPUEngine.h"
@@ -2616,7 +2619,36 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
         const SearchReflections lrefl = MakeSearchReflections(*merge.merged, cell, merge.space_group, d_llg,
                                                               result.d_max_search);
         plan_gpu(m);
-        const std::vector<double> rf = rotation_function(refl, cell, merge.space_group, rotations);
+#ifdef __APPLE__
+        // On the Apple GPU where there is one, from its copy of the transform, made for this stage and given back
+        // before the translation function: the seeds' few orientations later on are scored on the CPU. Where the
+        // GPU fails, on the CPU from there on - a Mac has no other GPU to try.
+        std::unique_ptr<MolecularReplacementMetal> metal;
+        const auto metal_failed = [&](const JFJochException &e) {
+            logger.Warning("Molecular replacement: the rotation function on the Apple GPU failed, on the CPU instead: {}",
+                           e.what());
+            metal.reset();
+        };
+        try {
+            if (MolecularReplacementMetal::Available(ft))
+                metal = std::make_unique<MolecularReplacementMetal>(ft, nthreads);
+        } catch (const JFJochException &e) {
+            metal_failed(e);
+        }
+#endif
+        const auto rotation_stage = [&](const std::vector<gemmi::Mat33> &r) {
+#ifdef __APPLE__
+            if (metal) {
+                try {
+                    return metal->RotationFunction(refl, cell, merge.space_group, r);
+                } catch (const JFJochException &e) {
+                    metal_failed(e);
+                }
+            }
+#endif
+            return rotation_function(refl, cell, merge.space_group, r);
+        };
+        const std::vector<double> rf = rotation_stage(rotations);
         double rf_mean = 0, rf_sd = 0;
         for (double x : rf)
             rf_mean += x;
@@ -2658,7 +2690,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                         for (int axis = 0; axis < 3; axis++)
                             for (double sign : {1.0, -1.0})
                                 trial.push_back(axis_rotation(axis, sign * step).multiply(peak_rotations[k]));
-                const std::vector<double> v = rotation_function(refl, cell, merge.space_group, trial);
+                const std::vector<double> v = rotation_stage(trial);
                 size_t next = 0;
                 for (size_t k = 0; k < peaks.size(); k++) {
                     if (!active[k])
@@ -2677,10 +2709,17 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
                 }
             }
         }
+#ifdef __APPLE__
+        const bool on_metal = metal != nullptr;
+        metal.reset();
+#else
+        const bool on_metal = false;
+#endif
         result.t_rotation += seconds_since(t0);
-        logger.Info("Molecular replacement: {} in {}: {} reflections {:.1f}-{:.1f} A, {} orientations, best RF Z {:.1f}",
+        logger.Info("Molecular replacement: {} in {}: {} reflections {:.1f}-{:.1f} A, {} orientations{}, best RF Z {:.1f}",
                     merge.name, merge.space_group.xhm(), refl.hkl.size(), result.d_max_search, result.d_min_search,
-                    rotations.size(), rf_sd > 0 ? (peak_rf.front() - rf_mean) / rf_sd : 0.0);
+                    rotations.size(), on_metal ? " (rotation function on the Apple GPU)" : "",
+                    rf_sd > 0 ? (peak_rf.front() - rf_mean) / rf_sd : 0.0);
 
         // --- translation, for every group and every peak orientation ---
         t0 = std::chrono::steady_clock::now();