leonarski_fandClaude Opus 5.5
b6d62d7309
axis_bricks() walked the 2 d + 1 points of a box and looked each one's brick up among those found so far. The wrapped points are one run from imod(c - d) up and, where they pass the cell's end, a second from 0 (2 d + 1 <= n, Supports()), so the bricks are now taken run by run, in the order the points reach them as before. The same bricks, counted and filled the same way: the grid is unchanged. RefinementGPU_Benchmark on an RTX 5080: count_bricks_kernel 19 -> 10 us, fill_bricks_kernel 32 -> 26 us. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37
457 lines
22 KiB
Plaintext
457 lines
22 KiB
Plaintext
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "ModelDensityGPU.h"
|
|
|
|
#include <algorithm>
|
|
#include <climits>
|
|
#include <cmath>
|
|
|
|
#include <cub/cub.cuh>
|
|
|
|
#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 GATHER_THREADS = BRICK * BRICK * BRICK / 2; // two points of the brick a thread
|
|
constexpr int MAX_BRICKS_PER_AXIS = 16; // an atom's box touches at most this many bricks along an axis
|
|
constexpr int THREADS = 256;
|
|
constexpr int SORT_ATOMS = 4096; // a brick's atoms are sorted in shared memory up to this many
|
|
|
|
__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<int>(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.
|
|
// The wrapped points are one run from imod(c - d) up, and a second from 0 where they pass the cell's end
|
|
// (2 d + 1 <= n, Supports()); the bricks are listed in the order the points reach them.
|
|
__device__ int axis_bricks(int c, int d, int n, int *out) {
|
|
const int lo = imod(c - d, n), hi = lo + 2 * d;
|
|
const int run_first[2] = {lo, 0}, run_last[2] = {min(hi, n - 1), hi - n};
|
|
int k = 0;
|
|
for (int r = 0; r < 2; r++)
|
|
for (int b = run_first[r] / BRICK; run_last[r] >= run_first[r] && b <= run_last[r] / BRICK; b++) {
|
|
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);
|
|
}
|
|
|
|
// How many atoms touch each brick: integer atomics, whose sum does not depend on their order.
|
|
__global__ void count_bricks_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw, int *brick_count) {
|
|
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;
|
|
for (int k = 0; k < c; k++)
|
|
for (int j = 0; j < b; j++)
|
|
for (int l = 0; l < a; l++)
|
|
atomicAdd(brick_count + bu[l] + nbu * (bv[j] + nbv * bw[k]), 1);
|
|
}
|
|
|
|
// Each atom into the segments of the bricks it touches, at whatever place in the segment the atomic hands out:
|
|
// the order within a brick is settled by order_bricks_kernel.
|
|
__global__ void fill_bricks_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw,
|
|
const int *brick_start, int *brick_fill, int *pair_atom) {
|
|
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;
|
|
for (int k = 0; k < c; k++)
|
|
for (int j = 0; j < b; j++)
|
|
for (int l = 0; l < a; l++) {
|
|
const int brick = bu[l] + nbu * (bv[j] + nbv * bw[k]);
|
|
pair_atom[brick_start[brick] + atomicAdd(brick_fill + brick, 1)] = i;
|
|
}
|
|
}
|
|
|
|
// Each brick's atoms into model order, one block per brick: a bitonic sort of the segment in shared memory,
|
|
// padded to a power of two, for the bricks of up to SORT_ATOMS atoms. A larger brick is ranked instead: an
|
|
// atom's place is the number of the brick's atoms with a lower index (an atom is in a brick once), counted
|
|
// against the segment staged a tile at a time, a pair staying with one thread through the tiles.
|
|
__global__ void order_bricks_kernel(const int *unordered, const int *brick_start, int *rank, int *pair_atom) {
|
|
const int brick = blockIdx.x;
|
|
const int s = brick_start[brick], e = brick_start[brick + 1], k = e - s;
|
|
__shared__ int keys[SORT_ATOMS];
|
|
if (k <= SORT_ATOMS) {
|
|
int n = 1;
|
|
while (n < k)
|
|
n <<= 1;
|
|
for (int t = threadIdx.x; t < n; t += THREADS)
|
|
keys[t] = t < k ? unordered[s + t] : INT_MAX;
|
|
__syncthreads();
|
|
for (int size = 2; size <= n; size <<= 1)
|
|
for (int stride = size >> 1; stride > 0; stride >>= 1) {
|
|
for (int t = threadIdx.x; t < n / 2; t += THREADS) {
|
|
const int i = 2 * t - (t & (stride - 1)), j = i + stride;
|
|
const int a = keys[i], b = keys[j];
|
|
if ((a > b) == ((i & size) == 0)) {
|
|
keys[i] = b;
|
|
keys[j] = a;
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
for (int t = threadIdx.x; t < k; t += THREADS)
|
|
pair_atom[s + t] = keys[t];
|
|
return;
|
|
}
|
|
for (int c = s; c < e; c += THREADS) {
|
|
const int m = min(THREADS, e - c);
|
|
__syncthreads();
|
|
if (threadIdx.x < m)
|
|
keys[threadIdx.x] = unordered[c + threadIdx.x];
|
|
__syncthreads();
|
|
for (int p = s + threadIdx.x; p < e; p += THREADS) {
|
|
const int atom = unordered[p];
|
|
int below = 0;
|
|
for (int q = 0; q < m; q++)
|
|
below += keys[q] < atom;
|
|
rank[p] = (c == s ? 0 : rank[p]) + below;
|
|
}
|
|
}
|
|
for (int p = s + threadIdx.x; p < e; p += THREADS)
|
|
pair_atom[s + rank[p]] = unordered[p];
|
|
}
|
|
|
|
// What the gather needs of one staged atom: its isotropic Gaussians (an anisotropic atom's matrices are read
|
|
// from the atoms in global memory), and its box as the brick sees it - a bit per point of the brick along each
|
|
// axis for whether the point is in the box, and for the points that are, what gemmi's box offset makes of them:
|
|
// the point's unwrapped box index is base + t, less n where t reaches wrap.
|
|
struct GatherAtom {
|
|
float a[5], b[5];
|
|
float occ, max_r2;
|
|
int index, aniso;
|
|
int masks; // bits 0-7: u, 8-15: v, 16-23: w
|
|
int base[3], wrap[3];
|
|
float nctr[3];
|
|
};
|
|
|
|
// The point p's offset in the box along an axis of n points, 0 to span where it is in it.
|
|
__device__ __forceinline__ int box_offset(int p, int first, int n) {
|
|
const int m = p - first;
|
|
return m < 0 ? m + n : m;
|
|
}
|
|
|
|
// One atom's density at a point of its box, f = centre - the point in the box's own (unwrapped) indices, as the
|
|
// CPU walks it; false where the point is beyond the radius.
|
|
__device__ __forceinline__ bool density_at(const GatherAtom &at, const ModelDensityAtom *atoms,
|
|
const ModelDensityGPU::Geometry &g, float fx, float fy, float fz,
|
|
float &density) {
|
|
const float z = g.m22 * fz;
|
|
const float y = g.m11 * fy + g.m12 * fz;
|
|
const float yz = y * y + z * z;
|
|
const float x = g.m00 * fx + (g.m01 * fy + g.m02 * fz);
|
|
const float r2 = yz + x * x;
|
|
if (r2 > at.max_r2)
|
|
return false;
|
|
density = 0.0f;
|
|
if (!at.aniso) {
|
|
for (int q = 0; q < 5; q++)
|
|
density += at.a[q] * expapprox(fmaxf(at.b[q] * r2, -88.f));
|
|
} else {
|
|
for (int q = 0; q < 5; q++) {
|
|
const float *bq = atoms[at.index].b[q];
|
|
// gemmi's SMat33::r_u_r(), in its order
|
|
const float rur = x * x * bq[0] + y * y * bq[1] + z * z * bq[2] +
|
|
2 * (x * y * bq[3] + x * z * bq[4] + y * z * bq[5]);
|
|
density += at.a[q] * expapprox(fmaxf(rur, -88.f));
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
// The density of the model on the grid, as UseDensityBoxesByPlane() and AtomDensity (ModelGrid.cpp) put
|
|
// it there, bit for bit: every point adds, in model order, each atom whose box holds it and whose radius
|
|
// reaches it, with the CPU's single-precision arithmetic in the CPU's order (no a*b+c is contracted:
|
|
// CMakeLists.txt compiles device code with --fmad=false). One block per brick; the brick's atoms are
|
|
// staged in shared memory. Supports() leaves only boxes narrower than the cell, so a point is in an atom's
|
|
// box at most once, at one index of the box's own unwrapped range.
|
|
// A thread takes two points, one above the other along w, and a warp a 4 x 4 x 4 corner of the brick; of each
|
|
// staged tile a warp takes only the atoms whose box reaches into its corner, in their order: the others hold none
|
|
// of its points.
|
|
__global__ void gather_kernel(const ModelDensityAtom *atoms, const float4 *pos, const int *box, 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 origin[3] = {(brick % nbu) * BRICK, (brick / nbu % nbv) * BRICK, (brick / (nbu * nbv)) * BRICK};
|
|
const int tid = threadIdx.x, lane = tid % 32, warp = tid / 32;
|
|
// the warp's corner starts at (cu, cv, cw) in the brick, the lane's points in it at (lane % 4, lane / 4 % 4) and
|
|
// w = 2 (lane / 16) and the one above
|
|
const int cu = (warp & 1) * 4, cv = (warp >> 1 & 1) * 4, cw = (warp >> 2) * 4;
|
|
const int tu = cu + lane % 4, tv = cv + lane / 4 % 4, tw = cw + 2 * (lane / 16);
|
|
const int u = origin[0] + tu, v = origin[1] + tv, w = origin[2] + tw;
|
|
const bool inside = u < g.nu && v < g.nv && w < g.nw; // the point above may still be outside
|
|
const int n3[3] = {g.nu, g.nv, g.nw};
|
|
|
|
__shared__ GatherAtom sh[GATHER_TILE];
|
|
__shared__ unsigned char reach[GATHER_THREADS / 32][GATHER_TILE]; // per warp, the tile's atoms it takes
|
|
float acc[2] = {0.0f, 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];
|
|
const ModelDensityAtom &at = atoms[a];
|
|
const float4 p = pos[a];
|
|
const float nctr[3] = {p.x, p.y, p.z};
|
|
GatherAtom &r = sh[tid];
|
|
for (int q = 0; q < 5; q++) {
|
|
r.a[q] = at.a[q];
|
|
r.b[q] = at.b[q][0];
|
|
}
|
|
r.occ = at.occ;
|
|
r.max_r2 = at.max_r2;
|
|
r.index = a;
|
|
r.aniso = at.aniso;
|
|
r.masks = 0;
|
|
for (int k = 0; k < 3; k++) {
|
|
const int d = box[3 * a + k];
|
|
const int lo = __float2int_rn(nctr[k]) - d, first = imod(lo, n3[k]), span = 2 * d;
|
|
const int m0 = box_offset(origin[k], first, n3[k]);
|
|
for (int t = 0; t < BRICK; t++) {
|
|
const int mt = m0 + t;
|
|
if ((mt < n3[k] ? mt : mt - n3[k]) <= span)
|
|
r.masks |= 1 << (8 * k + t);
|
|
}
|
|
r.base[k] = lo + m0;
|
|
r.wrap[k] = n3[k] - m0;
|
|
r.nctr[k] = nctr[k];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
int taken = 0;
|
|
for (int k0 = 0; k0 < m; k0 += 32) {
|
|
const int k = k0 + lane;
|
|
const int masks = k < m ? sh[k].masks : 0;
|
|
const bool hit = (masks >> cu & 0xF) && (masks >> (8 + cv) & 0xF) && (masks >> (16 + cw) & 0xF);
|
|
const unsigned ballot = __ballot_sync(0xffffffffu, hit);
|
|
if (hit)
|
|
reach[warp][taken + __popc(ballot & ((1u << lane) - 1))] = static_cast<unsigned char>(k);
|
|
taken += __popc(ballot);
|
|
}
|
|
__syncwarp();
|
|
if (!inside)
|
|
continue;
|
|
for (int j = 0; j < taken; j++) {
|
|
const GatherAtom &at = sh[reach[warp][j]];
|
|
const int in_w = at.masks >> (16 + tw) & 3; // bit 0: the lower point, bit 1: the one above
|
|
if (!((at.masks >> tu) & (at.masks >> (8 + tv)) & 1) || in_w == 0)
|
|
continue;
|
|
const int pu = at.base[0] + tu - (tu >= at.wrap[0] ? g.nu : 0);
|
|
const int pv = at.base[1] + tv - (tv >= at.wrap[1] ? g.nv : 0);
|
|
const float fx = at.nctr[0] - static_cast<float>(pu);
|
|
const float fy = at.nctr[1] - static_cast<float>(pv);
|
|
for (int k = 0; k < 2; k++) {
|
|
if (!(in_w >> k & 1))
|
|
continue;
|
|
const int pw = at.base[2] + tw + k - (tw + k >= at.wrap[2] ? g.nw : 0);
|
|
float density;
|
|
if (density_at(at, atoms, g, fx, fy, at.nctr[2] - static_cast<float>(pw), density))
|
|
acc[k] += at.occ * density;
|
|
}
|
|
}
|
|
}
|
|
for (int k = 0; k < 2; k++)
|
|
if (inside && w + k < g.nw)
|
|
grid[u + static_cast<size_t>(g.nu) * (v + static_cast<size_t>(g.nv) * (w + k))] = acc[k];
|
|
}
|
|
|
|
int blocks(size_t n) {
|
|
return static_cast<int>((n + THREADS - 1) / THREADS);
|
|
}
|
|
|
|
// cub's temporary storage for the scan over the bricks.
|
|
size_t CubBytes(size_t bricks) {
|
|
size_t scan = 0;
|
|
cuda_err(cub::DeviceScan::ExclusiveSum(nullptr, scan, static_cast<int *>(nullptr), static_cast<int *>(nullptr),
|
|
static_cast<int>(bricks + 1)));
|
|
return scan;
|
|
}
|
|
|
|
} // 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<size_t>((nu + BRICK - 1) / BRICK) * ((nv + BRICK - 1) / BRICK) * ((nw + BRICK - 1) / BRICK);
|
|
}
|
|
|
|
void ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const ModelDensityAtom *atoms, size_t n, int *box) {
|
|
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]));
|
|
for (size_t i = 0; i < n; i++)
|
|
for (int k = 0; k < 3; k++)
|
|
box[3 * i + k] = static_cast<int>(std::ceil(atoms[i].radius / spacing[k]));
|
|
}
|
|
|
|
std::vector<int> ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atoms) {
|
|
std::vector<int> box(3 * atoms.size());
|
|
Boxes(grid, atoms.data(), atoms.size(), box.data());
|
|
return box;
|
|
}
|
|
|
|
size_t ModelDensityGPU::PairBound(const ModelDensityGrid &grid, const int *box, size_t n) {
|
|
const int n3[3] = {grid.nu, grid.nv, grid.nw};
|
|
size_t pairs = 0;
|
|
for (size_t i = 0; i < n; i++) {
|
|
size_t product = 1;
|
|
for (int k = 0; k < 3; k++)
|
|
product *= AxisBrickBound(box[3 * i + k], n3[k]);
|
|
pairs += product;
|
|
}
|
|
return pairs;
|
|
}
|
|
|
|
size_t ModelDensityGPU::PairBound(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atoms) {
|
|
const std::vector<int> box = Boxes(grid, atoms);
|
|
return PairBound(grid, box.data(), atoms.size());
|
|
}
|
|
|
|
bool ModelDensityGPU::Supports(const ModelDensityGrid &grid, const int *box, size_t n) {
|
|
const int n3[3] = {grid.nu, grid.nv, grid.nw};
|
|
for (size_t i = 0; i < 3 * n; i++)
|
|
if (2 * box[i] + 1 > n3[i % 3] || AxisBrickBound(box[i], n3[i % 3]) > MAX_BRICKS_PER_AXIS)
|
|
return false;
|
|
return true;
|
|
}
|
|
|
|
bool ModelDensityGPU::Supports(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atoms) {
|
|
const std::vector<int> box = Boxes(grid, atoms);
|
|
return Supports(grid, box.data(), atoms.size());
|
|
}
|
|
|
|
size_t ModelDensityGPU::DeviceBytes(size_t max_atoms, size_t max_pairs, size_t max_bricks) {
|
|
const size_t na = max_atoms + 1, nb = max_bricks + 1;
|
|
return na * (sizeof(ModelDensityAtom) + 3 * sizeof(int)) + 3 * max_pairs * sizeof(int) + 3 * nb * sizeof(int) +
|
|
CubBytes(nb);
|
|
}
|
|
|
|
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<size_t>(max_atoms, 1), np = std::max<size_t>(max_pairs, 1), nb = max_bricks + 1;
|
|
atoms = CudaDevicePtr<ModelDensityAtom>(na, sync);
|
|
box = CudaDevicePtr<int>(3 * na, sync);
|
|
brick_count = CudaDevicePtr<int>(nb, sync);
|
|
brick_start = CudaDevicePtr<int>(nb, sync);
|
|
brick_fill = CudaDevicePtr<int>(nb, sync);
|
|
unordered = CudaDevicePtr<int>(np, sync);
|
|
rank = CudaDevicePtr<int>(np, sync);
|
|
pair_atom = CudaDevicePtr<int>(np, sync);
|
|
cub_bytes = CubBytes(nb);
|
|
cub_temp = CudaDevicePtr<char>(std::max<size_t>(cub_bytes, 1), sync);
|
|
host_box = CudaHostPtr<int>(3 * na);
|
|
}
|
|
|
|
void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atom_list) {
|
|
SetAtoms(grid, atom_list.data(), atom_list.size());
|
|
}
|
|
|
|
void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const ModelDensityAtom *atom_list, size_t n) {
|
|
if (n > max_atoms || Bricks(grid.nu, grid.nv, grid.nw) > max_bricks)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: grid or atoms over the reserve");
|
|
// The boxes into the pinned staging buffer, so that their copy and the atoms' (pinned where the caller keeps
|
|
// them so) go up by DMA instead of through the driver's staging of pageable memory.
|
|
Boxes(grid, atom_list, n, host_box);
|
|
// 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 box test does not; such a cell is left to the CPU.
|
|
if (!Supports(grid, host_box, n))
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: an atom wider than the cell");
|
|
if (PairBound(grid, host_box, n) > max_pairs)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: gather pairs over the reserve");
|
|
geom.nu = grid.nu;
|
|
geom.nv = grid.nv;
|
|
geom.nw = grid.nw;
|
|
// As gemmi's Grid sets orth_n (multiply_by_diagonal with 1/n) and ModelDensityGeometry() rounds it.
|
|
const double inv_n[3] = {1.0 / grid.nu, 1.0 / grid.nv, 1.0 / grid.nw};
|
|
geom.m00 = static_cast<float>(grid.orth[0] * inv_n[0]);
|
|
geom.m01 = static_cast<float>(grid.orth[1] * inv_n[1]);
|
|
geom.m02 = static_cast<float>(grid.orth[2] * inv_n[2]);
|
|
geom.m11 = static_cast<float>(grid.orth[4] * inv_n[1]);
|
|
geom.m12 = static_cast<float>(grid.orth[5] * inv_n[2]);
|
|
geom.m22 = static_cast<float>(grid.orth[8] * inv_n[2]);
|
|
n_atoms = static_cast<int>(n);
|
|
if (n_atoms > 0) {
|
|
cuda_err(cudaMemcpyAsync(atoms, atom_list, n * sizeof(ModelDensityAtom), cudaMemcpyHostToDevice, stream));
|
|
cuda_err(cudaMemcpyAsync(box, host_box, 3 * n * sizeof(int), cudaMemcpyHostToDevice, stream));
|
|
}
|
|
// The uploads read the caller's atoms, which it may change once this returns.
|
|
cuda_err(cudaStreamSynchronize(stream));
|
|
}
|
|
|
|
void ModelDensityGPU::Compute(const float4 *d_pos, float *d_grid) {
|
|
// Counted per brick, scanned, filled, and each brick put in model order: nothing comes back to the host, so
|
|
// the launches queue straight through (and record into a CUDA graph). The bricks' atoms end up as a stable
|
|
// sort of the atom-major pairs by brick would leave them.
|
|
const int nb = static_cast<int>(Bricks(geom.nu, geom.nv, geom.nw));
|
|
cuda_err(cudaMemsetAsync(brick_count, 0, (nb + 1) * sizeof(int), stream));
|
|
cuda_err(cudaMemsetAsync(brick_fill, 0, nb * sizeof(int), stream));
|
|
if (n_atoms > 0) {
|
|
count_bricks_kernel<<<blocks(n_atoms), THREADS, 0, stream>>>(d_pos, box, n_atoms, geom.nu, geom.nv, geom.nw,
|
|
brick_count);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
size_t temp = cub_bytes;
|
|
cuda_err(cub::DeviceScan::ExclusiveSum(cub_temp.get(), temp, brick_count.get(), brick_start.get(), nb + 1, stream));
|
|
if (n_atoms > 0) {
|
|
fill_bricks_kernel<<<blocks(n_atoms), THREADS, 0, stream>>>(d_pos, box, n_atoms, geom.nu, geom.nv, geom.nw,
|
|
brick_start, brick_fill, unordered);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
order_bricks_kernel<<<nb, THREADS, 0, stream>>>(unordered, brick_start, rank, pair_atom);
|
|
cuda_err(cudaGetLastError());
|
|
gather_kernel<<<nb, GATHER_THREADS, 0, stream>>>(atoms, d_pos, box, pair_atom, brick_start,
|
|
brick_start.get() + 1, geom, d_grid);
|
|
cuda_err(cudaGetLastError());
|
|
}
|