Files
leonarski_fandClaude Opus 5.5 868ad0c8f2 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
2026-10-18 17:04:17 +02:00

43 lines
2.2 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
// PutModelDensityOnGrid()'s density (ModelGrid.cpp) on an Apple GPU (macOS only), before the symmetrize: the same
// grid, bit for bit. A gather, as ModelDensityGPU.cu's under CUDA: the grid is cut into bricks of 4^3 points,
// one threadgroup each, and every point adds, in model order, each atom whose box holds it and whose radius
// reaches it - the points UseDensityBoxesByPlane() hands the atom, with the same single-precision arithmetic
// (ModelDensityMetal.metal spells it out), gemmi's float expapprox and the same order of the sums. Nothing is
// summed with an atomic. tests/ModelDensityMetalTest.cpp holds it to the CPU's grid.
#include <cstddef>
#include <vector>
// One atom, as ModelGrid.cpp's DensityBox and AtomDensity hold it: its box about the grid point c nearest its
// centre, c - d to c + d; gemmi's five Gaussians of the isotropic or the anisotropic density.
struct ModelDensityMetalAtom {
float nctr[3]; // the centre in grid units
int c[3], d[3];
float max_r2; // the radius squared
float occ;
int aniso; // 0: b[k][0] multiplies r^2; 1: b[k] is u11 u22 u33 u12 u13 u23
float a[5];
float b[5][6];
};
// The grid, u fastest, and gemmi's orth_n over its size in single precision (upper triangular).
struct ModelDensityMetalGrid {
int nu = 0, nv = 0, nw = 0;
float m00 = 0, m01 = 0, m02 = 0, m11 = 0, m12 = 0, m22 = 0;
};
// The density of `atoms` into `out` (nu nv nw floats, every point written) on the GPU. False where it is not
// taken there - no Metal device the process may use (rugnux --no-gpu hides it, see hide_gpus), or an atom whose
// box is wider than the cell, which reaches a point through two images;
// `out` is then untouched. `nthreads` CPU threads sort the atoms into the bricks. Throws JFJochException on a
// Metal failure.
bool ModelDensityMetal(const ModelDensityMetalGrid &grid, const std::vector<ModelDensityMetalAtom> &atoms, float *out,
size_t nthreads);
// How many grids the GPU has put the density on in this process.
size_t ModelDensityMetalGrids();