mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases 88 Wiki Activity

ModelDensityGPU: the gather culls a tile's atoms per warp and takes two points a thread

Browse Source 
A warp now covers a 4 x 4 x 4 corner of the brick and walks, of each staged tile, only the atoms whose
box reaches into that corner (a ballot over the tile's masks, in tile order), and each thread takes two
points one above the other along w. A point still adds every atom whose box holds it and whose radius
reaches it, in model order, with the same arithmetic: the grid stays bit-identical to the CPU's
(ModelDensityGPU_BitIdenticalToCPU).

RefinementGPU_Benchmark on an RTX 5080: gather_kernel 216 -> 183 us per evaluation.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-10-28 09:46:25 +01:00
1 parent 3a0c3774ff
commit 52dc7ff861
1 file changed
+66 -33
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/ModelDensityGPU.cu
+66 -33
View File
@@ -20,6 +20,7 @@ void cuda_err(cudaError_t 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
@@ -166,12 +167,43 @@ __device__ __forceinline__ int box_offset(int p, int first, int n) {
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) {
@@ -179,14 +211,18 @@ __global__ void gather_kernel(const ModelDensityAtom *atoms, const float4 *pos,
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 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;
const int tid = tu + BRICK * (tv + BRICK * 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];
float acc = 0.0f;
__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();
@@ -220,44 +256,41 @@ __global__ void gather_kernel(const ModelDensityAtom *atoms, const float4 *pos,
}
}
__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 k = 0; k < m; k++) {
const GatherAtom &at = sh[k];
if (!((at.masks >> tu) & (at.masks >> (8 + tv)) & (at.masks >> (16 + tw)) & 1))
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 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));
}
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;
}
acc += at.occ * density;
}
}
if (inside)
grid[u + static_cast<size_t>(g.nu) * (v + static_cast<size_t>(g.nv) * w)] = acc;
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) {
@@ -413,7 +446,7 @@ void ModelDensityGPU::Compute(const float4 *d_pos, float *d_grid) {
}
order_bricks_kernel<<<nb, THREADS, 0, stream>>>(unordered, brick_start, rank, pair_atom);
cuda_err(cudaGetLastError());
gather_kernel<<<nb, dim3(BRICK, BRICK, BRICK), 0, stream>>>(atoms, d_pos, box, pair_atom, brick_start,
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());
}