diff --git a/image_analysis/structure_refinement/ModelGrid.cpp b/image_analysis/structure_refinement/ModelGrid.cpp index 9323a2cfe..95cd5c355 100644 --- a/image_analysis/structure_refinement/ModelGrid.cpp +++ b/image_analysis/structure_refinement/ModelGrid.cpp @@ -224,7 +224,7 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator &dc, const gem // symmetry spread over the threads as above - the same mask, bit for bit. The island removal and the // shrink that follow are gemmi's own, except that the shrink is skipped where it cannot change a point: // it looks at the grid offsets within rshrink of each point (set_margin_around), and on a grid whose -// spacing is coarser than rshrink along all three axes there are none - which is every rigid-body zone - +// spacing is coarser than rshrink along all three axes there are none - which is most rigid-body zones - // yet it still walks the whole grid. void PutMaskOnGrid(gemmi::Grid &grid, const gemmi::Model &model, const std::vector &orbit_leaders, size_t nthreads) { diff --git a/image_analysis/structure_refinement/ModelMaskGPU.cu b/image_analysis/structure_refinement/ModelMaskGPU.cu index dd06d96dd..627da9aa8 100644 --- a/image_analysis/structure_refinement/ModelMaskGPU.cu +++ b/image_analysis/structure_refinement/ModelMaskGPU.cu @@ -1,11 +1,17 @@ // SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only +// The masking, the island removal and the shrink follow SolventMasker and set_margin_around() of GEMMI's +// solmask.hpp (https://github.com/project-gemmi/gemmi/blob/master/include/gemmi/solmask.hpp) +// (c) Global Phasing Ltd., Mozilla Public License Version 2.0 + #include "ModelMaskGPU.h" #include #include #include +#include +#include #include "../../common/JFJochException.h" @@ -172,6 +178,55 @@ __global__ void uf_remove(float *mask, const int *root, const int *size, int lim mask[i] = 0.0f; } +// gemmi's SolventMasker::shrink(), set_margin_around(mask, r_shrink, 1, -1) followed by -1 -> 1, in two +// passes. gemmi walks the solvent points and turns into solvent every macromolecule point at a stencil +// offset from one; the offsets are split at the grid spacing into near ones (stencil1) and far ones +// (stencil2), and the far ones are only taken from a solvent point that has a macromolecule point among +// its near ones - a solvent point at the edge. (With no far offsets gemmi walks the macromolecule points +// instead, asking whether any near offset is solvent, which is the same rule.) So: first mark the edge +// points, then turn every macromolecule point with an edge point at any stencil offset into solvent. +// The new solvent points are not edge points, so the second pass can write the mask in place. +__global__ void shrink_edges(const float *mask, int *edge, const int3 *stencil, int n_near, int nu, int nv, int nw) { + const size_t n = static_cast(nu) * nv * nw; + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) + return; + int is_edge = 0; + if (mask[i] >= 1.0f) { + const int u = static_cast(i % nu); + const int v = static_cast((i / nu) % nv); + const int w = static_cast(i / (static_cast(nu) * nv)); + for (int k = 0; k < n_near && !is_edge; k++) { + const size_t j = wrap(u + stencil[k].x, nu) + + static_cast(nu) * (wrap(v + stencil[k].y, nv) + + static_cast(nv) * wrap(w + stencil[k].z, nw)); + is_edge = mask[j] < 1.0f; + } + } + edge[i] = is_edge; +} + +// The stencil is symmetric (an offset and its negative are the same distance), so a macromolecule point +// that an edge point reaches is one that reaches an edge point. +__global__ void shrink_margin(float *mask, const int *edge, const int3 *stencil, int n_stencil, int nu, int nv, int nw) { + const size_t n = static_cast(nu) * nv * nw; + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n || mask[i] >= 1.0f) + return; + const int u = static_cast(i % nu); + const int v = static_cast((i / nu) % nv); + const int w = static_cast(i / (static_cast(nu) * nv)); + for (int k = 0; k < n_stencil; k++) { + const size_t j = wrap(u + stencil[k].x, nu) + + static_cast(nu) * (wrap(v + stencil[k].y, nv) + + static_cast(nv) * wrap(w + stencil[k].z, nw)); + if (edge[j]) { + mask[i] = 1.0f; + return; + } + } +} + unsigned int Blocks(size_t n) { return static_cast((n + THREADS - 1) / THREADS); } @@ -179,11 +234,12 @@ unsigned int Blocks(size_t n) { } // namespace size_t ModelMaskGPU::DeviceBytes(size_t max_points) { - return 2 * max_points * sizeof(int) + MAX_OPS * sizeof(ModelMaskOp); + return 2 * max_points * sizeof(int) + MAX_OPS * sizeof(ModelMaskOp) + MAX_STENCIL * sizeof(int3); } ModelMaskGPU::ModelMaskGPU(cudaStream_t stream, size_t max_points) - : stream(stream), max_points(max_points), ops_d(MAX_OPS), label(max_points), root(max_points) { + : stream(stream), max_points(max_points), ops_d(MAX_OPS), label(max_points), root(max_points), + stencil_d(MAX_STENCIL) { if (max_points > INT32_MAX) throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Solvent mask grid too large for int32 labels"); @@ -206,32 +262,6 @@ void GridSpacing(const ModelMaskGrid &grid, double spacing[3]) { } // namespace -bool ModelMaskGPU::ShrinkIsNoOp(const ModelMaskGrid &grid) { - const int n[3] = {grid.nu, grid.nv, grid.nw}; - double spacing[3]; - GridSpacing(grid, spacing); - const double *o = grid.orth; - // gemmi's shrink, set_margin_around(r_shrink), only ever looks at the grid offsets other than 0 - // within r_shrink, in the box of floor(r_shrink / spacing). No such offset, no change. - const int d[3] = {static_cast(std::floor(R_SHRINK / spacing[0])), - static_cast(std::floor(R_SHRINK / spacing[1])), - static_cast(std::floor(R_SHRINK / spacing[2]))}; - for (int w = -d[2]; w <= d[2]; w++) - for (int v = -d[1]; v <= d[1]; v++) - for (int u = -d[0]; u <= d[0]; u++) { - const double f[3] = {static_cast(u) / n[0], static_cast(v) / n[1], - static_cast(w) / n[2]}; - double r2 = 0; - for (int i = 0; i < 3; i++) { - const double x = o[3 * i] * f[0] + o[3 * i + 1] * f[1] + o[3 * i + 2] * f[2]; - r2 += x * x; - } - if (r2 <= R_SHRINK * R_SHRINK && r2 != 0.0) - return false; - } - return true; -} - void ModelMaskGPU::SetGrid(const ModelMaskGrid &grid, const std::vector &ops) { npoints = static_cast(grid.nu) * grid.nv * grid.nw; if (npoints > max_points) @@ -250,9 +280,7 @@ void ModelMaskGPU::SetGrid(const ModelMaskGrid &grid, const std::vector(grid.orth[3 * i + j] / n[j]); GridSpacing(grid, geom.spacing); - if (!ShrinkIsNoOp(grid)) - throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, - "Solvent mask grid too fine: the GPU mask has no shrink step"); + SetShrinkStencil(grid); // gemmi SolventMasker::remove_islands(), rounded the same way. island_limit = static_cast(static_cast(ISLAND_MIN_VOLUME * npoints / grid.volume)); @@ -270,6 +298,59 @@ void ModelMaskGPU::Compute(const ModelMaskAtom *d_atoms, int n_atoms, float *d_m cuda_err(cudaGetLastError()); } RemoveIslands(d_mask); + Shrink(d_mask); +} + +// gemmi's set_margin_around() stencil for r_shrink: every grid offset other than 0 within it, near +// (closer than the coarsest grid step) first and far after, in gemmi's order. +void ModelMaskGPU::SetShrinkStencil(const ModelMaskGrid &grid) { + const int n[3] = {grid.nu, grid.nv, grid.nw}; + const int d[3] = {static_cast(std::floor(R_SHRINK / geom.spacing[0])), + static_cast(std::floor(R_SHRINK / geom.spacing[1])), + static_cast(std::floor(R_SHRINK / geom.spacing[2]))}; + if (2 * d[0] >= n[0] || 2 * d[1] >= n[1] || 2 * d[2] >= n[2]) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "Solvent mask grid smaller than the shrink radius"); + const double *o = grid.orth; + // The cell edges a, b, c are the lengths of orth's columns. + double max_step = 0; + for (int k = 0; k < 3; k++) + max_step = std::max(max_step, std::sqrt(o[k] * o[k] + o[3 + k] * o[3 + k] + o[6 + k] * o[6 + k]) / n[k]); + const double max_spacing2 = max_step * max_step + 1e-6; + std::vector near, far; + for (int w = -d[2]; w <= d[2]; w++) + for (int v = -d[1]; v <= d[1]; v++) + for (int u = -d[0]; u <= d[0]; u++) { + const double f[3] = {static_cast(u) / n[0], static_cast(v) / n[1], + static_cast(w) / n[2]}; + double r2 = 0; + for (int i = 0; i < 3; i++) { + const double x = o[3 * i] * f[0] + o[3 * i + 1] * f[1] + o[3 * i + 2] * f[2]; + r2 += x * x; + } + if (r2 <= R_SHRINK * R_SHRINK && r2 != 0.0) + (r2 < max_spacing2 ? near : far).push_back(make_int3(u, v, w)); + } + n_near = static_cast(near.size()); + n_stencil = static_cast(near.size() + far.size()); + if (n_stencil > MAX_STENCIL) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Solvent mask grid too fine for the shrink"); + near.insert(near.end(), far.begin(), far.end()); + if (n_stencil > 0) { + cuda_err(cudaMemcpyAsync(stencil_d, near.data(), near.size() * sizeof(int3), cudaMemcpyHostToDevice, stream)); + cuda_err(cudaStreamSynchronize(stream)); + } +} + +void ModelMaskGPU::Shrink(float *d_mask) { + if (n_stencil == 0) + return; + int *edge = label; // free once the islands are gone + shrink_edges<<>>(d_mask, edge, stencil_d, n_near, geom.nu, geom.nv, geom.nw); + cuda_err(cudaGetLastError()); + shrink_margin<<>>(d_mask, edge, stencil_d, n_stencil, geom.nu, geom.nv, + geom.nw); + cuda_err(cudaGetLastError()); } // gemmi SolventMasker::remove_islands(): every component of solvent with at most island_limit points diff --git a/image_analysis/structure_refinement/ModelMaskGPU.h b/image_analysis/structure_refinement/ModelMaskGPU.h index 9d1968c63..537750822 100644 --- a/image_analysis/structure_refinement/ModelMaskGPU.h +++ b/image_analysis/structure_refinement/ModelMaskGPU.h @@ -14,8 +14,8 @@ // SolventMasker(AtomicRadiiSet::Refmac).put_mask_on_grid(), on the GPU. gemmi masks the atoms and then // symmetrizes the grid with the minimum; the operators are isometries, so that is the same as masking // every symmetry image of every atom, which is what is done here - no orbits, no symmetrize. The island -// removal is gemmi's (26-connected, periodic, the same size limit) as a union-find. The shrink is not -// implemented: it is a no-op on every rigid-body grid, and SetGrid() refuses a grid where it would not be. +// removal is gemmi's (26-connected, periodic, the same size limit) as a union-find, and the shrink is +// gemmi's set_margin_around() as two passes over the grid. // // Every write of the masking is the same idempotent store and the union-find partition is unique, so // the mask is bit-identical from run to run. It can differ from gemmi's at points lying exactly at an @@ -46,23 +46,25 @@ class ModelMaskGPU { public: // Enough for Fm-3m, 48 operators times 4 centring vectors. static constexpr int MAX_OPS = 192; + // Grid offsets within gemmi's r_shrink = 0.8 A: 26 at a 0.4 A spacing, 728 at 0.2 A. + static constexpr int MAX_STENCIL = 1024; static size_t DeviceBytes(size_t max_points); ModelMaskGPU(cudaStream_t stream, size_t max_points); - // Whether gemmi's shrink (r_shrink = 0.8 A) leaves this grid as it is - the only grids this mask is - // right on. Host code; asked before any engine is reserved, so a grid it fails goes to the CPU first. - static bool ShrinkIsNoOp(const ModelMaskGrid &grid); - // Per zone. Throws if the grid has more than max_points points, if there are more than MAX_OPS - // operators, or if ShrinkIsNoOp() fails on it. + // operators, or if the shrink's stencil is over MAX_STENCIL offsets. void SetGrid(const ModelMaskGrid &grid, const std::vector &ops); // d_atoms: hydrogens and unoccupied atoms already left out. d_mask: the whole grid, 1 = solvent, 0 = macromolecule. Queued on the stream. + // The masking, then RemoveIslands(), then Shrink(). void Compute(const ModelMaskAtom *d_atoms, int n_atoms, float *d_mask); - // The island removal alone, on a mask of 0 and 1 already on the grid. Compute() ends with it. + // The island removal alone, on a mask of 0 and 1 already on the grid. Compute() runs it before Shrink(). void RemoveIslands(float *d_mask); + // The shrink alone, on a mask whose islands are already removed. A no-op on a grid too coarse for + // any grid offset to lie within r_shrink. + void Shrink(float *d_mask); // What the masking kernel needs of the grid. struct Geometry { @@ -72,14 +74,18 @@ public: }; private: + void SetShrinkStencil(const ModelMaskGrid &grid); + cudaStream_t stream; size_t max_points; Geometry geom{}; size_t npoints = 0; int n_ops = 0; int island_limit = 0; + int n_near = 0, n_stencil = 0; // the shrink's offsets: the first n_near are the near ones CudaDevicePtr ops_d; CudaDevicePtr label; // union-find parent, then the component sizes CudaDevicePtr root; // each solvent point's component, -1 elsewhere + CudaDevicePtr stencil_d; }; diff --git a/image_analysis/structure_refinement/RigidBodyGPU.cpp b/image_analysis/structure_refinement/RigidBodyGPU.cpp index d85f696fc..3de8b7fd0 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.cpp +++ b/image_analysis/structure_refinement/RigidBodyGPU.cpp @@ -186,11 +186,6 @@ std::unique_ptr RigidBodyGPUPool::Create(const gemmi::Model &m "GPU's gridding at {:.1f} A", zone_d); return nullptr; } - if (!RigidBodyGPUEngine::MaskSupports(zone)) { - logger.Info("Model validation: the rigid body runs on the CPU - the solvent mask grid at {:.1f} A " - "is finer than the GPU's mask, which has no shrink step, can take", zone_d); - return nullptr; - } cap.atoms = std::max(cap.atoms, zone.atoms.size()); cap.grid_points = std::max(cap.grid_points, static_cast(zone.nu) * zone.nv * zone.nw); cap.complex_points = std::max(cap.complex_points, static_cast(zone.nu / 2 + 1) * zone.nv * zone.nw); diff --git a/image_analysis/structure_refinement/RigidBodyGPU.cu b/image_analysis/structure_refinement/RigidBodyGPU.cu index f7b283f3e..1db77013a 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.cu +++ b/image_analysis/structure_refinement/RigidBodyGPU.cu @@ -731,16 +731,6 @@ bool RigidBodyGPUEngine::Supports(const RigidBodyGPUZone &zone) { return true; } -bool RigidBodyGPUEngine::MaskSupports(const RigidBodyGPUZone &zone) { - ModelMaskGrid mask_grid; - mask_grid.nu = zone.nu; - mask_grid.nv = zone.nv; - mask_grid.nw = zone.nw; - std::copy(zone.orth, zone.orth + 9, mask_grid.orth); - mask_grid.volume = zone.volume; - return ModelMaskGPU::ShrinkIsNoOp(mask_grid); -} - void RigidBodyGPUEngine::MemoryInfo(size_t &free, size_t &total) { cuda_err(cudaMemGetInfo(&free, &total)); } diff --git a/image_analysis/structure_refinement/RigidBodyGPUEngine.h b/image_analysis/structure_refinement/RigidBodyGPUEngine.h index 648bc00f0..9b93f1554 100644 --- a/image_analysis/structure_refinement/RigidBodyGPUEngine.h +++ b/image_analysis/structure_refinement/RigidBodyGPUEngine.h @@ -101,8 +101,6 @@ public: // Whether the gather reproduces gemmi's box walk on this zone: every atom's box narrower than the cell, // so that no point is reached by two images of one atom. static bool Supports(const RigidBodyGPUZone &zone); - // Whether the GPU's solvent mask, which has no shrink step, is gemmi's on this zone's grid. - static bool MaskSupports(const RigidBodyGPUZone &zone); static void MemoryInfo(size_t &free, size_t &total); static int CurrentDevice(); diff --git a/tests/ModelMaskGPUTest.cpp b/tests/ModelMaskGPUTest.cpp index d6ca1e728..d10f5c92d 100644 --- a/tests/ModelMaskGPUTest.cpp +++ b/tests/ModelMaskGPUTest.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -202,21 +203,65 @@ TEST_CASE("ModelMaskGPU_MatchesGemmi", "[ModelValidation][gpu]") { CHECK(any_odd_size); } -// The shrink is not implemented: a grid fine enough for it to change anything is refused. -TEST_CASE("ModelMaskGPU_RefusesAGridTheShrinkWouldChange", "[ModelValidation][gpu]") { +// Grids fine enough for gemmi's shrink to act. The shrink alone, on gemmi's own mask after its island +// step, is exact; the whole mask may differ from gemmi's only where the masking does (points at an atom's +// radius) and at the neighbours the shrink then treats differently - measured, it does not differ at +// all here - and is the same bits twice. +TEST_CASE("ModelMaskGPU_ShrinkMatchesGemmi", "[ModelValidation][gpu]") { if (get_gpu_count() == 0) SKIP("No GPU"); CudaStream stream; - const gemmi::Structure st = ClusterModel(kCrysts[0]); - const std::vector ops = MaskOps(*st.find_spacegroup()); - const gemmi::Grid coarse = MaskGrid(st, 3.5); - const gemmi::Grid fine = MaskGrid(st, 2.0); - ModelMaskGPU mask(stream, fine.point_count()); - CHECK_NOTHROW(mask.SetGrid(MaskGridDescription(coarse), ops)); - CHECK_THROWS(mask.SetGrid(MaskGridDescription(fine), ops)); - // The same answer ahead of time, which is what sends such a grid to the CPU before any engine is made. - CHECK(ModelMaskGPU::ShrinkIsNoOp(MaskGridDescription(coarse))); - CHECK_FALSE(ModelMaskGPU::ShrinkIsNoOp(MaskGridDescription(fine))); + // The five groups at 1.5 A, and a strongly oblique monoclinic cell at 3.5 A, where the planes of the + // grid along a and c are closer than r_shrink although its steps along the axes are not. + std::vector> cases; + for (const char *cryst : kCrysts) + cases.emplace_back(cryst, 1.5); + cases.emplace_back("CRYST1 40.000 50.000 60.000 90.00 141.00 90.00 P 1 21 1 2\n", 3.5); + for (const auto &[cryst, d_min] : cases) { + const gemmi::Structure st = ClusterModel(cryst); + const gemmi::SpaceGroup *sg = st.find_spacegroup(); + REQUIRE(sg != nullptr); + const std::vector atoms = MaskAtoms(st); + const std::vector ops = MaskOps(*sg); + CudaDevicePtr atoms_d(atoms.size()); + REQUIRE(cudaMemcpyAsync(atoms_d, atoms.data(), atoms.size() * sizeof(ModelMaskAtom), cudaMemcpyHostToDevice, stream) == cudaSuccess); + + const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac); + gemmi::Grid islands = MaskGrid(st, d_min); + masker.clear(islands); + masker.mask_points(islands, st.models[0]); + masker.symmetrize(islands); + masker.remove_islands(islands); + gemmi::Grid ref = islands; + masker.shrink(ref); + const size_t n = ref.data.size(); + REQUIRE(std::memcmp(ref.data.data(), islands.data.data(), n * sizeof(float)) != 0); // the shrink acts + + ModelMaskGPU mask(stream, n); + mask.SetGrid(MaskGridDescription(ref), ops); + CudaDevicePtr mask_d(n); + std::vector out(n), first; + REQUIRE(cudaMemcpyAsync(mask_d, islands.data.data(), n * sizeof(float), cudaMemcpyHostToDevice, stream) == cudaSuccess); + mask.Shrink(mask_d); + REQUIRE(cudaMemcpyAsync(out.data(), mask_d, n * sizeof(float), cudaMemcpyDeviceToHost, stream) == cudaSuccess); + REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess); + CHECK(std::memcmp(out.data(), ref.data.data(), n * sizeof(float)) == 0); + + for (int repeat = 0; repeat < 2; repeat++) { + mask.Compute(atoms_d, static_cast(atoms.size()), mask_d); + REQUIRE(cudaMemcpyAsync(out.data(), mask_d, n * sizeof(float), cudaMemcpyDeviceToHost, stream) == cudaSuccess); + REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess); + if (repeat == 0) { + first = out; + size_t differ = 0; + for (size_t i = 0; i < n; i++) + differ += out[i] != ref.data[i]; + INFO(sg->hm << " grid " << ref.nu << "x" << ref.nv << "x" << ref.nw << ": " << differ << " points differ"); + CHECK(differ <= 27 * (15 * n / 10000000 + 1)); + } else + CHECK(std::memcmp(out.data(), first.data(), n * sizeof(float)) == 0); + } + } } #endif