From 868ad0c8f2247ce7230f4b98d8691af9ee0037b1 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 18 Oct 2026 17:04:17 +0200 Subject: [PATCH] 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 Claude-Session: https://claude.ai/code/session_018My7tqt8a4LyKKjP3FhScT --- .../structure_refinement/CMakeLists.txt | 13 ++ .../structure_refinement/ModelDensityMetal.h | 42 ++++ .../ModelDensityMetal.metal | 108 ++++++++++ .../structure_refinement/ModelDensityMetal.mm | 196 ++++++++++++++++++ .../structure_refinement/ModelGrid.cpp | 64 +++++- .../structure_refinement/ModelGrid.h | 6 +- tests/CMakeLists.txt | 1 + tests/ModelDensityMetalTest.cpp | 91 ++++++++ 8 files changed, 517 insertions(+), 4 deletions(-) create mode 100644 image_analysis/structure_refinement/ModelDensityMetal.h create mode 100644 image_analysis/structure_refinement/ModelDensityMetal.metal create mode 100644 image_analysis/structure_refinement/ModelDensityMetal.mm create mode 100644 tests/ModelDensityMetalTest.cpp diff --git a/image_analysis/structure_refinement/CMakeLists.txt b/image_analysis/structure_refinement/CMakeLists.txt index d4583917d..62ac58f77 100644 --- a/image_analysis/structure_refinement/CMakeLists.txt +++ b/image_analysis/structure_refinement/CMakeLists.txt @@ -39,5 +39,18 @@ IF (APPLE) MolecularReplacementMetal.metal ${CMAKE_CURRENT_BINARY_DIR}/MolecularReplacementMetalSource.cpp) SET_SOURCE_FILES_PROPERTIES(MolecularReplacementMetal.mm PROPERTIES COMPILE_OPTIONS -fobjc-arc) SET_SOURCE_FILES_PROPERTIES(MolecularReplacementMetal.metal PROPERTIES HEADER_FILE_ONLY ON) + # The model density's gather (ModelDensityMetal.h), carried in the same way. + FILE(READ ${CMAKE_CURRENT_SOURCE_DIR}/ModelDensityMetal.metal text) + SET_PROPERTY(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/ModelDensityMetal.metal) + FILE(WRITE ${CMAKE_CURRENT_BINARY_DIR}/ModelDensityMetalSource.cpp.tmp + "// Generated from image_analysis/structure_refinement by its CMakeLists.txt - do not edit.\n" + "extern const char *const model_density_metal_source;\n" + "const char *const model_density_metal_source = R\"DENSITY_METAL(\n${text})DENSITY_METAL\";\n") + CONFIGURE_FILE(${CMAKE_CURRENT_BINARY_DIR}/ModelDensityMetalSource.cpp.tmp + ${CMAKE_CURRENT_BINARY_DIR}/ModelDensityMetalSource.cpp COPYONLY) + TARGET_SOURCES(JFJochStructureRefinement PRIVATE ModelDensityMetal.mm ModelDensityMetal.h + ModelDensityMetal.metal ${CMAKE_CURRENT_BINARY_DIR}/ModelDensityMetalSource.cpp) + SET_SOURCE_FILES_PROPERTIES(ModelDensityMetal.mm PROPERTIES COMPILE_OPTIONS -fobjc-arc) + SET_SOURCE_FILES_PROPERTIES(ModelDensityMetal.metal PROPERTIES HEADER_FILE_ONLY ON) TARGET_LINK_LIBRARIES(JFJochStructureRefinement "-framework Metal" "-framework Foundation") ENDIF() diff --git a/image_analysis/structure_refinement/ModelDensityMetal.h b/image_analysis/structure_refinement/ModelDensityMetal.h new file mode 100644 index 000000000..6586a3d0d --- /dev/null +++ b/image_analysis/structure_refinement/ModelDensityMetal.h @@ -0,0 +1,42 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// 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 +#include + +// 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 &atoms, float *out, + size_t nthreads); +// How many grids the GPU has put the density on in this process. +size_t ModelDensityMetalGrids(); diff --git a/image_analysis/structure_refinement/ModelDensityMetal.metal b/image_analysis/structure_refinement/ModelDensityMetal.metal new file mode 100644 index 000000000..5a40eadfb --- /dev/null +++ b/image_analysis/structure_refinement/ModelDensityMetal.metal @@ -0,0 +1,108 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +// The model density's gather (ModelDensityMetal.h), compiled when it is first used from this source, which +// CMakeLists.txt copies into the binary - with fast math off and no a*b+c contracted, as on the host, so that +// every operation below rounds as the CPU's does. + +#include +using namespace metal; + +#pragma clang fp contract(off) + +// A brick's edge in points (ModelDensityMetal.mm says the same): every point of a brick tests each atom whose box +// reaches the brick, and with 4 rather than CUDA's 8 far fewer of them test atoms that miss them - the kernel +// took a half to a third of the time on an Apple M6. +constant int BRICK = 4; + +struct Atom { // ModelDensityMetalAtom + float nctr[3]; + int c[3], d[3]; + float max_r2; + float occ; + int aniso; + float a[5]; + float b[5][6]; +}; + +// Where an atom's box starts along each axis, wrapped into the grid, and its width less one (2 d): the test +// whether a point is in the box, without a division. +struct Box { + int first[3]; + int span[3]; +}; + +struct Grid { // ModelDensityMetalGrid + int nu, nv, nw; + float m00, m01, m02, m11, m12, m22; +}; + +// gemmi's unsafe_expapprox() (formfact.hpp, MPL-2.0; https://github.com/project-gemmi/gemmi), which is based on +// expapprox() from https://github.com/jhjourdan/SIMD-math-prims: the same operations, the bits moved with as_type +// where gemmi uses memcpy. +static float expapprox(float x) { + const float val = 12102203.1615614f * x + 1065353216.f; + const int vali = static_cast(val); + const int xu1 = vali & 0x7F800000; + const int xu2 = (vali & 0x7FFFFF) | 0x3F800000; + const float a = as_type(xu1); + const float b = as_type(xu2); + return a * (0.509871020f + b * (0.312146713f + b * (0.166617139f + b * + (-2.190619930e-3f + b * 1.3555747234e-2f)))); +} + +// The point p's offset in the box along an axis of n points, 0 to span where it is in it. +static int box_offset(int p, int first, int n) { + const int m = p - first; + return m < 0 ? m + n : m; +} + +// One threadgroup per brick of BRICK^3 points: each point adds the brick's atoms in model order, as +// UseDensityBoxesByPlane() and AtomDensity (ModelGrid.cpp) add them on the CPU. +kernel void model_density(device const Atom *atoms [[buffer(0)]], constant Grid &g [[buffer(1)]], + device const uint *pair_atom [[buffer(2)]], device const uint *brick_first [[buffer(3)]], + device float *out [[buffer(4)]], device const Box *boxes [[buffer(5)]], + uint3 group [[threadgroup_position_in_grid]], uint3 t [[thread_position_in_threadgroup]]) { + const uint brick = group.x; + const int nbu = (g.nu + BRICK - 1) / BRICK, nbv = (g.nv + BRICK - 1) / BRICK; + const int u = int(brick % uint(nbu)) * BRICK + int(t.x); + const int v = int(brick / uint(nbu) % uint(nbv)) * BRICK + int(t.y); + const int w = int(brick / uint(nbu * nbv)) * BRICK + int(t.z); + if (u >= g.nu || v >= g.nv || w >= g.nw) + return; + float acc = 0.0f; + for (uint k = brick_first[brick]; k < brick_first[brick + 1]; k++) { + const uint i = pair_atom[k]; + device const Box &box = boxes[i]; + const int mu = box_offset(u, box.first[0], g.nu), mv = box_offset(v, box.first[1], g.nv), + mw = box_offset(w, box.first[2], g.nw); + if (mu > box.span[0] || mv > box.span[1] || mw > box.span[2]) + continue; + device const Atom &at = atoms[i]; + // The point's index in the box's own (unwrapped) range, as the CPU walks it. + const int uu = at.c[0] - at.d[0] + mu, vv = at.c[1] - at.d[1] + mv, ww = at.c[2] - at.d[2] + mw; + const float fx = at.nctr[0] - float(uu), fy = at.nctr[1] - float(vv), fz = at.nctr[2] - float(ww); + const float z = g.m22 * fz; + const float y = g.m11 * fy + g.m12 * fz; + const float yz = y * y + z * z; + const float x = g.m00 * fx + (g.m01 * fy + g.m02 * fz); + const float r2 = yz + x * x; + if (r2 > at.max_r2) + continue; + float density = 0.0f; + if (at.aniso == 0) { + for (int q = 0; q < 5; q++) + density += at.a[q] * expapprox(max(at.b[q][0] * r2, -88.f)); + } else { + for (int q = 0; q < 5; q++) { + device const float *b = at.b[q]; + // gemmi's SMat33::r_u_r(), in its order + const float rur = x * x * b[0] + y * y * b[1] + z * z * b[2] + + 2 * (x * y * b[3] + x * z * b[4] + y * z * b[5]); + density += at.a[q] * expapprox(max(rur, -88.f)); + } + } + acc += at.occ * density; + } + out[uint(u) + uint(g.nu) * (uint(v) + uint(g.nv) * uint(w))] = acc; +} diff --git a/image_analysis/structure_refinement/ModelDensityMetal.mm b/image_analysis/structure_refinement/ModelDensityMetal.mm new file mode 100644 index 000000000..b1050260f --- /dev/null +++ b/image_analysis/structure_refinement/ModelDensityMetal.mm @@ -0,0 +1,196 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#import + +#include "ModelDensityMetal.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "../../common/CUDAWrapper.h" // gpus_are_hidden +#include "../../common/JFJochException.h" +#include "../../common/ParallelFor.h" + +// The kernel's source, ModelDensityMetal.metal, copied in at configure time (CMakeLists.txt) so the binary needs +// no file beside it. +extern const char *const model_density_metal_source; + +namespace { + +constexpr int BRICK = 4; // as in the kernel + +std::atomic grids_on_gpu{0}; + +static_assert(sizeof(ModelDensityMetalAtom) == 47 * 4, "the kernel's Atom is 47 words, none of them padding"); + +struct Kernel { + id device = nil; // nil: no Metal device + id pipeline = nil; + id queue = nil; +}; + +// Compiled once per process, the first time it is asked for. +const Kernel &GetKernel() { + static const Kernel kernel = [] { + Kernel k; + k.device = MTLCreateSystemDefaultDevice(); + if (k.device == nil) + return k; + MTLCompileOptions *options = [MTLCompileOptions new]; + // Fast math off: nothing reassociated or approximated, as on the CPU (MTLCompileOptions.mathMode says + // the same from macOS 15; the build targets 13). The language of macOS 13 on every macOS. + options.fastMathEnabled = NO; + options.languageVersion = MTLLanguageVersion3_0; + NSError *error = nil; + id library = [k.device newLibraryWithSource:@(model_density_metal_source) options:options + error:&error]; + if (library == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + std::string("model density kernel: ") + error.localizedDescription.UTF8String); + id function = [library newFunctionWithName:@"model_density"]; + k.pipeline = [k.device newComputePipelineStateWithFunction:function error:&error]; + if (k.pipeline == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + std::string("model_density: ") + error.localizedDescription.UTF8String); + k.queue = [k.device newCommandQueue]; + return k; + }(); + return kernel; +} + +id NewBuffer(id device, size_t bytes) { + id b = [device newBufferWithLength:std::max(bytes, 4) options:MTLResourceStorageModeShared]; + if (b == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, "model density: out of GPU memory"); + return b; +} + +// The bricks along an axis of n points that a box's points first .. first + span (wrapped, span < n, first in +// [0, n)) fall in, each once, into out; how many. +int AxisBricks(int first, int span, int n, int *out) { + int count = 0; + for (int p = first; p <= first + span;) { + const int q = p < n ? p : p - n; + const int b = q / BRICK; + if (std::find(out, out + count, b) == out + count) + out[count++] = b; + p += std::min(BRICK - q % BRICK, n - q); // to the next brick, or past the axis' last point + } + return count; +} + +} // namespace + +bool ModelDensityMetal(const ModelDensityMetalGrid &grid, const std::vector &atoms, + float *out, size_t nthreads) { + // In a pool of its own: the command buffer and its encoder are autoreleased, and the threads this runs on have + // none, so they - and every buffer they hold - would otherwise live as long as the thread. + @autoreleasepool { + const size_t points = static_cast(grid.nu) * grid.nv * grid.nw; + if (gpus_are_hidden() || points > std::numeric_limits::max()) + return false; + const int n3[3] = {grid.nu, grid.nv, grid.nw}; + for (const ModelDensityMetalAtom &a : atoms) + for (int k = 0; k < 3; k++) + if (2 * a.d[k] + 1 > n3[k]) + return false; // a box wider than the cell: gemmi's walk adds both images of the atom + const Kernel &k = GetKernel(); + if (k.device == nil) + return false; + + // Where each atom's box starts along each axis, wrapped into the grid, and its width less one. + id d_boxes = NewBuffer(k.device, atoms.size() * 6 * sizeof(int32_t)); + auto *boxes = static_cast(d_boxes.contents); + for (size_t i = 0; i < atoms.size(); i++) + for (int a = 0; a < 3; a++) { + const int first_point = (atoms[i].c[a] - atoms[i].d[a]) % n3[a]; + boxes[6 * i + a] = first_point < 0 ? first_point + n3[a] : first_point; + boxes[6 * i + 3 + a] = 2 * atoms[i].d[a]; + } + // Each brick's atoms in model order: the (brick, atom) pairs counted, then filled, each worker taking the + // bricks of a slab of w - a contiguous range of them - and the atoms in order. + const int nb3[3] = {(grid.nu + BRICK - 1) / BRICK, (grid.nv + BRICK - 1) / BRICK, + (grid.nw + BRICK - 1) / BRICK}; + const size_t bricks = static_cast(nb3[0]) * nb3[1] * nb3[2]; + std::vector first(bricks + 1, 0); + const auto pairs_of_slab = [&](int z0, int z1, const auto &pair) { + std::vector bu(nb3[0]), bv(nb3[1]), bw(nb3[2]); + for (size_t i = 0; i < atoms.size(); i++) { + const int32_t *box = boxes + 6 * i; + const int nw = AxisBricks(box[2], box[5], grid.nw, bw.data()); + int nv = -1, nu = 0; + for (int kw = 0; kw < nw; kw++) { + if (bw[kw] < z0 || bw[kw] >= z1) + continue; + if (nv < 0) { + nv = AxisBricks(box[1], box[4], grid.nv, bv.data()); + nu = AxisBricks(box[0], box[3], grid.nu, bu.data()); + } + for (int kv = 0; kv < nv; kv++) + for (int ku = 0; ku < nu; ku++) + pair(bu[ku] + static_cast(nb3[0]) * (bv[kv] + static_cast(nb3[1]) * bw[kw]), + i); + } + } + }; + ParallelChunks(nb3[2], nthreads, [&](int z0, int z1) { + pairs_of_slab(z0, z1, [&](size_t brick, size_t) { first[brick + 1]++; }); + }); + for (size_t b = 0; b < bricks; b++) + first[b + 1] += first[b]; + id d_atoms = NewBuffer(k.device, atoms.size() * sizeof(ModelDensityMetalAtom)); + id d_pairs = NewBuffer(k.device, first[bricks] * sizeof(uint32_t)); + id d_first = NewBuffer(k.device, first.size() * sizeof(uint32_t)); + std::memcpy(d_atoms.contents, atoms.data(), atoms.size() * sizeof(ModelDensityMetalAtom)); + std::memcpy(d_first.contents, first.data(), first.size() * sizeof(uint32_t)); + auto *pairs = static_cast(d_pairs.contents); + ParallelChunks(nb3[2], nthreads, [&](int z0, int z1) { + pairs_of_slab(z0, z1, [&](size_t brick, size_t i) { pairs[first[brick]++] = static_cast(i); }); + }); + // The fill moved every brick's start to the next one's: d_first keeps the starts. + + // The kernel writes into `out` itself where its pages can be handed to the GPU as they are; into a buffer + // of its own, copied out, otherwise. + const size_t page = static_cast(getpagesize()); + const size_t bytes = points * sizeof(float); + id d_out = nil; + if (reinterpret_cast(out) % page == 0) + d_out = [k.device newBufferWithBytesNoCopy:out length:(bytes + page - 1) / page * page + options:MTLResourceStorageModeShared deallocator:nil]; + const bool in_place = d_out != nil; + if (!in_place) + d_out = NewBuffer(k.device, bytes); + + id cb = [k.queue commandBuffer]; + id e = [cb computeCommandEncoder]; + [e setComputePipelineState:k.pipeline]; + [e setBuffer:d_atoms offset:0 atIndex:0]; + [e setBytes:&grid length:sizeof(grid) atIndex:1]; + [e setBuffer:d_pairs offset:0 atIndex:2]; + [e setBuffer:d_first offset:0 atIndex:3]; + [e setBuffer:d_out offset:0 atIndex:4]; + [e setBuffer:d_boxes offset:0 atIndex:5]; + [e dispatchThreadgroups:MTLSizeMake(bricks, 1, 1) threadsPerThreadgroup:MTLSizeMake(BRICK, BRICK, BRICK)]; + [e endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + if (cb.status != MTLCommandBufferStatusCompleted) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + cb.error ? cb.error.localizedDescription.UTF8String + : "model density: command buffer failed"); + if (!in_place) + std::memcpy(out, d_out.contents, bytes); + grids_on_gpu++; + return true; + } +} + +size_t ModelDensityMetalGrids() { + return grids_on_gpu; +} diff --git a/image_analysis/structure_refinement/ModelGrid.cpp b/image_analysis/structure_refinement/ModelGrid.cpp index 82345e07d..4e3177182 100644 --- a/image_analysis/structure_refinement/ModelGrid.cpp +++ b/image_analysis/structure_refinement/ModelGrid.cpp @@ -12,6 +12,13 @@ #include "gemmi/solmask.hpp" // SolventMasker #include "../../common/ParallelFor.h" +#ifdef __APPLE__ +#include + +#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 density_metal_failed{false}; + +bool PutDensityOnGridMetal(gemmi::Grid &grid, const std::vector &boxes, + const std::vector &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 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 OrbitLeaders(const gemmi::Grid &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 &dc, const gemmi::Model &model, - const std::vector &orbit_leaders, size_t nthreads) { + const std::vector &orbit_leaders, size_t nthreads, [[maybe_unused]] bool gpu) { using CReal = AtomDensity::CReal; dc.initialize_grid(); const std::vector atoms = ModelAtoms(model); @@ -401,7 +458,10 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator &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; }); } diff --git a/image_analysis/structure_refinement/ModelGrid.h b/image_analysis/structure_refinement/ModelGrid.h index 1ee4e9249..1eea50952 100644 --- a/image_analysis/structure_refinement/ModelGrid.h +++ b/image_analysis/structure_refinement/ModelGrid.h @@ -22,9 +22,11 @@ // (empty in P1). Both functions below take them, for a grid of the same size. std::vector OrbitLeaders(const gemmi::Grid &grid, size_t nthreads); -// `dc` set up as for put_model_density_on_grid(): d_min, rate, blur, the grid's cell and group. +// `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, float> &dc, const gemmi::Model &model, - const std::vector &orbit_leaders, size_t nthreads); + const std::vector &orbit_leaders, size_t nthreads, bool gpu = true); // `grid` sized, with its cell and group set. void PutMaskOnGrid(gemmi::Grid &grid, const gemmi::Model &model, const std::vector &orbit_leaders, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 7dd31e085..1895e035d 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -22,6 +22,7 @@ ADD_EXECUTABLE(jfjoch_portable_test EXCLUDE_FROM_ALL BraggIntegrationEngineCPUTest.cpp BraggIntegrationEngineMetalTest.cpp MolecularReplacementMetalTest.cpp + ModelDensityMetalTest.cpp WriteReflectionsTest.cpp LoadReferenceMtzFreeFlagsTest.cpp LoadReferenceSfCifTest.cpp diff --git a/tests/ModelDensityMetalTest.cpp b/tests/ModelDensityMetalTest.cpp new file mode 100644 index 000000000..7d9aee164 --- /dev/null +++ b/tests/ModelDensityMetalTest.cpp @@ -0,0 +1,91 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#ifdef __APPLE__ + +#include +#include +#include + +#include +#include + +#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 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(b(rng)); + a.occ = static_cast(occ(rng)); + if (i % 4 == 3) { + const float d1 = static_cast(u(rng)), d2 = static_cast(u(rng)), + d3 = static_cast(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, 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 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 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