ModelGrid: the model density on the Apple GPU (Metal), bit-identical to the CPU

On macOS, where a Metal device may be used (not under --no-gpu), PutModelDensityOnGrid puts the density
on the grid with a Metal gather (ModelDensityMetal), as ModelDensityGPU.cu does under CUDA: the grid in
bricks of 4^3 points, one threadgroup each; every point adds, in model order, each atom whose box holds
it and whose radius reaches it. No float atomics. The arithmetic is the CPU walk's of the previous commit,
spelled out in ModelDensityMetal.metal - the point's place in single precision in the same order, gemmi's
float expapprox ported as is, fast math and contraction off - so the grid is the CPU's bit for bit
(ModelDensityMetal_BitIdenticalToCPU, [portable][Metal]: a triclinic cell and a P 21 21 21 one, a quarter
of the atoms anisotropic, boxes wrapping at every face). The symmetrize stays on the CPU.

- Each brick's atoms: (brick, atom) pairs counted and filled on the CPU in parallel over slabs of w,
  each worker taking the atoms in order, so the lists keep model order; at most the atoms' boxes' bricks.
- The kernel writes straight into the grid FFTW reads where its pages can be handed to the GPU (large
  vectors are page-aligned on macOS), into its own buffer and copied otherwise.
- 4^3 bricks rather than CUDA's 8^3: every point tests each atom reaching its brick, and the smaller brick
  took a half to a third of the kernel time. The box test is a subtraction against a start the host wraps.
- Every grid goes to the GPU: giving it only grids of 256k points or more, or models of 8000 atoms or
  more, left 9s02's MR 0.8 s slower.
- A box wider than the cell (an atom reaching a point through two images) stays on the CPU, whose walk
  adds both. A Metal failure puts that grid and every later one on the CPU, with one warning.
- Each call runs in an autorelease pool of its own (see the rotation function's).

Apple M6 (16 GB), alternating with the rotation-function build (the previous commit but one), two each:
- 4nwv: wall 295.8/292.6 -> 275.2/277.8 s; model validation 48.9/49.0 -> 41.0/40.9 s; MR 130.2/130.3 ->
  121.4/122.5 s (rescoring 40.9/40.3 -> 35.6/36.5 s, null 31.7/32.1 -> 26.8/26.6 s); user CPU 1998/2011 ->
  1692/1681 s; peak footprint 17.8/18.5 -> 17.2/18.6 GB. The density's CPU time, sampled over a whole run,
  316 CPU-s (16% of the busy samples) -> none, the pairs' sort 31 CPU-s.
- 9s02: wall 76.7/76.5 -> 72.9/73.1 s; MR 16.9/16.8 -> 13.5/13.4 s.
- 9zlo: wall 70.4/70.7 -> 69.8/69.6 s; MR 3.1/3.1 -> 2.2/2.2 s.
p.hkl, the placed model and the maps are byte-identical to the previous commit's (density on the CPU) on
all three, and so are the MR and validation logs but for their timings.

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

image_analysis/structure_refinement/ModelGrid.cpp

@@ -12,6 +12,13 @@
 #include "gemmi/solmask.hpp"   // SolventMasker
 
 #include "../../common/ParallelFor.h"
+#ifdef __APPLE__
+#include <atomic>
+
+#include "../../common/JFJochException.h"
+#include "../../common/Logger.h"
+#include "ModelDensityMetal.h"
+#endif
 
 namespace {
 
@@ -337,6 +344,56 @@ struct AtomDensity {
     }
 };
 
+#ifdef __APPLE__
+// UseDensityBoxesByPlane() with AtomDensity on the Apple GPU (ModelDensityMetal.h): the same grid, bit for bit.
+// False where the GPU does not take it - and from the first time it fails on, with one warning, as the CPU
+// then does the same work.
+std::atomic<bool> density_metal_failed{false};
+
+bool PutDensityOnGridMetal(gemmi::Grid<float> &grid, const std::vector<DensityBox> &boxes,
+                           const std::vector<AtomDensity> &density, size_t nthreads) {
+    if (density_metal_failed)
+        return false;
+    static_assert(std::size(AtomDensity::IsoSum{}.a) == 5 && std::size(AtomDensity::AnisoSum{}.a) == 5);
+    const DensityGeometry m = ModelDensityGeometry(grid);
+    const ModelDensityMetalGrid g{grid.nu, grid.nv, grid.nw, m.m00, m.m01, m.m02, m.m11, m.m12, m.m22};
+    std::vector<ModelDensityMetalAtom> atoms(boxes.size());
+    for (size_t i = 0; i < boxes.size(); i++) {
+        ModelDensityMetalAtom &a = atoms[i];
+        const DensityBox &b = boxes[i];
+        const AtomDensity &d = density[i];
+        for (int k = 0; k < 3; k++) {
+            a.nctr[k] = b.nctr[k];
+            a.c[k] = b.c[k];
+            a.d[k] = b.d[k];
+        }
+        a.max_r2 = b.max_r2;
+        a.occ = d.occ;
+        a.aniso = d.is_aniso ? 1 : 0;
+        for (int q = 0; q < 5; q++) {
+            std::fill(std::begin(a.b[q]), std::end(a.b[q]), 0.0f);
+            if (!d.is_aniso) {
+                a.a[q] = d.iso.a[q];
+                a.b[q][0] = d.iso.b[q];
+            } else {
+                const auto &u = d.aniso.b[q];
+                a.a[q] = d.aniso.a[q];
+                const float row[6] = {u.u11, u.u22, u.u33, u.u12, u.u13, u.u23};
+                std::copy(std::begin(row), std::end(row), a.b[q]);
+            }
+        }
+    }
+    try {
+        return ModelDensityMetal(g, atoms, grid.data.data(), nthreads);
+    } catch (const JFJochException &e) {
+        if (!density_metal_failed.exchange(true))
+            Logger("Rugnux").Warning("Model density: the Apple GPU failed, on the CPU from here on: {}", e.what());
+        std::fill(grid.data.begin(), grid.data.end(), 0.0f);   // the CPU adds to a zero grid; the GPU may have written
+        return false;
+    }
+}
+#endif
+
 }  // namespace
 
 // The points gemmi's Grid::symmetrize() (grid.hpp) reduces the grid's orbits into, ascending. gemmi
@@ -372,7 +429,7 @@ std::vector<size_t> OrbitLeaders(const gemmi::Grid<float> &grid, size_t nthreads
 // and gemmi's but for the last bits of the single-precision walk (UseDensityBoxesByPlane). Each atom's density
 // coefficients and radius are worked out once, up front, rather than once per plane.
 void PutModelDensityOnGrid(gemmi::DensityCalculator<Table, float> &dc, const gemmi::Model &model,
-                           const std::vector<size_t> &orbit_leaders, size_t nthreads) {
+                           const std::vector<size_t> &orbit_leaders, size_t nthreads, [[maybe_unused]] bool gpu) {
     using CReal = AtomDensity::CReal;
     dc.initialize_grid();
     const std::vector<const gemmi::Atom *> atoms = ModelAtoms(model);
@@ -401,7 +458,10 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator<Table, float> &dc, const gem
             }
         }
     });
-    UseDensityBoxesByPlane(dc.grid, boxes, nthreads, [&](int i) { return density[i]; });
+#ifdef __APPLE__
+    if (!gpu || !PutDensityOnGridMetal(dc.grid, boxes, density, nthreads))
+#endif
+        UseDensityBoxesByPlane(dc.grid, boxes, nthreads, [&](int i) { return density[i]; });
     SymmetrizeOrbits(dc.grid, orbit_leaders, nthreads, [](float a, float b) { return a + b; });
 }