diff --git a/image_analysis/structure_refinement/ModelDensityGPU.cu b/image_analysis/structure_refinement/ModelDensityGPU.cu index a60879d06..ccdf501a2 100644 --- a/image_analysis/structure_refinement/ModelDensityGPU.cu +++ b/image_analysis/structure_refinement/ModelDensityGPU.cu @@ -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(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(pu); const float fy = at.nctr[1] - static_cast(pv); - const float fz = at.nctr[2] - static_cast(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(pw), density)) + acc[k] += at.occ * density; } - acc += at.occ * density; } } - if (inside) - grid[u + static_cast(g.nu) * (v + static_cast(g.nv) * w)] = acc; + for (int k = 0; k < 2; k++) + if (inside && w + k < g.nw) + grid[u + static_cast(g.nu) * (v + static_cast(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<<>>(unordered, brick_start, rank, pair_atom); cuda_err(cudaGetLastError()); - gather_kernel<<>>(atoms, d_pos, box, pair_atom, brick_start, + gather_kernel<<>>(atoms, d_pos, box, pair_atom, brick_start, brick_start.get() + 1, geom, d_grid); cuda_err(cudaGetLastError()); }