leonarski_fandClaude Fable 5.1
3189f71e8f
The GPU engine's waits on its stream (every evaluation, the refits' transforms) and the model density's wait for the (brick, atom) pair count were blocking event and stream synchronisations; where a weight trial has a card to itself, as on a box with a card per trial, the wake-up of the sleeping thread is on the critical path of every evaluation. Both now poll, and the pair count and overflow flag come back into pinned memory instead of a pageable int (a pageable copy back is staged and blocks the host for the whole stream). OwnSlotDevices places the slots beyond one per card on the cards the ladder's long trials (those refining B) leave free, so that on four cards the three B rungs have a card each and the three B-held rungs share the fourth, instead of two cards each carrying a B rung beside a B-held one. The outcome does not depend on the placement. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
387 lines
18 KiB
Plaintext
387 lines
18 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 <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 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<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.
|
|
__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;
|
|
}
|
|
|
|
// 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;
|
|
}
|
|
|
|
// 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.
|
|
__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 tu = threadIdx.x, tv = threadIdx.y, tw = threadIdx.z;
|
|
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;
|
|
const int tid = tu + BRICK * (tv + BRICK * tw);
|
|
const int n3[3] = {g.nu, g.nv, g.nw};
|
|
|
|
__shared__ GatherAtom sh[GATHER_TILE];
|
|
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];
|
|
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();
|
|
if (!inside)
|
|
continue;
|
|
for (int k = 0; k < m; k++) {
|
|
const GatherAtom &at = sh[k];
|
|
if (!((at.masks >> tu) & (at.masks >> (8 + tv)) & (at.masks >> (16 + tw)) & 1))
|
|
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 int pw = at.base[2] + tw - (tw >= at.wrap[2] ? g.nw : 0);
|
|
// f = centre - (u, v, w) in the box's own (unwrapped) indices, then in A, as the CPU walks it
|
|
const float fx = at.nctr[0] - static_cast<float>(pu);
|
|
const float fy = at.nctr[1] - static_cast<float>(pv);
|
|
const float fz = at.nctr[2] - static_cast<float>(pw);
|
|
const float z = g.m22 * fz;
|
|
const float y = g.m11 * fy + g.m12 * fz;
|
|
const float yz = y * y + z * z;
|
|
const float x = g.m00 * fx + (g.m01 * fy + g.m02 * fz);
|
|
const float r2 = yz + x * x;
|
|
if (r2 > at.max_r2)
|
|
continue;
|
|
float density = 0.0f;
|
|
if (!at.aniso) {
|
|
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));
|
|
}
|
|
}
|
|
acc += at.occ * density;
|
|
}
|
|
}
|
|
if (inside)
|
|
grid[u + static_cast<size_t>(g.nu) * (v + static_cast<size_t>(g.nv) * w)] = acc;
|
|
}
|
|
|
|
int blocks(size_t n) {
|
|
return static_cast<int>((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<int *>(nullptr), static_cast<int *>(nullptr),
|
|
static_cast<int>(atoms + 1)));
|
|
cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, sort, static_cast<int *>(nullptr), static_cast<int *>(nullptr),
|
|
static_cast<int *>(nullptr), static_cast<int *>(nullptr),
|
|
static_cast<int>(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<size_t>((nu + BRICK - 1) / BRICK) * ((nv + BRICK - 1) / BRICK) * ((nw + BRICK - 1) / BRICK);
|
|
}
|
|
|
|
std::vector<int> ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &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<int> 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<int>(std::ceil(atoms[i].radius / spacing[k]));
|
|
return box;
|
|
}
|
|
|
|
size_t ModelDensityGPU::PairBound(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &atoms) {
|
|
const std::vector<int> 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<ModelDensityAtom> &atoms) {
|
|
const std::vector<int> 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<size_t>(max_atoms, 1), np = std::max<size_t>(max_pairs, 1),
|
|
nb = std::max<size_t>(max_bricks, 1);
|
|
atoms = CudaDevicePtr<ModelDensityAtom>(na, sync);
|
|
box = CudaDevicePtr<int>(3 * na, sync);
|
|
overflow = CudaDevicePtr<int>(1, sync);
|
|
count = CudaDevicePtr<int>(na + 1, sync);
|
|
offset = CudaDevicePtr<int>(na + 1, sync);
|
|
key = CudaDevicePtr<int>(np, sync);
|
|
value = CudaDevicePtr<int>(np, sync);
|
|
key_sorted = CudaDevicePtr<int>(np, sync);
|
|
value_sorted = CudaDevicePtr<int>(np, sync);
|
|
brick_start = CudaDevicePtr<int>(nb, sync);
|
|
brick_end = CudaDevicePtr<int>(nb, sync);
|
|
host_count = CudaHostPtr<int>(2);
|
|
cub_bytes = CubBytes(na, np);
|
|
cub_temp = CudaDevicePtr<char>(std::max<size_t>(cub_bytes, 1), sync);
|
|
}
|
|
|
|
void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const std::vector<ModelDensityAtom> &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 box test 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;
|
|
// 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]);
|
|
const std::vector<int> b = Boxes(grid, atom_list);
|
|
n_atoms = static_cast<int>(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<<<blocks(n_atoms + 1), THREADS, 0, stream>>>(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));
|
|
// the count and the flag into pinned memory, the stream polled for them: this wait is on every evaluation
|
|
cuda_err(cudaMemcpyAsync(host_count.get(), offset.get() + n_atoms, sizeof(int), cudaMemcpyDeviceToHost, stream));
|
|
cuda_err(cudaMemcpyAsync(host_count.get() + 1, overflow.get(), sizeof(int), cudaMemcpyDeviceToHost, stream));
|
|
cudaError_t err;
|
|
while ((err = cudaStreamQuery(stream)) == cudaErrorNotReady) {}
|
|
cuda_err(err);
|
|
const int npairs = host_count[0];
|
|
const int over = host_count[1];
|
|
if (over)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: an atom's box over the bricks per axis");
|
|
if (static_cast<size_t>(npairs) > max_pairs)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: gather pairs over the reserve");
|
|
fill_pairs_kernel<<<blocks(n_atoms), THREADS, 0, stream>>>(d_pos, box, n_atoms, geom.nu, geom.nv, geom.nw,
|
|
offset, key, value);
|
|
cuda_err(cudaGetLastError());
|
|
const int nb = static_cast<int>(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<<<blocks(npairs), THREADS, 0, stream>>>(key_sorted, npairs, brick_start, brick_end);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
gather_kernel<<<nb, dim3(BRICK, BRICK, BRICK), 0, stream>>>(atoms, d_pos, box, value_sorted, brick_start, brick_end,
|
|
geom, d_grid);
|
|
cuda_err(cudaGetLastError());
|
|
}
|