leonarski_fandClaude Fable 5.1
128e3bb9ae
exp(2 pi i h.t) depends only on the integer h.t (in units of 1/den), so the engine tabulates it once over the range the rows and operators make, with the same op_phase() at the same double argument - the same bits. The compose kernel reads the table at h.t; the scatter terms carry their table entry, and their row directly, so neither hot kernel divides: the rotations go to the device in whole units (gemmi's are den times an integer matrix), the term's row no longer has to be split from its operator. The table is a few thousand entries, the scatter terms one int wider. Outputs are byte-identical; the X-ray evaluation counter drops 7-10% on the demo-sized sets. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
851 lines
39 KiB
Plaintext
851 lines
39 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) for h.t = ht / den, exact for the fractions a translation is made of.
|
|
__device__ __forceinline__ void op_phase(int ht, int den, double &c, double &s) {
|
|
sincospi(2.0 * ht / den, &s, &c);
|
|
}
|
|
|
|
// The phases of every h.t the rows and operators make, as a table over the integer h.t from `min`: the hot
|
|
// kernels read it in place of op_phase() - the same function at the same argument, so the same bits.
|
|
__global__ void phase_table_kernel(int min, int count, int den, double2 *phase) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= count)
|
|
return;
|
|
double s, c;
|
|
op_phase(min + i, den, c, s);
|
|
phase[i] = make_double2(c, s);
|
|
}
|
|
|
|
// 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. `ops` holds each operator's rotation in whole
|
|
// units and its translation in units of 1 / den.
|
|
__global__ void compose_target_kernel(const float2 *f1, float norm, int nu, int nv, int nw, const int *ops, int n_ops,
|
|
const double2 *phase, int phase_min, 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];
|
|
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;
|
|
const double2 p = phase[h[0] * tran[0] + h[1] * tran[1] + h[2] * tran[2] - phase_min];
|
|
sr += p.x * vr - p.y * vi;
|
|
si += p.x * vi + p.y * 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. A term is
|
|
// its row * 2 + conj, and its phase the table's entry at term_phase.
|
|
__global__ void scatter_kernel(const float2 *coef, const double2 *phase, const int *point, const int *start,
|
|
const int *term, const int *term_phase, 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;
|
|
const double2 p = phase[term_phase[e]];
|
|
const double ar = 0.5 * (coef[m].x * p.x - coef[m].y * p.y), ai = 0.5 * (coef[m].x * p.y + coef[m].y * p.x);
|
|
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};
|
|
// the box's first point, unwrapped and wrapped into the grid; a box is narrower than the cell (Supports())
|
|
int lo[3], first[3], span[3];
|
|
const int n3[3] = {g.nu, g.nv, g.nw};
|
|
for (int k = 0; k < 3; k++) {
|
|
lo[k] = __float2int_rn(nctr[k]) - box[3 * a + k];
|
|
first[k] = imod(lo[k], n3[k]);
|
|
span[k] = 2 * box[3 * a + k] + 1;
|
|
}
|
|
const int total = span[0] * span[1] * span[2];
|
|
// the lane's points t = lane, lane + WARP, ... as (iu, iv, iw) in the box, stepped on without dividing
|
|
int iu = lane % span[0], iv = lane / span[0] % span[1], iw = lane / (span[0] * span[1]);
|
|
const int step_u = WARP % span[0], step_v = WARP / span[0];
|
|
auto next = [&]() {
|
|
iu += step_u;
|
|
iv += step_v;
|
|
if (iu >= span[0]) {
|
|
iu -= span[0];
|
|
iv++;
|
|
}
|
|
while (iv >= span[1]) {
|
|
iv -= span[1];
|
|
iw++;
|
|
}
|
|
};
|
|
float gx = 0, gy = 0, gz = 0, gb = 0;
|
|
for (int t = lane; t < total; t += WARP, next()) {
|
|
const int u = lo[0] + iu, v = lo[1] + iv, w = lo[2] + iw;
|
|
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 int mu = first[0] + iu, mv = first[1] + iv, mw = first[2] + iw;
|
|
const float value = map[(mu < g.nu ? mu : mu - g.nu) +
|
|
static_cast<size_t>(g.nu) * ((mv < g.nv ? mv : mv - g.nv) +
|
|
static_cast<size_t>(g.nv) * (mw < g.nw ? mw : mw - 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: {
|
|
float gi = 0, gj = 0;
|
|
e = refine_math::BSimilarityTerm(xyzb[r.atom[0]].w, xyzb[r.atom[1]].w, r.sigma, gi, gj);
|
|
contrib[r.slot] = make_float4(0, 0, 0, gi);
|
|
contrib[r.slot + 1] = make_float4(0, 0, 0, gj);
|
|
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);
|
|
}
|
|
|
|
// h.t of a row and an operator, in units of 1 / den.
|
|
int HT(const std::array<int, 3> &h, const std::array<int, 12> &op) {
|
|
return h[0] * op[9] + h[1] * op[10] + h[2] * op[11];
|
|
}
|
|
|
|
// The range of h.t over the rows and operators: the phase table's.
|
|
void PhaseRange(const RefinementGPUSetup &setup, int &min, int &max) {
|
|
min = max = 0;
|
|
for (const auto &h : setup.rows)
|
|
for (const auto &op : setup.sym_ops) {
|
|
const int ht = HT(h, op);
|
|
min = std::min(min, ht);
|
|
max = std::max(max, ht);
|
|
}
|
|
}
|
|
|
|
// 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;
|
|
StreamWaiter waiter; // the engine's waits on the stream, asleep
|
|
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, scatter_phase;
|
|
CudaDevicePtr<double2> phase; // exp(2 pi i h.t) over the integer h.t from phase_min
|
|
int phase_min = 0;
|
|
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);
|
|
scatter_phase = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_term.size(), 1), sync);
|
|
int phase_max = 0;
|
|
PhaseRange(setup, phase_min, phase_max);
|
|
const int phase_count = phase_max - phase_min + 1;
|
|
phase = CudaDevicePtr<double2>(phase_count, 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);
|
|
}
|
|
|
|
// the rotations in whole units (gemmi's are den times an integer matrix), the translations as given
|
|
std::vector<int> flat_ops;
|
|
for (const auto &o : setup.sym_ops) {
|
|
for (int i = 0; i < 9; i++)
|
|
flat_ops.push_back(o[i] / setup.den);
|
|
flat_ops.insert(flat_ops.end(), o.begin() + 9, o.end());
|
|
}
|
|
std::vector<int> hkl;
|
|
for (const auto &h : setup.rows)
|
|
hkl.insert(hkl.end(), h.begin(), h.end());
|
|
// each scatter term as row * 2 + conj, and its phase's place in the table
|
|
const int n_ops = static_cast<int>(setup.sym_ops.size());
|
|
std::vector<int> term, term_phase;
|
|
for (const int t : setup.scatter_term) {
|
|
const int m = (t >> 1) / n_ops, o = (t >> 1) % n_ops;
|
|
term.push_back(m * 2 + (t & 1));
|
|
term_phase.push_back(HT(setup.rows[m], setup.sym_ops[o]) - phase_min);
|
|
}
|
|
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, term, stream);
|
|
Upload(scatter_phase, term_phase, stream);
|
|
phase_table_kernel<<<blocks(phase_count), THREADS, 0, stream>>>(phase_min, phase_count, setup.den, phase);
|
|
cuda_err(cudaGetLastError());
|
|
|
|
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()),
|
|
phase, phase_min, 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, phase, scatter_point, scatter_start,
|
|
scatter_term, scatter_phase, 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() + 2 * s.scatter_term.size()) * sizeof(int);
|
|
int phase_min = 0, phase_max = 0;
|
|
PhaseRange(s, phase_min, phase_max);
|
|
bytes += static_cast<size_t>(phase_max - phase_min + 1) * sizeof(double2);
|
|
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);
|
|
// the two restraint sets of SetRestraints(), each up to these
|
|
bytes += 2 * (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(e.waiter.Wait(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(e.waiter.Wait(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(e.waiter.Wait(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(e.waiter.Wait(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(e.waiter.Wait(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(e.waiter.Wait(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(e.waiter.Wait(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(e.waiter.Wait(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(e.waiter.Wait(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(e.waiter.Wait(e.stream));
|
|
if (grad != nullptr)
|
|
CopyOut(e.host_restraint_grad, r.atoms, *grad);
|
|
return e.host_sums[1];
|
|
}
|