From f74df573cf1d3c861e43c16fec34426a02c981e1 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 18 Oct 2026 17:04:17 +0200 Subject: [PATCH 1/3] MolecularReplacementMetal: each rotation-function call in an autorelease pool of its own The calls run on the search's worker threads, which have no autorelease pool, so each call's autoreleased command buffers and every buffer they held lived until the thread exited: about 7 MB per call. The Bragg engine already runs its calls in a pool. On an Apple M6 this removes a 4-5 GB rise of 4nwv's peak footprint seen with more Metal calls, and probably the rotation-function build's unexplained +1.7 GB; a 20-call worker-thread test grew 0.41 -> 0.55 GB before and stays flat after. Output unchanged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_018My7tqt8a4LyKKjP3FhScT --- .../MolecularReplacementMetal.mm | 178 +++++++++--------- 1 file changed, 92 insertions(+), 86 deletions(-) diff --git a/image_analysis/structure_refinement/MolecularReplacementMetal.mm b/image_analysis/structure_refinement/MolecularReplacementMetal.mm index 768fdd143..3e84c77a5 100644 --- a/image_analysis/structure_refinement/MolecularReplacementMetal.mm +++ b/image_analysis/structure_refinement/MolecularReplacementMetal.mm @@ -109,7 +109,8 @@ MolecularReplacementMetal::MolecularReplacementMetal(const MolecularTransform &f for (int l = 2 * bl; l < std::min(2 * bl + 2, g.nw); l++) for (int j = 0; j < g.nv; j++) for (int i = 0; i < g.nu; i++) { - const size_t brick = (i >> 1) + static_cast(bu) * ((j >> 1) + static_cast(bv) * (l >> 1)); + const size_t brick = + (i >> 1) + static_cast(bu) * ((j >> 1) + static_cast(bv) * (l >> 1)); out[brick * 8 + (i & 1) + 2 * (j & 1) + 4 * (l & 1)] = g.data[i + static_cast(g.nu) * (j + static_cast(g.nv) * l)]; } @@ -122,94 +123,99 @@ std::vector MolecularReplacementMetal::RotationFunction(const SearchRefl const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, const std::vector &rotations) const { - const Kernel &k = GetKernel(); - const size_t n = refl.hkl.size(); - const std::vector s = CopyVectors(refl, cell, sg); - const size_t nops = s.size() / n; - double w_mean = 0; - for (double e2 : refl.e2) - w_mean += e2; - w_mean /= static_cast(n); - double w_ss = 0; - for (double e2 : refl.e2) - w_ss += (e2 - w_mean) * (e2 - w_mean); + // 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 Kernel &k = GetKernel(); + const size_t n = refl.hkl.size(); + const std::vector s = CopyVectors(refl, cell, sg); + const size_t nops = s.size() / n; + double w_mean = 0; + for (double e2 : refl.e2) + w_mean += e2; + w_mean /= static_cast(n); + double w_ss = 0; + for (double e2 : refl.e2) + w_ss += (e2 - w_mean) * (e2 - w_mean); - // The reflections in the order of their bins, each bin's from bin_first[b] to bin_first[b + 1]. - std::vector order(n); - std::iota(order.begin(), order.end(), 0); - std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { return refl.bin[a] < refl.bin[b]; }); - std::vector bin_first(refl.bins + 1, 0); - for (size_t i = 0; i < n; i++) - bin_first[refl.bin[i] + 1]++; - for (int b = 0; b < refl.bins; b++) - bin_first[b + 1] += bin_first[b]; - std::vector sf(3 * n * nops), w(n); - for (size_t i = 0; i < n; i++) { - for (size_t j = 0; j < nops; j++) { - const gemmi::Vec3 &v = s[order[i] * nops + j]; - sf[3 * (i * nops + j)] = static_cast(v.x); - sf[3 * (i * nops + j) + 1] = static_cast(v.y); - sf[3 * (i * nops + j) + 2] = static_cast(v.z); + // The reflections in the order of their bins, each bin's from bin_first[b] to bin_first[b + 1]. + std::vector order(n); + std::iota(order.begin(), order.end(), 0); + std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { return refl.bin[a] < refl.bin[b]; }); + std::vector bin_first(refl.bins + 1, 0); + for (size_t i = 0; i < n; i++) + bin_first[refl.bin[i] + 1]++; + for (int b = 0; b < refl.bins; b++) + bin_first[b + 1] += bin_first[b]; + std::vector sf(3 * n * nops), w(n); + for (size_t i = 0; i < n; i++) { + for (size_t j = 0; j < nops; j++) { + const gemmi::Vec3 &v = s[order[i] * nops + j]; + sf[3 * (i * nops + j)] = static_cast(v.x); + sf[3 * (i * nops + j) + 1] = static_cast(v.y); + sf[3 * (i * nops + j) + 2] = static_cast(v.z); + } + w[i] = static_cast(refl.e2[order[i]] - w_mean); } - w[i] = static_cast(refl.e2[order[i]] - w_mean); - } - // R^T as Metal's float3x3: three columns of four floats. - std::vector rot(12 * rotations.size(), 0.0f); - for (size_t r = 0; r < rotations.size(); r++) - for (int col = 0; col < 3; col++) - for (int row = 0; row < 3; row++) - rot[12 * r + 4 * col + row] = static_cast(rotations[r].a[col][row]); + // R^T as Metal's float3x3: three columns of four floats. + std::vector rot(12 * rotations.size(), 0.0f); + for (size_t r = 0; r < rotations.size(); r++) + for (int col = 0; col < 3; col++) + for (int row = 0; row < 3; row++) + rot[12 * r + 4 * col + row] = static_cast(rotations[r].a[col][row]); - Params params = state_->params; - params.nops = static_cast(nops); - params.bins = static_cast(refl.bins); - id d_s = Upload(k.device, sf.data(), sf.size() * sizeof(float)); - id d_w = Upload(k.device, w.data(), w.size() * sizeof(float)); - id d_bin_first = Upload(k.device, bin_first.data(), bin_first.size() * sizeof(uint32_t)); - id d_rot = Upload(k.device, rot.data(), rot.size() * sizeof(float)); - id d_out = [k.device newBufferWithLength:std::max(rotations.size() * refl.bins * 3 * sizeof(float), 4) - options:MTLResourceStorageModeShared]; - if (d_out == nil) - throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory"); - std::vector> running; - for (size_t first = 0; first < rotations.size(); first += BATCH) { - const size_t count = std::min(BATCH, rotations.size() - first); - id cb = [state_->queue commandBuffer]; - id e = [cb computeCommandEncoder]; - [e setComputePipelineState:k.pipeline]; - [e setBuffer:state_->grid offset:0 atIndex:0]; - [e setBytes:¶ms length:sizeof(params) atIndex:1]; - [e setBuffer:d_rot offset:first * 12 * sizeof(float) atIndex:2]; - [e setBuffer:d_s offset:0 atIndex:3]; - [e setBuffer:d_w offset:0 atIndex:4]; - [e setBuffer:d_bin_first offset:0 atIndex:5]; - [e setBuffer:d_out offset:first * refl.bins * 3 * sizeof(float) atIndex:6]; - [e dispatchThreadgroups:MTLSizeMake(count, 1, 1) threadsPerThreadgroup:MTLSizeMake(THREADS, 1, 1)]; - [e endEncoding]; - [cb commit]; - running.push_back(cb); - } - for (id cb : running) { - [cb waitUntilCompleted]; - if (cb.status != MTLCommandBufferStatusCompleted) - throw JFJochException(JFJochExceptionCategory::GPUMetalError, - cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed"); - } - - // The correlation of E^2 - 1 with ic over its bin's mean, from the bins' sums: x = ic / mean_b. - const auto *sums = static_cast(d_out.contents); - std::vector out(rotations.size()); - for (size_t r = 0; r < rotations.size(); r++) { - double x_sum = 0, x_ss = 0, xy = 0; - for (int b = 0; b < refl.bins; b++) { - const float *o = sums + (r * refl.bins + b) * 3; - const double mean = o[0] / static_cast(bin_first[b + 1] - bin_first[b]); - x_sum += o[0] / mean; - x_ss += o[1] / (mean * mean); - xy += o[2] / mean; + Params params = state_->params; + params.nops = static_cast(nops); + params.bins = static_cast(refl.bins); + id d_s = Upload(k.device, sf.data(), sf.size() * sizeof(float)); + id d_w = Upload(k.device, w.data(), w.size() * sizeof(float)); + id d_bin_first = Upload(k.device, bin_first.data(), bin_first.size() * sizeof(uint32_t)); + id d_rot = Upload(k.device, rot.data(), rot.size() * sizeof(float)); + const size_t out_bytes = rotations.size() * refl.bins * 3 * sizeof(float); + id d_out = [k.device newBufferWithLength:std::max(out_bytes, 4) + options:MTLResourceStorageModeShared]; + if (d_out == nil) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory"); + std::vector> running; + for (size_t first = 0; first < rotations.size(); first += BATCH) { + const size_t count = std::min(BATCH, rotations.size() - first); + id cb = [state_->queue commandBuffer]; + id e = [cb computeCommandEncoder]; + [e setComputePipelineState:k.pipeline]; + [e setBuffer:state_->grid offset:0 atIndex:0]; + [e setBytes:¶ms length:sizeof(params) atIndex:1]; + [e setBuffer:d_rot offset:first * 12 * sizeof(float) atIndex:2]; + [e setBuffer:d_s offset:0 atIndex:3]; + [e setBuffer:d_w offset:0 atIndex:4]; + [e setBuffer:d_bin_first offset:0 atIndex:5]; + [e setBuffer:d_out offset:first * refl.bins * 3 * sizeof(float) atIndex:6]; + [e dispatchThreadgroups:MTLSizeMake(count, 1, 1) threadsPerThreadgroup:MTLSizeMake(THREADS, 1, 1)]; + [e endEncoding]; + [cb commit]; + running.push_back(cb); } - const double sxx = x_ss - x_sum * x_sum / static_cast(n); - out[r] = sxx > 0 && w_ss > 0 ? xy / std::sqrt(sxx * w_ss) : 0.0; + for (id cb : running) { + [cb waitUntilCompleted]; + if (cb.status != MTLCommandBufferStatusCompleted) + throw JFJochException(JFJochExceptionCategory::GPUMetalError, + cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed"); + } + + // The correlation of E^2 - 1 with ic over its bin's mean, from the bins' sums: x = ic / mean_b. + const auto *sums = static_cast(d_out.contents); + std::vector out(rotations.size()); + for (size_t r = 0; r < rotations.size(); r++) { + double x_sum = 0, x_ss = 0, xy = 0; + for (int b = 0; b < refl.bins; b++) { + const float *o = sums + (r * refl.bins + b) * 3; + const double mean = o[0] / static_cast(bin_first[b + 1] - bin_first[b]); + x_sum += o[0] / mean; + x_ss += o[1] / (mean * mean); + xy += o[2] / mean; + } + const double sxx = x_ss - x_sum * x_sum / static_cast(n); + out[r] = sxx > 0 && w_ss > 0 ? xy / std::sqrt(sxx * w_ss) : 0.0; + } + return out; } - return out; } From fe58c5041ce0e7ffb966cfbc8ce87358b4a9567d Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 18 Oct 2026 17:04:17 +0200 Subject: [PATCH 2/3] RESULT-CHANGING (ulp level): the model density's walk places its points in single precision PutModelDensityOnGrid walked each atom's box as gemmi does, with every point's place relative to the atom and its test against the atom's radius in double. The Apple GPU has no double, so a Metal gather could only come to the same grid to rounding. The density now has its own walk, UseDensityBoxesByPlane: the same planes, atoms, box and order as before, with the point's place in single precision in one fixed order of operations (f = centre - (u, v, w); x = m00 fx + (m01 fy + m02 fz), y = m11 fy + m12 fz, z = m22 fz with gemmi's upper-triangular orth_n; r2 = (y y + z z) + x x, taken where not above the radius squared). A gather that spells the same arithmetic out then finds the same points with the same bits, and gemmi's float expapprox and the Gaussian sums are already single precision. The solvent mask keeps the double walk: its points are 0 or 1, and a point at the edge would flip. The density differs from gemmi's in the last bits, so everything computed from it moves at that level. 9s02 on an Apple M6 against the previous commit: p.hkl unchanged; the MR decision (AGREES, P 21 21 2) and every R-work/R-free within 0.001; validation R(all) 0.3719 -> 0.3718; the validation rigid body took 29 evaluations instead of 22 to the same place (rotation 0.003 deg); Z of rejected candidates in other groups by up to 0.7; the model and maps differ in their last digits. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_018My7tqt8a4LyKKjP3FhScT --- .../structure_refinement/ModelGrid.cpp | 139 ++++++++++++++++-- .../structure_refinement/ModelGrid.h | 4 +- 2 files changed, 133 insertions(+), 10 deletions(-) diff --git a/image_analysis/structure_refinement/ModelGrid.cpp b/image_analysis/structure_refinement/ModelGrid.cpp index c10ea95fd..82345e07d 100644 --- a/image_analysis/structure_refinement/ModelGrid.cpp +++ b/image_analysis/structure_refinement/ModelGrid.cpp @@ -130,6 +130,126 @@ void UseAtomBoxesByPlane(gemmi::Grid &grid, const std::vector &b }); } +// The density's walk: UseAtomBoxesByPlane()'s box and order, with each point's place relative to the atom and +// the test against its radius taken in single precision, in a fixed order of operations, so that the Apple +// GPU's gather (ModelDensityMetal.metal) - which has no double precision - finds the same points with the same +// bits, spelling the same arithmetic out. gemmi takes them in double; the density differs from its own in the +// last bits. + +// An atom's box: about the grid point nearest its centre, c - d to c + d along each axis. +struct DensityBox { + float nctr[3]; // the centre in grid units, (fractional position) x (grid size) + int c[3], d[3]; + float max_r2; // the radius squared +}; + +DensityBox MakeDensityBox(const gemmi::Grid &grid, const gemmi::Fractional &fpos, double radius) { + const AtomBox b = MakeAtomBox(grid, fpos, radius); + const double n[3] = {fpos.x * grid.nu, fpos.y * grid.nv, fpos.z * grid.nw}; + DensityBox out; + for (int k = 0; k < 3; k++) { + out.nctr[k] = static_cast(n[k]); + out.c[k] = gemmi::iround(n[k]); + } + out.d[0] = b.du; + out.d[1] = b.dv; + out.d[2] = b.dw; + out.max_r2 = static_cast(radius * radius); + return out; +} + +// The orthogonalization over the grid's size (gemmi's orth_n, upper triangular) in single precision: a point's +// place along v and w does not depend on u. +struct DensityGeometry { + float m00, m01, m02, m11, m12, m22; +}; + +DensityGeometry ModelDensityGeometry(const gemmi::Grid &grid) { + const gemmi::UpperTriangularMat33 &m = grid.orth_n; + return {static_cast(m.a11), static_cast(m.a12), static_cast(m.a13), + static_cast(m.a22), static_cast(m.a23), static_cast(m.a33)}; +} + +// The points of one atom's box on one plane within its radius, in the order gemmi visits them: their index +// in the grid, squared distance and delta (the atom's centre minus the point, orthogonal). +struct DensityPoints { + int n = 0; + std::vector idx; + std::vector r2, dx, dy, dz; + void Clear(size_t capacity) { + n = 0; + if (idx.size() < capacity) { + idx.resize(capacity); + r2.resize(capacity); + dx.resize(capacity); + dy.resize(capacity); + dz.resize(capacity); + } + } +}; + +// UseAtomBoxesByPlane() for the density's boxes: the same planes, atoms and walk. A point at (u, v, w), the +// box's own (unwrapped) indices, is f = centre - (u, v, w) from the atom, and in A +// x = m00 fx + (m01 fy + m02 fz), y = m11 fy + m12 fz, z = m22 fz, r2 = (y y + z z) + x x, +// and is taken where r2 is not above the radius squared. +template +void UseDensityBoxesByPlane(gemmi::Grid &grid, const std::vector &boxes, size_t nthreads, + FuncFor func_for) { + const int nu = grid.nu, nv = grid.nv, nw = grid.nw; + const DensityGeometry m = ModelDensityGeometry(grid); + std::vector> plane_atoms(nw); + for (int i = 0; i < static_cast(boxes.size()); i++) + for (int w = boxes[i].c[2] - boxes[i].d[2]; w <= boxes[i].c[2] + boxes[i].d[2]; w++) { + std::vector &atoms = plane_atoms[gemmi::modulo(w, nw)]; + if (atoms.empty() || atoms.back() != i) + atoms.push_back(i); + } + ParallelFor(nw, nthreads, [&](int plane) { + DensityPoints points; + for (int i : plane_atoms[plane]) { + const DensityBox &b = boxes[i]; + const auto func = func_for(i); + const int u_lo = b.c[0] - b.d[0], u_hi = b.c[0] + b.d[0]; + const int v_lo = b.c[1] - b.d[1], v_hi = b.c[1] + b.d[1]; + const int w_lo = b.c[2] - b.d[2], w_hi = b.c[2] + b.d[2]; + const int u_0 = gemmi::modulo(u_lo, nu), v_0 = gemmi::modulo(v_lo, nv); + points.Clear(static_cast(2 * b.d[0] + 1) * (2 * b.d[1] + 1) * (1 + 2 * b.d[2] / nw)); + int n = 0; + for (int w = w_lo + gemmi::modulo(plane - w_lo, nw); w <= w_hi; w += nw) { + const float fz = b.nctr[2] - static_cast(w); + const float z = m.m22 * fz; + for (int v = v_lo, v_ = v_0; v <= v_hi; ++v, v_ = (v_ + 1 == nv ? 0 : v_ + 1)) { + const float fy = b.nctr[1] - static_cast(v); + const float y = m.m11 * fy + m.m12 * fz; + const float yz = y * y + z * z; + if (yz > b.max_r2) + continue; // so is every point of the row: x x only adds + const float x_yz = m.m01 * fy + m.m02 * fz; + const size_t row = grid.index_q(0, v_, plane); + for (int u = u_lo, u_ = u_0;;) { + const float x = m.m00 * (b.nctr[0] - static_cast(u)) + x_yz; + const float r2 = yz + x * x; + if (!(r2 > b.max_r2)) { + points.idx[n] = row + u_; + points.r2[n] = r2; + points.dx[n] = x; + points.dy[n] = y; + points.dz[n] = z; + ++n; + } + if (u >= u_hi) + break; + ++u; + u_ = (u_ + 1 == nu ? 0 : u_ + 1); + } + } + } + points.n = n; + func(grid.data.data(), points); + } + }); +} + // gemmi's Grid::symmetrize() in parallel over the orbits found by OrbitLeaders(): each orbit is reduced // by one thread, into its leader, with gemmi's operands in gemmi's order, and the value written to its // leader and mates as gemmi writes it. Orbits share no points, so the grid is gemmi's, bit for bit. @@ -183,7 +303,7 @@ struct AtomDensity { // point still sums its Gaussians in gemmi's order, starting from 0, and the loops over the points // vectorise where gemmi's loop over one point's few Gaussians does not. The points are then added // to the grid in the order given. - void operator()(float *data, const BoxPoints &p) const { + void operator()(float *data, const DensityPoints &p) const { constexpr int BLOCK = 256; const int n = p.n; for (int k0 = 0; k0 < n; k0 += BLOCK) { @@ -194,7 +314,7 @@ struct AtomDensity { if (!is_aniso) { float r2[BLOCK]; for (int k = 0; k < m; ++k) - r2[k] = (CReal) p.d2[k0 + k]; + r2[k] = p.r2[k0 + k]; for (size_t i = 0; i < std::size(iso.a); ++i) { for (int k = 0; k < m; ++k) x[k] = std::max(iso.b[i] * r2[k], -88.f); @@ -204,7 +324,7 @@ struct AtomDensity { } else { for (size_t i = 0; i < std::size(aniso.a); ++i) { for (int k = 0; k < m; ++k) { - const gemmi::Vec3f r((float) p.dx[k0 + k], (float) p.dy[k0 + k], (float) p.dz[k0 + k]); + const gemmi::Vec3f r(p.dx[k0 + k], p.dy[k0 + k], p.dz[k0 + k]); x[k] = std::max(aniso.b[i].r_u_r(r), -88.f); } for (int k = 0; k < m; ++k) @@ -248,15 +368,16 @@ std::vector OrbitLeaders(const gemmi::Grid &grid, size_t nthreads } // DensityCalculator::put_model_density_on_grid() (gemmi dencalc.hpp, do_add_atom_density_to_grid) -// with the atoms and the symmetry spread over the threads as above - the same grid, bit for bit. Each -// atom's density coefficients and radius are worked out once, up front, rather than once per plane. +// with the atoms and the symmetry spread over the threads as above - the same grid on any number of threads, +// 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) { using CReal = AtomDensity::CReal; dc.initialize_grid(); const std::vector atoms = ModelAtoms(model); const int n = static_cast(atoms.size()); - std::vector boxes(n); + std::vector boxes(n); std::vector density(n); ParallelChunks(n, nthreads, [&](int lo, int hi) { for (int i = lo; i < hi; i++) { @@ -270,17 +391,17 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator &dc, const gem if (!d.is_aniso) { const CReal b = static_cast(atom.b_iso + dc.blur); d.iso = coef.precalculate_density_iso(b, addend); - boxes[i] = MakeAtomBox(dc.grid, fpos, dc.estimate_radius(d.iso, b)); + boxes[i] = MakeDensityBox(dc.grid, fpos, dc.estimate_radius(d.iso, b)); } else { const auto aniso_b = atom.aniso.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(dc.blur)); const CReal b_max = std::max(std::max(aniso_b.u11, aniso_b.u22), aniso_b.u33); const double radius = dc.estimate_radius(coef.precalculate_density_iso(b_max, addend), b_max); d.aniso = coef.precalculate_density_aniso_b(aniso_b, addend); - boxes[i] = MakeAtomBox(dc.grid, fpos, radius); + boxes[i] = MakeDensityBox(dc.grid, fpos, radius); } } }); - UseAtomBoxesByPlane(dc.grid, boxes, nthreads, [&](int i) { return density[i]; }); + 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 b6075a999..1ee4e9249 100644 --- a/image_analysis/structure_refinement/ModelGrid.h +++ b/image_analysis/structure_refinement/ModelGrid.h @@ -14,7 +14,9 @@ // 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 gemmi's, bit for bit, on any number of threads. +// 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. From 868ad0c8f2247ce7230f4b98d8691af9ee0037b1 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 18 Oct 2026 17:04:17 +0200 Subject: [PATCH 3/3] 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