Files
Jungfraujoch/tests/ModelDensityMetalTest.cpp
T
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

92 lines
4.0 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include <catch2/catch_all.hpp>
#ifdef __APPLE__
#include <cstring>
#include <random>
#include <vector>
#include <gemmi/dencalc.hpp>
#include <gemmi/it92.hpp>
#include "../image_analysis/structure_refinement/ModelDensityMetal.h"
#include "../image_analysis/structure_refinement/ModelGrid.h"
// The model density on the Apple GPU is the CPU's, bit for bit (ModelDensityMetal.h), so the two grids are compared
// as bit patterns. A Mac without a Metal device fails rather than skips - the macOS CI job is the only place this
// can be checked, as in BraggIntegrationEngineMetalTest.cpp.
namespace {
// Atoms scattered over the whole cell - so that boxes wrap at every face - of a few elements, a quarter of them
// anisotropic, with occupancies below one.
gemmi::Structure random_crystal(const gemmi::UnitCell &cell, const char *group) {
std::mt19937 rng(11);
std::uniform_real_distribution<double> frac(0.0, 1.0), b(8.0, 60.0), occ(0.3, 1.0), u(0.05, 0.4);
const gemmi::El elements[] = {gemmi::El::C, gemmi::El::N, gemmi::El::O, gemmi::El::S};
gemmi::Structure st;
st.cell = cell;
st.spacegroup_hm = group;
st.setup_cell_images();
gemmi::Model model(1);
gemmi::Chain chain("A");
gemmi::Residue res;
res.name = "UNK";
res.seqid = gemmi::SeqId(1, ' ');
for (int i = 0; i < 3000; i++) {
gemmi::Atom a;
a.name = "X";
a.element = elements[i % 4];
a.pos = cell.orthogonalize(gemmi::Fractional(frac(rng), frac(rng), frac(rng)));
a.b_iso = static_cast<float>(b(rng));
a.occ = static_cast<float>(occ(rng));
if (i % 4 == 3) {
const float d1 = static_cast<float>(u(rng)), d2 = static_cast<float>(u(rng)),
d3 = static_cast<float>(u(rng));
a.aniso = {d1, d2, d3, 0.3f * std::min(d1, d2), -0.2f * std::min(d1, d3), 0.1f * std::min(d2, d3)};
}
res.atoms.push_back(a);
}
chain.residues.push_back(res);
model.chains.push_back(chain);
st.models.push_back(model);
return st;
}
void check_same_grid(const gemmi::UnitCell &cell, const char *group, double d_min) {
const gemmi::Structure st = random_crystal(cell, group);
gemmi::DensityCalculator<gemmi::IT92<float>, float> cpu, gpu;
for (auto *dc : {&cpu, &gpu}) {
dc->d_min = d_min;
dc->rate = 1.5;
dc->set_grid_cell_and_spacegroup(st);
dc->set_refmac_compatible_blur(st.models[0]);
}
// The orbits of the grid initialize_grid() sizes, as ModelValidation.cpp finds them.
gemmi::Grid<float> sized;
sized.unit_cell = cpu.grid.unit_cell;
sized.spacegroup = cpu.grid.spacegroup;
sized.set_size_from_spacing(cpu.requested_grid_spacing(), gemmi::GridSizeRounding::Up);
const std::vector<size_t> leaders = OrbitLeaders(sized, 4);
PutModelDensityOnGrid(cpu, st.models[0], leaders, 4, false);
const size_t before = ModelDensityMetalGrids();
PutModelDensityOnGrid(gpu, st.models[0], leaders, 4);
REQUIRE(ModelDensityMetalGrids() == before + 1); // the GPU put it on the grid
REQUIRE(gpu.grid.data.size() == cpu.grid.data.size());
size_t different = 0;
for (size_t i = 0; i < cpu.grid.data.size(); i++)
different += std::memcmp(&cpu.grid.data[i], &gpu.grid.data[i], sizeof(float)) != 0;
CHECK(different == 0);
}
}
TEST_CASE("ModelDensityMetal_BitIdenticalToCPU", "[portable][Metal]") {
// Triclinic, so that every term of the orthogonalization is there; then a group with symmetry to reduce.
check_same_grid(gemmi::UnitCell(61, 68, 77, 81, 96, 107), "P 1", 2.0);
check_same_grid(gemmi::UnitCell(70, 80, 90, 90, 90, 90), "P 21 21 21", 2.0);
}
#endif