leonarski_f
e3782cabe4
Every evaluation uploaded both copies of the atoms' densities (the gather's and the gradient's) and their boxes, although in a positions-only minimisation no B changes between two Refits. The engine now keeps one copy (the gradient kernel reads ModelDensityGPU's atoms and boxes), uploads it only where some atom's density was made again, and sends the positions from a pinned buffer without waiting. The densities that are made again (every atom in a B minimisation) and the grid positions are worked out on the refinement's threads. Same bits. 9yl4: upload 3.29 -> 0.56 s, host side 1.29 -> 0.61 s, refinement 22.3 -> 17.7 s; 5reo 4.8 -> 4.4 s (RTX 5080, -N 8, loaded box). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 46c3b67c91aee0aec3473b0fb1a4945c00bda070) (cherry picked from commit 0aebc98110c8155932af647d6204891aec662294)
695 lines
32 KiB
Plaintext
695 lines
32 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};
|
|
}
|
|
|
|
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
|
|
CudaDevicePtr<ModelMaskAtom> mask_atoms;
|
|
CudaDevicePtr<float> grid;
|
|
CudaDevicePtr<float2> spectrum;
|
|
CudaDevicePtr<char> fft_work;
|
|
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;
|
|
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);
|
|
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.grid));
|
|
ModelMaskGrid mask_grid;
|
|
mask_grid.nu = setup.grid.nu;
|
|
mask_grid.nv = setup.grid.nv;
|
|
mask_grid.nw = setup.grid.nw;
|
|
std::copy(setup.grid.orth, setup.grid.orth + 9, mask_grid.orth);
|
|
mask_grid.volume = setup.volume;
|
|
mask->SetGrid(mask_grid, setup.mask_ops);
|
|
|
|
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 (forward.WorkBytes() > work_bytes || inverse.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))); }
|
|
|
|
// 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());
|
|
}
|
|
};
|
|
|
|
size_t RefinementGPUEngine::FFTWorkBytes(int device, const ModelDensityGrid &grid) {
|
|
DeviceGuard guard(device);
|
|
return std::max(CudaFFTPlan::WorkBytes(ForwardShape(grid)), CudaFFTPlan::WorkBytes(InverseShape(grid)));
|
|
}
|
|
|
|
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.grid));
|
|
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));
|
|
e.mask->Compute(e.mask_atoms, static_cast<int>(mask_atoms.size()), e.grid);
|
|
cufft_err(e.forward.R2C(e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
|
|
const int rows = static_cast<int>(e.setup.rows.size());
|
|
f.resize(rows);
|
|
if (rows > 0) {
|
|
fmask_kernel<<<blocks(rows), THREADS, 0, e.stream>>>(e.spectrum, e.Norm(), e.row_hkl, rows, e.setup.grid.nu,
|
|
e.setup.grid.nv, e.setup.grid.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);
|
|
const int rows = static_cast<int>(e.setup.rows.size());
|
|
e.Transform();
|
|
e.Compose(true);
|
|
double sum = 0;
|
|
if (rows > 0) {
|
|
size_t temp = e.reduce_bytes;
|
|
cuda_err(cub::DeviceReduce::Sum(e.reduce_temp.get(), temp, e.nll.get(), e.nll_sum.get(), rows, e.stream));
|
|
cuda_err(cudaMemcpyAsync(&sum, e.nll_sum, sizeof(double), cudaMemcpyDeviceToHost, e.stream));
|
|
}
|
|
if (grad != nullptr) {
|
|
const ModelDensityGrid &g = e.setup.grid;
|
|
cuda_err(cudaMemsetAsync(e.spectrum, 0, ComplexPoints(g) * sizeof(float2), e.stream));
|
|
const int points = static_cast<int>(e.setup.scatter_point.size());
|
|
if (points > 0) {
|
|
scatter_kernel<<<blocks(points), THREADS, 0, e.stream>>>(e.coef, e.row_hkl, e.ops,
|
|
static_cast<int>(e.setup.sym_ops.size()), e.setup.den,
|
|
e.scatter_point, e.scatter_start, e.scatter_term,
|
|
points, e.spectrum);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
cufft_err(e.inverse.C2R(reinterpret_cast<cufftComplex *>(e.spectrum.get()), e.grid));
|
|
if (e.n_atoms > 0) {
|
|
const size_t threads = static_cast<size_t>(e.n_atoms) * WARP;
|
|
atom_gradient_kernel<<<blocks(threads), THREADS, 0, e.stream>>>(e.density->DeviceAtoms(), e.pos,
|
|
e.density->DeviceBoxes(), e.n_atoms, e.geom,
|
|
e.grid, e.Norm(), e.grad);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
grad->resize(e.n_atoms);
|
|
if (e.n_atoms > 0)
|
|
cuda_err(cudaMemcpyAsync(grad->data(), e.grad, e.n_atoms * sizeof(float4), cudaMemcpyDeviceToHost, e.stream));
|
|
}
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
return sum;
|
|
}
|
|
|
|
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");
|
|
if (r.atoms > 0)
|
|
cuda_err(cudaMemcpyAsync(r.xyzb, xyzb.data(), xyzb.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream));
|
|
double sum = 0;
|
|
if (r.records > 0) {
|
|
restraint_kernel<<<blocks(r.records), THREADS, 0, e.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, e.stream));
|
|
cuda_err(cudaMemcpyAsync(&sum, r.sum, sizeof(double), cudaMemcpyDeviceToHost, e.stream));
|
|
}
|
|
if (grad != nullptr) {
|
|
grad->resize(r.atoms);
|
|
if (r.atoms > 0) {
|
|
if (r.records > 0) {
|
|
slot_gather_kernel<<<blocks(r.atoms), THREADS, 0, e.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), e.stream));
|
|
cuda_err(cudaMemcpyAsync(grad->data(), r.grad, r.atoms * sizeof(float4), cudaMemcpyDeviceToHost, e.stream));
|
|
}
|
|
}
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
return sum;
|
|
}
|