leonarski_f
8718065a7c
The score draws the mask on a grid of its own (min(d_min/3, 0.4 A), BulkSolventMask) on the host; on a GPU engine that was a CPU mask and transform every cycle, per trial slot on one thread. Now the engine masks on that grid itself: - RefinementGPUSetup::mask_grid (BulkSolventMask::MaskGrid, the grid the host mask uses); where it is not the structure factors' grid the engine reserves a grid, a spectrum and an R2C plan for it (counted in DeviceBytes and the FFT work area), and Fmask masks (ModelMaskGPU), transforms and reads the rows there. Refit takes the engine's Fmask on a GPU, BulkSolventMask on the CPU. - ModelMaskGPU's shrink: each edge point marks its stencil as solvent instead of every macromolecule point searching its stencil for an edge - the same idempotent stores to the same points, so the same mask to the bit (ModelMaskGPU_* tests unchanged), and a thin shell of points visited instead of half the grid. - BulkSolventMask finds its symmetry orbits at its first mask: a GPU refinement never draws one on the host (1.1 s of trial-slot setup on the 3.6 A set). The card's mask is gemmi's but at points exactly on an atom's radius (ModelMaskGPU.h): R-work/R-free of the 14 parity sets on the GPU are the CPU engine's to 4 digits. Refinement deterministic run to run on the GPU (7 timing sets, refined model and report identical over 5 runs of successive builds that did not change results). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 89b5d00907c9ad99e8b389f6272a3724f981cd10)
778 lines
36 KiB
Plaintext
778 lines
36 KiB
Plaintext
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "RefinementGPUEngine.h"
|
|
|
|
#include <algorithm>
|
|
#include <string>
|
|
|
|
#include <cub/cub.cuh>
|
|
#include <cufft.h>
|
|
|
|
#include "../../common/JFJochException.h"
|
|
#include "../indexing/CUDAMemHelpers.h"
|
|
|
|
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<int>(val)));
|
|
}
|
|
|
|
constexpr int THREADS = 256;
|
|
constexpr int WARP = 32;
|
|
constexpr double FOUR_PI2 = 4 * refine_math::PI * refine_math::PI;
|
|
|
|
int blocks(size_t n) {
|
|
return static_cast<int>((n + THREADS - 1) / THREADS);
|
|
}
|
|
|
|
__device__ __forceinline__ int imod(int a, int n) {
|
|
const int r = a % n;
|
|
return r < 0 ? r + n : r;
|
|
}
|
|
|
|
// gemmi's unsafe_expapprox(), as ModelDensityGPU.cu has it.
|
|
__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))));
|
|
}
|
|
|
|
// exp(+2 pi i h.t) of row h and operator o, exact for the fractions a translation is made of.
|
|
__device__ __forceinline__ void op_phase(const int *h, const int *tran, int den, double &c, double &s) {
|
|
sincospi(2.0 * (h[0] * tran[0] + h[1] * tran[1] + h[2] * tran[2]) / den, &s, &c);
|
|
}
|
|
|
|
// F of the atoms at each row, composed from the half-u transform of one copy (ModelStructureFactorsGPU.cu's
|
|
// compose_kernel), and with `evaluate` the target's term of a working row and the coefficient
|
|
// conj(dT/dF) n_cen unblur the gradient's map is made of.
|
|
__global__ void compose_target_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, bool evaluate,
|
|
const float *k_total, const float2 *fmask, const float *ml_d, const float *ml_var,
|
|
const float *fo, const char *work, const char *centric, float2 *f, double *nll,
|
|
float2 *coef) {
|
|
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;
|
|
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<size_t>(nu / 2 + 1) * (imod(k[1], nv) + static_cast<size_t>(nv) * imod(k[2], nw))];
|
|
const double vr = out.x * norm, vi = (conj ? out.y : -out.y) * norm;
|
|
double s, c;
|
|
op_phase(h, tran, den, c, s);
|
|
sr += c * vr - s * vi;
|
|
si += c * vi + s * vr;
|
|
}
|
|
const double scale = row_scale[m];
|
|
sr *= scale;
|
|
si *= scale;
|
|
f[m] = make_float2(static_cast<float>(sr), static_cast<float>(si));
|
|
if (!evaluate)
|
|
return;
|
|
double term = 0, cr = 0, ci = 0;
|
|
if (work[m]) {
|
|
const double k = k_total[m];
|
|
const double mr = k * (sr + fmask[m].x), mi = k * (si + fmask[m].y);
|
|
const double fc = sqrt(mr * mr + mi * mi);
|
|
double dfc = 0;
|
|
term = refine_math::AmplitudeNegLogLikelihood<double>(fo[m], fc, ml_d[m], ml_var[m], centric[m], dfc);
|
|
if (fc > 0) {
|
|
// dT/dF = k dT/d|F_model| F_model / |F_model|; the coefficient is its conjugate times the scale
|
|
const double g = k * dfc / fc * scale;
|
|
cr = g * mr;
|
|
ci = -g * mi;
|
|
}
|
|
}
|
|
nll[m] = term;
|
|
coef[m] = make_float2(static_cast<float>(cr), static_cast<float>(ci));
|
|
}
|
|
|
|
// The gradient's coefficients on the half-u spectrum: each point the sum, in order, of the terms
|
|
// (coef[row] exp(2 pi i h.t_op)) / 2 that land on it, conjugated where they land as a Friedel mate.
|
|
__global__ void scatter_kernel(const float2 *coef, const int *row_hkl, const int *ops, int n_ops, int den,
|
|
const int *point, const int *start, const int *term, int points, float2 *spectrum) {
|
|
const int s = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (s >= points)
|
|
return;
|
|
double sr = 0, si = 0;
|
|
for (int e = start[s]; e < start[s + 1]; e++) {
|
|
const int t = term[e];
|
|
const bool conj = t & 1;
|
|
const int m = (t >> 1) / n_ops, o = (t >> 1) % n_ops;
|
|
double c, sn;
|
|
op_phase(row_hkl + 3 * m, ops + 12 * o + 9, den, c, sn);
|
|
const double ar = 0.5 * (coef[m].x * c - coef[m].y * sn), ai = 0.5 * (coef[m].x * sn + coef[m].y * c);
|
|
sr += ar;
|
|
si += conj ? -ai : ai;
|
|
}
|
|
spectrum[point[s]] = make_float2(static_cast<float>(sr), static_cast<float>(si));
|
|
}
|
|
|
|
// F_mask at each row, read at h directly (ModelStructureFactorsGPU.cu's fmask_kernel).
|
|
__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<size_t>(nu / 2 + 1) * (imod(k, nv) + static_cast<size_t>(nv) * imod(l, nw))];
|
|
const float re = out.x * norm, im = -out.y * norm;
|
|
fmask[m] = make_float2(re, conj ? -im : im);
|
|
}
|
|
|
|
// dT/d(x, y, z, B) of every atom: the derivative map summed against the derivatives of the atom's density over
|
|
// its box, one warp per atom, the lanes' sums added in a fixed tree. The box, the points within the radius and
|
|
// their offsets from the centre are the density gather's (ModelDensityGPU.cu).
|
|
__global__ void atom_gradient_kernel(const ModelDensityAtom *atoms, const float4 *pos, const int *box, int n,
|
|
ModelDensityGPU::Geometry g, const float *map, float map_scale, float4 *grad) {
|
|
const int a = (blockIdx.x * blockDim.x + threadIdx.x) / WARP;
|
|
const int lane = threadIdx.x % WARP;
|
|
if (a >= n)
|
|
return;
|
|
const ModelDensityAtom at = atoms[a];
|
|
const float4 p = pos[a];
|
|
const float nctr[3] = {p.x, p.y, p.z};
|
|
int lo[3], span[3];
|
|
for (int k = 0; k < 3; k++) {
|
|
lo[k] = __float2int_rn(nctr[k]) - box[3 * a + k];
|
|
span[k] = 2 * box[3 * a + k] + 1;
|
|
}
|
|
const int total = span[0] * span[1] * span[2];
|
|
float gx = 0, gy = 0, gz = 0, gb = 0;
|
|
for (int t = lane; t < total; t += WARP) {
|
|
const int u = lo[0] + t % span[0], v = lo[1] + t / span[0] % span[1], w = lo[2] + t / (span[0] * span[1]);
|
|
const float fx = nctr[0] - static_cast<float>(u);
|
|
const float fy = nctr[1] - static_cast<float>(v);
|
|
const float fz = nctr[2] - static_cast<float>(w);
|
|
const float z = g.m22 * fz;
|
|
const float y = g.m11 * fy + g.m12 * fz;
|
|
const float x = g.m00 * fx + (g.m01 * fy + g.m02 * fz);
|
|
const float r2 = (y * y + z * z) + x * x;
|
|
if (r2 > at.max_r2)
|
|
continue;
|
|
const float value = map[imod(u, g.nu) + static_cast<size_t>(g.nu) * (imod(v, g.nv) + static_cast<size_t>(g.nv) * imod(w, g.nw))];
|
|
float dr2 = 0, db = 0;
|
|
for (int q = 0; q < 5; q++) {
|
|
const float b = at.b[q][0];
|
|
const float e = at.a[q] * expapprox(fmaxf(b * r2, -88.f));
|
|
dr2 += e * b;
|
|
db += e * b * (1.5f + b * r2);
|
|
}
|
|
gx += value * dr2 * x;
|
|
gy += value * dr2 * y;
|
|
gz += value * dr2 * z;
|
|
gb += value * db;
|
|
}
|
|
for (int off = WARP / 2; off > 0; off /= 2) {
|
|
gx += __shfl_down_sync(0xffffffff, gx, off);
|
|
gy += __shfl_down_sync(0xffffffff, gy, off);
|
|
gz += __shfl_down_sync(0xffffffff, gz, off);
|
|
gb += __shfl_down_sync(0xffffffff, gb, off);
|
|
}
|
|
if (lane == 0) {
|
|
const float s = 2.0f * at.occ * map_scale;
|
|
grad[a] = make_float4(s * gx, s * gy, s * gz, at.occ * map_scale * gb / static_cast<float>(FOUR_PI2));
|
|
}
|
|
}
|
|
|
|
using V3f = refine_math::V3<float>;
|
|
|
|
__device__ V3f xyz(const float4 *p, int i) {
|
|
const float4 q = p[i];
|
|
return {q.x, q.y, q.z};
|
|
}
|
|
|
|
// Every restraint's E, and its gradient into its slots, one thread per restraint.
|
|
__global__ void restraint_kernel(const RestraintRecord *records, int n, const int *plane_atoms, const float4 *xyzb,
|
|
float4 *contrib, double *energy) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n)
|
|
return;
|
|
const RestraintRecord &r = records[i];
|
|
V3f g[4] = {};
|
|
int na = 0;
|
|
float e = 0;
|
|
switch (static_cast<RestraintKind>(r.kind)) {
|
|
case RestraintKind::Bond:
|
|
e = refine_math::BondTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), r.ideal, r.sigma, g);
|
|
na = 2;
|
|
break;
|
|
case RestraintKind::Angle:
|
|
e = refine_math::AngleTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), xyz(xyzb, r.atom[2]), r.ideal,
|
|
r.sigma, g);
|
|
na = 3;
|
|
break;
|
|
case RestraintKind::Torsion:
|
|
e = refine_math::TorsionTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), xyz(xyzb, r.atom[2]),
|
|
xyz(xyzb, r.atom[3]), r.ideal, r.sigma, r.period, g);
|
|
na = 4;
|
|
break;
|
|
case RestraintKind::Chiral:
|
|
e = refine_math::ChiralTerm(xyz(xyzb, r.atom[0]), xyz(xyzb, r.atom[1]), xyz(xyzb, r.atom[2]),
|
|
xyz(xyzb, r.atom[3]), r.ideal, r.sigma, g);
|
|
na = 4;
|
|
break;
|
|
case RestraintKind::Repulsion: {
|
|
const V3f q = xyz(xyzb, r.atom[1]);
|
|
const V3f image{r.rot[0] * q.x + r.rot[1] * q.y + r.rot[2] * q.z + r.tran[0],
|
|
r.rot[3] * q.x + r.rot[4] * q.y + r.rot[5] * q.z + r.tran[1],
|
|
r.rot[6] * q.x + r.rot[7] * q.y + r.rot[8] * q.z + r.tran[2]};
|
|
e = refine_math::RepulsionTerm(xyz(xyzb, r.atom[0]), image, r.ideal, r.sigma, g);
|
|
const V3f gi = g[1];
|
|
g[1] = {r.rot[0] * gi.x + r.rot[3] * gi.y + r.rot[6] * gi.z,
|
|
r.rot[1] * gi.x + r.rot[4] * gi.y + r.rot[7] * gi.z,
|
|
r.rot[2] * gi.x + r.rot[5] * gi.y + r.rot[8] * gi.z};
|
|
na = 2;
|
|
break;
|
|
}
|
|
case RestraintKind::Plane: {
|
|
V3f px[RefinementGPUEngine::MAX_PLANE_ATOMS], pg[RefinementGPUEngine::MAX_PLANE_ATOMS];
|
|
const int m = r.atom[1];
|
|
for (int k = 0; k < m; k++)
|
|
px[k] = xyz(xyzb, plane_atoms[r.atom[0] + k]);
|
|
e = refine_math::PlaneTerm(px, m, r.sigma, pg);
|
|
for (int k = 0; k < m; k++)
|
|
contrib[r.slot + k] = make_float4(pg[k].x, pg[k].y, pg[k].z, 0.0f);
|
|
break;
|
|
}
|
|
case RestraintKind::BSimilarity: {
|
|
const float rr = (xyzb[r.atom[0]].w - xyzb[r.atom[1]].w) / r.sigma;
|
|
e = rr * rr;
|
|
contrib[r.slot] = make_float4(0, 0, 0, 2 * rr / r.sigma);
|
|
contrib[r.slot + 1] = make_float4(0, 0, 0, -2 * rr / r.sigma);
|
|
break;
|
|
}
|
|
}
|
|
for (int k = 0; k < na; k++)
|
|
contrib[r.slot + k] = make_float4(g[k].x, g[k].y, g[k].z, 0.0f);
|
|
energy[i] = e;
|
|
}
|
|
|
|
// Each atom's restraint gradient: its slots summed in order.
|
|
__global__ void slot_gather_kernel(const float4 *contrib, const int *start, const int *index, int n, float4 *grad) {
|
|
const int a = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (a >= n)
|
|
return;
|
|
float4 s = make_float4(0, 0, 0, 0);
|
|
for (int k = start[a]; k < start[a + 1]; k++) {
|
|
const float4 c = contrib[index[k]];
|
|
s.x += c.x;
|
|
s.y += c.y;
|
|
s.z += c.z;
|
|
s.w += c.w;
|
|
}
|
|
grad[a] = s;
|
|
}
|
|
|
|
CudaFFTShape ForwardShape(const ModelDensityGrid &g) {
|
|
return {.n = {g.nw, g.nv, g.nu}, .idist = g.nu * g.nv * g.nw, .odist = (g.nu / 2 + 1) * g.nv * g.nw,
|
|
.type = CUFFT_R2C};
|
|
}
|
|
|
|
bool SameGrid(const ModelDensityGrid &a, const ModelDensityGrid &b) {
|
|
return a.nu == b.nu && a.nv == b.nv && a.nw == b.nw;
|
|
}
|
|
|
|
CudaFFTShape InverseShape(const ModelDensityGrid &g) {
|
|
return {.n = {g.nw, g.nv, g.nu}, .idist = (g.nu / 2 + 1) * g.nv * g.nw, .odist = g.nu * g.nv * g.nw,
|
|
.type = CUFFT_C2R};
|
|
}
|
|
|
|
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<size_t>(g.nu) * g.nv * g.nw;
|
|
}
|
|
|
|
size_t ComplexPoints(const ModelDensityGrid &g) {
|
|
return static_cast<size_t>(g.nu / 2 + 1) * g.nv * g.nw;
|
|
}
|
|
|
|
template <typename T>
|
|
void Upload(CudaDevicePtr<T> &dst, const std::vector<T> &src, cudaStream_t stream) {
|
|
if (!src.empty())
|
|
cuda_err(cudaMemcpyAsync(dst, src.data(), src.size() * sizeof(T), cudaMemcpyHostToDevice, stream));
|
|
}
|
|
|
|
// A device buffer of at least n elements, kept between calls and grown when it is too small.
|
|
template <typename T>
|
|
void Ensure(CudaDevicePtr<T> &p, size_t &capacity, size_t n) {
|
|
if (n <= capacity && p)
|
|
return;
|
|
capacity = std::max<size_t>(n, 1);
|
|
p = CudaDevicePtr<T>(capacity, CudaAlloc::Synchronous);
|
|
}
|
|
|
|
// cub's temporary storage to reduce n doubles.
|
|
size_t ReduceBytes(size_t n) {
|
|
size_t bytes = 0;
|
|
cuda_err(cub::DeviceReduce::Sum(nullptr, bytes, static_cast<double *>(nullptr), static_cast<double *>(nullptr),
|
|
static_cast<int>(std::max<size_t>(n, 1))));
|
|
return bytes;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
struct RefinementGPUEngineImpl {
|
|
struct Restraints {
|
|
int records = 0, atoms = 0;
|
|
size_t cap_records = 0, cap_planes = 0, cap_start = 0, cap_index = 0, cap_slots = 0, cap_reduce = 0,
|
|
cap_atoms = 0, cap_xyzb = 0, reduce_bytes = 0;
|
|
CudaDevicePtr<RestraintRecord> rec;
|
|
CudaDevicePtr<int> plane_atoms, slot_start, slot_index;
|
|
CudaDevicePtr<float4> contrib, xyzb, grad;
|
|
CudaDevicePtr<double> energy, sum;
|
|
CudaDevicePtr<char> reduce_temp;
|
|
};
|
|
|
|
RefinementGPUSetup setup;
|
|
size_t work_bytes;
|
|
int device = 0;
|
|
CudaStream stream;
|
|
int n_atoms = 0;
|
|
|
|
CudaDevicePtr<float4> pos;
|
|
CudaHostPtr<float4> host_pos; // the positions' staging buffer, pinned
|
|
// What comes back, pinned, so that a copy back queues without waiting: the X-ray target and the restraints'
|
|
// sums, their gradients, and the restraints' positions going up.
|
|
CudaHostPtr<double> host_sums;
|
|
CudaHostPtr<float4> host_grad, host_xyzb, host_restraint_grad;
|
|
CudaDevicePtr<ModelMaskAtom> mask_atoms;
|
|
CudaDevicePtr<float> grid;
|
|
CudaDevicePtr<float2> spectrum;
|
|
CudaDevicePtr<char> fft_work;
|
|
// The mask's grid, spectrum and transform where its grid is not the structure factors' (grid, spectrum,
|
|
// forward otherwise).
|
|
CudaDevicePtr<float> mask_grid;
|
|
CudaDevicePtr<float2> mask_spectrum;
|
|
CudaDevicePtr<int> ops, row_hkl;
|
|
CudaDevicePtr<double> row_scale;
|
|
CudaDevicePtr<float> k_total, ml_d, ml_var, fo;
|
|
CudaDevicePtr<float2> fmask_scaled, f, fm, coef;
|
|
CudaDevicePtr<char> work, centric;
|
|
CudaDevicePtr<double> nll, nll_sum;
|
|
CudaDevicePtr<char> reduce_temp;
|
|
size_t reduce_bytes = 0;
|
|
CudaDevicePtr<int> scatter_point, scatter_start, scatter_term;
|
|
CudaDevicePtr<float4> grad;
|
|
std::unique_ptr<ModelDensityGPU> density;
|
|
std::unique_ptr<ModelMaskGPU> mask;
|
|
CudaFFTPlan forward, inverse, mask_forward;
|
|
ModelDensityGPU::Geometry geom{};
|
|
Restraints restraints[2];
|
|
|
|
RefinementGPUEngineImpl(int dev, const RefinementGPUSetup &s, size_t fft_work_bytes) : setup(s), work_bytes(fft_work_bytes), device(dev) {
|
|
const auto sync = CudaAlloc::Synchronous;
|
|
const size_t na = std::max<size_t>(setup.max_atoms, 1), nr = std::max<size_t>(setup.rows.size(), 1);
|
|
pos = CudaDevicePtr<float4>(na, sync);
|
|
host_pos = CudaHostPtr<float4>(na);
|
|
host_sums = CudaHostPtr<double>(2);
|
|
host_grad = CudaHostPtr<float4>(na);
|
|
host_xyzb = CudaHostPtr<float4>(na);
|
|
host_restraint_grad = CudaHostPtr<float4>(na);
|
|
mask_atoms = CudaDevicePtr<ModelMaskAtom>(na, sync);
|
|
grad = CudaDevicePtr<float4>(na, sync);
|
|
grid = CudaDevicePtr<float>(GridPoints(setup.grid), sync);
|
|
spectrum = CudaDevicePtr<float2>(ComplexPoints(setup.grid), sync);
|
|
fft_work = CudaDevicePtr<char>(std::max<size_t>(work_bytes, 1), sync);
|
|
ops = CudaDevicePtr<int>(std::max<size_t>(12 * setup.sym_ops.size(), 1), sync);
|
|
row_hkl = CudaDevicePtr<int>(3 * nr, sync);
|
|
row_scale = CudaDevicePtr<double>(nr, sync);
|
|
k_total = CudaDevicePtr<float>(nr, sync);
|
|
ml_d = CudaDevicePtr<float>(nr, sync);
|
|
ml_var = CudaDevicePtr<float>(nr, sync);
|
|
fo = CudaDevicePtr<float>(nr, sync);
|
|
fmask_scaled = CudaDevicePtr<float2>(nr, sync);
|
|
f = CudaDevicePtr<float2>(nr, sync);
|
|
fm = CudaDevicePtr<float2>(nr, sync);
|
|
coef = CudaDevicePtr<float2>(nr, sync);
|
|
work = CudaDevicePtr<char>(nr, sync);
|
|
centric = CudaDevicePtr<char>(nr, sync);
|
|
nll = CudaDevicePtr<double>(nr, sync);
|
|
nll_sum = CudaDevicePtr<double>(1, sync);
|
|
reduce_bytes = ReduceBytes(nr);
|
|
reduce_temp = CudaDevicePtr<char>(std::max<size_t>(reduce_bytes, 1), sync);
|
|
scatter_point = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_point.size(), 1), sync);
|
|
scatter_start = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_start.size(), 1), sync);
|
|
scatter_term = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_term.size(), 1), sync);
|
|
density = std::make_unique<ModelDensityGPU>(stream, setup.max_atoms, setup.max_pairs,
|
|
ModelDensityGPU::Bricks(setup.grid.nu, setup.grid.nv, setup.grid.nw));
|
|
mask = std::make_unique<ModelMaskGPU>(stream, GridPoints(setup.mask_grid));
|
|
ModelMaskGrid mg;
|
|
mg.nu = setup.mask_grid.nu;
|
|
mg.nv = setup.mask_grid.nv;
|
|
mg.nw = setup.mask_grid.nw;
|
|
std::copy(setup.mask_grid.orth, setup.mask_grid.orth + 9, mg.orth);
|
|
mg.volume = setup.volume;
|
|
mask->SetGrid(mg, setup.mask_ops);
|
|
if (!SameGrid(setup.mask_grid, setup.grid)) {
|
|
mask_grid = CudaDevicePtr<float>(GridPoints(setup.mask_grid), sync);
|
|
mask_spectrum = CudaDevicePtr<float2>(ComplexPoints(setup.mask_grid), sync);
|
|
}
|
|
|
|
std::vector<int> flat_ops;
|
|
for (const auto &o : setup.sym_ops)
|
|
flat_ops.insert(flat_ops.end(), o.begin(), o.end());
|
|
std::vector<int> hkl;
|
|
for (const auto &h : setup.rows)
|
|
hkl.insert(hkl.end(), h.begin(), h.end());
|
|
Upload(ops, flat_ops, stream);
|
|
Upload(row_hkl, hkl, stream);
|
|
Upload(fo, setup.fo, stream);
|
|
Upload(work, setup.work, stream);
|
|
Upload(centric, setup.centric, stream);
|
|
Upload(scatter_point, setup.scatter_point, stream);
|
|
Upload(scatter_start, setup.scatter_start, stream);
|
|
Upload(scatter_term, setup.scatter_term, stream);
|
|
|
|
const double inv_n[3] = {1.0 / setup.grid.nu, 1.0 / setup.grid.nv, 1.0 / setup.grid.nw};
|
|
geom.nu = setup.grid.nu;
|
|
geom.nv = setup.grid.nv;
|
|
geom.nw = setup.grid.nw;
|
|
geom.m00 = static_cast<float>(setup.grid.orth[0] * inv_n[0]);
|
|
geom.m01 = static_cast<float>(setup.grid.orth[1] * inv_n[1]);
|
|
geom.m02 = static_cast<float>(setup.grid.orth[2] * inv_n[2]);
|
|
geom.m11 = static_cast<float>(setup.grid.orth[4] * inv_n[1]);
|
|
geom.m12 = static_cast<float>(setup.grid.orth[5] * inv_n[2]);
|
|
geom.m22 = static_cast<float>(setup.grid.orth[8] * inv_n[2]);
|
|
|
|
forward = CudaFFTPlan(ForwardShape(setup.grid), stream, fft_work);
|
|
inverse = CudaFFTPlan(InverseShape(setup.grid), stream, fft_work);
|
|
if (!SameGrid(setup.mask_grid, setup.grid))
|
|
mask_forward = CudaFFTPlan(ForwardShape(setup.mask_grid), stream, fft_work);
|
|
if (forward.WorkBytes() > work_bytes || inverse.WorkBytes() > work_bytes ||
|
|
(mask_grid.get() != nullptr && mask_forward.WorkBytes() > work_bytes))
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: cuFFT work area over the reserve");
|
|
cuda_err(cudaStreamSynchronize(stream));
|
|
}
|
|
|
|
~RefinementGPUEngineImpl() { cudaStreamSynchronize(stream); }
|
|
|
|
float Norm() const { return static_cast<float>(setup.volume / static_cast<double>(GridPoints(setup.grid))); }
|
|
float MaskNorm() const {
|
|
return static_cast<float>(setup.volume / static_cast<double>(GridPoints(setup.mask_grid)));
|
|
}
|
|
|
|
// The atoms' density on the grid and its transform into the spectrum.
|
|
void Transform() {
|
|
density->Compute(pos, grid);
|
|
cufft_err(forward.R2C(grid, reinterpret_cast<cufftComplex *>(spectrum.get())));
|
|
}
|
|
|
|
void Compose(bool evaluate) {
|
|
const int rows = static_cast<int>(setup.rows.size());
|
|
if (rows == 0)
|
|
return;
|
|
compose_target_kernel<<<blocks(rows), THREADS, 0, stream>>>(
|
|
spectrum, Norm(), setup.grid.nu, setup.grid.nv, setup.grid.nw, ops, static_cast<int>(setup.sym_ops.size()),
|
|
setup.den, row_hkl, row_scale, rows, evaluate, k_total, fmask_scaled, ml_d, ml_var, fo, work, centric, f,
|
|
nll, coef);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
|
|
// The X-ray target into host_sums[0] and, with `grad`, its gradient into host_grad: queued, not waited for.
|
|
void QueueEvaluate(bool grad) {
|
|
const int rows = static_cast<int>(setup.rows.size());
|
|
Transform();
|
|
Compose(true);
|
|
host_sums[0] = 0;
|
|
if (rows > 0) {
|
|
size_t temp = reduce_bytes;
|
|
cuda_err(cub::DeviceReduce::Sum(reduce_temp.get(), temp, nll.get(), nll_sum.get(), rows, stream));
|
|
cuda_err(cudaMemcpyAsync(host_sums.get(), nll_sum, sizeof(double), cudaMemcpyDeviceToHost, stream));
|
|
}
|
|
if (grad) {
|
|
const ModelDensityGrid &g = setup.grid;
|
|
cuda_err(cudaMemsetAsync(spectrum, 0, ComplexPoints(g) * sizeof(float2), stream));
|
|
const int points = static_cast<int>(setup.scatter_point.size());
|
|
if (points > 0) {
|
|
scatter_kernel<<<blocks(points), THREADS, 0, stream>>>(coef, row_hkl, ops,
|
|
static_cast<int>(setup.sym_ops.size()), setup.den,
|
|
scatter_point, scatter_start, scatter_term,
|
|
points, spectrum);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
cufft_err(inverse.C2R(reinterpret_cast<cufftComplex *>(spectrum.get()), this->grid));
|
|
if (n_atoms > 0) {
|
|
const size_t threads = static_cast<size_t>(n_atoms) * WARP;
|
|
atom_gradient_kernel<<<blocks(threads), THREADS, 0, stream>>>(density->DeviceAtoms(), pos,
|
|
density->DeviceBoxes(), n_atoms, geom,
|
|
this->grid, Norm(), this->grad);
|
|
cuda_err(cudaGetLastError());
|
|
cuda_err(cudaMemcpyAsync(host_grad.get(), this->grad, n_atoms * sizeof(float4), cudaMemcpyDeviceToHost,
|
|
stream));
|
|
}
|
|
}
|
|
}
|
|
|
|
// Restraint set `which` at host_xyzb into host_sums[1] and, with `grad`, its gradient into
|
|
// host_restraint_grad: queued, not waited for.
|
|
void QueueRestraints(int which, bool grad) {
|
|
Restraints &r = restraints[which];
|
|
if (r.atoms > 0)
|
|
cuda_err(cudaMemcpyAsync(r.xyzb, host_xyzb.get(), r.atoms * sizeof(float4), cudaMemcpyHostToDevice, stream));
|
|
host_sums[1] = 0;
|
|
if (r.records > 0) {
|
|
restraint_kernel<<<blocks(r.records), THREADS, 0, stream>>>(r.rec, r.records, r.plane_atoms, r.xyzb, r.contrib,
|
|
r.energy);
|
|
cuda_err(cudaGetLastError());
|
|
size_t temp = r.reduce_bytes;
|
|
cuda_err(cub::DeviceReduce::Sum(r.reduce_temp.get(), temp, r.energy.get(), r.sum.get(), r.records, stream));
|
|
cuda_err(cudaMemcpyAsync(host_sums.get() + 1, r.sum, sizeof(double), cudaMemcpyDeviceToHost, stream));
|
|
}
|
|
if (grad && r.atoms > 0) {
|
|
if (r.records > 0) {
|
|
slot_gather_kernel<<<blocks(r.atoms), THREADS, 0, stream>>>(r.contrib, r.slot_start, r.slot_index,
|
|
r.atoms, r.grad);
|
|
cuda_err(cudaGetLastError());
|
|
} else
|
|
cuda_err(cudaMemsetAsync(r.grad, 0, r.atoms * sizeof(float4), stream));
|
|
cuda_err(cudaMemcpyAsync(host_restraint_grad.get(), r.grad, r.atoms * sizeof(float4), cudaMemcpyDeviceToHost,
|
|
stream));
|
|
}
|
|
}
|
|
};
|
|
|
|
namespace {
|
|
|
|
void CopyOut(const float4 *from, int n, std::vector<std::array<float, 4>> &to) {
|
|
to.resize(n);
|
|
std::copy(reinterpret_cast<const std::array<float, 4> *>(from), reinterpret_cast<const std::array<float, 4> *>(from) + n,
|
|
to.begin());
|
|
}
|
|
|
|
} // namespace
|
|
|
|
size_t RefinementGPUEngine::FFTWorkBytes(int device, const RefinementGPUSetup &setup) {
|
|
DeviceGuard guard(device);
|
|
size_t bytes = std::max(CudaFFTPlan::WorkBytes(ForwardShape(setup.grid)),
|
|
CudaFFTPlan::WorkBytes(InverseShape(setup.grid)));
|
|
if (!SameGrid(setup.mask_grid, setup.grid))
|
|
bytes = std::max(bytes, CudaFFTPlan::WorkBytes(ForwardShape(setup.mask_grid)));
|
|
return bytes;
|
|
}
|
|
|
|
size_t RefinementGPUEngine::DeviceBytes(const RefinementGPUSetup &s, size_t fft_work_bytes, size_t max_records,
|
|
size_t max_slots) {
|
|
const size_t na = s.max_atoms + 1, nr = s.rows.size() + 1;
|
|
size_t bytes = na * (2 * sizeof(float4) + 3 * sizeof(int) + sizeof(ModelMaskAtom) + sizeof(ModelDensityAtom));
|
|
bytes += GridPoints(s.grid) * sizeof(float) + ComplexPoints(s.grid) * sizeof(float2) + fft_work_bytes;
|
|
bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 4 * sizeof(float) + 4 * sizeof(float2) + 2);
|
|
bytes += (s.scatter_point.size() + s.scatter_start.size() + s.scatter_term.size()) * sizeof(int);
|
|
bytes += ModelDensityGPU::DeviceBytes(s.max_atoms, s.max_pairs, ModelDensityGPU::Bricks(s.grid.nu, s.grid.nv, s.grid.nw));
|
|
bytes += ModelMaskGPU::DeviceBytes(GridPoints(s.mask_grid));
|
|
if (!SameGrid(s.mask_grid, s.grid))
|
|
bytes += GridPoints(s.mask_grid) * sizeof(float) + ComplexPoints(s.mask_grid) * sizeof(float2);
|
|
bytes += max_records * (sizeof(RestraintRecord) + sizeof(double)) + max_slots * (sizeof(float4) + sizeof(int)) +
|
|
na * (2 * sizeof(float4) + sizeof(int));
|
|
return bytes;
|
|
}
|
|
|
|
RefinementGPUEngine::RefinementGPUEngine(int device, const RefinementGPUSetup &setup, size_t fft_work_bytes) {
|
|
DeviceGuard guard(device);
|
|
impl_ = std::make_unique<RefinementGPUEngineImpl>(device, setup, fft_work_bytes);
|
|
}
|
|
|
|
RefinementGPUEngine::~RefinementGPUEngine() {
|
|
if (impl_) {
|
|
DeviceGuard guard(impl_->device);
|
|
impl_.reset();
|
|
}
|
|
}
|
|
|
|
void RefinementGPUEngine::SetBlur(double blur) {
|
|
RefinementGPUEngineImpl &e = *impl_;
|
|
DeviceGuard guard(e.device);
|
|
std::vector<double> scale(e.setup.rows.size());
|
|
for (size_t m = 0; m < scale.size(); m++)
|
|
scale[m] = e.setup.centring * std::exp(blur * 0.25 * e.setup.inv_d2[m]);
|
|
Upload(e.row_scale, scale, e.stream);
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
}
|
|
|
|
void RefinementGPUEngine::SetRowModel(const std::vector<float> &k_total, const std::vector<std::array<float, 2>> &fmask,
|
|
const std::vector<float> &d, const std::vector<float> &var) {
|
|
RefinementGPUEngineImpl &e = *impl_;
|
|
DeviceGuard guard(e.device);
|
|
Upload(e.k_total, k_total, e.stream);
|
|
if (!fmask.empty())
|
|
cuda_err(cudaMemcpyAsync(e.fmask_scaled, fmask.data(), fmask.size() * sizeof(float2), cudaMemcpyHostToDevice,
|
|
e.stream));
|
|
Upload(e.ml_d, d, e.stream);
|
|
Upload(e.ml_var, var, e.stream);
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
}
|
|
|
|
void RefinementGPUEngine::SetDensities(const std::vector<ModelDensityAtom> &atoms) {
|
|
RefinementGPUEngineImpl &e = *impl_;
|
|
DeviceGuard guard(e.device);
|
|
if (atoms.size() > e.setup.max_atoms)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: atoms over the reserve");
|
|
e.density->SetAtoms(e.setup.grid, atoms);
|
|
e.n_atoms = static_cast<int>(atoms.size());
|
|
}
|
|
|
|
void RefinementGPUEngine::SetPositions(const std::vector<std::array<float, 4>> &pos) {
|
|
RefinementGPUEngineImpl &e = *impl_;
|
|
DeviceGuard guard(e.device);
|
|
if (static_cast<int>(pos.size()) != e.n_atoms)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: positions do not match the atoms");
|
|
// the staging buffer free again: the last copy out of it done
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
std::copy(pos.begin(), pos.end(), reinterpret_cast<std::array<float, 4> *>(e.host_pos.get()));
|
|
if (!pos.empty())
|
|
cuda_err(cudaMemcpyAsync(e.pos, e.host_pos, pos.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream));
|
|
}
|
|
|
|
void RefinementGPUEngine::Fcalc(std::vector<std::array<float, 2>> &f) {
|
|
RefinementGPUEngineImpl &e = *impl_;
|
|
DeviceGuard guard(e.device);
|
|
e.Transform();
|
|
e.Compose(false);
|
|
f.resize(e.setup.rows.size());
|
|
if (!f.empty())
|
|
cuda_err(cudaMemcpyAsync(f.data(), e.f, f.size() * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
}
|
|
|
|
void RefinementGPUEngine::Fmask(const std::vector<ModelMaskAtom> &mask_atoms, std::vector<std::array<float, 2>> &f) {
|
|
RefinementGPUEngineImpl &e = *impl_;
|
|
DeviceGuard guard(e.device);
|
|
if (mask_atoms.size() > e.setup.max_atoms)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: mask atoms over the reserve");
|
|
if (!mask_atoms.empty())
|
|
cuda_err(cudaMemcpyAsync(e.mask_atoms, mask_atoms.data(), mask_atoms.size() * sizeof(ModelMaskAtom),
|
|
cudaMemcpyHostToDevice, e.stream));
|
|
const bool own = e.mask_grid.get() != nullptr;
|
|
float *grid = own ? e.mask_grid.get() : e.grid.get();
|
|
float2 *spectrum = own ? e.mask_spectrum.get() : e.spectrum.get();
|
|
e.mask->Compute(e.mask_atoms, static_cast<int>(mask_atoms.size()), grid);
|
|
cufft_err((own ? e.mask_forward : e.forward).R2C(grid, reinterpret_cast<cufftComplex *>(spectrum)));
|
|
const int rows = static_cast<int>(e.setup.rows.size());
|
|
f.resize(rows);
|
|
if (rows > 0) {
|
|
const ModelDensityGrid &g = e.setup.mask_grid;
|
|
fmask_kernel<<<blocks(rows), THREADS, 0, e.stream>>>(spectrum, e.MaskNorm(), e.row_hkl, rows, g.nu, g.nv, g.nw,
|
|
e.fm);
|
|
cuda_err(cudaGetLastError());
|
|
cuda_err(cudaMemcpyAsync(f.data(), e.fm, rows * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
|
|
}
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
}
|
|
|
|
double RefinementGPUEngine::Evaluate(std::vector<std::array<float, 4>> *grad) {
|
|
RefinementGPUEngineImpl &e = *impl_;
|
|
DeviceGuard guard(e.device);
|
|
e.QueueEvaluate(grad != nullptr);
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
if (grad != nullptr)
|
|
CopyOut(e.host_grad, e.n_atoms, *grad);
|
|
return e.host_sums[0];
|
|
}
|
|
|
|
double RefinementGPUEngine::EvaluateWithRestraints(std::vector<std::array<float, 4>> *grad, int which,
|
|
const std::function<void(std::array<float, 4> *)> &xyzb,
|
|
std::vector<std::array<float, 4>> *restraint_grad,
|
|
double &restraints) {
|
|
RefinementGPUEngineImpl &e = *impl_;
|
|
DeviceGuard guard(e.device);
|
|
RefinementGPUEngineImpl::Restraints &r = e.restraints[which];
|
|
if (r.atoms != e.n_atoms)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: restraint atoms do not match");
|
|
e.QueueEvaluate(grad != nullptr);
|
|
xyzb(reinterpret_cast<std::array<float, 4> *>(e.host_xyzb.get()));
|
|
e.QueueRestraints(which, restraint_grad != nullptr);
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
if (grad != nullptr)
|
|
CopyOut(e.host_grad, e.n_atoms, *grad);
|
|
if (restraint_grad != nullptr)
|
|
CopyOut(e.host_restraint_grad, r.atoms, *restraint_grad);
|
|
restraints = e.host_sums[1];
|
|
return e.host_sums[0];
|
|
}
|
|
|
|
void RefinementGPUEngine::SetRestraints(int which, const std::vector<RestraintRecord> &records,
|
|
const std::vector<int> &plane_atoms, const std::vector<int> &slot_start,
|
|
const std::vector<int> &slot_index, int n_slots) {
|
|
RefinementGPUEngineImpl &e = *impl_;
|
|
DeviceGuard guard(e.device);
|
|
RefinementGPUEngineImpl::Restraints &r = e.restraints[which];
|
|
for (const RestraintRecord &rec : records)
|
|
if (rec.kind == static_cast<int>(RestraintKind::Plane) && rec.atom[1] > MAX_PLANE_ATOMS)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: a plane of too many atoms");
|
|
r.records = static_cast<int>(records.size());
|
|
r.atoms = static_cast<int>(slot_start.empty() ? 0 : slot_start.size() - 1);
|
|
Ensure(r.rec, r.cap_records, records.size());
|
|
Ensure(r.energy, r.cap_reduce, records.size());
|
|
Ensure(r.plane_atoms, r.cap_planes, plane_atoms.size());
|
|
Ensure(r.slot_start, r.cap_start, slot_start.size());
|
|
Ensure(r.slot_index, r.cap_index, slot_index.size());
|
|
Ensure(r.contrib, r.cap_slots, static_cast<size_t>(n_slots));
|
|
Ensure(r.xyzb, r.cap_xyzb, r.atoms);
|
|
Ensure(r.grad, r.cap_atoms, r.atoms);
|
|
if (!r.sum)
|
|
r.sum = CudaDevicePtr<double>(1, CudaAlloc::Synchronous);
|
|
const size_t bytes = ReduceBytes(r.cap_reduce);
|
|
if (!r.reduce_temp || bytes > r.reduce_bytes) {
|
|
r.reduce_bytes = bytes;
|
|
r.reduce_temp = CudaDevicePtr<char>(std::max<size_t>(bytes, 1), CudaAlloc::Synchronous);
|
|
}
|
|
Upload(r.rec, records, e.stream);
|
|
Upload(r.plane_atoms, plane_atoms, e.stream);
|
|
Upload(r.slot_start, slot_start, e.stream);
|
|
Upload(r.slot_index, slot_index, e.stream);
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
}
|
|
|
|
double RefinementGPUEngine::EvaluateRestraints(int which, const std::vector<std::array<float, 4>> &xyzb,
|
|
std::vector<std::array<float, 4>> *grad) {
|
|
RefinementGPUEngineImpl &e = *impl_;
|
|
DeviceGuard guard(e.device);
|
|
RefinementGPUEngineImpl::Restraints &r = e.restraints[which];
|
|
if (static_cast<int>(xyzb.size()) != r.atoms)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: restraint atoms do not match");
|
|
// the staging buffer free again
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
std::copy(xyzb.begin(), xyzb.end(), reinterpret_cast<std::array<float, 4> *>(e.host_xyzb.get()));
|
|
e.QueueRestraints(which, grad != nullptr);
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
if (grad != nullptr)
|
|
CopyOut(e.host_restraint_grad, r.atoms, *grad);
|
|
return e.host_sums[1];
|
|
}
|