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