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

46 lines
2.7 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
#include <cstddef>
#include <vector>
#include "gemmi/dencalc.hpp" // DensityCalculator
#include "gemmi/grid.hpp"
#include "gemmi/it92.hpp" // IT92 x-ray form factors
#include "gemmi/model.hpp"
// gemmi's gridding of a model - the density of DensityCalculator::put_model_density_on_grid() and the
// bulk-solvent mask of SolventMasker(AtomicRadiiSet::Refmac).put_mask_on_grid() - spread over
// `nthreads` threads of ParallelFor's pool. Every grid point receives the same operations in the same
// order as in gemmi's serial loops, so the grids are the same on any number of threads: the mask is gemmi's
// bit for bit, the density gemmi's but for the last bits (its points are placed in single precision, as the
// Apple GPU places them).
// The points gemmi's Grid::symmetrize() reduces the orbits of this grid's size and space group into
// (empty in P1). Both functions below take them, for a grid of the same size.
std::vector<size_t> OrbitLeaders(const gemmi::Grid<float> &grid, size_t nthreads);
// `dc` set up as for put_model_density_on_grid(): d_min, rate, blur, the grid's cell and group. On macOS the
// Apple GPU puts the density on the grid (ModelDensityMetal.h), to the same bits; `gpu` false keeps it on the
// CPU.
void PutModelDensityOnGrid(gemmi::DensityCalculator<gemmi::IT92<float>, float> &dc, const gemmi::Model &model,
const std::vector<size_t> &orbit_leaders, size_t nthreads, bool gpu = true);
// `grid` sized, with its cell and group set.
void PutMaskOnGrid(gemmi::Grid<float> &grid, const gemmi::Model &model, const std::vector<size_t> &orbit_leaders,
size_t nthreads);
// The number of points of the grid a model's structure factors to d_min are computed on at DensityCalculator's
// `rate` (a spacing of d_min / (2 rate)) in this cell and space group: the grid's size, without making it.
size_t ModelGridPoints(const gemmi::UnitCell &cell, const gemmi::SpaceGroup *sg, double d_min, double rate = 1.5);
// What computing a model's structure factors on the CPU holds at its peak, per point of that grid: the density
// and the bulk-solvent mask grids, and for each of them its transform's input copy, spectrum and coefficients,
// with the two made at once (ModelValidation's compute_model_factors).
constexpr size_t MODEL_FACTORS_BYTES_PER_GRID_POINT = 32;
// What a molecular-replacement placement's structure factors and rigid body hold on the CPU at their peak, per
// point of that grid: measured at under half of the above.
constexpr size_t MODEL_PLACEMENT_BYTES_PER_GRID_POINT = 16;