mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 81 Wiki Activity
Files
rc175
Jungfraujoch/image_analysis/structure_refinement/RigidBodyGPU.cu
T
leonarski_fandClaude Opus 5.5 dcb01100e7 RefineRigidBody: a body refined beside fixed copies of it 
An optional `fixed` model: other atoms of the crystal held where they are -
the other copies of a molecular-replacement search - whose Fcalc is added to
the body's at every evaluation and whose atoms the bulk-solvent mask covers
together with the body's (computed once per zone, with the body where the
zone starts, and held). With them the origin is no longer free, so the gauge
projection is skipped. On the CPU the fixed part is gridded once per zone; on
the GPU the fixed atoms must be copies of the body, and each one's placement
is read off by Kabsch superposition and composed on the device with the
engine's own density, transform and mask (SetFixedCopies), so a fit beside
three copies of a 90 kDa model costs what a fit alone does.

Without `fixed` nothing changes (the composition kernel's extra term is null).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
2026-10-12 00:26:49 +02:00

763 lines
34 KiB
Plaintext
Raw Permalink Blame History

// 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 <cufft.h>
#include "ModelDensityGPU.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 THREADS = 256;
constexpr int MAX_SCALE_PARAMS = 7; // k_overall + up to six B* constraints
// The grid of a zone.
struct GridGeom {
int nu, nv, nw;
};
// 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;
}
// 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];
}
}
// 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, const double2 *fixed, 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] = fixed != nullptr ? make_double2(scale * sr + fixed[m].x, scale * si + fixed[m].y)
: 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]);
}
}
__global__ void add_kernel(const double2 *a, int n, double2 *sum) {
const int m = blockIdx.x * blockDim.x + threadIdx.x;
if (m < n)
sum[m] = make_double2(sum[m].x + a[m].x, sum[m].y + a[m].y);
}
// 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);
}
// 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;
}
int GreatestStreamPriority() {
int least = 0, greatest = 0;
cuda_err(cudaDeviceGetStreamPriorityRange(&least, &greatest));
return greatest;
}
} // namespace
// The zone's grid as the density gather takes it.
static ModelDensityGrid DensityGrid(const RigidBodyGPUZone &zone) {
ModelDensityGrid g;
g.nu = zone.nu;
g.nv = zone.nv;
g.nw = zone.nw;
std::copy(zone.orth, zone.orth + 9, g.orth);
std::copy(zone.frac, zone.frac + 9, g.frac);
return g;
}
size_t RigidBodyGPUEngine::AxisBrickBound(int d, int n) {
return ModelDensityGPU::AxisBrickBound(d, n);
}
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<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<double> copies_mask_atoms; // the same for the body and fixed copies of it, grown as needed
size_t copies_mask_capacity = 0;
CudaDevicePtr<float> grid; // 3 grids back to back
CudaDevicePtr<float2> spectrum; // their 3 transforms
CudaDevicePtr<char> fft_work;
CudaDevicePtr<RigidBodyGPUTerm> terms;
CudaDevicePtr<int> row_hkl;
CudaDevicePtr<double> row_scale, row_stol2;
CudaDevicePtr<double2> f, df_dt, shifted, fixed;
bool has_fixed = false; // the zone's fixed Fcalc is in `fixed`
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<ModelDensityGPU> density;
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);
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);
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);
fixed = CudaDevicePtr<double2>(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);
density = std::make_unique<ModelDensityGPU>(stream, cap.atoms, cap.pairs, cap.bricks);
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());
density->Compute(pos, grid.get() + slot * GridPoints());
}
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),
has_fixed ? fixed.get() : nullptr, 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 ModelDensityGPU::Bricks(nu, nv, nw);
}
size_t RigidBodyGPUEngine::PairBound(const RigidBodyGPUZone &zone) {
return ModelDensityGPU::PairBound(DensityGrid(zone), zone.atoms);
}
bool RigidBodyGPUEngine::Supports(const RigidBodyGPUZone &zone) {
return ModelDensityGPU::Supports(DensityGrid(zone), zone.atoms);
}
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 * (7 * sizeof(double) + sizeof(float4) + sizeof(int) + sizeof(float));
bytes += 3 * c.grid_points * sizeof(float) + 3 * c.complex_points * sizeof(float2) + c.fft_work_bytes;
bytes += ModelDensityGPU::DeviceBytes(c.atoms, c.pairs, c.bricks);
bytes += c.terms * sizeof(RigidBodyGPUTerm);
bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 8 * sizeof(double2) + sizeof(float2));
bytes += no * (2 * sizeof(int) + sizeof(float) + (1 + 6 + MAX_SCALE_PARAMS) * sizeof(double) + 2 * sizeof(float2));
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);
}
// Released on its own card, which need not be the one the destroying thread is on.
RigidBodyGPUEngine::~RigidBodyGPUEngine() {
int current = 0;
cudaGetDevice(&current);
cudaSetDevice(impl_->device);
impl_.reset();
cudaSetDevice(current);
}
int RigidBodyGPUEngine::Device() const {
return impl_->device;
}
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;
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;
e.has_fixed = false;
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.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.density->SetAtoms(DensityGrid(zone), zone.atoms);
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::SetFixedCopies(const std::vector<std::array<double, 12>> &fixed,
const std::array<double, 12> &moving) {
RigidBodyGPUEngineImpl &e = *impl_;
cuda_err(cudaSetDevice(e.device));
const size_t copies = fixed.size() + 1;
if (e.copies_mask_capacity < copies * e.n_mask_atoms) {
e.copies_mask_capacity = copies * e.n_mask_atoms;
e.copies_mask_atoms = CudaDevicePtr<double>(4 * std::max<size_t>(e.copies_mask_capacity, 1), CudaAlloc::Synchronous);
}
cuda_err(cudaMemsetAsync(e.fixed, 0, e.rows * sizeof(double2), e.stream));
for (size_t c = 0; c < copies; c++) {
const std::array<double, 12> &x = c < fixed.size() ? fixed[c] : moving;
Placement p{};
std::copy(x.begin(), x.begin() + 9, p.r);
std::copy(x.begin() + 9, x.end(), p.t);
std::copy(e.frac, e.frac + 9, p.frac);
place_kernel<<<blocks(e.n_atoms), THREADS, 0, e.stream>>>(e.rel, e.n_atoms, p, e.geom.nu, e.geom.nv, e.geom.nw,
e.mask_slot, e.mask_radius, e.pos,
e.copies_mask_atoms.get() + 4 * c * e.n_mask_atoms);
cuda_err(cudaGetLastError());
if (c < fixed.size()) {
e.density->Compute(e.pos, e.grid.get());
cufft_err(cufftExecR2C(e.Plan(1), e.grid, reinterpret_cast<cufftComplex *>(e.spectrum.get())));
e.has_fixed = false;
e.Compose(0, e.shifted, nullptr);
add_kernel<<<blocks(e.rows), THREADS, 0, e.stream>>>(e.shifted, static_cast<int>(e.rows), e.fixed);
cuda_err(cudaGetLastError());
}
}
float *mask_grid = e.grid.get() + e.GridPoints();
cufftComplex *mask_spectrum = reinterpret_cast<cufftComplex *>(e.spectrum.get() + e.ComplexPoints());
e.mask->Compute(reinterpret_cast<const ModelMaskAtom *>(e.copies_mask_atoms.get()),
static_cast<int>(copies * 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());
cuda_err(cudaStreamSynchronize(e.stream));
e.has_fixed = true;
}
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 b_hi, 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, b_hi);
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));
}
Reference in New Issue View Git Blame Copy Permalink