diff --git a/image_analysis/structure_refinement/CMakeLists.txt b/image_analysis/structure_refinement/CMakeLists.txt index 34c5ac3e1..4f6c6b3d1 100644 --- a/image_analysis/structure_refinement/CMakeLists.txt +++ b/image_analysis/structure_refinement/CMakeLists.txt @@ -7,7 +7,11 @@ ADD_LIBRARY(JFJochStructureRefinement STATIC $<$:RigidBodyGPU.cpp> $<$:RigidBodyGPU.cu> RigidBodyGPU.h RigidBodyGPUEngine.h + $<$:ModelDensityGPU.cu> ModelDensityGPU.h $<$:ModelMaskGPU.cu> ModelMaskGPU.h + $<$:ModelStructureFactorsGPU.cpp> + $<$:ModelStructureFactorsGPU.cu> + ModelStructureFactorsGPU.h ModelStructureFactorsGPUEngine.h $<$:ModelScaleGPU.cu> ModelScaleGPU.h SigmaA.cpp SigmaA.h) diff --git a/image_analysis/structure_refinement/ModelDensityGPU.cu b/image_analysis/structure_refinement/ModelDensityGPU.cu new file mode 100644 index 000000000..d85f9e642 --- /dev/null +++ b/image_analysis/structure_refinement/ModelDensityGPU.cu @@ -0,0 +1,350 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "ModelDensityGPU.h" + +#include +#include + +#include + +#include "../../common/JFJochException.h" + +namespace { + +void cuda_err(cudaError_t val) { + if (val != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); +} + +constexpr int BRICK = 8; // the gather's bricks are BRICK^3 grid points, one block each +constexpr int GATHER_TILE = 128; // atoms staged in shared memory at a time +constexpr int MAX_BRICKS_PER_AXIS = 16; // an atom's box touches at most this many bricks along an axis +constexpr int THREADS = 256; + +__device__ __forceinline__ int imod(int a, int n) { + const int r = a % n; + return r < 0 ? r + n : r; +} + +// gemmi's unsafe_expapprox() (formfact.hpp), the exponential the density is computed with. +__device__ __forceinline__ float expapprox(float x) { + const float val = 12102203.1615614f * x + 1065353216.f; + const int vali = static_cast(val); + const float a = __int_as_float(vali & 0x7F800000); + const float b = __int_as_float((vali & 0x7FFFFF) | 0x3F800000); + return a * (0.509871020f + b * (0.312146713f + b * (0.166617139f + b * (-2.190619930e-3f + b * 1.3555747234e-2f)))); +} + +// The distinct bricks the points c - d ... c + d of one axis fall in, wrapped into the cell. A grid size +// that is not a multiple of BRICK leaves the last brick partial, which is why this is done on wrapped +// points and not in brick coordinates. -1 where there are more than MAX_BRICKS_PER_AXIS of them. +__device__ int axis_bricks(int c, int d, int n, int *out) { + int k = 0; + for (int p = c - d; p <= c + d; p++) { + const int b = imod(p, n) / BRICK; + bool seen = false; + for (int j = 0; j < k; j++) + seen = seen || out[j] == b; + if (seen) + continue; + if (k == MAX_BRICKS_PER_AXIS) + return -1; + out[k++] = b; + } + return k; +} + +// The bricks atom i's box touches, per axis. The box is gemmi's: `box` points either side of the +// point nearest the atom. +__device__ void atom_bricks(const float4 &pos, const int *box, int nu, int nv, int nw, int *bu, int &nbu, int *bv, + int &nbv, int *bw, int &nbw) { + nbu = axis_bricks(__float2int_rn(pos.x), box[0], nu, bu); + nbv = axis_bricks(__float2int_rn(pos.y), box[1], nv, bv); + nbw = axis_bricks(__float2int_rn(pos.z), box[2], nw, bw); +} + +__global__ void count_pairs_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw, int *count, + int *overflow) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i > n) + return; + if (i == n) { + count[n] = 0; // so the exclusive scan's last entry is the total + return; + } + int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c; + atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c); + if (a < 0 || b < 0 || c < 0) { + *overflow = 1; + a = b = c = 0; + } + count[i] = a * b * c; +} + +// The (brick, atom) pairs, atom-major: sorted stably by brick, each brick's atoms stay in model order. +__global__ void fill_pairs_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw, + const int *offset, int *key, int *value) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) + return; + int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c; + atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c); + const int nbu = (nu + BRICK - 1) / BRICK, nbv = (nv + BRICK - 1) / BRICK; + int o = offset[i]; + for (int k = 0; k < c; k++) + for (int j = 0; j < b; j++) + for (int l = 0; l < a; l++) { + key[o] = bu[l] + nbu * (bv[j] + nbv * bw[k]); + value[o] = i; + o++; + } +} + +__global__ void brick_ranges_kernel(const int *key, int npairs, int *start, int *end) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= npairs) + return; + const int k = key[i]; + if (i == 0 || key[i - 1] != k) + start[k] = i; + if (i == npairs - 1 || key[i + 1] != k) + end[k] = i + 1; +} + +// The density of the model on the grid, as gemmi's do_add_atom_density_to_grid() puts it there: every +// point adds, in model order, each atom whose sphere it is inside. One block per brick; the brick's atoms +// are staged in shared memory with the position of their image nearest to the brick, relative to the +// brick's first point. Where the cell is wider than an atom's box plus a brick, that is the only image +// that reaches any point of the brick; on a narrower cell each point finds its own nearest image. +__global__ void gather_kernel(const ModelDensityAtom *atoms, const float4 *pos, const int *pair_atom, + const int *brick_start, const int *brick_end, ModelDensityGPU::Geometry g, + float *grid) { + const int nbu = (g.nu + BRICK - 1) / BRICK, nbv = (g.nv + BRICK - 1) / BRICK; + const int brick = blockIdx.x; + const int s = brick_start[brick], e = brick_end[brick]; // s == e: an empty brick, written as zeros + const int u0 = (brick % nbu) * BRICK, v0 = (brick / nbu % nbv) * BRICK, w0 = (brick / (nbu * nbv)) * BRICK; + const int u = u0 + threadIdx.x, v = v0 + threadIdx.y, w = w0 + threadIdx.z; + const bool inside = u < g.nu && v < g.nv && w < g.nw; + const int tid = threadIdx.x + BRICK * (threadIdx.y + BRICK * threadIdx.z); + // This point relative to the brick's first one, in Cartesian coordinates. + const float tx = g.orth_n[0] * threadIdx.x + g.orth_n[1] * threadIdx.y + g.orth_n[2] * threadIdx.z; + const float ty = g.orth_n[3] * threadIdx.x + g.orth_n[4] * threadIdx.y + g.orth_n[5] * threadIdx.z; + const float tz = g.orth_n[6] * threadIdx.x + g.orth_n[7] * threadIdx.y + g.orth_n[8] * threadIdx.z; + + constexpr int WORDS = sizeof(ModelDensityAtom) / sizeof(int); + __shared__ int sh_index[GATHER_TILE]; + __shared__ ModelDensityAtom sh_atom[GATHER_TILE]; + __shared__ float3 sh_centre[GATHER_TILE]; // Cartesian + __shared__ float3 sh_grid[GATHER_TILE]; // the same, in grid units + float acc = 0.0f; + for (int c = s; c < e; c += GATHER_TILE) { + const int m = min(GATHER_TILE, e - c); + __syncthreads(); + if (tid < m) { + const int a = pair_atom[c + tid]; + sh_index[tid] = a; + // The atom's image nearest to the brick's centre, in grid units from its first point. + const float4 p = pos[a]; + float fx = p.x - u0, fy = p.y - v0, fz = p.z - w0; + fx -= g.nu * rintf((fx - 0.5f * (BRICK - 1)) / g.nu); + fy -= g.nv * rintf((fy - 0.5f * (BRICK - 1)) / g.nv); + fz -= g.nw * rintf((fz - 0.5f * (BRICK - 1)) / g.nw); + sh_grid[tid] = make_float3(fx, fy, fz); + sh_centre[tid] = make_float3(g.orth_n[0] * fx + g.orth_n[1] * fy + g.orth_n[2] * fz, + g.orth_n[3] * fx + g.orth_n[4] * fy + g.orth_n[5] * fz, + g.orth_n[6] * fx + g.orth_n[7] * fy + g.orth_n[8] * fz); + } + __syncthreads(); + for (int i = tid; i < m * WORDS; i += BRICK * BRICK * BRICK) + reinterpret_cast(sh_atom)[i] = reinterpret_cast(atoms + sh_index[i / WORDS])[i % WORDS]; + __syncthreads(); + if (!inside) + continue; + for (int k = 0; k < m; k++) { + const ModelDensityAtom &at = sh_atom[k]; + float x = sh_centre[k].x - tx, y = sh_centre[k].y - ty, z = sh_centre[k].z - tz; + if (g.narrow) { + float fx = sh_grid[k].x - threadIdx.x, fy = sh_grid[k].y - threadIdx.y, fz = sh_grid[k].z - threadIdx.z; + fx -= g.nu * rintf(fx / g.nu); + fy -= g.nv * rintf(fy / g.nv); + fz -= g.nw * rintf(fz / g.nw); + x = g.orth_n[0] * fx + g.orth_n[1] * fy + g.orth_n[2] * fz; + y = g.orth_n[3] * fx + g.orth_n[4] * fy + g.orth_n[5] * fz; + z = g.orth_n[6] * fx + g.orth_n[7] * fy + g.orth_n[8] * fz; + } + const float r2 = x * x + y * y + z * z; + if (r2 > at.radius * at.radius) + continue; + float density = 0.0f; + if (!at.aniso) { + for (int q = 0; q < 5; q++) + density += at.a[q] * expapprox(fmaxf(at.b[q][0] * r2, -88.f)); + } else { + for (int q = 0; q < 5; q++) { + const float *b = at.b[q]; + 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(fmaxf(rur, -88.f)); + } + } + acc += at.occ * density; + } + } + if (inside) + grid[u + static_cast(g.nu) * (v + static_cast(g.nv) * w)] = acc; +} + +int blocks(size_t n) { + return static_cast((n + THREADS - 1) / THREADS); +} + +// cub's temporary storage for the pair scan and sort at the capacity. +size_t CubBytes(size_t atoms, size_t pairs) { + size_t scan = 0, sort = 0; + cuda_err(cub::DeviceScan::ExclusiveSum(nullptr, scan, static_cast(nullptr), static_cast(nullptr), + static_cast(atoms + 1))); + cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, sort, static_cast(nullptr), static_cast(nullptr), + static_cast(nullptr), static_cast(nullptr), + static_cast(pairs))); + return std::max(scan, sort); +} + +} // namespace + +// The most bricks the 2 d + 1 wrapped points of a box can fall in along an axis of n points: one more +// than the points span for where they start in a brick, and one more again where they wrap past a last +// brick that is partial (n = 17: the points 15, 16, 0 fall in bricks 1, 2 and 0). +size_t ModelDensityGPU::AxisBrickBound(int d, int n) { + const int nb = (n + BRICK - 1) / BRICK; + return std::min(nb, (2 * d + 1 + BRICK - 1) / BRICK + 2); +} + +size_t ModelDensityGPU::Bricks(int nu, int nv, int nw) { + return static_cast((nu + BRICK - 1) / BRICK) * ((nv + BRICK - 1) / BRICK) * ((nw + BRICK - 1) / BRICK); +} + +std::vector ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const std::vector &atoms) { + const int n3[3] = {grid.nu, grid.nv, grid.nw}; + double spacing[3]; + for (int k = 0; k < 3; k++) + spacing[k] = 1.0 / (n3[k] * std::sqrt(grid.frac[3 * k] * grid.frac[3 * k] + grid.frac[3 * k + 1] * grid.frac[3 * k + 1] + + grid.frac[3 * k + 2] * grid.frac[3 * k + 2])); + std::vector box(3 * atoms.size()); + for (size_t i = 0; i < atoms.size(); i++) + for (int k = 0; k < 3; k++) + box[3 * i + k] = static_cast(std::ceil(atoms[i].radius / spacing[k])); + return box; +} + +size_t ModelDensityGPU::PairBound(const ModelDensityGrid &grid, const std::vector &atoms) { + const std::vector box = Boxes(grid, atoms); + const int n3[3] = {grid.nu, grid.nv, grid.nw}; + size_t pairs = 0; + for (size_t i = 0; i < atoms.size(); i++) { + size_t product = 1; + for (int k = 0; k < 3; k++) + product *= AxisBrickBound(box[3 * i + k], n3[k]); + pairs += product; + } + return pairs; +} + +bool ModelDensityGPU::Supports(const ModelDensityGrid &grid, const std::vector &atoms) { + const std::vector box = Boxes(grid, atoms); + const int n3[3] = {grid.nu, grid.nv, grid.nw}; + for (size_t i = 0; i < box.size(); i++) + if (2 * box[i] + 1 > n3[i % 3] || AxisBrickBound(box[i], n3[i % 3]) > MAX_BRICKS_PER_AXIS) + return false; + return true; +} + +size_t ModelDensityGPU::DeviceBytes(size_t max_atoms, size_t max_pairs, size_t max_bricks) { + const size_t na = max_atoms + 1; + return na * (sizeof(ModelDensityAtom) + 3 * sizeof(int) + 2 * sizeof(int)) + 4 * max_pairs * sizeof(int) + + 2 * max_bricks * sizeof(int) + CubBytes(na, max_pairs + 1); +} + +ModelDensityGPU::ModelDensityGPU(cudaStream_t stream, size_t max_atoms_, size_t max_pairs_, size_t max_bricks_) + : stream(stream), max_atoms(max_atoms_), max_pairs(max_pairs_), max_bricks(max_bricks_) { + const auto sync = CudaAlloc::Synchronous; + const size_t na = std::max(max_atoms, 1), np = std::max(max_pairs, 1), + nb = std::max(max_bricks, 1); + atoms = CudaDevicePtr(na, sync); + box = CudaDevicePtr(3 * na, sync); + overflow = CudaDevicePtr(1, sync); + count = CudaDevicePtr(na + 1, sync); + offset = CudaDevicePtr(na + 1, sync); + key = CudaDevicePtr(np, sync); + value = CudaDevicePtr(np, sync); + key_sorted = CudaDevicePtr(np, sync); + value_sorted = CudaDevicePtr(np, sync); + brick_start = CudaDevicePtr(nb, sync); + brick_end = CudaDevicePtr(nb, sync); + cub_bytes = CubBytes(na, np); + cub_temp = CudaDevicePtr(std::max(cub_bytes, 1), sync); +} + +void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const std::vector &atom_list) { + if (atom_list.size() > max_atoms || Bricks(grid.nu, grid.nv, grid.nw) > max_bricks) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: grid or atoms over the reserve"); + // An atom whose box is wider than the cell would reach a point through more than one image, which + // gemmi's box walk adds and the gather's nearest image does not; such a cell is left to the CPU. + if (!Supports(grid, atom_list)) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: an atom wider than the cell"); + geom.nu = grid.nu; + geom.nv = grid.nv; + geom.nw = grid.nw; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) + geom.orth_n[3 * i + j] = static_cast(grid.orth[3 * i + j] / (j == 0 ? grid.nu : j == 1 ? grid.nv : grid.nw)); + const std::vector b = Boxes(grid, atom_list); + geom.narrow = 0; + for (size_t i = 0; i < b.size(); i++) + if (2 * b[i] + BRICK > (i % 3 == 0 ? grid.nu : i % 3 == 1 ? grid.nv : grid.nw)) + geom.narrow = 1; + n_atoms = static_cast(atom_list.size()); + if (n_atoms > 0) { + cuda_err(cudaMemcpyAsync(atoms, atom_list.data(), atom_list.size() * sizeof(ModelDensityAtom), + cudaMemcpyHostToDevice, stream)); + cuda_err(cudaMemcpyAsync(box, b.data(), b.size() * sizeof(int), cudaMemcpyHostToDevice, stream)); + } + // The uploads read host vectors that end with this call. + cuda_err(cudaStreamSynchronize(stream)); +} + +void ModelDensityGPU::Compute(const float4 *d_pos, float *d_grid) { + cuda_err(cudaMemsetAsync(overflow, 0, sizeof(int), stream)); + count_pairs_kernel<<>>(d_pos, box, n_atoms, geom.nu, geom.nv, geom.nw, + count, overflow); + cuda_err(cudaGetLastError()); + size_t temp = cub_bytes; + cuda_err(cub::DeviceScan::ExclusiveSum(cub_temp.get(), temp, count.get(), offset.get(), n_atoms + 1, stream)); + int npairs = 0; + int over = 0; + cuda_err(cudaMemcpyAsync(&npairs, offset.get() + n_atoms, sizeof(int), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(&over, overflow.get(), sizeof(int), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaStreamSynchronize(stream)); + if (over) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: an atom's box over the bricks per axis"); + if (static_cast(npairs) > max_pairs) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: gather pairs over the reserve"); + fill_pairs_kernel<<>>(d_pos, box, n_atoms, geom.nu, geom.nv, geom.nw, + offset, key, value); + cuda_err(cudaGetLastError()); + const int nb = static_cast(Bricks(geom.nu, geom.nv, geom.nw)); + int bits = 1; + while ((1 << bits) < nb) + bits++; + temp = cub_bytes; + cuda_err(cub::DeviceRadixSort::SortPairs(cub_temp.get(), temp, key.get(), key_sorted.get(), value.get(), + value_sorted.get(), npairs, 0, bits, stream)); + cuda_err(cudaMemsetAsync(brick_start, 0, nb * sizeof(int), stream)); + cuda_err(cudaMemsetAsync(brick_end, 0, nb * sizeof(int), stream)); + if (npairs > 0) { + brick_ranges_kernel<<>>(key_sorted, npairs, brick_start, brick_end); + cuda_err(cudaGetLastError()); + } + gather_kernel<<>>(atoms, d_pos, value_sorted, brick_start, brick_end, + geom, d_grid); + cuda_err(cudaGetLastError()); +} diff --git a/image_analysis/structure_refinement/ModelDensityGPU.h b/image_analysis/structure_refinement/ModelDensityGPU.h new file mode 100644 index 000000000..d58734d9c --- /dev/null +++ b/image_analysis/structure_refinement/ModelDensityGPU.h @@ -0,0 +1,83 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// The density of a model's atoms on a grid, on the GPU (CUDA builds only): gemmi's +// DensityCalculator::put_model_density_on_grid() without the symmetrize - ONE copy of the content. The +// rigid body (RigidBodyGPU.cu) and the model's structure factors (ModelStructureFactorsGPU.cu) both grid +// with it, and compose the crystal's symmetry in reciprocal space (SymmetryComposition, RigidBodyRefine.h). +// +// A gather, not a scatter: the grid is cut into bricks of BRICK^3 points, each brick is one block, and +// every point adds, in model order, each atom whose sphere it is inside. Nothing is summed with a +// floating-point atomic, so the grid is the same, bit for bit, on every run. + +#include +#include + +#include + +#include "../indexing/CUDAMemHelpers.h" + +// One atom's density on a grid, as PutModelDensityOnGrid() (ModelGrid.cpp) sets it up: gemmi's +// precalculated five-Gaussian sum and the radius it cuts the sum at. +struct ModelDensityAtom { + float a[5]; + float b[5][6]; // isotropic: b[k][0] multiplies r^2; anisotropic: the matrix, u11 u22 u33 u12 u13 u23 + float occ; + float radius; + int aniso; +}; + +// The grid the density is put on, u fastest: index = u + nu * (v + nv * w). +struct ModelDensityGrid { + int nu = 0, nv = 0, nw = 0; + double orth[9] = {}, frac[9] = {}; // the cell's, row-major +}; + +class ModelDensityGPU { +public: + // The bytes an instance of this capacity reserves on the device. + static size_t DeviceBytes(size_t max_atoms, size_t max_pairs, size_t max_bricks); + static size_t Bricks(int nu, int nv, int nw); + // The most gather bricks the 2 d + 1 points of an atom's box can fall in along an axis of n points. + static size_t AxisBrickBound(int d, int n); + // gemmi's box around each atom (MakeAtomBox, ModelGrid.cpp): the points within ceil(radius / spacing) + // of the nearest one along each axis, three per atom. + static std::vector Boxes(const ModelDensityGrid &grid, const std::vector &atoms); + // (brick, atom) pairs of the gather over at most. + static size_t PairBound(const ModelDensityGrid &grid, const std::vector &atoms); + // Whether the gather reproduces gemmi's box walk on this grid: every atom's box narrower than the cell, + // so that no point is reached by two images of one atom. + static bool Supports(const ModelDensityGrid &grid, const std::vector &atoms); + + ModelDensityGPU(cudaStream_t stream, size_t max_atoms, size_t max_pairs, size_t max_bricks); + + // The grid and the atoms' densities. Throws if they are over the capacity or if !Supports(). + void SetAtoms(const ModelDensityGrid &grid, const std::vector &atoms); + + // The density into d_grid (nu * nv * nw floats, every point written). d_pos: each atom's position in + // grid units - fractional times the grid size, wrapped into [0, n) - in the order SetAtoms() got them; + // w is not read. Waits on the stream once, for the number of (brick, atom) pairs. + void Compute(const float4 *d_pos, float *d_grid); + + // What the gather kernel needs of the grid. + struct Geometry { + int nu, nv, nw; + float orth_n[9]; // orth * diag(1/nu, 1/nv, 1/nw), row-major: grid offset -> Cartesian + int narrow; // the cell is too small for one image of an atom to serve a whole brick + }; + +private: + cudaStream_t stream; + size_t max_atoms, max_pairs, max_bricks; + Geometry geom{}; + int n_atoms = 0; + + CudaDevicePtr atoms; + CudaDevicePtr box; + CudaDevicePtr overflow; // an atom's box over MAX_BRICKS_PER_AXIS + CudaDevicePtr count, offset, key, value, key_sorted, value_sorted, brick_start, brick_end; + CudaDevicePtr cub_temp; + size_t cub_bytes = 0; +}; diff --git a/image_analysis/structure_refinement/ModelStructureFactorsGPU.cpp b/image_analysis/structure_refinement/ModelStructureFactorsGPU.cpp new file mode 100644 index 000000000..02b8bfe3e --- /dev/null +++ b/image_analysis/structure_refinement/ModelStructureFactorsGPU.cpp @@ -0,0 +1,254 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +// The reflection list follows ReciprocalGrid::prepare_asu_data() of GEMMI's recgrid.hpp +// (https://github.com/project-gemmi/gemmi/blob/master/include/gemmi/recgrid.hpp) +// (c) Global Phasing Ltd., Mozilla Public License Version 2.0 + +#include "ModelStructureFactorsGPU.h" + +#include +#include + +#include "gemmi/fourier.hpp" // get_size_for_hkl, get_f_phi_on_grid +#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent + +#include "RigidBodyRefine.h" // RigidBodyZoneGrid + +namespace { + +using Table = gemmi::IT92; + +// The density calculator model validation grids F_calc with: d_min, rate 1.5 and the blur gemmi's +// set_refmac_compatible_blur() gives this model. +gemmi::DensityCalculator Calculator(const gemmi::Model &model, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, double d_min) { + gemmi::DensityCalculator dc; + dc.d_min = d_min; + dc.rate = 1.5; + dc.grid.unit_cell = cell; + dc.grid.spacegroup = &sg; + dc.set_refmac_compatible_blur(model); + return dc; +} + +// The reflections prepare_asu_data(d_min) lists for the half-l transform of an (nu, nv, nw) grid, in its +// order (h, then k, then l ascending): in the group's reciprocal ASU, strictly inside d_min, inside the +// grid's Nyquist box, not systematically absent and not 000. +std::vector AsuReflections(const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, + int nu, int nv, int nw) { + const gemmi::Miller lim = cell.get_hkl_limits(d_min); + const int max_h = std::min((nu - 1) / 2, lim[0]); + const int max_k = std::min((nv - 1) / 2, lim[1]); + const int max_l = std::min(nw / 2, lim[2]); + const double max_1_d2 = 1.0 / (d_min * d_min); + const gemmi::ReciprocalAsu asu(&sg); + const gemmi::GroupOps gops = sg.operations(); + std::vector rows; + gemmi::Miller hkl; + for (hkl[0] = -max_h; hkl[0] <= max_h; ++hkl[0]) + for (hkl[1] = -max_k; hkl[1] <= max_k; ++hkl[1]) + for (hkl[2] = -max_l; hkl[2] <= max_l; ++hkl[2]) + if (asu.is_in(hkl) && cell.calculate_1_d2(hkl) < max_1_d2 && !gops.is_systematically_absent(hkl) && + !(hkl[0] == 0 && hkl[1] == 0 && hkl[2] == 0)) + rows.push_back(hkl); + return rows; +} + +} // namespace + +void ModelDensityAtoms(const gemmi::Model &model, const gemmi::DensityCalculator &dc, + std::vector &atoms, std::vector &mask_atom, + std::vector &mask_radius) { + atoms.clear(); + mask_atom.clear(); + mask_radius.clear(); + const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac); + int index = 0; + for (const gemmi::Chain &ch : model.chains) + for (const gemmi::Residue &r : ch.residues) + for (const gemmi::Atom &atom : r.atoms) { + using CReal = Table::Coef::coef_type; + const auto &coef = Table::get(atom.element, atom.charge, atom.serial); + const float addend = dc.addends.get(atom.element); + ModelDensityAtom a{}; + a.occ = atom.occ; + a.aniso = atom.aniso.nonzero(); + if (!a.aniso) { + const CReal b = static_cast(atom.b_iso + dc.blur); + const auto precal = coef.precalculate_density_iso(b, addend); + a.radius = dc.estimate_radius(precal, b); + for (int k = 0; k < 5; k++) { + a.a[k] = precal.a[k]; + a.b[k][0] = precal.b[k]; + } + } 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); + a.radius = static_cast(dc.estimate_radius(coef.precalculate_density_iso(b_max, addend), b_max)); + const auto precal = coef.precalculate_density_aniso_b(aniso_b, addend); + for (int k = 0; k < 5; k++) { + a.a[k] = precal.a[k]; + const gemmi::SMat33 &m = precal.b[k]; + const float e[6] = {m.u11, m.u22, m.u33, m.u12, m.u13, m.u23}; + std::copy(e, e + 6, a.b[k]); + } + } + atoms.push_back(a); + if (!((masker.ignore_hydrogen && atom.is_hydrogen()) || + (masker.ignore_zero_occupancy_atoms && atom.occ <= 0))) { + mask_atom.push_back(index); + mask_radius.push_back(static_cast( + masker.constant_r + masker.rprobe + gemmi::refmac_radius_for_bulk_solvent(atom.element.elem))); + } + ++index; + } +} + +ModelStructureFactorsGPU::ModelStructureFactorsGPU(int device, const gemmi::Model &model, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, double d_min) + : device_(device), cell_(cell), sg_(&sg), d_min_(d_min), dc_(Calculator(model, cell, sg, d_min)) { + const gemmi::Grid grid = RigidBodyZoneGrid(cell, sg, d_min); // DensityCalculator's for d_min + setup_.grid.nu = grid.nu; + setup_.grid.nv = grid.nv; + setup_.grid.nw = grid.nw; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) { + setup_.grid.orth[3 * i + j] = cell.orth.mat[i][j]; + setup_.grid.frac[3 * i + j] = cell.frac.mat[i][j]; + } + setup_.volume = cell.volume; + + const gemmi::GroupOps gops = sg.operations(); + setup_.den = gemmi::Op::DEN; + for (const gemmi::Op &op : gops.sym_ops) { + std::array o{}; + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) + o[3 * i + j] = op.rot[i][j]; + o[9 + i] = op.tran[i]; + } + setup_.sym_ops.push_back(o); + } + for (const gemmi::Op::Tran &cen : gops.cen_ops) + for (const gemmi::Op &op : gops.sym_ops) { + ModelMaskOp m{}; + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) + m.rot[3 * i + j] = static_cast(op.rot[i][j]) / gemmi::Op::DEN; + m.tran[i] = static_cast(op.tran[i] + cen[i]) / gemmi::Op::DEN; + } + setup_.mask_ops.push_back(m); + } + + rows_ = AsuReflections(cell, sg, d_min, grid.nu, grid.nv, grid.nw); + const double n_cen = static_cast(gops.cen_ops.size()); + for (const gemmi::Miller &h : rows_) { + setup_.rows.push_back({h[0], h[1], h[2]}); + // prepare_asu_data()'s unblur, exp(B_blur |s|^2 / 4) + setup_.row_scale.push_back(n_cen * std::exp(dc_.blur * 0.25 * cell.calculate_1_d2(h))); + } + + std::vector atoms; + std::vector mask_atom; + std::vector mask_radius; + ModelDensityAtoms(model, dc_, atoms, mask_atom, mask_radius); + setup_.max_atoms = atoms.size(); + setup_.max_pairs = ModelDensityGPU::PairBound(setup_.grid, atoms); + supported_ = !rows_.empty() && ModelDensityGPU::Supports(setup_.grid, atoms) && + setup_.mask_ops.size() <= ModelMaskGPU::MAX_OPS; + if (!supported_) + return; + + // The largest map: Map()'s coefficients are on a subset of these reflections, and get_size_for_hkl() + // only grows with the indices and the resolution it is given. + gemmi::AsuData> all; + all.unit_cell_ = cell; + all.spacegroup_ = &sg; + for (const gemmi::Miller &h : rows_) + all.v.push_back({h, {1.0f, 0.0f}}); + map_size_ = gemmi::get_size_for_hkl(all, {{0, 0, 0}}, 3.0); + // MapFromFPhi() transforms a half-l grid of nw / 2 + 1 planes back to 2 (nw / 2) of them. + const int map_nw = 2 * (map_size_[2] / 2); + setup_.map_points = static_cast(map_size_[0]) * map_size_[1] * map_nw; + setup_.map_complex_points = static_cast(map_size_[0]) * map_size_[1] * (map_nw / 2 + 1); + fft_work_bytes_ = ModelStructureFactorsGPUEngine::FFTWorkBytes(device_, setup_.grid, map_size_[0], map_size_[1], map_nw); + device_bytes_ = ModelStructureFactorsGPUEngine::DeviceBytes(setup_, fft_work_bytes_); +} + +ModelStructureFactorsGPU::~ModelStructureFactorsGPU() = default; + +void ModelStructureFactorsGPU::Reserve() { + std::lock_guard lock(m_); + if (!engine_) + engine_ = std::make_unique(device_, setup_, fft_work_bytes_); +} + +void ModelStructureFactorsGPU::Compute(const gemmi::Model &model, gemmi::AsuData> &fcalc, + gemmi::AsuData> &fmask) { + std::vector atoms; + std::vector mask_atom; + std::vector mask_radius; + ModelDensityAtoms(model, dc_, atoms, mask_atom, mask_radius); + + // Each atom in grid units and, for the mask, fractional - both wrapped into the cell. + const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw}; + std::vector> pos; + std::vector> frac; + for (const gemmi::Chain &ch : model.chains) + for (const gemmi::Residue &r : ch.residues) + for (const gemmi::Atom &a : r.atoms) { + const gemmi::Fractional f0 = cell_.fractionalize(a.pos); + std::array f{f0.x, f0.y, f0.z}; + std::array g{}; + for (int k = 0; k < 3; k++) { + f[k] -= std::floor(f[k]); + if (f[k] >= 1.0) // a tiny negative coordinate wraps to 1 - epsilon, which rounds to 1 + f[k] = 0.0; + g[k] = static_cast(f[k] * n3[k]); + if (g[k] >= n3[k]) + g[k] -= n3[k]; + } + frac.push_back(f); + pos.push_back(g); + } + std::vector mask_atoms; + for (size_t j = 0; j < mask_atom.size(); j++) { + const std::array &f = frac[mask_atom[j]]; + mask_atoms.push_back({f[0], f[1], f[2], mask_radius[j]}); + } + + std::vector> fc, fm; + { + std::lock_guard lock(m_); + engine_->Compute(atoms, pos, mask_atoms, fc, fm); + } + fcalc.v.clear(); + fmask.v.clear(); + fcalc.v.reserve(rows_.size()); + fmask.v.reserve(rows_.size()); + for (size_t i = 0; i < rows_.size(); i++) { + fcalc.v.push_back({rows_[i], {fc[i][0], fc[i][1]}}); + fmask.v.push_back({rows_[i], {fm[i][0], fm[i][1]}}); + } + fcalc.unit_cell_ = fmask.unit_cell_ = cell_; + fcalc.spacegroup_ = fmask.spacegroup_ = sg_; +} + +gemmi::Grid ModelStructureFactorsGPU::Map(gemmi::AsuData> &coef) { + coef.ensure_sorted(); + const std::array size = gemmi::get_size_for_hkl(coef, {{0, 0, 0}}, 3.0); + // ZYX: l fastest and halved, the layout a c2r transform takes, where gemmi's XYZ grid halves the slowest + // axis. The coefficients written are the same. + const gemmi::FPhiGrid hkl = gemmi::get_f_phi_on_grid(coef, size, true, gemmi::AxisOrder::ZYX); + gemmi::Grid map; + map.spacegroup = coef.spacegroup_; + map.unit_cell = coef.unit_cell_; + // As MapFromFPhi(): 2 (nw - 1) planes back from nw stored ones. + map.set_size(size[0], size[1], 2 * (hkl.nu - 1)); + map.axis_order = gemmi::AxisOrder::XYZ; + std::lock_guard lock(m_); + engine_->Map(map.nu, map.nv, map.nw, map.unit_cell.volume, reinterpret_cast(hkl.data.data()), + map.data.data()); + return map; +} diff --git a/image_analysis/structure_refinement/ModelStructureFactorsGPU.cu b/image_analysis/structure_refinement/ModelStructureFactorsGPU.cu new file mode 100644 index 000000000..ce4b7eb94 --- /dev/null +++ b/image_analysis/structure_refinement/ModelStructureFactorsGPU.cu @@ -0,0 +1,365 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "ModelStructureFactorsGPUEngine.h" + +#include +#include + +#include + +#include "../../common/JFJochException.h" +#include "../indexing/CUDAMemHelpers.h" + +// Every kernel below is deterministic: the density is ModelDensityGPU's gather, the mask ModelMaskGPU's, +// and each reflection is composed by one thread over the operators in order. The same atoms give the same +// structure factors, bit for bit, on every run. + +namespace { + +void cuda_err(cudaError_t val) { + if (val != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); +} + +void cufft_err(cufftResult val) { + if (val != CUFFT_SUCCESS) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, + "cuFFT error " + std::to_string(static_cast(val))); +} + +constexpr int THREADS = 256; + +int blocks(size_t n) { + return static_cast((n + THREADS - 1) / THREADS); +} + +__device__ __forceinline__ int imod(int a, int n) { + const int r = a % n; + return r < 0 ? r + n : r; +} + +// F of the crystal at each row, composed from the transform of one copy of its content (SymmetryComposition, +// RigidBodyRefine.h): F(h) = n_cen * unblur * sum over the operators x -> Rx + t of exp(+2 pi i h.t) F1(hR). +// `f1` is cuFFT's r2c transform of the one copy, u halved, which carries exp(-2 pi i h.x): conjugated and +// scaled by V/N it is gemmi's F1. An operator is gemmi's Op: rot row-major and tran, both in units of 1/den, +// and hR is its apply_to_hkl(). +__global__ void compose_kernel(const float2 *f1, float norm, int nu, int nv, int nw, const int *ops, int n_ops, + int den, const int *row_hkl, const double *row_scale, int rows, float2 *f) { + const int m = blockIdx.x * blockDim.x + threadIdx.x; + if (m >= rows) + return; + const int h[3] = {row_hkl[3 * m], row_hkl[3 * m + 1], row_hkl[3 * m + 2]}; + double sr = 0, si = 0; + for (int o = 0; o < n_ops; o++) { + const int *rot = ops + 12 * o, *tran = rot + 9; + int k[3]; + for (int i = 0; i < 3; i++) + k[i] = (rot[i] * h[0] + rot[3 + i] * h[1] + rot[6 + i] * h[2]) / den; + // The half-u transform holds k_u >= 0, and F1(-k) = conj F1(k) for a real map. + const bool conj = k[0] < 0; + if (conj) + for (int i = 0; i < 3; i++) + k[i] = -k[i]; + const float2 out = f1[k[0] + static_cast(nu / 2 + 1) * (imod(k[1], nv) + static_cast(nv) * imod(k[2], nw))]; + const double vr = out.x * norm, vi = (conj ? out.y : -out.y) * norm; + // exp(+2 pi i h.t), exact for the twelfths and quarters a translation is made of + double s, c; + sincospi(2.0 * (h[0] * tran[0] + h[1] * tran[1] + h[2] * tran[2]) / den, &s, &c); + sr += c * vr - s * vi; + si += c * vi + s * vr; + } + const double scale = row_scale[m]; + f[m] = make_float2(static_cast(scale * sr), static_cast(scale * si)); +} + +// F_mask at each row, read at h directly (the mask is symmetric): gemmi's prepare_asu_data() of the mask's +// transform, taking the Friedel mate where h is outside the stored half. +__global__ void fmask_kernel(const float2 *fm, float norm, const int *row_hkl, int rows, int nu, int nv, int nw, + float2 *fmask) { + const int m = blockIdx.x * blockDim.x + threadIdx.x; + if (m >= rows) + return; + int h = row_hkl[3 * m], k = row_hkl[3 * m + 1], l = row_hkl[3 * m + 2]; + const bool conj = h < 0; + if (conj) { + h = -h; + k = -k; + l = -l; + } + const float2 out = fm[h + static_cast(nu / 2 + 1) * (imod(k, nv) + static_cast(nv) * imod(l, nw))]; + const float re = out.x * norm, im = -out.y * norm; + fmask[m] = make_float2(re, conj ? -im : im); +} + +// MapFromFPhi()'s input step: conjugated, as gemmi does it (rho = 1/V sum F exp(-2 pi i h.x), and the c2r +// transform carries exp(+2 pi i h.x)), with a missing coefficient (NaN) read as zero. +__global__ void conjugate_kernel(float2 *c, size_t n) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) + return; + const float2 x = c[i]; + c[i] = isnan(x.y) ? make_float2(0.0f, 0.0f) : make_float2(x.x, -x.y); +} + +// The c2r output, z fastest and each z row padded to 2 (nw / 2 + 1) floats, into x fastest, times 1/V. +__global__ void transpose_kernel(const float *in, int nu, int nv, int nw, float scale, float *out) { + 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; + const int x = static_cast(i % nu); + const int y = static_cast(i / nu % nv); + const int z = static_cast(i / (static_cast(nu) * nv)); + out[i] = in[z + static_cast(2 * (nw / 2 + 1)) * (y + static_cast(nv) * x)] * scale; +} + +// The structure factors' transform: r2c of an (nu, nv, nw) grid stored u fastest, which halves u. +cufftHandle ForwardPlan(const ModelDensityGrid &g, size_t &work) { + cufftHandle plan; + cufft_err(cufftCreate(&plan)); + cufft_err(cufftSetAutoAllocation(plan, 0)); + int n[3] = {g.nw, g.nv, g.nu}; + const int idist = g.nu * g.nv * g.nw, odist = (g.nu / 2 + 1) * g.nv * g.nw; + const cufftResult r = cufftMakePlanMany(plan, 3, n, nullptr, 1, idist, nullptr, 1, odist, CUFFT_R2C, 1, &work); + if (r != CUFFT_SUCCESS) + cufftDestroy(plan); + cufft_err(r); + return plan; +} + +// A map's transform: in-place c2r of an (nu, nv, nw) grid stored as gemmi's ZYX half-l grid, w fastest and +// halved, the real output padded to 2 (nw / 2 + 1) per row as cuFFT does in place. +cufftHandle MapPlan(int nu, int nv, int nw, size_t &work) { + cufftHandle plan; + cufft_err(cufftCreate(&plan)); + cufft_err(cufftSetAutoAllocation(plan, 0)); + int n[3] = {nu, nv, nw}; + int inembed[3] = {nu, nv, nw / 2 + 1}; + int onembed[3] = {nu, nv, 2 * (nw / 2 + 1)}; + const int idist = nu * nv * (nw / 2 + 1), odist = 2 * idist; + const cufftResult r = cufftMakePlanMany(plan, 3, n, inembed, 1, idist, onembed, 1, odist, CUFFT_C2R, 1, &work); + if (r != CUFFT_SUCCESS) + cufftDestroy(plan); + cufft_err(r); + return plan; +} + +// The calling thread on `device` for as long as this lives, and back on its own device after. +class DeviceGuard { +public: + explicit DeviceGuard(int device) { + cuda_err(cudaGetDevice(&previous_)); + cuda_err(cudaSetDevice(device)); + } + ~DeviceGuard() { cudaSetDevice(previous_); } + DeviceGuard(const DeviceGuard &) = delete; + DeviceGuard &operator=(const DeviceGuard &) = delete; + +private: + int previous_ = 0; +}; + +size_t GridPoints(const ModelDensityGrid &g) { + return static_cast(g.nu) * g.nv * g.nw; +} + +size_t ComplexPoints(const ModelDensityGrid &g) { + return static_cast(g.nu / 2 + 1) * g.nv * g.nw; +} + +} // namespace + +struct ModelStructureFactorsGPUEngineImpl { + ModelStructureFactorsGPUSetup setup; + size_t work_bytes; + int device = 0; + CudaStream stream; + + CudaDevicePtr pos; + CudaDevicePtr mask_atoms; + CudaDevicePtr grid; // the density, then the mask; a map's real grid + CudaDevicePtr spectrum; // their transforms; a map's coefficients, transformed in place + CudaDevicePtr fft_work; + CudaDevicePtr ops, row_hkl; + CudaDevicePtr row_scale; + CudaDevicePtr f, fmask; + std::unique_ptr density; + std::unique_ptr mask; + cufftHandle forward = 0; + + // On `dev`, which the caller has made current. + ModelStructureFactorsGPUEngineImpl(int dev, const ModelStructureFactorsGPUSetup &s, size_t work) + : setup(s), work_bytes(work), device(dev) { + const auto sync = CudaAlloc::Synchronous; + const size_t na = std::max(setup.max_atoms, 1), nr = std::max(setup.rows.size(), 1); + pos = CudaDevicePtr(na, sync); + mask_atoms = CudaDevicePtr(na, sync); + grid = CudaDevicePtr(std::max(GridPoints(setup.grid), setup.map_points), sync); + spectrum = CudaDevicePtr(std::max(ComplexPoints(setup.grid), setup.map_complex_points), sync); + fft_work = CudaDevicePtr(std::max(work_bytes, 1), sync); + ops = CudaDevicePtr(std::max(12 * setup.sym_ops.size(), 1), sync); + row_hkl = CudaDevicePtr(3 * nr, sync); + row_scale = CudaDevicePtr(nr, sync); + f = CudaDevicePtr(nr, sync); + fmask = CudaDevicePtr(nr, sync); + density = std::make_unique(stream, setup.max_atoms, setup.max_pairs, + ModelDensityGPU::Bricks(setup.grid.nu, setup.grid.nv, setup.grid.nw)); + mask = std::make_unique(stream, GridPoints(setup.grid)); + + ModelMaskGrid mask_grid; + mask_grid.nu = setup.grid.nu; + mask_grid.nv = setup.grid.nv; + mask_grid.nw = setup.grid.nw; + std::copy(setup.grid.orth, setup.grid.orth + 9, mask_grid.orth); + mask_grid.volume = setup.volume; + mask->SetGrid(mask_grid, setup.mask_ops); + + auto up = [&](void *dst, const void *src, size_t bytes) { + if (bytes > 0) + cuda_err(cudaMemcpyAsync(dst, src, bytes, cudaMemcpyHostToDevice, stream)); + }; + up(ops, setup.sym_ops.data(), setup.sym_ops.size() * 12 * sizeof(int)); + up(row_hkl, setup.rows.data(), setup.rows.size() * 3 * sizeof(int)); + up(row_scale, setup.row_scale.data(), setup.row_scale.size() * sizeof(double)); + + size_t plan_work = 0; + forward = ForwardPlan(setup.grid, plan_work); + if (plan_work > work_bytes) { + cufftDestroy(forward); + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "structure factors: cuFFT work area over the reserve"); + } + cufft_err(cufftSetWorkArea(forward, fft_work)); + cufft_err(cufftSetStream(forward, stream)); + cuda_err(cudaStreamSynchronize(stream)); + } + + ~ModelStructureFactorsGPUEngineImpl() { + cudaStreamSynchronize(stream); + cufftDestroy(forward); + } + + float Norm() const { + return static_cast(setup.volume / static_cast(GridPoints(setup.grid))); + } +}; + +size_t ModelStructureFactorsGPUEngine::TotalMemory(int device) { + DeviceGuard guard(device); + size_t free = 0, total = 0; + cuda_err(cudaMemGetInfo(&free, &total)); + return total; +} + +size_t ModelStructureFactorsGPUEngine::FFTWorkBytes(int device, const ModelDensityGrid &grid, int map_nu, int map_nv, + int map_nw) { + DeviceGuard guard(device); + size_t forward = 0, inverse = 0; + cufftDestroy(ForwardPlan(grid, forward)); + cufftDestroy(MapPlan(map_nu, map_nv, map_nw, inverse)); + return std::max(forward, inverse); +} + +size_t ModelStructureFactorsGPUEngine::DeviceBytes(const ModelStructureFactorsGPUSetup &s, size_t fft_work_bytes) { + const size_t na = s.max_atoms + 1, nr = s.rows.size() + 1; + const size_t points = GridPoints(s.grid); + size_t bytes = na * (sizeof(float4) + sizeof(ModelMaskAtom)); + bytes += std::max(points, s.map_points) * sizeof(float); + bytes += std::max(ComplexPoints(s.grid), s.map_complex_points) * sizeof(float2); + bytes += fft_work_bytes; + bytes += 12 * s.sym_ops.size() * sizeof(int) + nr * (3 * sizeof(int) + sizeof(double) + 2 * sizeof(float2)); + bytes += ModelDensityGPU::DeviceBytes(s.max_atoms, s.max_pairs, ModelDensityGPU::Bricks(s.grid.nu, s.grid.nv, s.grid.nw)); + bytes += ModelMaskGPU::DeviceBytes(points); + return bytes; +} + +ModelStructureFactorsGPUEngine::ModelStructureFactorsGPUEngine(int device, const ModelStructureFactorsGPUSetup &setup, + size_t fft_work_bytes) { + DeviceGuard guard(device); + impl_ = std::make_unique(device, setup, fft_work_bytes); +} + +ModelStructureFactorsGPUEngine::~ModelStructureFactorsGPUEngine() { + if (impl_) { + DeviceGuard guard(impl_->device); + impl_.reset(); + } +} + +void ModelStructureFactorsGPUEngine::Compute(const std::vector &atoms, + const std::vector> &pos, + const std::vector &mask_atoms, + std::vector> &fcalc, + std::vector> &fmask) { + ModelStructureFactorsGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + if (pos.size() != atoms.size() || mask_atoms.size() > e.setup.max_atoms) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "structure factors: atoms over the reserve"); + const ModelDensityGrid &g = e.setup.grid; + const int rows = static_cast(e.setup.rows.size()); + + e.density->SetAtoms(g, atoms); + if (!pos.empty()) + cuda_err(cudaMemcpyAsync(e.pos, pos.data(), pos.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream)); + if (!mask_atoms.empty()) + cuda_err(cudaMemcpyAsync(e.mask_atoms, mask_atoms.data(), mask_atoms.size() * sizeof(ModelMaskAtom), + cudaMemcpyHostToDevice, e.stream)); + + e.density->Compute(e.pos, e.grid); + cufft_err(cufftExecR2C(e.forward, e.grid, reinterpret_cast(e.spectrum.get()))); + if (rows > 0) { + compose_kernel<<>>(e.spectrum, e.Norm(), g.nu, g.nv, g.nw, e.ops, + static_cast(e.setup.sym_ops.size()), e.setup.den, + e.row_hkl, e.row_scale, rows, e.f); + cuda_err(cudaGetLastError()); + } + + e.mask->Compute(e.mask_atoms, static_cast(mask_atoms.size()), e.grid); + cufft_err(cufftExecR2C(e.forward, e.grid, reinterpret_cast(e.spectrum.get()))); + if (rows > 0) { + fmask_kernel<<>>(e.spectrum, e.Norm(), e.row_hkl, rows, g.nu, g.nv, g.nw, + e.fmask); + cuda_err(cudaGetLastError()); + } + + fcalc.resize(rows); + fmask.resize(rows); + if (rows > 0) { + cuda_err(cudaMemcpyAsync(fcalc.data(), e.f, rows * sizeof(float2), cudaMemcpyDeviceToHost, e.stream)); + cuda_err(cudaMemcpyAsync(fmask.data(), e.fmask, rows * sizeof(float2), cudaMemcpyDeviceToHost, e.stream)); + } + cuda_err(cudaStreamSynchronize(e.stream)); +} + +void ModelStructureFactorsGPUEngine::Map(int nu, int nv, int nw, double volume, const float *coefficients, float *map) { + ModelStructureFactorsGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + const size_t points = static_cast(nu) * nv * nw; + const size_t complex_points = static_cast(nu) * nv * (nw / 2 + 1); + if (points > e.setup.map_points || complex_points > e.setup.map_complex_points) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "map: grid over the reserve"); + size_t plan_work = 0; + const cufftHandle plan = MapPlan(nu, nv, nw, plan_work); + try { + if (plan_work > e.work_bytes) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "map: cuFFT work area over the reserve"); + cufft_err(cufftSetWorkArea(plan, e.fft_work)); + cufft_err(cufftSetStream(plan, e.stream)); + cuda_err(cudaMemcpyAsync(e.spectrum, coefficients, complex_points * sizeof(float2), cudaMemcpyHostToDevice, + e.stream)); + conjugate_kernel<<>>(e.spectrum, complex_points); + cuda_err(cudaGetLastError()); + cufft_err(cufftExecC2R(plan, reinterpret_cast(e.spectrum.get()), + reinterpret_cast(e.spectrum.get()))); + transpose_kernel<<>>(reinterpret_cast(e.spectrum.get()), + nu, nv, nw, static_cast(1.0 / volume), e.grid); + cuda_err(cudaGetLastError()); + cuda_err(cudaMemcpyAsync(map, e.grid, points * sizeof(float), cudaMemcpyDeviceToHost, e.stream)); + cuda_err(cudaStreamSynchronize(e.stream)); + } catch (...) { + cudaStreamSynchronize(e.stream); + cufftDestroy(plan); + throw; + } + cufftDestroy(plan); +} diff --git a/image_analysis/structure_refinement/ModelStructureFactorsGPU.h b/image_analysis/structure_refinement/ModelStructureFactorsGPU.h new file mode 100644 index 000000000..b206997b5 --- /dev/null +++ b/image_analysis/structure_refinement/ModelStructureFactorsGPU.h @@ -0,0 +1,95 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// A model's structure factors, and the maps made from coefficients on its reflections, on the GPU (CUDA +// builds only). It is what model validation computes on the CPU - F_calc from the model's density on a +// grid (gemmi's DensityCalculator, IT92, the Refmac-compatible blur and its unblur) and F_mask from the +// bulk-solvent mask (gemmi's SolventMasker with the Refmac radii), each transformed and read off as +// prepare_asu_data() lists them; and the inverse transform of a set of map coefficients, as +// get_f_phi_on_grid() and MapFromFPhi() make it - moved to the device. The two agree to rounding, not +// bit for bit: float distances and cuFFT for FFTW, and the crystal's symmetry is composed in reciprocal +// space (SymmetryComposition, RigidBodyRefine.h) instead of by symmetrizing the grid. The GPU is +// deterministic on its own. +// +// Made once for a cell, a group, a resolution and a model's atoms, then evaluated as often as the +// coordinates change: everything that depends only on the first four is worked out here, once, and +// the device buffers are reserved once. The atoms' B factors and occupancies are read at every +// evaluation, but the blur is the one the model had when the engine was made. + +#include +#include +#include +#include +#include + +#include "gemmi/asudata.hpp" +#include "gemmi/dencalc.hpp" // DensityCalculator +#include "gemmi/grid.hpp" +#include "gemmi/it92.hpp" +#include "gemmi/model.hpp" +#include "gemmi/symmetry.hpp" +#include "gemmi/unitcell.hpp" + +#include "ModelStructureFactorsGPUEngine.h" + +// Each atom's density as PutModelDensityOnGrid() (ModelGrid.cpp) sets it up for `dc` - its d_min, rate +// and blur - in model order; and the atoms of the bulk-solvent mask as PutMaskOnGrid() takes them: each +// one's index in the model and its radius, probe included. +void ModelDensityAtoms(const gemmi::Model &model, const gemmi::DensityCalculator, float> &dc, + std::vector &atoms, std::vector &mask_atom, + std::vector &mask_radius); + +class ModelStructureFactorsGPU { +public: + // For `device`. Host work only: the grid, the reflections and what the device will need for them. + // Nothing is reserved until Reserve(), so DeviceBytes() can decide first whether to. Every call works + // on `device` and leaves the calling thread's current device as it found it; which card it is does not + // change a number. + ModelStructureFactorsGPU(int device, const gemmi::Model &model, const gemmi::UnitCell &cell, + const gemmi::SpaceGroup &sg, double d_min); + ~ModelStructureFactorsGPU(); + + // Whether the GPU reproduces this case: the gather needs every atom's box narrower than the cell. + bool Supported() const { return supported_; } + // The device memory Reserve() takes: the grid and its transform, the solvent mask's labels, the cuFFT + // work area and the reflections. The largest map Map() can be given shares the same buffers. + size_t DeviceBytes() const { return device_bytes_; } + double DMin() const { return d_min_; } + int Device() const { return device_; } + // The device's total memory, which DeviceBytes() is to be judged against. + size_t DeviceTotalMemory() const { return ModelStructureFactorsGPUEngine::TotalMemory(device_); } + // The grid the structure factors are computed on. + std::array GridSize() const { return {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw}; } + std::array MapSizeBound() const { return map_size_; } + + // Reserves the device memory. Throws JFJochException on a CUDA failure. + void Reserve(); + + // F_calc and F_mask of `model` - the model this was made for, its atoms anywhere - as model validation's + // compute_model_factors() makes them: prepare_asu_data(d_min, blur) of the density and + // prepare_asu_data(d_min) of the mask. Throws JFJochException on a CUDA failure. Safe to call from + // several threads; they take turns. + void Compute(const gemmi::Model &model, gemmi::AsuData> &fcalc, + gemmi::AsuData> &fmask); + + // The real-space map of ASU coefficients on this engine's reflections, as model validation's + // map_from_coefficients() makes it: on the grid get_size_for_hkl(coef, {0, 0, 0}, 3.0) sizes. Sorts + // `coef`. Throws JFJochException on a CUDA failure. Safe to call from several threads; they take turns. + gemmi::Grid Map(gemmi::AsuData> &coef); + +private: + int device_; + gemmi::UnitCell cell_; + const gemmi::SpaceGroup *sg_; + double d_min_; + gemmi::DensityCalculator, float> dc_; + ModelStructureFactorsGPUSetup setup_; + std::vector rows_; + std::array map_size_{}; + bool supported_ = false; + size_t fft_work_bytes_ = 0, device_bytes_ = 0; + std::unique_ptr engine_; + std::mutex m_; +}; diff --git a/image_analysis/structure_refinement/ModelStructureFactorsGPUEngine.h b/image_analysis/structure_refinement/ModelStructureFactorsGPUEngine.h new file mode 100644 index 000000000..06555704c --- /dev/null +++ b/image_analysis/structure_refinement/ModelStructureFactorsGPUEngine.h @@ -0,0 +1,64 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +// The device half of ModelStructureFactorsGPU (ModelStructureFactorsGPU.h). The host half works out with +// gemmi - which stays out of nvcc - everything that depends on the cell, the group and the resolution, and +// per evaluation the atoms' densities and positions; this half grids, transforms and composes. CUDA builds +// only. + +#include +#include +#include +#include + +#include "ModelDensityGPU.h" +#include "ModelMaskGPU.h" + +// Everything an engine needs that does not depend on the model's coordinates. +struct ModelStructureFactorsGPUSetup { + ModelDensityGrid grid; // the structure-factor grid, u fastest + double volume = 0; + // The group's operators as gemmi holds them: rot (row-major) and tran, in units of 1 / den. + std::vector> sym_ops; + int den = 24; + std::vector mask_ops; // every operator with every centring vector + std::vector> rows; // the reflections computed, in the order they come back + std::vector row_scale; // n_cen * prepare_asu_data()'s unblur, per row + size_t max_atoms = 0, max_pairs = 0; + // The largest map MapFromFPhi() will be given: points of the real grid, and of its half-l transform. + size_t map_points = 0, map_complex_points = 0; +}; + +struct ModelStructureFactorsGPUEngineImpl; + +class ModelStructureFactorsGPUEngine { +public: + // Every call names its device and leaves the calling thread's current device as it found it. + + // The device's total memory. + static size_t TotalMemory(int device); + // The cuFFT work area the structure factors' r2c, and a map of (nu, nv, nw)'s c2r, take on `device`. + static size_t FFTWorkBytes(int device, const ModelDensityGrid &grid, int map_nu, int map_nv, int map_nw); + // The bytes an engine for `setup` reserves on the device, given its cuFFT work area. + static size_t DeviceBytes(const ModelStructureFactorsGPUSetup &setup, size_t fft_work_bytes); + + ModelStructureFactorsGPUEngine(int device, const ModelStructureFactorsGPUSetup &setup, size_t fft_work_bytes); + ~ModelStructureFactorsGPUEngine(); + + // F_calc and F_mask at the rows for one set of atoms: each atom's density, its grid position + // (fractional times the grid size, wrapped into [0, n); the 4th number is not read) and the atoms of + // the bulk-solvent mask. Both come back as (re, im) per row. + void Compute(const std::vector &atoms, const std::vector> &pos, + const std::vector &mask_atoms, std::vector> &fcalc, + std::vector> &fmask); + + // The real-space map rho(x) = 1/V sum F exp(-2 pi i h.x) of a half-l coefficient grid in gemmi's ZYX + // order - h slowest, l fastest and halved: nu * nv * (nw / 2 + 1) complex numbers, (re, im) + // interleaved, NaN read as zero - into `map`, nu * nv * nw floats with x fastest. + void Map(int nu, int nv, int nw, double volume, const float *coefficients, float *map); + +private: + std::unique_ptr impl_; +}; diff --git a/image_analysis/structure_refinement/ModelValidation.cpp b/image_analysis/structure_refinement/ModelValidation.cpp index 3673483dc..304484ffb 100644 --- a/image_analysis/structure_refinement/ModelValidation.cpp +++ b/image_analysis/structure_refinement/ModelValidation.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -38,8 +39,13 @@ #include "RigidBodyRefine.h" #include "SigmaA.h" #ifdef JFJOCH_USE_CUDA +#include "ModelStructureFactorsGPU.h" #include "RigidBodyGPU.h" +#include "RigidBodyGPUEngine.h" // RigidBodyGPUEngine::CurrentDevice #include "../../common/CUDAWrapper.h" +#include "../../common/JFJochException.h" +#else +class ModelStructureFactorsGPU; #endif namespace { @@ -51,8 +57,19 @@ long hkl_key(const gemmi::Miller &h) { return (h[0] + 512L) * 1048576 + (h[1] + 512L) * 1024 + (h[2] + 512L); } -// FFT ASU map coefficients into a real-space map. -gemmi::Grid map_from_coefficients(gemmi::AsuData> &coef) { +// FFT ASU map coefficients into a real-space map, on `gpu` where the validation has one. +gemmi::Grid map_from_coefficients(gemmi::AsuData> &coef, ModelStructureFactorsGPU *gpu) { +#ifdef JFJOCH_USE_CUDA + if (gpu != nullptr) { + try { + return gpu->Map(coef); + } catch (const JFJochException &e) { + throw RigidBodyGPUFailure(e.what()); + } + } +#else + (void) gpu; +#endif coef.ensure_sorted(); std::array size = gemmi::get_size_for_hkl(coef, {{0, 0, 0}}, 3.0); return MapFromFPhi(gemmi::get_f_phi_on_grid(coef, size, true)); @@ -335,13 +352,53 @@ double frame_probe_r(const gemmi::Structure &st, const gemmi::SpaceGroup *sg, return scaling.calculate_r_factor(); } +#ifdef JFJOCH_USE_CUDA +// The model's structure factors to the data's resolution, and the maps, on the GPU - or not, decided here, +// once, before any of them is computed, and never revisited: nothing is moved to the CPU part way through. +// Decided on what the GPU path needs against the card's TOTAL memory and not against what happens to be +// free at the moment, so the same input on the same machine always takes the same path, whatever else is +// running beside it. The structure factors may take half the card; the rigid body's engines take at most +// a quarter beside them (RigidBodyGPUPool::Create), and the last quarter is left to everything else. +std::unique_ptr StructureFactorsGPU(int device, const gemmi::Model &model, + const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, + double d_min, Logger &logger) { + auto sf = std::make_unique(device, model, cell, sg, d_min); + const std::array n = sf->GridSize(); + if (!sf->Supported()) { + logger.Info("Model validation: structure factors and maps on the CPU - the cell is too small for the " + "GPU's gridding"); + return nullptr; + } + size_t total = 0; + try { + total = sf->DeviceTotalMemory(); + } catch (const JFJochException &e) { + throw RigidBodyGPUFailure(e.what()); + } + if (sf->DeviceBytes() > total / 2) { + logger.Info("Model validation: structure factors and maps on the CPU - on a {}x{}x{} grid they need " + "{:.2f} GB of GPU memory, over half of the card's {:.2f} GB", n[0], n[1], n[2], + sf->DeviceBytes() / 1e9, total / 1e9); + return nullptr; + } + try { + sf->Reserve(); + } catch (const JFJochException &e) { + throw RigidBodyGPUFailure(e.what()); + } + logger.Info("Model validation: structure factors and maps on the GPU - a {}x{}x{} grid, {:.2f} GB of the " + "card's {:.2f} GB", n[0], n[1], n[2], sf->DeviceBytes() / 1e9, total / 1e9); + return sf; +} +#endif } // namespace namespace { -// ValidateAgainstModel() with the rigid body on the GPU where `rigid_body_gpu` allows it and a card is -// there to take it, on the CPU otherwise - for the whole validation either way. +// ValidateAgainstModel() with the rigid body, the structure factors and the maps on the GPU where +// `rigid_body_gpu` allows it and a card is there to take them, on the CPU otherwise - for the whole +// validation either way. ModelValidationResult Validate(const std::vector &merged, const UnitCell &cell, const std::string &model_path, @@ -528,10 +585,33 @@ ModelValidationResult Validate(const std::vector &merged, logger.Info("Model validation: {} atoms, cell a={:.2f} b={:.2f} c={:.2f}, sg {}, to {:.2f} A", gemmi::count_atom_sites(st.models[0]), ucell.a, ucell.b, ucell.c, sg->hm, d_min); + // The structure factors to d_min and the maps on the GPU, where `rigid_body_gpu` allows it and the + // card can take them; decided here, before any is computed. Anything to a coarser limit - the null's, + // the indexing probe's - stays on the CPU, where it is cheap, so that every replicate of the null and + // the real model's side of it are computed the same way. + ModelStructureFactorsGPU *sf_gpu = nullptr; +#ifdef JFJOCH_USE_CUDA + std::unique_ptr sf_engine; + // On the current device, as the rigid body's engines take it: one card for the whole validation. + if (rigid_body_gpu && get_gpu_count() > 0) + sf_engine = StructureFactorsGPU(RigidBodyGPUEngine::CurrentDevice(), st.models[0], ucell, *sg, d_min, logger); + sf_gpu = sf_engine.get(); +#endif + // --- Fcalc (atomic) via electron density on a grid + FFT, plus a flat bulk-solvent mask -> Fmask. // A lambda because the rigid-body step below moves the model and then needs both again, and it // takes the state to work on so that a null replicate can run it on its own copy. --- auto compute_model_factors = [&](ModelState &ms, double to_d) { +#ifdef JFJOCH_USE_CUDA + if (sf_gpu != nullptr && to_d == sf_gpu->DMin()) { + try { + sf_gpu->Compute(ms.st.models[0], ms.fcalc, ms.fmask); + } catch (const JFJochException &e) { + throw RigidBodyGPUFailure(e.what()); + } + return; + } +#endif gemmi::DensityCalculator dc; dc.d_min = to_d; dc.rate = 1.5; @@ -1265,7 +1345,7 @@ ModelValidationResult Validate(const std::vector &merged, // The three maps - this one, the difference map and the anomalous map - share nothing but what // they read, so the other two are made and written on threads of their own beside this one. std::future fofc_map = std::async(std::launch::async, [&] { - write_ccp4(map_from_coefficients(mapfofc), output_prefix + "_fofc.ccp4"); + write_ccp4(map_from_coefficients(mapfofc, sf_gpu), output_prefix + "_fofc.ccp4"); }); // --- anomalous difference map, where the merge kept the Bijvoet split --- @@ -1312,7 +1392,7 @@ ModelValidationResult Validate(const std::vector &merged, result.anomalous_pairs = static_cast(mapanom.v.size()); if (!mapanom.v.empty()) { - const gemmi::Grid grid = map_from_coefficients(mapanom); + const gemmi::Grid grid = map_from_coefficients(mapanom, sf_gpu); const double rms = write_ccp4(grid, output_prefix + "_anom.ccp4"); std::vector sites; double scatterer_sum = 0; @@ -1355,7 +1435,7 @@ ModelValidationResult Validate(const std::vector &merged, } }); - const gemmi::Grid grid2fofc = map_from_coefficients(map2fofc); + const gemmi::Grid grid2fofc = map_from_coefficients(map2fofc, sf_gpu); const double rms2 = write_ccp4(grid2fofc, output_prefix + "_2fofc.ccp4"); { double s = 0; int n = 0; @@ -1490,7 +1570,7 @@ ModelValidationResult ValidateAgainstModel(const std::vector & double wavelength_A, const std::vector &report_shell_d_min) { #ifdef JFJOCH_USE_CUDA - // A CUDA failure in the rigid body does not end the run: the validation is started again from the + // A CUDA failure on the GPU path does not end the run: the validation is started again from the // model as read, on the CPU throughout. It is re-runnable, and a dead model check should not take a // finished merge with it. try { diff --git a/image_analysis/structure_refinement/RigidBodyGPU.cpp b/image_analysis/structure_refinement/RigidBodyGPU.cpp index 3de8b7fd0..2c0b5eb16 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.cpp +++ b/image_analysis/structure_refinement/RigidBodyGPU.cpp @@ -12,9 +12,9 @@ #include "gemmi/dencalc.hpp" // DensityCalculator #include "gemmi/it92.hpp" // IT92 x-ray form factors #include "gemmi/scaling.hpp" // Scaling -#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent #include "ModelScaling.h" // FitModelScale +#include "ModelStructureFactorsGPU.h" // ModelDensityAtoms #include "RigidBodyGPUEngine.h" #include "../../common/CUDAWrapper.h" #include "../../common/JFJochException.h" @@ -40,7 +40,6 @@ RigidBodyGPUZone ModelZone(const gemmi::Model &model, const gemmi::UnitCell &cel } zone.volume = cell.volume; - // The density of each atom as PutModelDensityOnGrid() (ModelGrid.cpp) precalculates it. gemmi::DensityCalculator dc; dc.d_min = d_min; dc.rate = 1.5; @@ -48,47 +47,7 @@ RigidBodyGPUZone ModelZone(const gemmi::Model &model, const gemmi::UnitCell &cel dc.grid.spacegroup = &sg; dc.set_refmac_compatible_blur(model); zone.blur = dc.blur; - const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac); - int index = 0; - for (const gemmi::Chain &ch : model.chains) - for (const gemmi::Residue &r : ch.residues) - for (const gemmi::Atom &atom : r.atoms) { - using CReal = Table::Coef::coef_type; - const auto &coef = Table::get(atom.element, atom.charge, atom.serial); - const float addend = dc.addends.get(atom.element); - RigidBodyGPUAtom a{}; - a.occ = atom.occ; - a.aniso = atom.aniso.nonzero(); - if (!a.aniso) { - const CReal b = static_cast(atom.b_iso + dc.blur); - const auto precal = coef.precalculate_density_iso(b, addend); - a.radius = dc.estimate_radius(precal, b); - for (int k = 0; k < 5; k++) { - a.a[k] = precal.a[k]; - a.b[k][0] = precal.b[k]; - } - } 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); - a.radius = static_cast(dc.estimate_radius(coef.precalculate_density_iso(b_max, addend), b_max)); - const auto precal = coef.precalculate_density_aniso_b(aniso_b, addend); - for (int k = 0; k < 5; k++) { - a.a[k] = precal.a[k]; - const gemmi::SMat33 &m = precal.b[k]; - const float e[6] = {m.u11, m.u22, m.u33, m.u12, m.u13, m.u23}; - std::copy(e, e + 6, a.b[k]); - } - } - zone.atoms.push_back(a); - // The bulk-solvent mask's atoms, as PutMaskOnGrid() takes them. - if (!((masker.ignore_hydrogen && atom.is_hydrogen()) || - (masker.ignore_zero_occupancy_atoms && atom.occ <= 0))) { - zone.mask_atom.push_back(index); - zone.mask_radius.push_back(static_cast( - masker.constant_r + masker.rprobe + gemmi::refmac_radius_for_bulk_solvent(atom.element.elem))); - } - ++index; - } + ModelDensityAtoms(model, dc, zone.atoms, zone.mask_atom, zone.mask_radius); const gemmi::GroupOps gops = sg.operations(); for (const gemmi::Op::Tran &cen : gops.cen_ops) diff --git a/image_analysis/structure_refinement/RigidBodyGPU.cu b/image_analysis/structure_refinement/RigidBodyGPU.cu index 1db77013a..689ab5513 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.cu +++ b/image_analysis/structure_refinement/RigidBodyGPU.cu @@ -8,9 +8,9 @@ #include #include -#include #include +#include "ModelDensityGPU.h" #include "ModelMaskGPU.h" #include "ModelScaleGPU.h" #include "../../common/JFJochException.h" @@ -33,17 +33,12 @@ void cufft_err(cufftResult val) { "cuFFT error " + std::to_string(static_cast(val))); } -constexpr int BRICK = 8; // the gather's bricks are BRICK^3 grid points, one block each -constexpr int GATHER_TILE = 128; // atoms staged in shared memory at a time -constexpr int MAX_BRICKS_PER_AXIS = 8; // an atom's box touches at most this many bricks along an axis constexpr int THREADS = 256; constexpr int MAX_SCALE_PARAMS = 7; // k_overall + up to six B* constraints -// The grid of a zone as the kernels see it. +// The grid of a zone. struct GridGeom { int nu, nv, nw; - float orth_n[9]; // orth * diag(1/nu, 1/nv, 1/nw), row-major: grid offset -> Cartesian - int narrow; // the cell is too small for one image of an atom to serve a whole brick }; // Where the body is: x -> R x_rel + t, and the fractionalization. @@ -67,15 +62,6 @@ __device__ __forceinline__ int imod(int a, int n) { return r < 0 ? r + n : r; } -// gemmi's unsafe_expapprox() (formfact.hpp), the exponential the density is computed with. -__device__ __forceinline__ float expapprox(float x) { - const float val = 12102203.1615614f * x + 1065353216.f; - const int vali = static_cast(val); - const float a = __int_as_float(vali & 0x7F800000); - const float b = __int_as_float((vali & 0x7FFFFF) | 0x3F800000); - return a * (0.509871020f + b * (0.312146713f + b * (0.166617139f + b * (-2.190619930e-3f + b * 1.3555747234e-2f)))); -} - // b_star.r_u_r(h), as gemmi's SMat33. __device__ __forceinline__ double r_u_r(const double *u, const int *h) { const double x = h[0], y = h[1], z = h[2]; @@ -116,163 +102,6 @@ __global__ void place_kernel(const double *rel, int n, Placement p, int nu, int } } -// The distinct bricks the points c - d ... c + d of one axis fall in, wrapped into the cell. A grid size -// that is not a multiple of BRICK leaves the last brick partial, which is why this is done on wrapped -// points and not in brick coordinates. -1 where there are more than MAX_BRICKS_PER_AXIS of them. -__device__ int axis_bricks(int c, int d, int n, int *out) { - int k = 0; - for (int p = c - d; p <= c + d; p++) { - const int b = imod(p, n) / BRICK; - bool seen = false; - for (int j = 0; j < k; j++) - seen = seen || out[j] == b; - if (seen) - continue; - if (k == MAX_BRICKS_PER_AXIS) - return -1; - out[k++] = b; - } - return k; -} - -// The bricks atom i's box touches, per axis. The box is gemmi's: `box` points either side of the -// point nearest the atom. -__device__ void atom_bricks(const float4 &pos, const int *box, int nu, int nv, int nw, int *bu, int &nbu, int *bv, - int &nbv, int *bw, int &nbw) { - nbu = axis_bricks(__float2int_rn(pos.x), box[0], nu, bu); - nbv = axis_bricks(__float2int_rn(pos.y), box[1], nv, bv); - nbw = axis_bricks(__float2int_rn(pos.z), box[2], nw, bw); -} - -__global__ void count_pairs_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw, int *count, - int *overflow) { - const int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i > n) - return; - if (i == n) { - count[n] = 0; // so the exclusive scan's last entry is the total - return; - } - int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c; - atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c); - if (a < 0 || b < 0 || c < 0) { - *overflow = 1; - a = b = c = 0; - } - count[i] = a * b * c; -} - -// The (brick, atom) pairs, atom-major: sorted stably by brick, each brick's atoms stay in model order. -__global__ void fill_pairs_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw, - const int *offset, int *key, int *value) { - const int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) - return; - int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c; - atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c); - const int nbu = (nu + BRICK - 1) / BRICK, nbv = (nv + BRICK - 1) / BRICK; - int o = offset[i]; - for (int k = 0; k < c; k++) - for (int j = 0; j < b; j++) - for (int l = 0; l < a; l++) { - key[o] = bu[l] + nbu * (bv[j] + nbv * bw[k]); - value[o] = i; - o++; - } -} - -__global__ void brick_ranges_kernel(const int *key, int npairs, int *start, int *end) { - const int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= npairs) - return; - const int k = key[i]; - if (i == 0 || key[i - 1] != k) - start[k] = i; - if (i == npairs - 1 || key[i + 1] != k) - end[k] = i + 1; -} - -// The density of the model on the grid, as gemmi's do_add_atom_density_to_grid() puts it there: every -// point adds, in model order, each atom whose sphere it is inside. One block per brick; the brick's atoms -// are staged in shared memory with the position of their image nearest to the brick, relative to the -// brick's first point. Where the cell is wider than an atom's box plus a brick, that is the only image -// that reaches any point of the brick; on a narrower cell each point finds its own nearest image. -__global__ void gather_kernel(const RigidBodyGPUAtom *atoms, const float4 *pos, const int *pair_atom, - const int *brick_start, const int *brick_end, GridGeom g, float *grid) { - const int nbu = (g.nu + BRICK - 1) / BRICK, nbv = (g.nv + BRICK - 1) / BRICK; - const int brick = blockIdx.x; - const int s = brick_start[brick], e = brick_end[brick]; // s == e: an empty brick, written as zeros - const int u0 = (brick % nbu) * BRICK, v0 = (brick / nbu % nbv) * BRICK, w0 = (brick / (nbu * nbv)) * BRICK; - const int u = u0 + threadIdx.x, v = v0 + threadIdx.y, w = w0 + threadIdx.z; - const bool inside = u < g.nu && v < g.nv && w < g.nw; - const int tid = threadIdx.x + BRICK * (threadIdx.y + BRICK * threadIdx.z); - // This point relative to the brick's first one, in Cartesian coordinates. - const float tx = g.orth_n[0] * threadIdx.x + g.orth_n[1] * threadIdx.y + g.orth_n[2] * threadIdx.z; - const float ty = g.orth_n[3] * threadIdx.x + g.orth_n[4] * threadIdx.y + g.orth_n[5] * threadIdx.z; - const float tz = g.orth_n[6] * threadIdx.x + g.orth_n[7] * threadIdx.y + g.orth_n[8] * threadIdx.z; - - constexpr int WORDS = sizeof(RigidBodyGPUAtom) / sizeof(int); - __shared__ int sh_index[GATHER_TILE]; - __shared__ RigidBodyGPUAtom sh_atom[GATHER_TILE]; - __shared__ float3 sh_centre[GATHER_TILE]; // Cartesian - __shared__ float3 sh_grid[GATHER_TILE]; // the same, in grid units - float acc = 0.0f; - for (int c = s; c < e; c += GATHER_TILE) { - const int m = min(GATHER_TILE, e - c); - __syncthreads(); - if (tid < m) { - const int a = pair_atom[c + tid]; - sh_index[tid] = a; - // The atom's image nearest to the brick's centre, in grid units from its first point. - const float4 p = pos[a]; - float fx = p.x - u0, fy = p.y - v0, fz = p.z - w0; - fx -= g.nu * rintf((fx - 0.5f * (BRICK - 1)) / g.nu); - fy -= g.nv * rintf((fy - 0.5f * (BRICK - 1)) / g.nv); - fz -= g.nw * rintf((fz - 0.5f * (BRICK - 1)) / g.nw); - sh_grid[tid] = make_float3(fx, fy, fz); - sh_centre[tid] = make_float3(g.orth_n[0] * fx + g.orth_n[1] * fy + g.orth_n[2] * fz, - g.orth_n[3] * fx + g.orth_n[4] * fy + g.orth_n[5] * fz, - g.orth_n[6] * fx + g.orth_n[7] * fy + g.orth_n[8] * fz); - } - __syncthreads(); - for (int i = tid; i < m * WORDS; i += BRICK * BRICK * BRICK) - reinterpret_cast(sh_atom)[i] = reinterpret_cast(atoms + sh_index[i / WORDS])[i % WORDS]; - __syncthreads(); - if (!inside) - continue; - for (int k = 0; k < m; k++) { - const RigidBodyGPUAtom &at = sh_atom[k]; - float x = sh_centre[k].x - tx, y = sh_centre[k].y - ty, z = sh_centre[k].z - tz; - if (g.narrow) { - float fx = sh_grid[k].x - threadIdx.x, fy = sh_grid[k].y - threadIdx.y, fz = sh_grid[k].z - threadIdx.z; - fx -= g.nu * rintf(fx / g.nu); - fy -= g.nv * rintf(fy / g.nv); - fz -= g.nw * rintf(fz / g.nw); - x = g.orth_n[0] * fx + g.orth_n[1] * fy + g.orth_n[2] * fz; - y = g.orth_n[3] * fx + g.orth_n[4] * fy + g.orth_n[5] * fz; - z = g.orth_n[6] * fx + g.orth_n[7] * fy + g.orth_n[8] * fz; - } - const float r2 = x * x + y * y + z * z; - if (r2 > at.radius * at.radius) - continue; - float density = 0.0f; - if (!at.aniso) { - for (int q = 0; q < 5; q++) - density += at.a[q] * expapprox(fmaxf(at.b[q][0] * r2, -88.f)); - } else { - for (int q = 0; q < 5; q++) { - const float *b = at.b[q]; - 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(fmaxf(rur, -88.f)); - } - } - acc += at.occ * density; - } - } - if (inside) - grid[u + static_cast(g.nu) * (v + static_cast(g.nv) * w)] = acc; -} - // F at the zone's rows composed from the transform of one copy (SymmetryComposition::Compose), with // dF/dt where `df_dt` is not null. `f1` is cuFFT's forward transform, which carries exp(-2 pi i h.x): // conjugated and scaled by V/N it is gemmi's F1. @@ -445,17 +274,6 @@ int blocks(size_t n) { return static_cast((n + THREADS - 1) / THREADS); } -// cub's temporary storage for the pair scan and sort at the capacity. -size_t CubBytes(size_t atoms, size_t pairs) { - size_t scan = 0, sort = 0; - cuda_err(cub::DeviceScan::ExclusiveSum(nullptr, scan, static_cast(nullptr), static_cast(nullptr), - static_cast(atoms + 1))); - cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, sort, static_cast(nullptr), static_cast(nullptr), - static_cast(nullptr), static_cast(nullptr), - static_cast(pairs))); - return std::max(scan, sort); -} - // A plan for an (nu, nv, nw) R2C transform, u fastest, `batch` grids back to back, whose work area the // caller provides; its size in `work`. cufftHandle MakePlan(int nu, int nv, int nw, int batch, size_t &work) { @@ -471,21 +289,6 @@ cufftHandle MakePlan(int nu, int nv, int nw, int batch, size_t &work) { return plan; } -// gemmi's box around each atom (MakeAtomBox, ModelGrid.cpp): the points within ceil(radius / spacing) -// of the nearest one along each axis, spacing being the distance between the grid's lattice planes. -std::vector AtomBoxes(const RigidBodyGPUZone &zone) { - const int n3[3] = {zone.nu, zone.nv, zone.nw}; - double spacing[3]; - for (int k = 0; k < 3; k++) - spacing[k] = 1.0 / (n3[k] * std::sqrt(zone.frac[3 * k] * zone.frac[3 * k] + zone.frac[3 * k + 1] * zone.frac[3 * k + 1] + - zone.frac[3 * k + 2] * zone.frac[3 * k + 2])); - std::vector box(3 * zone.atoms.size()); - for (size_t i = 0; i < zone.atoms.size(); i++) - for (int k = 0; k < 3; k++) - box[3 * i + k] = static_cast(std::ceil(zone.atoms[i].radius / spacing[k])); - return box; -} - int GreatestStreamPriority() { int least = 0, greatest = 0; cuda_err(cudaDeviceGetStreamPriorityRange(&least, &greatest)); @@ -494,12 +297,19 @@ int GreatestStreamPriority() { } // namespace -// The most bricks the 2 d + 1 wrapped points of a box can fall in along an axis of n points: one more -// than the points span for where they start in a brick, and one more again where they wrap past a last -// brick that is partial (n = 17: the points 15, 16, 0 fall in bricks 1, 2 and 0). +// The zone's grid as the density gather takes it. +static ModelDensityGrid DensityGrid(const RigidBodyGPUZone &zone) { + ModelDensityGrid g; + g.nu = zone.nu; + g.nv = zone.nv; + g.nw = zone.nw; + std::copy(zone.orth, zone.orth + 9, g.orth); + std::copy(zone.frac, zone.frac + 9, g.frac); + return g; +} + size_t RigidBodyGPUEngine::AxisBrickBound(int d, int n) { - const int nb = (n + BRICK - 1) / BRICK; - return std::min(nb, (2 * d + 1 + BRICK - 1) / BRICK + 2); + return ModelDensityGPU::AxisBrickBound(d, n); } struct RigidBodyGPUEngineImpl { @@ -509,8 +319,6 @@ struct RigidBodyGPUEngineImpl { // otherwise fill the card ahead of a fit's short ones, one evaluation at a time. CudaStream stream{cudaStreamNonBlocking, GreatestStreamPriority()}; - CudaDevicePtr atoms; - CudaDevicePtr box; CudaDevicePtr rel; CudaDevicePtr pos; CudaDevicePtr mask_slot; @@ -521,11 +329,6 @@ struct RigidBodyGPUEngineImpl { CudaDevicePtr spectrum; // their 3 transforms CudaDevicePtr fft_work; - CudaDevicePtr overflow; // set by count_pairs_kernel where an atom's box is over MAX_BRICKS_PER_AXIS - CudaDevicePtr count, offset, key, value, key_sorted, value_sorted, brick_start, brick_end; - CudaDevicePtr cub_temp; - size_t cub_bytes = 0; - CudaDevicePtr terms; CudaDevicePtr row_hkl; CudaDevicePtr row_scale, row_stol2; @@ -539,6 +342,7 @@ struct RigidBodyGPUEngineImpl { CudaDevicePtr products, x; std::map, cufftHandle> plans; + std::unique_ptr density; std::unique_ptr mask; std::unique_ptr scale; @@ -555,8 +359,6 @@ struct RigidBodyGPUEngineImpl { : device(dev), cap(c) { const auto sync = CudaAlloc::Synchronous; const size_t na = std::max(cap.atoms, 1); - atoms = CudaDevicePtr(na, sync); - box = CudaDevicePtr(3 * na, sync); rel = CudaDevicePtr(3 * na, sync); pos = CudaDevicePtr(na, sync); mask_slot = CudaDevicePtr(na, sync); @@ -565,18 +367,6 @@ struct RigidBodyGPUEngineImpl { grid = CudaDevicePtr(3 * cap.grid_points, sync); spectrum = CudaDevicePtr(3 * cap.complex_points, sync); fft_work = CudaDevicePtr(std::max(cap.fft_work_bytes, 1), sync); - overflow = CudaDevicePtr(1, sync); - count = CudaDevicePtr(na + 1, sync); - offset = CudaDevicePtr(na + 1, sync); - const size_t np = std::max(cap.pairs, 1); - key = CudaDevicePtr(np, sync); - value = CudaDevicePtr(np, sync); - key_sorted = CudaDevicePtr(np, sync); - value_sorted = CudaDevicePtr(np, sync); - brick_start = CudaDevicePtr(std::max(cap.bricks, 1), sync); - brick_end = CudaDevicePtr(std::max(cap.bricks, 1), sync); - cub_bytes = CubBytes(na, np); - cub_temp = CudaDevicePtr(std::max(cub_bytes, 1), sync); const size_t nr = std::max(cap.rows, 1), no = std::max(cap.observations, 1); terms = CudaDevicePtr(std::max(cap.terms, 1), sync); row_hkl = CudaDevicePtr(3 * nr, sync); @@ -596,6 +386,7 @@ struct RigidBodyGPUEngineImpl { point_fmask = CudaDevicePtr(no, sync); products = CudaDevicePtr(MAX_SCALE_PARAMS * (MAX_SCALE_PARAMS + 6), sync); x = CudaDevicePtr(MAX_SCALE_PARAMS * 6, sync); + density = std::make_unique(stream, cap.atoms, cap.pairs, cap.bricks); mask = std::make_unique(stream, std::max(cap.grid_points, 1)); scale = std::make_unique(stream, no); } @@ -646,42 +437,7 @@ struct RigidBodyGPUEngineImpl { place_kernel<<>>(rel, n_atoms, p, geom.nu, geom.nv, geom.nw, mask_slot, mask_radius, pos, mask_atoms); cuda_err(cudaGetLastError()); - cuda_err(cudaMemsetAsync(overflow, 0, sizeof(int), stream)); - count_pairs_kernel<<>>(pos, box, n_atoms, geom.nu, geom.nv, geom.nw, - count, overflow); - cuda_err(cudaGetLastError()); - size_t temp = cub_bytes; - cuda_err(cub::DeviceScan::ExclusiveSum(cub_temp.get(), temp, count.get(), offset.get(), n_atoms + 1, stream)); - int npairs = 0; - int over = 0; - cuda_err(cudaMemcpyAsync(&npairs, offset.get() + n_atoms, sizeof(int), cudaMemcpyDeviceToHost, stream)); - cuda_err(cudaMemcpyAsync(&over, overflow.get(), sizeof(int), cudaMemcpyDeviceToHost, stream)); - cuda_err(cudaStreamSynchronize(stream)); - if (over) - throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: an atom's box over the bricks per axis"); - if (static_cast(npairs) > cap.pairs) - throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: gather pairs over the reserve"); - fill_pairs_kernel<<>>(pos, box, n_atoms, geom.nu, geom.nv, geom.nw, - offset, key, value); - cuda_err(cudaGetLastError()); - const int nbu = (geom.nu + BRICK - 1) / BRICK, nbv = (geom.nv + BRICK - 1) / BRICK, - nbw = (geom.nw + BRICK - 1) / BRICK; - const int nb = nbu * nbv * nbw; - int bits = 1; - while ((1 << bits) < nb) - bits++; - temp = cub_bytes; - cuda_err(cub::DeviceRadixSort::SortPairs(cub_temp.get(), temp, key.get(), key_sorted.get(), value.get(), - value_sorted.get(), npairs, 0, bits, stream)); - cuda_err(cudaMemsetAsync(brick_start, 0, nb * sizeof(int), stream)); - cuda_err(cudaMemsetAsync(brick_end, 0, nb * sizeof(int), stream)); - if (npairs > 0) { - brick_ranges_kernel<<>>(key_sorted, npairs, brick_start, brick_end); - cuda_err(cudaGetLastError()); - } - gather_kernel<<>>(atoms, pos, value_sorted, brick_start, brick_end, - geom, grid.get() + slot * GridPoints()); - cuda_err(cudaGetLastError()); + density->Compute(pos, grid.get() + slot * GridPoints()); } void Compose(int slot, double2 *out, double2 *derivative) { @@ -706,29 +462,15 @@ struct RigidBodyGPUEngineImpl { }; size_t RigidBodyGPUEngine::Bricks(int nu, int nv, int nw) { - return static_cast((nu + BRICK - 1) / BRICK) * ((nv + BRICK - 1) / BRICK) * ((nw + BRICK - 1) / BRICK); + return ModelDensityGPU::Bricks(nu, nv, nw); } size_t RigidBodyGPUEngine::PairBound(const RigidBodyGPUZone &zone) { - const std::vector box = AtomBoxes(zone); - const int n3[3] = {zone.nu, zone.nv, zone.nw}; - size_t pairs = 0; - for (size_t i = 0; i < zone.atoms.size(); i++) { - size_t product = 1; - for (int k = 0; k < 3; k++) - product *= AxisBrickBound(box[3 * i + k], n3[k]); - pairs += product; - } - return pairs; + return ModelDensityGPU::PairBound(DensityGrid(zone), zone.atoms); } bool RigidBodyGPUEngine::Supports(const RigidBodyGPUZone &zone) { - const std::vector box = AtomBoxes(zone); - const int n3[3] = {zone.nu, zone.nv, zone.nw}; - for (size_t i = 0; i < box.size(); i++) - if (2 * box[i] + 1 > n3[i % 3] || AxisBrickBound(box[i], n3[i % 3]) > MAX_BRICKS_PER_AXIS) - return false; - return true; + return ModelDensityGPU::Supports(DensityGrid(zone), zone.atoms); } void RigidBodyGPUEngine::MemoryInfo(size_t &free, size_t &total) { @@ -754,14 +496,12 @@ size_t RigidBodyGPUEngine::FFTWorkBytes(int nu, int nv, int nw) { size_t RigidBodyGPUEngine::DeviceBytes(const RigidBodyGPUCapacity &c) { const size_t na = c.atoms + 1, no = c.observations + 1, nr = c.rows + 1; - size_t bytes = na * (sizeof(RigidBodyGPUAtom) + 3 * sizeof(int) + 7 * sizeof(double) + sizeof(float4) + - sizeof(int) + sizeof(float) + 2 * sizeof(int)); + size_t bytes = na * (7 * sizeof(double) + sizeof(float4) + sizeof(int) + sizeof(float)); bytes += 3 * c.grid_points * sizeof(float) + 3 * c.complex_points * sizeof(float2) + c.fft_work_bytes; - bytes += 4 * c.pairs * sizeof(int) + 2 * c.bricks * sizeof(int); + bytes += ModelDensityGPU::DeviceBytes(c.atoms, c.pairs, c.bricks); bytes += c.terms * sizeof(RigidBodyGPUTerm); bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 7 * sizeof(double2) + sizeof(float2)); bytes += no * (2 * sizeof(int) + sizeof(float) + (1 + 6 + MAX_SCALE_PARAMS) * sizeof(double) + 2 * sizeof(float2)); - bytes += CubBytes(na, c.pairs + 1); bytes += ModelMaskGPU::DeviceBytes(c.grid_points) + ModelScaleGPU::DeviceBytes(no); return bytes; } @@ -798,9 +538,6 @@ void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) { e.geom.nu = zone.nu; e.geom.nv = zone.nv; e.geom.nw = zone.nw; - for (int i = 0; i < 3; i++) - for (int j = 0; j < 3; j++) - e.geom.orth_n[3 * i + j] = static_cast(zone.orth[3 * i + j] / (j == 0 ? zone.nu : j == 1 ? zone.nv : zone.nw)); std::copy(zone.frac, zone.frac + 9, e.frac); e.volume = zone.volume; e.rows = zone.row_hkl.size(); @@ -812,16 +549,6 @@ void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) { e.constraints = zone.constraints; e.f_mean = zone.f_mean; - // An atom whose box is wider than the cell would reach a point through more than one image, which - // gemmi's box walk adds and the gather's nearest image does not; such a cell is left to the CPU. - if (!Supports(zone)) - throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: an atom wider than the cell"); - const std::vector box = AtomBoxes(zone); - e.geom.narrow = 0; - for (size_t i = 0; i < box.size(); i++) - if (2 * box[i] + BRICK > (i % 3 == 0 ? zone.nu : i % 3 == 1 ? zone.nv : zone.nw)) - e.geom.narrow = 1; - std::vector slot(zone.atoms.size(), -1); std::vector radius(zone.atoms.size(), 0.0f); for (size_t j = 0; j < zone.mask_atom.size(); j++) { @@ -833,8 +560,6 @@ void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) { if (bytes > 0) cuda_err(cudaMemcpyAsync(dst, src, bytes, cudaMemcpyHostToDevice, e.stream)); }; - up(e.atoms, zone.atoms.data(), zone.atoms.size() * sizeof(RigidBodyGPUAtom)); - up(e.box, box.data(), box.size() * sizeof(int)); up(e.mask_slot, slot.data(), slot.size() * sizeof(int)); up(e.mask_radius, radius.data(), radius.size() * sizeof(float)); up(e.terms, zone.terms.data(), zone.terms.size() * sizeof(RigidBodyGPUTerm)); @@ -856,6 +581,7 @@ void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) { std::copy(zone.images[i].begin() + 9, zone.images[i].end(), images[i].tran); } e.mask->SetGrid(mask_grid, images); + e.density->SetAtoms(DensityGrid(zone), zone.atoms); e.scale->SetPoints(zone.point_hkl, zone.point_stol2, zone.point_fobs, zone.point_sigma, zone.constraints, zone.frac); e.Plan(1); e.Plan(3); diff --git a/image_analysis/structure_refinement/RigidBodyGPU.h b/image_analysis/structure_refinement/RigidBodyGPU.h index 7f2651fb5..bf312db45 100644 --- a/image_analysis/structure_refinement/RigidBodyGPU.h +++ b/image_analysis/structure_refinement/RigidBodyGPU.h @@ -24,7 +24,8 @@ class Logger; class RigidBodyGPUEngine; struct RigidBodyGPUZone; -// A CUDA failure inside the rigid body. Model validation catches it and starts again on the CPU. +// A CUDA failure inside the rigid body, or in model validation's structure factors and maps +// (ModelStructureFactorsGPU). Model validation catches it and starts again on the CPU. class RigidBodyGPUFailure : public std::runtime_error { public: explicit RigidBodyGPUFailure(const std::string &what) : std::runtime_error(what) {} diff --git a/image_analysis/structure_refinement/RigidBodyGPUEngine.h b/image_analysis/structure_refinement/RigidBodyGPUEngine.h index 9b93f1554..9a1919858 100644 --- a/image_analysis/structure_refinement/RigidBodyGPUEngine.h +++ b/image_analysis/structure_refinement/RigidBodyGPUEngine.h @@ -15,15 +15,7 @@ #include #include -// One atom's density on one zone's grid, as PutModelDensityOnGrid() (ModelGrid.cpp) sets it up: gemmi's -// precalculated five-Gaussian sum and the radius it cuts the sum at. -struct RigidBodyGPUAtom { - float a[5]; - float b[5][6]; // isotropic: b[k][0] multiplies r^2; anisotropic: the matrix, u11 u22 u33 u12 u13 u23 - float occ; - float radius; - int aniso; -}; +#include "ModelDensityGPU.h" // ModelDensityAtom // One term of SymmetryComposition: F1 at k = hR, times the phase of the operator's translation. struct RigidBodyGPUTerm { @@ -39,7 +31,7 @@ struct RigidBodyGPUZone { double orth[9] = {}, frac[9] = {}; // row-major double volume = 0; double blur = 0; // DensityCalculator's, which the rows' unblur undoes - std::vector atoms; // model order + std::vector atoms; // model order // The bulk-solvent mask's atoms: the model's index of each, and its radius (probe included). std::vector mask_atom; diff --git a/tests/ModelValidationTest.cpp b/tests/ModelValidationTest.cpp index baf75aff2..a7e0eb588 100644 --- a/tests/ModelValidationTest.cpp +++ b/tests/ModelValidationTest.cpp @@ -23,6 +23,7 @@ #include "../image_analysis/structure_refinement/ModelValidation.h" #include "../image_analysis/structure_refinement/RigidBodyRefine.h" #ifdef JFJOCH_USE_CUDA +#include "../image_analysis/structure_refinement/ModelStructureFactorsGPU.h" #include "../image_analysis/structure_refinement/RigidBodyGPU.h" #include "../image_analysis/structure_refinement/RigidBodyGPUEngine.h" #include "../common/CUDAWrapper.h" @@ -1414,4 +1415,98 @@ TEST_CASE("RigidBodyGPU_Deterministic", "[ModelValidation][gpu]") { } } +// The model's structure factors on the GPU against the CPU path model validation takes - gemmi's density and +// mask on the symmetrized grid, transformed by FFTW and read off by prepare_asu_data() - in the five groups, +// at a resolution where the mask's shrink step does nothing and at one where it does. The same reflections +// in the same order; F_calc to rounding; F_mask to the few grid points at an atom's mask radius that float +// and double distances put on different sides. The same atoms twice give the same bits. +TEST_CASE("ModelStructureFactorsGPU_MatchesCPU", "[ModelValidation][gpu]") { + if (get_gpu_count() == 0) + SKIP("No GPU"); + for (const char *cryst : kRigidBodyCrysts) { + const gemmi::Structure st = AnisoCluster(cryst); + const gemmi::Model &model = st.models[0]; + const gemmi::SpaceGroup *sg = st.find_spacegroup(); + for (double d_min : {3.5, 1.5}) { + auto dc = ZoneDensity(st, model, d_min); + dc.put_model_density_on_grid(model); + const auto ref_fc = MapToFPhi(dc.grid).prepare_asu_data(d_min, dc.blur, false, false, false); + gemmi::Grid mask; + mask.unit_cell = st.cell; + mask.spacegroup = sg; + mask.set_size_from_spacing(dc.requested_grid_spacing(), gemmi::GridSizeRounding::Up); + gemmi::SolventMasker(gemmi::AtomicRadiiSet::Refmac).put_mask_on_grid(mask, model); + const auto ref_fm = MapToFPhi(mask).prepare_asu_data(d_min, 0); + + ModelStructureFactorsGPU gpu(0, model, st.cell, *sg, d_min); + REQUIRE(gpu.Supported()); + gpu.Reserve(); + gemmi::AsuData> fc, fm, fc2, fm2; + gpu.Compute(model, fc, fm); + gpu.Compute(model, fc2, fm2); + REQUIRE(fc.v.size() == ref_fc.v.size()); + REQUIRE(fm.v.size() == ref_fm.v.size()); + double fc_mean = 0, fc_worst = 0, fm_norm = 0, fm_diff = 0; + bool repeat_same = true; + for (size_t i = 0; i < fc.v.size(); i++) { + REQUIRE(fc.v[i].hkl == ref_fc.v[i].hkl); + REQUIRE(fm.v[i].hkl == ref_fm.v[i].hkl); + fc_mean += std::abs(ref_fc.v[i].value) / static_cast(fc.v.size()); + fc_worst = std::max(fc_worst, static_cast(std::abs(fc.v[i].value - ref_fc.v[i].value))); + fm_norm += std::norm(ref_fm.v[i].value); + fm_diff += std::norm(fm.v[i].value - ref_fm.v[i].value); + repeat_same = repeat_same && fc2.v[i].value == fc.v[i].value && fm2.v[i].value == fm.v[i].value; + } + INFO(cryst << " at " << d_min << " A: F_calc worst " << fc_worst / fc_mean << " of the mean |F|, F_mask " + << std::sqrt(fm_diff / fm_norm) << " relative rms"); + CHECK(fc_worst <= 1e-4 * fc_mean); + CHECK(std::sqrt(fm_diff) <= 1e-3 * std::sqrt(fm_norm)); + CHECK(repeat_same); + } + } +} + +// A map from coefficients on the GPU against MapFromFPhi() of gemmi's get_f_phi_on_grid(): the same grid, +// the same values to rounding, the same bits twice. +TEST_CASE("ModelStructureFactorsGPU_MapMatchesCPU", "[ModelValidation][gpu]") { + if (get_gpu_count() == 0) + SKIP("No GPU"); + for (const char *cryst : kRigidBodyCrysts) { + const gemmi::Structure st = AnisoCluster(cryst); + const gemmi::Model &model = st.models[0]; + const gemmi::SpaceGroup *sg = st.find_spacegroup(); + const double d_min = 2.0; + ModelStructureFactorsGPU gpu(0, model, st.cell, *sg, d_min); + REQUIRE(gpu.Supported()); + gpu.Reserve(); + gemmi::AsuData> coef, fm; + gpu.Compute(model, coef, fm); + // Coefficients on every second reflection only, as a map's are on the observed ones. + gemmi::AsuData> some = coef; + some.v.clear(); + for (size_t i = 0; i < coef.v.size(); i += 2) + some.v.push_back(coef.v[i]); + + gemmi::AsuData> cpu_coef = some; + cpu_coef.ensure_sorted(); + const auto size = gemmi::get_size_for_hkl(cpu_coef, {{0, 0, 0}}, 3.0); + const gemmi::Grid cpu = MapFromFPhi(gemmi::get_f_phi_on_grid(cpu_coef, size, true)); + gemmi::AsuData> gpu_coef = some; + const gemmi::Grid map = gpu.Map(gpu_coef); + const gemmi::Grid map2 = gpu.Map(gpu_coef); + REQUIRE(map.nu == cpu.nu); + REQUIRE(map.nv == cpu.nv); + REQUIRE(map.nw == cpu.nw); + double rms = 0, worst = 0; + for (size_t i = 0; i < cpu.data.size(); i++) { + rms += gemmi::sq(cpu.data[i]) / static_cast(cpu.data.size()); + worst = std::max(worst, static_cast(std::fabs(map.data[i] - cpu.data[i]))); + } + INFO(cryst << ": map differs by at most " << worst / std::sqrt(rms) << " of its rms"); + CHECK(worst <= 1e-4 * std::sqrt(rms)); + CHECK(std::memcmp(map.data.data(), map2.data.data(), map.data.size() * sizeof(float)) == 0); + } +} + #endif +