leonarski_fandClaude Opus 5.5
1227d4168f
The P1 cross-check (<prefix>_P1.mtz) and --no-p1-crosscheck are gone. The tail waited for that merge after the reflection files were written: 12-18 s on a large rotation set (7200 frames), ~2 s on a small protein. A user who wants a P1 merge re-runs with -S 1. --finalist-ledger keeps working: it now starts a P1 merge of its own beside the analyses (made only when the flag is given) and folds the ledger from it. The twin-immune zone evidence, which was read on the cross-check merge, moves with it and is reported only with --finalist-ledger. The unmerged MTZ was built in the shadow of the cross-check merge; it is now built beside the reflection files, from the point where the model validation has settled the indexing. WriteMtzReflections loses its amplitudes flag (only the cross-check wrote an intensity-only MTZ). The merged outputs (p.mtz, .hkl, _unmerged.mtz) are byte-identical to the base on the gate matrix. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
972 lines
45 KiB
Plaintext
972 lines
45 KiB
Plaintext
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "RigidBodyGPUEngine.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <map>
|
|
#include <string>
|
|
|
|
#include <cub/cub.cuh>
|
|
#include <cufft.h>
|
|
|
|
#include "ModelMaskGPU.h"
|
|
#include "ModelScaleGPU.h"
|
|
#include "../../common/JFJochException.h"
|
|
#include "../indexing/CUDAMemHelpers.h"
|
|
|
|
// Every kernel below is deterministic: nothing is summed with a floating-point atomic, each grid point
|
|
// adds its atoms in model order, and every reduction runs in a fixed order. The same placement gives the
|
|
// same residuals, bit for bit, on every run and whichever engine runs it.
|
|
|
|
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 BRICK = 8; // the gather's bricks are BRICK^3 grid points, one block each
|
|
constexpr int GATHER_TILE = 128; // atoms staged in shared memory at a time
|
|
constexpr int MAX_BRICKS_PER_AXIS = 8; // an atom's box touches at most this many bricks along an axis
|
|
constexpr int THREADS = 256;
|
|
constexpr int MAX_SCALE_PARAMS = 7; // k_overall + up to six B* constraints
|
|
|
|
// The grid of a zone as the kernels see it.
|
|
struct GridGeom {
|
|
int nu, nv, nw;
|
|
float orth_n[9]; // orth * diag(1/nu, 1/nv, 1/nw), row-major: grid offset -> Cartesian
|
|
int narrow; // the cell is too small for one image of an atom to serve a whole brick
|
|
};
|
|
|
|
// Where the body is: x -> R x_rel + t, and the fractionalization.
|
|
struct Placement {
|
|
double r[9];
|
|
double t[3];
|
|
double frac[9];
|
|
};
|
|
|
|
// The scale a residual or a Jacobian row is taken at.
|
|
struct Scale {
|
|
double k_overall;
|
|
double b_star[6]; // u11 u22 u33 u12 u13 u23
|
|
double k_sol, b_sol;
|
|
int n_constraints;
|
|
double constraints[6][6];
|
|
};
|
|
|
|
__device__ __forceinline__ int imod(int a, int n) {
|
|
const int r = a % n;
|
|
return r < 0 ? r + n : r;
|
|
}
|
|
|
|
// gemmi's unsafe_expapprox() (formfact.hpp), the exponential the density is computed with.
|
|
__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))));
|
|
}
|
|
|
|
// b_star.r_u_r(h), as gemmi's SMat33.
|
|
__device__ __forceinline__ double r_u_r(const double *u, const int *h) {
|
|
const double x = h[0], y = h[1], z = h[2];
|
|
return x * x * u[0] + y * y * u[1] + z * z * u[2] + 2 * (x * y * u[3] + x * z * u[4] + y * z * u[5]);
|
|
}
|
|
|
|
// Each atom placed: its grid coordinates (fractional times the grid size, wrapped into the cell) for
|
|
// the gather and, for a mask atom, its wrapped fractional coordinates and mask radius for the mask.
|
|
__global__ void place_kernel(const double *rel, int n, Placement p, int nu, int nv, int nw, const int *mask_slot,
|
|
const float *mask_radius, float4 *pos, double *mask_atoms) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n)
|
|
return;
|
|
const double x = rel[3 * i], y = rel[3 * i + 1], z = rel[3 * i + 2];
|
|
const double c[3] = {p.r[0] * x + p.r[1] * y + p.r[2] * z + p.t[0],
|
|
p.r[3] * x + p.r[4] * y + p.r[5] * z + p.t[1],
|
|
p.r[6] * x + p.r[7] * y + p.r[8] * z + p.t[2]};
|
|
double f[3];
|
|
for (int k = 0; k < 3; k++) {
|
|
f[k] = p.frac[3 * k] * c[0] + p.frac[3 * k + 1] * c[1] + p.frac[3 * k + 2] * c[2];
|
|
f[k] -= floor(f[k]);
|
|
if (f[k] >= 1.0) // a tiny negative coordinate wraps to 1 - epsilon, which rounds to 1
|
|
f[k] = 0.0;
|
|
}
|
|
const int n3[3] = {nu, nv, nw};
|
|
float g[3];
|
|
for (int k = 0; k < 3; k++) {
|
|
g[k] = static_cast<float>(f[k] * n3[k]);
|
|
if (g[k] >= n3[k])
|
|
g[k] -= n3[k];
|
|
}
|
|
pos[i] = make_float4(g[0], g[1], g[2], 0.0f);
|
|
const int slot = mask_slot[i];
|
|
if (slot >= 0) {
|
|
for (int k = 0; k < 3; k++)
|
|
mask_atoms[4 * slot + k] = f[k];
|
|
mask_atoms[4 * slot + 3] = mask_radius[i];
|
|
}
|
|
}
|
|
|
|
// The distinct bricks the points c - d ... c + d of one axis fall in, wrapped into the cell. A grid size
|
|
// that is not a multiple of BRICK leaves the last brick partial, which is why this is done on wrapped
|
|
// points and not in brick coordinates. -1 where there are more than MAX_BRICKS_PER_AXIS of them.
|
|
__device__ int axis_bricks(int c, int d, int n, int *out) {
|
|
int k = 0;
|
|
for (int p = c - d; p <= c + d; p++) {
|
|
const int b = imod(p, n) / BRICK;
|
|
bool seen = false;
|
|
for (int j = 0; j < k; j++)
|
|
seen = seen || out[j] == b;
|
|
if (seen)
|
|
continue;
|
|
if (k == MAX_BRICKS_PER_AXIS)
|
|
return -1;
|
|
out[k++] = b;
|
|
}
|
|
return k;
|
|
}
|
|
|
|
// The bricks atom i's box touches, per axis. The box is gemmi's: `box` points either side of the
|
|
// point nearest the atom.
|
|
__device__ void atom_bricks(const float4 &pos, const int *box, int nu, int nv, int nw, int *bu, int &nbu, int *bv,
|
|
int &nbv, int *bw, int &nbw) {
|
|
nbu = axis_bricks(__float2int_rn(pos.x), box[0], nu, bu);
|
|
nbv = axis_bricks(__float2int_rn(pos.y), box[1], nv, bv);
|
|
nbw = axis_bricks(__float2int_rn(pos.z), box[2], nw, bw);
|
|
}
|
|
|
|
__global__ void count_pairs_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw, int *count,
|
|
int *overflow) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i > n)
|
|
return;
|
|
if (i == n) {
|
|
count[n] = 0; // so the exclusive scan's last entry is the total
|
|
return;
|
|
}
|
|
int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c;
|
|
atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c);
|
|
if (a < 0 || b < 0 || c < 0) {
|
|
*overflow = 1;
|
|
a = b = c = 0;
|
|
}
|
|
count[i] = a * b * c;
|
|
}
|
|
|
|
// The (brick, atom) pairs, atom-major: sorted stably by brick, each brick's atoms stay in model order.
|
|
__global__ void fill_pairs_kernel(const float4 *pos, const int *box, int n, int nu, int nv, int nw,
|
|
const int *offset, int *key, int *value) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n)
|
|
return;
|
|
int bu[MAX_BRICKS_PER_AXIS], bv[MAX_BRICKS_PER_AXIS], bw[MAX_BRICKS_PER_AXIS], a, b, c;
|
|
atom_bricks(pos[i], box + 3 * i, nu, nv, nw, bu, a, bv, b, bw, c);
|
|
const int nbu = (nu + BRICK - 1) / BRICK, nbv = (nv + BRICK - 1) / BRICK;
|
|
int o = offset[i];
|
|
for (int k = 0; k < c; k++)
|
|
for (int j = 0; j < b; j++)
|
|
for (int l = 0; l < a; l++) {
|
|
key[o] = bu[l] + nbu * (bv[j] + nbv * bw[k]);
|
|
value[o] = i;
|
|
o++;
|
|
}
|
|
}
|
|
|
|
__global__ void brick_ranges_kernel(const int *key, int npairs, int *start, int *end) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= npairs)
|
|
return;
|
|
const int k = key[i];
|
|
if (i == 0 || key[i - 1] != k)
|
|
start[k] = i;
|
|
if (i == npairs - 1 || key[i + 1] != k)
|
|
end[k] = i + 1;
|
|
}
|
|
|
|
// The density of the model on the grid, as gemmi's do_add_atom_density_to_grid() puts it there: every
|
|
// point adds, in model order, each atom whose sphere it is inside. One block per brick; the brick's atoms
|
|
// are staged in shared memory with the position of their image nearest to the brick, relative to the
|
|
// brick's first point. Where the cell is wider than an atom's box plus a brick, that is the only image
|
|
// that reaches any point of the brick; on a narrower cell each point finds its own nearest image.
|
|
__global__ void gather_kernel(const RigidBodyGPUAtom *atoms, const float4 *pos, const int *pair_atom,
|
|
const int *brick_start, const int *brick_end, GridGeom g, float *grid) {
|
|
const int nbu = (g.nu + BRICK - 1) / BRICK, nbv = (g.nv + BRICK - 1) / BRICK;
|
|
const int brick = blockIdx.x;
|
|
const int s = brick_start[brick], e = brick_end[brick]; // s == e: an empty brick, written as zeros
|
|
const int u0 = (brick % nbu) * BRICK, v0 = (brick / nbu % nbv) * BRICK, w0 = (brick / (nbu * nbv)) * BRICK;
|
|
const int u = u0 + threadIdx.x, v = v0 + threadIdx.y, w = w0 + threadIdx.z;
|
|
const bool inside = u < g.nu && v < g.nv && w < g.nw;
|
|
const int tid = threadIdx.x + BRICK * (threadIdx.y + BRICK * threadIdx.z);
|
|
// This point relative to the brick's first one, in Cartesian coordinates.
|
|
const float tx = g.orth_n[0] * threadIdx.x + g.orth_n[1] * threadIdx.y + g.orth_n[2] * threadIdx.z;
|
|
const float ty = g.orth_n[3] * threadIdx.x + g.orth_n[4] * threadIdx.y + g.orth_n[5] * threadIdx.z;
|
|
const float tz = g.orth_n[6] * threadIdx.x + g.orth_n[7] * threadIdx.y + g.orth_n[8] * threadIdx.z;
|
|
|
|
constexpr int WORDS = sizeof(RigidBodyGPUAtom) / sizeof(int);
|
|
__shared__ int sh_index[GATHER_TILE];
|
|
__shared__ RigidBodyGPUAtom sh_atom[GATHER_TILE];
|
|
__shared__ float3 sh_centre[GATHER_TILE]; // Cartesian
|
|
__shared__ float3 sh_grid[GATHER_TILE]; // the same, in grid units
|
|
float acc = 0.0f;
|
|
for (int c = s; c < e; c += GATHER_TILE) {
|
|
const int m = min(GATHER_TILE, e - c);
|
|
__syncthreads();
|
|
if (tid < m) {
|
|
const int a = pair_atom[c + tid];
|
|
sh_index[tid] = a;
|
|
// The atom's image nearest to the brick's centre, in grid units from its first point.
|
|
const float4 p = pos[a];
|
|
float fx = p.x - u0, fy = p.y - v0, fz = p.z - w0;
|
|
fx -= g.nu * rintf((fx - 0.5f * (BRICK - 1)) / g.nu);
|
|
fy -= g.nv * rintf((fy - 0.5f * (BRICK - 1)) / g.nv);
|
|
fz -= g.nw * rintf((fz - 0.5f * (BRICK - 1)) / g.nw);
|
|
sh_grid[tid] = make_float3(fx, fy, fz);
|
|
sh_centre[tid] = make_float3(g.orth_n[0] * fx + g.orth_n[1] * fy + g.orth_n[2] * fz,
|
|
g.orth_n[3] * fx + g.orth_n[4] * fy + g.orth_n[5] * fz,
|
|
g.orth_n[6] * fx + g.orth_n[7] * fy + g.orth_n[8] * fz);
|
|
}
|
|
__syncthreads();
|
|
for (int i = tid; i < m * WORDS; i += BRICK * BRICK * BRICK)
|
|
reinterpret_cast<int *>(sh_atom)[i] = reinterpret_cast<const int *>(atoms + sh_index[i / WORDS])[i % WORDS];
|
|
__syncthreads();
|
|
if (!inside)
|
|
continue;
|
|
for (int k = 0; k < m; k++) {
|
|
const RigidBodyGPUAtom &at = sh_atom[k];
|
|
float x = sh_centre[k].x - tx, y = sh_centre[k].y - ty, z = sh_centre[k].z - tz;
|
|
if (g.narrow) {
|
|
float fx = sh_grid[k].x - threadIdx.x, fy = sh_grid[k].y - threadIdx.y, fz = sh_grid[k].z - threadIdx.z;
|
|
fx -= g.nu * rintf(fx / g.nu);
|
|
fy -= g.nv * rintf(fy / g.nv);
|
|
fz -= g.nw * rintf(fz / g.nw);
|
|
x = g.orth_n[0] * fx + g.orth_n[1] * fy + g.orth_n[2] * fz;
|
|
y = g.orth_n[3] * fx + g.orth_n[4] * fy + g.orth_n[5] * fz;
|
|
z = g.orth_n[6] * fx + g.orth_n[7] * fy + g.orth_n[8] * fz;
|
|
}
|
|
const float r2 = x * x + y * y + z * z;
|
|
if (r2 > at.radius * at.radius)
|
|
continue;
|
|
float density = 0.0f;
|
|
if (!at.aniso) {
|
|
for (int q = 0; q < 5; q++)
|
|
density += at.a[q] * expapprox(fmaxf(at.b[q][0] * r2, -88.f));
|
|
} else {
|
|
for (int q = 0; q < 5; q++) {
|
|
const float *b = at.b[q];
|
|
const float rur = x * x * b[0] + y * y * b[1] + z * z * b[2] + 2 * (x * y * b[3] + x * z * b[4] + y * z * b[5]);
|
|
density += at.a[q] * expapprox(fmaxf(rur, -88.f));
|
|
}
|
|
}
|
|
acc += at.occ * density;
|
|
}
|
|
}
|
|
if (inside)
|
|
grid[u + static_cast<size_t>(g.nu) * (v + static_cast<size_t>(g.nv) * w)] = acc;
|
|
}
|
|
|
|
// F at the zone's rows composed from the transform of one copy (SymmetryComposition::Compose), with
|
|
// dF/dt where `df_dt` is not null. `f1` is cuFFT's forward transform, which carries exp(-2 pi i h.x):
|
|
// conjugated and scaled by V/N it is gemmi's F1.
|
|
__global__ void compose_kernel(const float2 *f1, float norm, const RigidBodyGPUTerm *terms, const double *row_scale,
|
|
int rows, int ops, double2 *f, double2 *df_dt) {
|
|
const int m = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (m >= rows)
|
|
return;
|
|
double sr = 0, si = 0, d[3][2] = {};
|
|
for (int o = 0; o < ops; o++) {
|
|
const RigidBodyGPUTerm t = terms[static_cast<size_t>(m) * ops + o];
|
|
const float2 out = f1[t.index];
|
|
double vr = out.x * norm, vi = -out.y * norm; // F1 at the index read
|
|
if (t.conj)
|
|
vi = -vi;
|
|
const double tr = t.phase[0] * vr - t.phase[1] * vi, ti = t.phase[0] * vi + t.phase[1] * vr;
|
|
sr += tr;
|
|
si += ti;
|
|
for (int k = 0; k < 3; k++) {
|
|
d[k][0] += tr * t.s[k];
|
|
d[k][1] += ti * t.s[k];
|
|
}
|
|
}
|
|
const double scale = row_scale[m];
|
|
f[m] = make_double2(scale * sr, scale * si);
|
|
if (df_dt != nullptr) {
|
|
// times 2 pi i
|
|
const double two_pi = 2 * 3.14159265358979323846;
|
|
for (int k = 0; k < 3; k++)
|
|
df_dt[3 * static_cast<size_t>(m) + k] = make_double2(-scale * two_pi * d[k][1], scale * two_pi * d[k][0]);
|
|
}
|
|
}
|
|
|
|
// Fmask at each row, read at h directly (the mask is symmetric): gemmi's get_value_by_hkl() on the
|
|
// transform of the mask, taking the Friedel mate where h is outside the stored half.
|
|
__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);
|
|
}
|
|
|
|
// Scaling<float>::Point's fcmol and fmask of each point.
|
|
__global__ void pack_points_kernel(const int *point_obs, const int *obs_row, int n, const double2 *f,
|
|
const float2 *fmask, float2 *fcmol_p, float2 *fmask_p) {
|
|
const int p = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (p >= n)
|
|
return;
|
|
const int m = obs_row[point_obs[p]];
|
|
fcmol_p[p] = make_float2(static_cast<float>(f[m].x), static_cast<float>(f[m].y));
|
|
fmask_p[p] = fmask[m];
|
|
}
|
|
|
|
// gemmi's Scaling::get_fcalc() of a row: fcmol + (float) solvent scale * fmask, in float.
|
|
__device__ float2 solvent_fcalc(const double2 &f, const float2 &fm, double stol2, const Scale &sc) {
|
|
const float solvent = static_cast<float>(sc.k_sol * exp(-sc.b_sol * stol2));
|
|
return make_float2(static_cast<float>(f.x) + solvent * fm.x, static_cast<float>(f.y) + solvent * fm.y);
|
|
}
|
|
|
|
// The residual of each observation: scale_data() applied to the row, then (Fobs - |F|) / <Fobs>.
|
|
__global__ void residual_kernel(const int *obs_row, const float *obs_fobs, int n, const double2 *f,
|
|
const float2 *fmask, const int *row_hkl, const double *row_stol2, Scale sc,
|
|
double f_mean, double *residual) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n)
|
|
return;
|
|
const int m = obs_row[i];
|
|
if (m < 0) {
|
|
residual[i] = 0.0;
|
|
return;
|
|
}
|
|
float2 value = solvent_fcalc(f[m], fmask[m], row_stol2[m], sc);
|
|
const float k = static_cast<float>(sc.k_overall * exp(-0.25 * r_u_r(sc.b_star, row_hkl + 3 * m)));
|
|
value.x *= k;
|
|
value.y *= k;
|
|
residual[i] = (obs_fobs[i] - hypotf(value.x, value.y)) / f_mean;
|
|
}
|
|
|
|
// One row of the fixed-scale Jacobian and of the scale parameters' own columns (RigidBodyTarget::Jacobian).
|
|
__global__ void jacobian_kernel(const int *obs_row, int n, const double2 *f, const double2 *df_dt,
|
|
const double2 *shifted, size_t rows, const float2 *fmask, const int *row_hkl,
|
|
const double *row_stol2, Scale sc, double f_mean, double step, double *jacobian,
|
|
double *jk) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n)
|
|
return;
|
|
const int p = 1 + sc.n_constraints;
|
|
const int m = obs_row[i];
|
|
if (m < 0) {
|
|
for (int j = 0; j < 6; j++)
|
|
jacobian[6 * static_cast<size_t>(i) + j] = 0.0;
|
|
for (int a = 0; a < p; a++)
|
|
jk[static_cast<size_t>(i) * p + a] = 0.0;
|
|
return;
|
|
}
|
|
const int *h = row_hkl + 3 * m;
|
|
// compute_value_and_derivatives() with k_sol and b_sol fixed
|
|
const float2 fcalc = solvent_fcalc(f[m], fmask[m], row_stol2[m], sc);
|
|
const double fcalc_abs = hypotf(fcalc.x, fcalc.y);
|
|
const double kaniso = exp(-0.25 * r_u_r(sc.b_star, h));
|
|
const double fe = fcalc_abs * kaniso;
|
|
const double y = sc.k_overall * fe;
|
|
const double x = h[0], yy = h[1], z = h[2];
|
|
const double du[6] = {-0.25 * y * (x * x), -0.25 * y * (yy * yy), -0.25 * y * (z * z),
|
|
-0.5 * y * (x * yy), -0.5 * y * (x * z), -0.5 * y * (yy * z)};
|
|
jk[static_cast<size_t>(i) * p] = -fe / f_mean;
|
|
for (int c = 0; c < sc.n_constraints; c++) {
|
|
double dot = 0;
|
|
for (int q = 0; q < 6; q++)
|
|
dot += sc.constraints[c][q] * du[q];
|
|
jk[static_cast<size_t>(i) * p + 1 + c] = -dot / f_mean;
|
|
}
|
|
// d|F_scaled|/dq = K Re(conj(F_total) dFcalc/dq) / |F_total|
|
|
const double k = sc.k_overall * kaniso / (fcalc_abs * f_mean);
|
|
for (int j = 0; j < 6; j++) {
|
|
double2 d;
|
|
if (j < 3) {
|
|
const double2 s = shifted[j * rows + m];
|
|
d = make_double2((s.x - f[m].x) / step, (s.y - f[m].y) / step);
|
|
} else {
|
|
d = df_dt[3 * static_cast<size_t>(m) + j - 3];
|
|
}
|
|
jacobian[6 * static_cast<size_t>(i) + j] = -k * (fcalc.x * d.x + fcalc.y * d.y);
|
|
}
|
|
}
|
|
|
|
// out[b] = sum over i of a[i * lda + ia] * c[i * ldc + ic] for the pair b = (ia, ic): one block per pair,
|
|
// a fixed stride per thread and a fixed shared-memory tree, so the sum is the same on every run.
|
|
__global__ void products_kernel(const double *a, int lda, const double *c, int ldc, int n, int nc, double *out) {
|
|
const int ia = blockIdx.x / nc, ic = blockIdx.x % nc;
|
|
double sum = 0;
|
|
for (int i = threadIdx.x; i < n; i += THREADS)
|
|
sum += a[static_cast<size_t>(i) * lda + ia] * c[static_cast<size_t>(i) * ldc + ic];
|
|
__shared__ double sh[THREADS];
|
|
sh[threadIdx.x] = sum;
|
|
__syncthreads();
|
|
for (int s = THREADS / 2; s > 0; s >>= 1) {
|
|
if (threadIdx.x < s)
|
|
sh[threadIdx.x] += sh[threadIdx.x + s];
|
|
__syncthreads();
|
|
}
|
|
if (threadIdx.x == 0)
|
|
out[blockIdx.x] = sh[0];
|
|
}
|
|
|
|
// J - J_k x, the scale re-fit projected out.
|
|
__global__ void project_kernel(double *jacobian, const double *jk, int p, const double *x, int n) {
|
|
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n)
|
|
return;
|
|
for (int j = 0; j < 6; j++) {
|
|
double v = jacobian[6 * static_cast<size_t>(i) + j];
|
|
for (int a = 0; a < p; a++)
|
|
v -= jk[static_cast<size_t>(i) * p + a] * x[a * 6 + j];
|
|
jacobian[6 * static_cast<size_t>(i) + j] = v;
|
|
}
|
|
}
|
|
|
|
int blocks(size_t n) {
|
|
return static_cast<int>((n + THREADS - 1) / THREADS);
|
|
}
|
|
|
|
// cub's temporary storage for the pair scan and sort at the capacity.
|
|
size_t CubBytes(size_t atoms, size_t pairs) {
|
|
size_t scan = 0, sort = 0;
|
|
cuda_err(cub::DeviceScan::ExclusiveSum(nullptr, scan, static_cast<int *>(nullptr), static_cast<int *>(nullptr),
|
|
static_cast<int>(atoms + 1)));
|
|
cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, sort, static_cast<int *>(nullptr), static_cast<int *>(nullptr),
|
|
static_cast<int *>(nullptr), static_cast<int *>(nullptr),
|
|
static_cast<int>(pairs)));
|
|
return std::max(scan, sort);
|
|
}
|
|
|
|
// A plan for an (nu, nv, nw) R2C transform, u fastest, `batch` grids back to back, whose work area the
|
|
// caller provides; its size in `work`.
|
|
cufftHandle MakePlan(int nu, int nv, int nw, int batch, size_t &work) {
|
|
cufftHandle plan;
|
|
cufft_err(cufftCreate(&plan));
|
|
cufft_err(cufftSetAutoAllocation(plan, 0));
|
|
int n[3] = {nw, nv, nu};
|
|
const int idist = nu * nv * nw, odist = (nu / 2 + 1) * nv * nw;
|
|
const cufftResult r = cufftMakePlanMany(plan, 3, n, nullptr, 1, idist, nullptr, 1, odist, CUFFT_R2C, batch, &work);
|
|
if (r != CUFFT_SUCCESS)
|
|
cufftDestroy(plan);
|
|
cufft_err(r);
|
|
return plan;
|
|
}
|
|
|
|
// gemmi's box around each atom (MakeAtomBox, ModelGrid.cpp): the points within ceil(radius / spacing)
|
|
// of the nearest one along each axis, spacing being the distance between the grid's lattice planes.
|
|
std::vector<int> AtomBoxes(const RigidBodyGPUZone &zone) {
|
|
const int n3[3] = {zone.nu, zone.nv, zone.nw};
|
|
double spacing[3];
|
|
for (int k = 0; k < 3; k++)
|
|
spacing[k] = 1.0 / (n3[k] * std::sqrt(zone.frac[3 * k] * zone.frac[3 * k] + zone.frac[3 * k + 1] * zone.frac[3 * k + 1] +
|
|
zone.frac[3 * k + 2] * zone.frac[3 * k + 2]));
|
|
std::vector<int> box(3 * zone.atoms.size());
|
|
for (size_t i = 0; i < zone.atoms.size(); i++)
|
|
for (int k = 0; k < 3; k++)
|
|
box[3 * i + k] = static_cast<int>(std::ceil(zone.atoms[i].radius / spacing[k]));
|
|
return box;
|
|
}
|
|
|
|
int GreatestStreamPriority() {
|
|
int least = 0, greatest = 0;
|
|
cuda_err(cudaDeviceGetStreamPriorityRange(&least, &greatest));
|
|
return greatest;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
// The most bricks the 2 d + 1 wrapped points of a box can fall in along an axis of n points: one more
|
|
// than the points span for where they start in a brick, and one more again where they wrap past a last
|
|
// brick that is partial (n = 17: the points 15, 16, 0 fall in bricks 1, 2 and 0).
|
|
size_t RigidBodyGPUEngine::AxisBrickBound(int d, int n) {
|
|
const int nb = (n + BRICK - 1) / BRICK;
|
|
return std::min(nb, (2 * d + 1 + BRICK - 1) / BRICK + 2);
|
|
}
|
|
|
|
struct RigidBodyGPUEngineImpl {
|
|
int device = 0;
|
|
RigidBodyGPUCapacity cap;
|
|
// At the device's highest priority: a validation can run beside a merge, whose kernels would
|
|
// otherwise fill the card ahead of a fit's short ones, one evaluation at a time.
|
|
CudaStream stream{cudaStreamNonBlocking, GreatestStreamPriority()};
|
|
|
|
CudaDevicePtr<RigidBodyGPUAtom> atoms;
|
|
CudaDevicePtr<int> box;
|
|
CudaDevicePtr<double> rel;
|
|
CudaDevicePtr<float4> pos;
|
|
CudaDevicePtr<int> mask_slot;
|
|
CudaDevicePtr<float> mask_radius;
|
|
CudaDevicePtr<double> mask_atoms; // 4 per mask atom: fractional x, y, z in [0,1), radius
|
|
|
|
CudaDevicePtr<float> grid; // 3 grids back to back
|
|
CudaDevicePtr<float2> spectrum; // their 3 transforms
|
|
CudaDevicePtr<char> fft_work;
|
|
|
|
CudaDevicePtr<int> overflow; // set by count_pairs_kernel where an atom's box is over MAX_BRICKS_PER_AXIS
|
|
CudaDevicePtr<int> count, offset, key, value, key_sorted, value_sorted, brick_start, brick_end;
|
|
CudaDevicePtr<char> cub_temp;
|
|
size_t cub_bytes = 0;
|
|
|
|
CudaDevicePtr<RigidBodyGPUTerm> terms;
|
|
CudaDevicePtr<int> row_hkl;
|
|
CudaDevicePtr<double> row_scale, row_stol2;
|
|
CudaDevicePtr<double2> f, df_dt, shifted;
|
|
CudaDevicePtr<float2> fmask;
|
|
|
|
CudaDevicePtr<int> obs_row, point_obs;
|
|
CudaDevicePtr<float> obs_fobs;
|
|
CudaDevicePtr<double> residual, jacobian, jk;
|
|
CudaDevicePtr<float2> point_fcmol, point_fmask;
|
|
CudaDevicePtr<double> products, x;
|
|
|
|
std::map<std::array<int, 4>, cufftHandle> plans;
|
|
std::unique_ptr<ModelMaskGPU> mask;
|
|
std::unique_ptr<ModelScaleGPU> scale;
|
|
|
|
// The zone.
|
|
int n_atoms = 0;
|
|
GridGeom geom{};
|
|
double frac[9] = {};
|
|
double volume = 0;
|
|
size_t rows = 0, ops = 1, n_obs = 0, n_points = 0, n_constraints = 0, n_mask_atoms = 0;
|
|
std::vector<std::array<double, 6>> constraints;
|
|
double f_mean = 1;
|
|
|
|
explicit RigidBodyGPUEngineImpl(const RigidBodyGPUCapacity &c, int dev)
|
|
: device(dev), cap(c) {
|
|
const auto sync = CudaAlloc::Synchronous;
|
|
const size_t na = std::max<size_t>(cap.atoms, 1);
|
|
atoms = CudaDevicePtr<RigidBodyGPUAtom>(na, sync);
|
|
box = CudaDevicePtr<int>(3 * na, sync);
|
|
rel = CudaDevicePtr<double>(3 * na, sync);
|
|
pos = CudaDevicePtr<float4>(na, sync);
|
|
mask_slot = CudaDevicePtr<int>(na, sync);
|
|
mask_radius = CudaDevicePtr<float>(na, sync);
|
|
mask_atoms = CudaDevicePtr<double>(4 * na, sync);
|
|
grid = CudaDevicePtr<float>(3 * cap.grid_points, sync);
|
|
spectrum = CudaDevicePtr<float2>(3 * cap.complex_points, sync);
|
|
fft_work = CudaDevicePtr<char>(std::max<size_t>(cap.fft_work_bytes, 1), sync);
|
|
overflow = CudaDevicePtr<int>(1, sync);
|
|
count = CudaDevicePtr<int>(na + 1, sync);
|
|
offset = CudaDevicePtr<int>(na + 1, sync);
|
|
const size_t np = std::max<size_t>(cap.pairs, 1);
|
|
key = CudaDevicePtr<int>(np, sync);
|
|
value = CudaDevicePtr<int>(np, sync);
|
|
key_sorted = CudaDevicePtr<int>(np, sync);
|
|
value_sorted = CudaDevicePtr<int>(np, sync);
|
|
brick_start = CudaDevicePtr<int>(std::max<size_t>(cap.bricks, 1), sync);
|
|
brick_end = CudaDevicePtr<int>(std::max<size_t>(cap.bricks, 1), sync);
|
|
cub_bytes = CubBytes(na, np);
|
|
cub_temp = CudaDevicePtr<char>(std::max<size_t>(cub_bytes, 1), sync);
|
|
const size_t nr = std::max<size_t>(cap.rows, 1), no = std::max<size_t>(cap.observations, 1);
|
|
terms = CudaDevicePtr<RigidBodyGPUTerm>(std::max<size_t>(cap.terms, 1), sync);
|
|
row_hkl = CudaDevicePtr<int>(3 * nr, sync);
|
|
row_scale = CudaDevicePtr<double>(nr, sync);
|
|
row_stol2 = CudaDevicePtr<double>(nr, sync);
|
|
f = CudaDevicePtr<double2>(nr, sync);
|
|
df_dt = CudaDevicePtr<double2>(3 * nr, sync);
|
|
shifted = CudaDevicePtr<double2>(3 * nr, sync);
|
|
fmask = CudaDevicePtr<float2>(nr, sync);
|
|
obs_row = CudaDevicePtr<int>(no, sync);
|
|
point_obs = CudaDevicePtr<int>(no, sync);
|
|
obs_fobs = CudaDevicePtr<float>(no, sync);
|
|
residual = CudaDevicePtr<double>(no, sync);
|
|
jacobian = CudaDevicePtr<double>(6 * no, sync);
|
|
jk = CudaDevicePtr<double>(MAX_SCALE_PARAMS * no, sync);
|
|
point_fcmol = CudaDevicePtr<float2>(no, sync);
|
|
point_fmask = CudaDevicePtr<float2>(no, sync);
|
|
products = CudaDevicePtr<double>(MAX_SCALE_PARAMS * (MAX_SCALE_PARAMS + 6), sync);
|
|
x = CudaDevicePtr<double>(MAX_SCALE_PARAMS * 6, sync);
|
|
mask = std::make_unique<ModelMaskGPU>(stream, std::max<size_t>(cap.grid_points, 1));
|
|
scale = std::make_unique<ModelScaleGPU>(stream, no);
|
|
}
|
|
|
|
// The points' fcmol and fmask at this evaluation.
|
|
void PackPoints() {
|
|
if (n_points == 0)
|
|
return;
|
|
pack_points_kernel<<<blocks(n_points), THREADS, 0, stream>>>(point_obs, obs_row, static_cast<int>(n_points),
|
|
f, fmask, point_fcmol, point_fmask);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
|
|
~RigidBodyGPUEngineImpl() {
|
|
cudaSetDevice(device);
|
|
cudaStreamSynchronize(stream);
|
|
for (auto &[k, plan] : plans)
|
|
cufftDestroy(plan);
|
|
}
|
|
|
|
cufftHandle Plan(int batch) {
|
|
const std::array<int, 4> k{geom.nu, geom.nv, geom.nw, batch};
|
|
auto it = plans.find(k);
|
|
if (it != plans.end())
|
|
return it->second;
|
|
size_t work = 0;
|
|
const cufftHandle plan = MakePlan(geom.nu, geom.nv, geom.nw, batch, work);
|
|
if (work > cap.fft_work_bytes) {
|
|
cufftDestroy(plan);
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: cuFFT work area over the reserve");
|
|
}
|
|
cufft_err(cufftSetWorkArea(plan, fft_work));
|
|
cufft_err(cufftSetStream(plan, stream));
|
|
plans.emplace(k, plan);
|
|
return plan;
|
|
}
|
|
|
|
size_t GridPoints() const { return static_cast<size_t>(geom.nu) * geom.nv * geom.nw; }
|
|
size_t ComplexPoints() const { return static_cast<size_t>(geom.nu / 2 + 1) * geom.nv * geom.nw; }
|
|
float Norm() const { return static_cast<float>(volume / static_cast<double>(GridPoints())); }
|
|
|
|
// The body placed at (R, t) and its density gathered into grid `slot`.
|
|
void Density(const double r[9], const double t[3], int slot) {
|
|
Placement p{};
|
|
std::copy(r, r + 9, p.r);
|
|
std::copy(t, t + 3, p.t);
|
|
std::copy(frac, frac + 9, p.frac);
|
|
place_kernel<<<blocks(n_atoms), THREADS, 0, stream>>>(rel, n_atoms, p, geom.nu, geom.nv, geom.nw, mask_slot,
|
|
mask_radius, pos, mask_atoms);
|
|
cuda_err(cudaGetLastError());
|
|
cuda_err(cudaMemsetAsync(overflow, 0, sizeof(int), stream));
|
|
count_pairs_kernel<<<blocks(n_atoms + 1), THREADS, 0, stream>>>(pos, box, n_atoms, geom.nu, geom.nv, geom.nw,
|
|
count, overflow);
|
|
cuda_err(cudaGetLastError());
|
|
size_t temp = cub_bytes;
|
|
cuda_err(cub::DeviceScan::ExclusiveSum(cub_temp.get(), temp, count.get(), offset.get(), n_atoms + 1, stream));
|
|
int npairs = 0;
|
|
int over = 0;
|
|
cuda_err(cudaMemcpyAsync(&npairs, offset.get() + n_atoms, sizeof(int), cudaMemcpyDeviceToHost, stream));
|
|
cuda_err(cudaMemcpyAsync(&over, overflow.get(), sizeof(int), cudaMemcpyDeviceToHost, stream));
|
|
cuda_err(cudaStreamSynchronize(stream));
|
|
if (over)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: an atom's box over the bricks per axis");
|
|
if (static_cast<size_t>(npairs) > cap.pairs)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: gather pairs over the reserve");
|
|
fill_pairs_kernel<<<blocks(n_atoms), THREADS, 0, stream>>>(pos, box, n_atoms, geom.nu, geom.nv, geom.nw,
|
|
offset, key, value);
|
|
cuda_err(cudaGetLastError());
|
|
const int nbu = (geom.nu + BRICK - 1) / BRICK, nbv = (geom.nv + BRICK - 1) / BRICK,
|
|
nbw = (geom.nw + BRICK - 1) / BRICK;
|
|
const int nb = nbu * nbv * nbw;
|
|
int bits = 1;
|
|
while ((1 << bits) < nb)
|
|
bits++;
|
|
temp = cub_bytes;
|
|
cuda_err(cub::DeviceRadixSort::SortPairs(cub_temp.get(), temp, key.get(), key_sorted.get(), value.get(),
|
|
value_sorted.get(), npairs, 0, bits, stream));
|
|
cuda_err(cudaMemsetAsync(brick_start, 0, nb * sizeof(int), stream));
|
|
cuda_err(cudaMemsetAsync(brick_end, 0, nb * sizeof(int), stream));
|
|
if (npairs > 0) {
|
|
brick_ranges_kernel<<<blocks(npairs), THREADS, 0, stream>>>(key_sorted, npairs, brick_start, brick_end);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
gather_kernel<<<nb, dim3(BRICK, BRICK, BRICK), 0, stream>>>(atoms, pos, value_sorted, brick_start, brick_end,
|
|
geom, grid.get() + slot * GridPoints());
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
|
|
void Compose(int slot, double2 *out, double2 *derivative) {
|
|
compose_kernel<<<blocks(rows), THREADS, 0, stream>>>(spectrum.get() + slot * ComplexPoints(), Norm(), terms,
|
|
row_scale, static_cast<int>(rows), static_cast<int>(ops),
|
|
out, derivative);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
|
|
Scale MakeScale(double k_overall, const double b_star[6], double k_sol, double b_sol) const {
|
|
Scale s{};
|
|
s.k_overall = k_overall;
|
|
std::copy(b_star, b_star + 6, s.b_star);
|
|
s.k_sol = k_sol;
|
|
s.b_sol = b_sol;
|
|
s.n_constraints = static_cast<int>(n_constraints);
|
|
for (size_t c = 0; c < n_constraints; c++)
|
|
for (int q = 0; q < 6; q++)
|
|
s.constraints[c][q] = constraints[c][q];
|
|
return s;
|
|
}
|
|
};
|
|
|
|
size_t RigidBodyGPUEngine::Bricks(int nu, int nv, int nw) {
|
|
return static_cast<size_t>((nu + BRICK - 1) / BRICK) * ((nv + BRICK - 1) / BRICK) * ((nw + BRICK - 1) / BRICK);
|
|
}
|
|
|
|
size_t RigidBodyGPUEngine::PairBound(const RigidBodyGPUZone &zone) {
|
|
const std::vector<int> box = AtomBoxes(zone);
|
|
const int n3[3] = {zone.nu, zone.nv, zone.nw};
|
|
size_t pairs = 0;
|
|
for (size_t i = 0; i < zone.atoms.size(); i++) {
|
|
size_t product = 1;
|
|
for (int k = 0; k < 3; k++)
|
|
product *= AxisBrickBound(box[3 * i + k], n3[k]);
|
|
pairs += product;
|
|
}
|
|
return pairs;
|
|
}
|
|
|
|
bool RigidBodyGPUEngine::Supports(const RigidBodyGPUZone &zone) {
|
|
const std::vector<int> box = AtomBoxes(zone);
|
|
const int n3[3] = {zone.nu, zone.nv, zone.nw};
|
|
for (size_t i = 0; i < box.size(); i++)
|
|
if (2 * box[i] + 1 > n3[i % 3] || AxisBrickBound(box[i], n3[i % 3]) > MAX_BRICKS_PER_AXIS)
|
|
return false;
|
|
return true;
|
|
}
|
|
|
|
void RigidBodyGPUEngine::MemoryInfo(size_t &free, size_t &total) {
|
|
cuda_err(cudaMemGetInfo(&free, &total));
|
|
}
|
|
|
|
int RigidBodyGPUEngine::CurrentDevice() {
|
|
int device = 0;
|
|
cuda_err(cudaGetDevice(&device));
|
|
return device;
|
|
}
|
|
|
|
size_t RigidBodyGPUEngine::FFTWorkBytes(int nu, int nv, int nw) {
|
|
size_t largest = 0;
|
|
for (int batch : {1, 3}) {
|
|
size_t work = 0;
|
|
const cufftHandle plan = MakePlan(nu, nv, nw, batch, work);
|
|
cufftDestroy(plan);
|
|
largest = std::max(largest, work);
|
|
}
|
|
return largest;
|
|
}
|
|
|
|
size_t RigidBodyGPUEngine::DeviceBytes(const RigidBodyGPUCapacity &c) {
|
|
const size_t na = c.atoms + 1, no = c.observations + 1, nr = c.rows + 1;
|
|
size_t bytes = na * (sizeof(RigidBodyGPUAtom) + 3 * sizeof(int) + 7 * sizeof(double) + sizeof(float4) +
|
|
sizeof(int) + sizeof(float) + 2 * sizeof(int));
|
|
bytes += 3 * c.grid_points * sizeof(float) + 3 * c.complex_points * sizeof(float2) + c.fft_work_bytes;
|
|
bytes += 4 * c.pairs * sizeof(int) + 2 * c.bricks * sizeof(int);
|
|
bytes += c.terms * sizeof(RigidBodyGPUTerm);
|
|
bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 7 * sizeof(double2) + sizeof(float2));
|
|
bytes += no * (2 * sizeof(int) + sizeof(float) + (1 + 6 + MAX_SCALE_PARAMS) * sizeof(double) + 2 * sizeof(float2));
|
|
bytes += CubBytes(na, c.pairs + 1);
|
|
bytes += ModelMaskGPU::DeviceBytes(c.grid_points) + ModelScaleGPU::DeviceBytes(no);
|
|
return bytes;
|
|
}
|
|
|
|
RigidBodyGPUEngine::RigidBodyGPUEngine(const RigidBodyGPUCapacity &capacity, int device) {
|
|
cuda_err(cudaSetDevice(device));
|
|
impl_ = std::make_unique<RigidBodyGPUEngineImpl>(capacity, device);
|
|
}
|
|
|
|
RigidBodyGPUEngine::~RigidBodyGPUEngine() = default;
|
|
|
|
void RigidBodyGPUEngine::SetBody(const std::vector<std::array<double, 3>> &relative) {
|
|
RigidBodyGPUEngineImpl &e = *impl_;
|
|
cuda_err(cudaSetDevice(e.device));
|
|
if (relative.size() > e.cap.atoms)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: more atoms than the reserve");
|
|
e.n_atoms = static_cast<int>(relative.size());
|
|
cuda_err(cudaMemcpyAsync(e.rel, relative.data(), relative.size() * 3 * sizeof(double), cudaMemcpyHostToDevice,
|
|
e.stream));
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
}
|
|
|
|
void RigidBodyGPUEngine::SetZone(const RigidBodyGPUZone &zone) {
|
|
RigidBodyGPUEngineImpl &e = *impl_;
|
|
cuda_err(cudaSetDevice(e.device));
|
|
const size_t np = static_cast<size_t>(zone.nu) * zone.nv * zone.nw;
|
|
const size_t nc = static_cast<size_t>(zone.nu / 2 + 1) * zone.nv * zone.nw;
|
|
if (static_cast<int>(zone.atoms.size()) != e.n_atoms || np > e.cap.grid_points || nc > e.cap.complex_points ||
|
|
Bricks(zone.nu, zone.nv, zone.nw) > e.cap.bricks || zone.row_hkl.size() > e.cap.rows ||
|
|
zone.terms.size() > e.cap.terms || zone.obs_row.size() > e.cap.observations ||
|
|
zone.constraints.size() + 1 > MAX_SCALE_PARAMS)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: zone over the engine's reserve");
|
|
|
|
e.geom.nu = zone.nu;
|
|
e.geom.nv = zone.nv;
|
|
e.geom.nw = zone.nw;
|
|
for (int i = 0; i < 3; i++)
|
|
for (int j = 0; j < 3; j++)
|
|
e.geom.orth_n[3 * i + j] = static_cast<float>(zone.orth[3 * i + j] / (j == 0 ? zone.nu : j == 1 ? zone.nv : zone.nw));
|
|
std::copy(zone.frac, zone.frac + 9, e.frac);
|
|
e.volume = zone.volume;
|
|
e.rows = zone.row_hkl.size();
|
|
e.ops = zone.ops;
|
|
e.n_obs = zone.obs_row.size();
|
|
e.n_points = zone.point_obs.size();
|
|
e.n_mask_atoms = zone.mask_atom.size();
|
|
e.n_constraints = zone.constraints.size();
|
|
e.constraints = zone.constraints;
|
|
e.f_mean = zone.f_mean;
|
|
|
|
// An atom whose box is wider than the cell would reach a point through more than one image, which
|
|
// gemmi's box walk adds and the gather's nearest image does not; such a cell is left to the CPU.
|
|
if (!Supports(zone))
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "rigid body: an atom wider than the cell");
|
|
const std::vector<int> box = AtomBoxes(zone);
|
|
e.geom.narrow = 0;
|
|
for (size_t i = 0; i < box.size(); i++)
|
|
if (2 * box[i] + BRICK > (i % 3 == 0 ? zone.nu : i % 3 == 1 ? zone.nv : zone.nw))
|
|
e.geom.narrow = 1;
|
|
|
|
std::vector<int> slot(zone.atoms.size(), -1);
|
|
std::vector<float> radius(zone.atoms.size(), 0.0f);
|
|
for (size_t j = 0; j < zone.mask_atom.size(); j++) {
|
|
slot[zone.mask_atom[j]] = static_cast<int>(j);
|
|
radius[zone.mask_atom[j]] = zone.mask_radius[j];
|
|
}
|
|
|
|
auto up = [&](void *dst, const void *src, size_t bytes) {
|
|
if (bytes > 0)
|
|
cuda_err(cudaMemcpyAsync(dst, src, bytes, cudaMemcpyHostToDevice, e.stream));
|
|
};
|
|
up(e.atoms, zone.atoms.data(), zone.atoms.size() * sizeof(RigidBodyGPUAtom));
|
|
up(e.box, box.data(), box.size() * sizeof(int));
|
|
up(e.mask_slot, slot.data(), slot.size() * sizeof(int));
|
|
up(e.mask_radius, radius.data(), radius.size() * sizeof(float));
|
|
up(e.terms, zone.terms.data(), zone.terms.size() * sizeof(RigidBodyGPUTerm));
|
|
up(e.row_hkl, zone.row_hkl.data(), zone.row_hkl.size() * 3 * sizeof(int));
|
|
up(e.row_scale, zone.row_scale.data(), zone.row_scale.size() * sizeof(double));
|
|
up(e.row_stol2, zone.row_stol2.data(), zone.row_stol2.size() * sizeof(double));
|
|
up(e.obs_row, zone.obs_row.data(), zone.obs_row.size() * sizeof(int));
|
|
up(e.obs_fobs, zone.obs_fobs.data(), zone.obs_fobs.size() * sizeof(float));
|
|
up(e.point_obs, zone.point_obs.data(), zone.point_obs.size() * sizeof(int));
|
|
ModelMaskGrid mask_grid;
|
|
mask_grid.nu = zone.nu;
|
|
mask_grid.nv = zone.nv;
|
|
mask_grid.nw = zone.nw;
|
|
std::copy(zone.orth, zone.orth + 9, mask_grid.orth);
|
|
mask_grid.volume = zone.volume;
|
|
std::vector<ModelMaskOp> images(zone.images.size());
|
|
for (size_t i = 0; i < images.size(); i++) {
|
|
std::copy(zone.images[i].begin(), zone.images[i].begin() + 9, images[i].rot);
|
|
std::copy(zone.images[i].begin() + 9, zone.images[i].end(), images[i].tran);
|
|
}
|
|
e.mask->SetGrid(mask_grid, images);
|
|
e.scale->SetPoints(zone.point_hkl, zone.point_stol2, zone.point_fobs, zone.point_sigma, zone.constraints, zone.frac);
|
|
e.Plan(1);
|
|
e.Plan(3);
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
}
|
|
|
|
void RigidBodyGPUEngine::Fcalc(const double rotation[9], const double translation[3]) {
|
|
RigidBodyGPUEngineImpl &e = *impl_;
|
|
cuda_err(cudaSetDevice(e.device));
|
|
e.Density(rotation, translation, 0);
|
|
cufft_err(cufftExecR2C(e.Plan(1), e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
|
|
e.Compose(0, e.f, e.df_dt);
|
|
}
|
|
|
|
void RigidBodyGPUEngine::Fmask() {
|
|
RigidBodyGPUEngineImpl &e = *impl_;
|
|
cuda_err(cudaSetDevice(e.device));
|
|
float *mask_grid = e.grid.get() + e.GridPoints();
|
|
cufftComplex *mask_spectrum = reinterpret_cast<cufftComplex *>(e.spectrum.get() + e.ComplexPoints());
|
|
// The mask atoms Fcalc() placed, four doubles each, as ModelMaskAtom lays them out.
|
|
e.mask->Compute(reinterpret_cast<const ModelMaskAtom *>(e.mask_atoms.get()), static_cast<int>(e.n_mask_atoms),
|
|
mask_grid);
|
|
cufft_err(cufftExecR2C(e.Plan(1), mask_grid, mask_spectrum));
|
|
fmask_kernel<<<blocks(e.rows), THREADS, 0, e.stream>>>(reinterpret_cast<float2 *>(mask_spectrum), e.Norm(),
|
|
e.row_hkl, static_cast<int>(e.rows), e.geom.nu, e.geom.nv,
|
|
e.geom.nw, e.fmask);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
|
|
void RigidBodyGPUEngine::DownloadPoints(std::vector<std::array<float, 2>> &fcmol,
|
|
std::vector<std::array<float, 2>> &fmask) {
|
|
RigidBodyGPUEngineImpl &e = *impl_;
|
|
cuda_err(cudaSetDevice(e.device));
|
|
fcmol.resize(e.n_points);
|
|
fmask.resize(e.n_points);
|
|
if (e.n_points == 0)
|
|
return;
|
|
e.PackPoints();
|
|
cuda_err(cudaMemcpyAsync(fcmol.data(), e.point_fcmol, e.n_points * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
|
|
cuda_err(cudaMemcpyAsync(fmask.data(), e.point_fmask, e.n_points * sizeof(float2), cudaMemcpyDeviceToHost, e.stream));
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
}
|
|
|
|
void RigidBodyGPUEngine::FitSolvent(double &k_sol, double &b_sol) {
|
|
RigidBodyGPUEngineImpl &e = *impl_;
|
|
cuda_err(cudaSetDevice(e.device));
|
|
e.PackPoints();
|
|
const ModelSolventFit fit = e.scale->FitSolvent(e.point_fcmol, e.point_fmask);
|
|
k_sol = fit.k_sol;
|
|
b_sol = fit.b_sol;
|
|
}
|
|
|
|
void RigidBodyGPUEngine::FitScale(double k_sol, double b_sol, double &k_overall, double b_star[6]) {
|
|
RigidBodyGPUEngineImpl &e = *impl_;
|
|
cuda_err(cudaSetDevice(e.device));
|
|
e.PackPoints();
|
|
const ModelScaleParams p = e.scale->Fit(e.point_fcmol, e.point_fmask, k_sol, b_sol);
|
|
k_overall = p.k_overall;
|
|
std::copy(p.b_star, p.b_star + 6, b_star);
|
|
}
|
|
|
|
void RigidBodyGPUEngine::Residuals(double k_overall, const double b_star[6], double k_sol, double b_sol,
|
|
double *residuals) {
|
|
RigidBodyGPUEngineImpl &e = *impl_;
|
|
cuda_err(cudaSetDevice(e.device));
|
|
const Scale sc = e.MakeScale(k_overall, b_star, k_sol, b_sol);
|
|
residual_kernel<<<blocks(e.n_obs), THREADS, 0, e.stream>>>(e.obs_row, e.obs_fobs, static_cast<int>(e.n_obs), e.f,
|
|
e.fmask, e.row_hkl, e.row_stol2, sc, e.f_mean,
|
|
e.residual);
|
|
cuda_err(cudaGetLastError());
|
|
cuda_err(cudaMemcpyAsync(residuals, e.residual, e.n_obs * sizeof(double), cudaMemcpyDeviceToHost, e.stream));
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
}
|
|
|
|
void RigidBodyGPUEngine::Jacobian(const double rotation[3][9], const double translation[3][3], double step,
|
|
double k_overall, const double b_star[6], double k_sol, double b_sol,
|
|
std::vector<double> &jtj, std::vector<double> &jtq) {
|
|
RigidBodyGPUEngineImpl &e = *impl_;
|
|
cuda_err(cudaSetDevice(e.device));
|
|
for (int j = 0; j < 3; j++)
|
|
e.Density(rotation[j], translation[j], j);
|
|
cufft_err(cufftExecR2C(e.Plan(3), e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
|
|
for (int j = 0; j < 3; j++)
|
|
e.Compose(j, e.shifted.get() + j * e.rows, nullptr);
|
|
|
|
const Scale sc = e.MakeScale(k_overall, b_star, k_sol, b_sol);
|
|
const int p = 1 + static_cast<int>(e.n_constraints);
|
|
const int n = static_cast<int>(e.n_obs);
|
|
jacobian_kernel<<<blocks(e.n_obs), THREADS, 0, e.stream>>>(e.obs_row, n, e.f, e.df_dt, e.shifted, e.rows, e.fmask,
|
|
e.row_hkl, e.row_stol2, sc, e.f_mean, step, e.jacobian,
|
|
e.jk);
|
|
cuda_err(cudaGetLastError());
|
|
products_kernel<<<p * p, THREADS, 0, e.stream>>>(e.jk, p, e.jk, p, n, p, e.products);
|
|
cuda_err(cudaGetLastError());
|
|
products_kernel<<<p * 6, THREADS, 0, e.stream>>>(e.jk, p, e.jacobian, 6, n, 6, e.products.get() + p * p);
|
|
cuda_err(cudaGetLastError());
|
|
std::vector<double> out(p * p + p * 6);
|
|
cuda_err(cudaMemcpyAsync(out.data(), e.products, out.size() * sizeof(double), cudaMemcpyDeviceToHost, e.stream));
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
jtj.assign(out.begin(), out.begin() + p * p);
|
|
jtq.assign(out.begin() + p * p, out.end());
|
|
}
|
|
|
|
void RigidBodyGPUEngine::ProjectJacobian(const std::vector<double> &x, double *jacobian) {
|
|
RigidBodyGPUEngineImpl &e = *impl_;
|
|
cuda_err(cudaSetDevice(e.device));
|
|
const int p = 1 + static_cast<int>(e.n_constraints);
|
|
cuda_err(cudaMemcpyAsync(e.x, x.data(), x.size() * sizeof(double), cudaMemcpyHostToDevice, e.stream));
|
|
project_kernel<<<blocks(e.n_obs), THREADS, 0, e.stream>>>(e.jacobian, e.jk, p, e.x, static_cast<int>(e.n_obs));
|
|
cuda_err(cudaGetLastError());
|
|
cuda_err(cudaMemcpyAsync(jacobian, e.jacobian, e.n_obs * 6 * sizeof(double), cudaMemcpyDeviceToHost, e.stream));
|
|
cuda_err(cudaStreamSynchronize(e.stream));
|
|
}
|