leonarski_fandClaude Opus 5.5
97ef0936b2
Every cuFFT plan has the driver compile cuFFT's PTX kernels: cufftMakePlanMany -> cuModuleLoadData -> the PTX JIT, under cuFFT's process-wide lock. A second plan of the same shape pays again, even while the first is alive. The CUDA JIT cache makes that ~1 ms. Without one (CUDA_CACHE_DISABLE, no writable $HOME) a plan costs 20-250 ms. That holds for every size (64^3 and 1D 256 included), every transform type and batch, on sm_75 and sm_120, with cuFFT 11.2, 12.0 and 12.4. No CUFFT_* setting changes it, so a choice of grid sizes cannot avoid it; only the number of plans made can be cut. Before this change, a plan was lent to one user at a time, so each concurrent user of a shape on a card made its own: 9-24 rigid-body engines, the translation slices and the score engines. Now a plan is made once per card and shape and shared. Each user keeps its own stream and work area, and CudaFFTPlan::R2C/C2R set both and launch under the plan's mutex. Shared plans give the same bits as private ones: 8 threads on one plan with their own work areas, 9 shapes x 160 transforms, 0 mismatches on RTX 5080 and T4. A shared work area, the negative control, gives 2-13 mismatches per shape where the plan has a work area. The rigid-body engine also dropped its batch-3 plan: the Jacobian transforms its three grids one at a time on the batch-1 plan. That gives the same bits (10 zone-grid shapes, sm_75 and sm_120) and one plan less per zone and card. The cytochrome C and thaumatin MR test sets, AF models (MR step, s): 4x T4 cytc cache on 2.22 -> 2.22 cache off 8.41 -> 4.00 plans 86 -> 24 4x T4 thau cache on 1.58 -> 1.61 cache off 11.13 -> 3.08 plans 131 -> 23 1x T4 cytc cache on 2.87 -> 2.81 cache off 8.82 -> 3.84 plans 80 -> 12 1x T4 thau cache on 2.41 -> 2.39 cache off 10.42 -> 3.17 plans 104 -> 11 RTX5080 cytc cache on 1.41 -> 1.12 cache off 6.5-10.7 -> 2.02 RTX5080 thau cache on 0.90 -> 0.85 cache off 8.03 -> 1.60 Outputs are bit-identical before and after, cache on or off: model, merged hkl, maps, and every MR and validation log line. The mtz/cif differ only in the command line they record. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4
714 lines
33 KiB
Plaintext
714 lines
33 KiB
Plaintext
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "MolecularReplacementGPU.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <mutex>
|
|
#include <string>
|
|
|
|
#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 MAX_BINS = 16; // SEARCH_BINS (MolecularReplacement.cpp) is 10
|
|
// The rotation function's per-reflection scratch is made for this many orientations at a time.
|
|
constexpr size_t ROTATION_BATCH = 2048;
|
|
|
|
// The eight corners of an interpolation in the half-l grid lie in four rows of u. Packed, every point (i, j, l)
|
|
// holds the four (v, w) corners about it, so that two neighbouring records hold all eight: the same values,
|
|
// read from two places instead of four, for four times the memory per point. The packed grid is cut to the indices
|
|
// the plain one is read at in u as it is in v (2 reach + 1, so nv), and has the conjugate plane l = -1 in front.
|
|
struct Corners {
|
|
float2 v[4]; // (j, l), (j, l + 1), (j + 1, l), (j + 1, l + 1)
|
|
};
|
|
|
|
struct Transform {
|
|
const float2 *data; // the plain half-l grid, or null
|
|
const Corners *corners; // the packed one, or null: nv x nv x (nwh + 1), the plane l = -1 first
|
|
int nu, nv, nwh;
|
|
float box, blur;
|
|
};
|
|
|
|
struct Rot {
|
|
float m[9]; // R^T, row-major: the model turned by R has F_R(s) = F_model(R^T s)
|
|
};
|
|
|
|
__device__ __forceinline__ float2 transform_at(const Transform &g, int i, int j, int l) {
|
|
bool conj = false;
|
|
if (l < 0) { // the half-l grid holds l >= 0; F(-h) = conj F(h)
|
|
i = -i; j = -j; l = -l;
|
|
conj = true;
|
|
}
|
|
// Wrapped into the grid; the division only for an index outside it.
|
|
if (i < 0 || i >= g.nu) { i %= g.nu; if (i < 0) i += g.nu; }
|
|
if (j < 0 || j >= g.nv) { j %= g.nv; if (j < 0) j += g.nv; }
|
|
float2 v = g.data[i + static_cast<size_t>(g.nu) * (j + static_cast<size_t>(g.nv) * l)];
|
|
if (conj)
|
|
v.y = -v.y;
|
|
return v;
|
|
}
|
|
|
|
// MolecularTransform::At on the device.
|
|
__device__ float2 transform_interpolate(const Transform &g, const Rot &r, float3 s0) {
|
|
const float s[3] = {r.m[0] * s0.x + r.m[1] * s0.y + r.m[2] * s0.z,
|
|
r.m[3] * s0.x + r.m[4] * s0.y + r.m[5] * s0.z,
|
|
r.m[6] * s0.x + r.m[7] * s0.y + r.m[8] * s0.z};
|
|
float f[3];
|
|
int i0[3];
|
|
for (int a = 0; a < 3; a++) {
|
|
const float u = s[a] * g.box;
|
|
i0[a] = static_cast<int>(floorf(u));
|
|
f[a] = u - i0[a];
|
|
}
|
|
float2 out = make_float2(0, 0);
|
|
if (g.corners) {
|
|
// Below l = -1 every corner is the conjugate of one about the opposite point, -i0 - 1: read there,
|
|
// in the mirrored order.
|
|
const bool mirror = i0[2] < -1;
|
|
int i = mirror ? -i0[0] - 1 : i0[0], j = mirror ? -i0[1] - 1 : i0[1];
|
|
const int l = mirror ? -i0[2] - 1 : i0[2];
|
|
if (i < 0 || i >= g.nv) { i %= g.nv; if (i < 0) i += g.nv; }
|
|
if (j < 0 || j >= g.nv) { j %= g.nv; if (j < 0) j += g.nv; }
|
|
const size_t row = static_cast<size_t>(g.nv) * (j + static_cast<size_t>(g.nv) * (l + 1));
|
|
const Corners c[2] = {g.corners[i + row], g.corners[(i + 1 < g.nv ? i + 1 : 0) + row]};
|
|
for (int dx = 0; dx < 2; dx++)
|
|
for (int dy = 0; dy < 2; dy++)
|
|
for (int dz = 0; dz < 2; dz++) {
|
|
const float w = (dx ? f[0] : 1 - f[0]) * (dy ? f[1] : 1 - f[1]) * (dz ? f[2] : 1 - f[2]);
|
|
float2 v = mirror ? c[1 - dx].v[2 * (1 - dy) + 1 - dz] : c[dx].v[2 * dy + dz];
|
|
if (mirror)
|
|
v.y = -v.y;
|
|
out.x += w * v.x;
|
|
out.y += w * v.y;
|
|
}
|
|
} else {
|
|
for (int dx = 0; dx < 2; dx++)
|
|
for (int dy = 0; dy < 2; dy++)
|
|
for (int dz = 0; dz < 2; dz++) {
|
|
const float w = (dx ? f[0] : 1 - f[0]) * (dy ? f[1] : 1 - f[1]) * (dz ? f[2] : 1 - f[2]);
|
|
const float2 v = transform_at(g, i0[0] + dx, i0[1] + dy, i0[2] + dz);
|
|
out.x += w * v.x;
|
|
out.y += w * v.y;
|
|
}
|
|
}
|
|
const float unblur = expf(g.blur * (s[0] * s[0] + s[1] * s[1] + s[2] * s[2]) / 4.0f);
|
|
return make_float2(out.x * unblur, out.y * unblur);
|
|
}
|
|
|
|
// The packed grid from the plain one: record (i, j, l) holds transform_at() of its four (v, w) corners, i the
|
|
// index of u's frequency in a grid of nv as j is of v's; zero past the last plane, which is never read.
|
|
__global__ void pack_corners(Transform g, Corners *out) {
|
|
const size_t n = static_cast<size_t>(g.nv) * g.nv * (g.nwh + 1);
|
|
const size_t r = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
|
if (r >= n)
|
|
return;
|
|
const int i = static_cast<int>(r % g.nv), j = static_cast<int>((r / g.nv) % g.nv);
|
|
const int l = static_cast<int>(r / (static_cast<size_t>(g.nv) * g.nv)) - 1;
|
|
const int u = i <= g.nv / 2 ? i : i - g.nv;
|
|
Corners c;
|
|
for (int dy = 0; dy < 2; dy++)
|
|
for (int dz = 0; dz < 2; dz++)
|
|
c.v[2 * dy + dz] = l + dz < g.nwh ? transform_at(g, u, j + dy, l + dz) : make_float2(0, 0);
|
|
out[r] = c;
|
|
}
|
|
|
|
// A block-wide sum in a fixed order: each thread's value into shared memory, then a tree.
|
|
__device__ double block_sum(double v, double *sh) {
|
|
sh[threadIdx.x] = v;
|
|
__syncthreads();
|
|
for (int stride = THREADS / 2; stride > 0; stride /= 2) {
|
|
if (threadIdx.x < stride)
|
|
sh[threadIdx.x] += sh[threadIdx.x + stride];
|
|
__syncthreads();
|
|
}
|
|
const double out = sh[0];
|
|
__syncthreads();
|
|
return out;
|
|
}
|
|
|
|
// One block per orientation: |F|^2 summed over the copies for every reflection, binned, normalised by
|
|
// its bin's mean, and correlated with E^2 - 1.
|
|
__global__ void rotation_function(Transform g, const Rot *rot, const float3 *s, const float *e2, const int *bin,
|
|
const int *bin_n, int n, int nops, int bins, double w_mean, double w_ss,
|
|
float *scratch, double *out) {
|
|
__shared__ double sh[THREADS];
|
|
__shared__ double bin_mean[MAX_BINS];
|
|
const Rot r = rot[blockIdx.x];
|
|
float *ic = scratch + static_cast<size_t>(blockIdx.x) * n;
|
|
double part[MAX_BINS];
|
|
for (int b = 0; b < bins; b++)
|
|
part[b] = 0;
|
|
for (int i = threadIdx.x; i < n; i += THREADS) {
|
|
float sum = 0;
|
|
for (int j = 0; j < nops; j++) {
|
|
const float2 f = transform_interpolate(g, r, s[static_cast<size_t>(i) * nops + j]);
|
|
sum += f.x * f.x + f.y * f.y;
|
|
}
|
|
ic[i] = sum;
|
|
part[bin[i]] += sum;
|
|
}
|
|
for (int b = 0; b < bins; b++) {
|
|
const double total = block_sum(part[b], sh);
|
|
if (threadIdx.x == 0)
|
|
bin_mean[b] = total / bin_n[b];
|
|
}
|
|
__syncthreads();
|
|
double x_sum = 0, x_ss = 0, xy = 0;
|
|
for (int i = threadIdx.x; i < n; i += THREADS) {
|
|
const double x = ic[i] / bin_mean[bin[i]];
|
|
x_sum += x;
|
|
x_ss += x * x;
|
|
xy += x * (e2[i] - w_mean);
|
|
}
|
|
x_sum = block_sum(x_sum, sh);
|
|
x_ss = block_sum(x_ss, sh);
|
|
xy = block_sum(xy, sh);
|
|
if (threadIdx.x == 0) {
|
|
const double sxx = x_ss - x_sum * x_sum / n;
|
|
out[blockIdx.x] = sxx > 0 && w_ss > 0 ? xy / sqrt(sxx * w_ss) : 0.0;
|
|
}
|
|
}
|
|
|
|
// g_j = G_j(h) exp(2 pi i h.t_j) for every term and operator, and each term's sum of |g_j|^2.
|
|
__global__ void translation_copies(Transform g, Rot r, const float3 *s, const float2 *phase, int terms, int nops,
|
|
float2 *copies, float *term_sum) {
|
|
const int t = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (t >= terms)
|
|
return;
|
|
float sum = 0;
|
|
for (int j = 0; j < nops; j++) {
|
|
const size_t e = static_cast<size_t>(t) * nops + j;
|
|
const float2 f = transform_interpolate(g, r, s[e]);
|
|
const float2 p = phase[e];
|
|
const float2 v = make_float2(f.x * p.x - f.y * p.y, f.x * p.y + f.y * p.x);
|
|
copies[e] = v;
|
|
sum += v.x * v.x + v.y * v.y;
|
|
}
|
|
term_sum[t] = sum;
|
|
}
|
|
|
|
// Each term's weight over its bin's mean |g|^2 sum, the bins summed in a fixed order by one block.
|
|
__global__ void translation_weights(const float *term_sum, const float *weight, const int *bin, const int *bin_n,
|
|
int terms, int bins, float *scaled) {
|
|
__shared__ double sh[THREADS];
|
|
__shared__ double bin_mean[MAX_BINS];
|
|
double part[MAX_BINS];
|
|
for (int b = 0; b < bins; b++)
|
|
part[b] = 0;
|
|
for (int t = threadIdx.x; t < terms; t += THREADS)
|
|
part[bin[t]] += term_sum[t];
|
|
for (int b = 0; b < bins; b++) {
|
|
const double total = block_sum(part[b], sh);
|
|
if (threadIdx.x == 0)
|
|
bin_mean[b] = total / bin_n[b];
|
|
}
|
|
__syncthreads();
|
|
for (int t = threadIdx.x; t < terms; t += THREADS)
|
|
scaled[t] = static_cast<float>(weight[t] / bin_mean[bin[t]]);
|
|
}
|
|
|
|
// Every coefficient of the series, gathered: a grid point sums the (term, pair) entries that land on it,
|
|
// in the order they were listed.
|
|
__global__ void translation_gather(const int *offset, const int *entry, const float2 *copies, const float *scaled,
|
|
const int *pair_j, const int *pair_k, int nops, int npairs, size_t cells,
|
|
float2 *coef) {
|
|
const size_t c = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
|
if (c >= cells)
|
|
return;
|
|
float2 sum = make_float2(0, 0);
|
|
for (int e = offset[c]; e < offset[c + 1]; e++) {
|
|
const int t = entry[e] / npairs, p = entry[e] % npairs;
|
|
const float2 a = copies[static_cast<size_t>(t) * nops + pair_j[p]];
|
|
const float2 b = copies[static_cast<size_t>(t) * nops + pair_k[p]];
|
|
const float w = scaled[t];
|
|
sum.x += w * (a.x * b.x + a.y * b.y); // a conj(b)
|
|
sum.y += w * (a.y * b.x - a.x * b.y);
|
|
}
|
|
coef[c] = sum;
|
|
}
|
|
|
|
// The cross terms with a known partial structure, added to the coefficients: w conj(F_known) g_j, or its
|
|
// conjugate, for every entry landing on a grid point, in the order they were listed.
|
|
__global__ void translation_gather_cross(const int *offset, const int *entry, const float2 *copies,
|
|
const float *scaled, const float2 *known, const int *cross_term,
|
|
const int *cross_op, const char *cross_conj, int nops, size_t cells,
|
|
float2 *coef) {
|
|
const size_t c = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
|
if (c >= cells)
|
|
return;
|
|
float2 sum = coef[c];
|
|
for (int e = offset[c]; e < offset[c + 1]; e++) {
|
|
const int x = entry[e], t = cross_term[x];
|
|
const float2 k = known[t];
|
|
const float2 g = copies[static_cast<size_t>(t) * nops + cross_op[x]];
|
|
const float w = scaled[t];
|
|
const float re = k.x * g.x + k.y * g.y, im = k.x * g.y - k.y * g.x; // conj(k) g
|
|
sum.x += w * re;
|
|
sum.y += w * (cross_conj[x] ? -im : im);
|
|
}
|
|
coef[c] = sum;
|
|
}
|
|
|
|
// Per block: the sum and the sum of squares of the function, and its highest point (the lowest index
|
|
// on a tie).
|
|
__global__ void translation_stats(const float *map, size_t n, double *sum, double *sum2, float *best,
|
|
long long *best_index) {
|
|
__shared__ double sh[THREADS];
|
|
__shared__ float sh_best[THREADS];
|
|
__shared__ long long sh_index[THREADS];
|
|
double s = 0, s2 = 0;
|
|
float b = -INFINITY;
|
|
long long bi = -1;
|
|
for (size_t i = static_cast<size_t>(blockIdx.x) * THREADS + threadIdx.x; i < n;
|
|
i += static_cast<size_t>(gridDim.x) * THREADS) {
|
|
const float v = map[i];
|
|
s += v;
|
|
s2 += static_cast<double>(v) * v;
|
|
if (v > b) {
|
|
b = v;
|
|
bi = static_cast<long long>(i);
|
|
}
|
|
}
|
|
s = block_sum(s, sh);
|
|
s2 = block_sum(s2, sh);
|
|
sh_best[threadIdx.x] = b;
|
|
sh_index[threadIdx.x] = bi;
|
|
__syncthreads();
|
|
for (int stride = THREADS / 2; stride > 0; stride /= 2) {
|
|
if (threadIdx.x < stride) {
|
|
const float o = sh_best[threadIdx.x + stride];
|
|
const long long oi = sh_index[threadIdx.x + stride];
|
|
if (o > sh_best[threadIdx.x] || (o == sh_best[threadIdx.x] && oi >= 0 && oi < sh_index[threadIdx.x])) {
|
|
sh_best[threadIdx.x] = o;
|
|
sh_index[threadIdx.x] = oi;
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
if (threadIdx.x == 0) {
|
|
sum[blockIdx.x] = s;
|
|
sum2[blockIdx.x] = s2;
|
|
best[blockIdx.x] = sh_best[0];
|
|
best_index[blockIdx.x] = sh_index[0];
|
|
}
|
|
}
|
|
|
|
// cuFFT's r2c output of a map stored u fastest ([w][v][u/2+1], u halved) taken to MapToFPhi's half-l
|
|
// layout (u fastest, w halved), conjugated and scaled as gemmi does, and cropped as CropHalfL crops it.
|
|
__global__ void half_u_to_half_w(const cufftComplex *in, int nu, int nv, int nw, int reach, float norm, float2 *out) {
|
|
const int nuh = nu / 2 + 1, nk = 2 * reach + 1;
|
|
const size_t n = static_cast<size_t>(nu) * nk * (reach + 1);
|
|
const size_t i = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
|
if (i >= n)
|
|
return;
|
|
const int u = static_cast<int>(i % nu), k = static_cast<int>((i / nu) % nk), l = static_cast<int>(i / (static_cast<size_t>(nu) * nk));
|
|
const int v = k <= reach ? k : k - nk + nv;
|
|
cufftComplex x;
|
|
if (u < nuh) {
|
|
x = in[u + static_cast<size_t>(nuh) * (v + static_cast<size_t>(nv) * l)];
|
|
} else { // F(-h) = conj F(h)
|
|
x = in[(nu - u) + static_cast<size_t>(nuh) * ((nv - v) % nv + static_cast<size_t>(nv) * ((nw - l) % nw))];
|
|
x.y = -x.y;
|
|
}
|
|
out[i] = make_float2(x.x * norm, -x.y * norm);
|
|
}
|
|
|
|
Rot transposed(const gemmi::Mat33 &m) {
|
|
Rot r;
|
|
for (int i = 0; i < 3; i++)
|
|
for (int j = 0; j < 3; j++)
|
|
r.m[3 * i + j] = static_cast<float>(m.a[j][i]);
|
|
return r;
|
|
}
|
|
|
|
template <typename T>
|
|
CudaDevicePtr<T> upload(const std::vector<T> &v, cudaStream_t stream) {
|
|
CudaDevicePtr<T> d(std::max<size_t>(v.size(), 1));
|
|
if (!v.empty())
|
|
cuda_err(cudaMemcpyAsync(d.get(), v.data(), v.size() * sizeof(T), cudaMemcpyHostToDevice, stream));
|
|
return d;
|
|
}
|
|
|
|
std::vector<float3> float3s(const std::vector<gemmi::Vec3> &v) {
|
|
std::vector<float3> out(v.size());
|
|
for (size_t i = 0; i < v.size(); i++)
|
|
out[i] = make_float3(static_cast<float>(v[i].x), static_cast<float>(v[i].y), static_cast<float>(v[i].z));
|
|
return out;
|
|
}
|
|
|
|
std::vector<int> bin_counts(const std::vector<int> &bin, int bins) {
|
|
std::vector<int> n(bins, 0);
|
|
for (int b : bin)
|
|
n[b]++;
|
|
return n;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
size_t MapToFPhiGPUBytes(const gemmi::Grid<float> &map, int reach) {
|
|
const size_t points = static_cast<size_t>(map.nu) * map.nv * map.nw;
|
|
const size_t half_u = static_cast<size_t>(map.nu / 2 + 1) * map.nv * map.nw;
|
|
const size_t cropped = static_cast<size_t>(map.nu) * (2 * reach + 1) * (reach + 1);
|
|
return points * sizeof(float) + (half_u + cropped) * sizeof(cufftComplex)
|
|
+ CudaFFTPlan::WorkBytes(CudaFFTShape{.n = {map.nw, map.nv, map.nu}, .type = CUFFT_R2C});
|
|
}
|
|
|
|
gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &map, int reach) {
|
|
// One at a time: concurrent searches (the PAE domains') would each hold a large grid on the card.
|
|
static std::mutex m;
|
|
std::lock_guard lock(m);
|
|
gemmi::FPhiGrid<float> hkl;
|
|
hkl.unit_cell = map.unit_cell;
|
|
hkl.spacegroup = map.spacegroup;
|
|
hkl.axis_order = map.axis_order;
|
|
hkl.half_l = true;
|
|
hkl.set_size_without_checking(map.nu, 2 * reach + 1, reach + 1);
|
|
const float norm = static_cast<float>(map.unit_cell.volume / map.point_count());
|
|
CudaStream stream;
|
|
auto d_map = upload(map.data, stream);
|
|
CudaDevicePtr<cufftComplex> d_half(static_cast<size_t>(map.nu / 2 + 1) * map.nv * map.nw);
|
|
CudaDevicePtr<float2> d_out(hkl.data.size());
|
|
CudaFFTPlan plan(CudaFFTShape{.n = {map.nw, map.nv, map.nu}, .type = CUFFT_R2C}, stream);
|
|
cufft_err(plan.R2C(d_map.get(), d_half.get()));
|
|
half_u_to_half_w<<<static_cast<unsigned>((hkl.data.size() + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
|
|
d_half, map.nu, map.nv, map.nw, reach, norm, d_out);
|
|
cuda_err(cudaGetLastError());
|
|
static_assert(sizeof(std::complex<float>) == sizeof(float2));
|
|
cuda_err(cudaMemcpyAsync(hkl.data.data(), d_out.get(), hkl.data.size() * sizeof(float2),
|
|
cudaMemcpyDeviceToHost, stream));
|
|
cuda_err(cudaStreamSynchronize(stream));
|
|
return hkl;
|
|
}
|
|
|
|
// Every call runs on the card the transform was uploaded to, whichever card the calling thread is on: its buffers,
|
|
// its plans and its kernels beside the transform and the stream.
|
|
struct MolecularReplacementGPUState {
|
|
int device = 0;
|
|
CudaStream stream;
|
|
CudaDevicePtr<float2> grid;
|
|
CudaDevicePtr<Corners> corners;
|
|
Transform transform{};
|
|
};
|
|
|
|
size_t MolecularReplacementGPU::TransformBytes(const MolecularTransform &ft, bool packed) {
|
|
const gemmi::FPhiGrid<float> &g = ft.Grid();
|
|
if (packed)
|
|
return static_cast<size_t>(g.nv) * g.nv * (g.nw + 1) * sizeof(Corners);
|
|
return g.data.size() * sizeof(float2);
|
|
}
|
|
|
|
MolecularReplacementGPU::MolecularReplacementGPU(const MolecularTransform &ft, bool packed)
|
|
: state_(std::make_unique<MolecularReplacementGPUState>()) {
|
|
cuda_err(cudaGetDevice(&state_->device));
|
|
const gemmi::FPhiGrid<float> &g = ft.Grid();
|
|
static_assert(sizeof(std::complex<float>) == sizeof(float2));
|
|
state_->grid = CudaDevicePtr<float2>(std::max<size_t>(g.data.size(), 1));
|
|
cuda_err(cudaMemcpyAsync(state_->grid.get(), g.data.data(), g.data.size() * sizeof(float2),
|
|
cudaMemcpyHostToDevice, state_->stream));
|
|
state_->transform = {state_->grid.get(), nullptr, g.nu, g.nv, g.nw, static_cast<float>(ft.Box()),
|
|
static_cast<float>(ft.Blur())};
|
|
if (packed) {
|
|
// Made from the plain grid on the device, which is then given back.
|
|
const size_t n = TransformBytes(ft, true) / sizeof(Corners);
|
|
state_->corners = CudaDevicePtr<Corners>(n);
|
|
pack_corners<<<static_cast<unsigned>((n + THREADS - 1) / THREADS), THREADS, 0, state_->stream>>>(
|
|
state_->transform, state_->corners);
|
|
cuda_err(cudaGetLastError());
|
|
cuda_err(cudaStreamSynchronize(state_->stream));
|
|
state_->grid = CudaDevicePtr<float2>();
|
|
state_->transform.data = nullptr;
|
|
state_->transform.corners = state_->corners.get();
|
|
}
|
|
cuda_err(cudaStreamSynchronize(state_->stream));
|
|
}
|
|
|
|
MolecularReplacementGPU::~MolecularReplacementGPU() {
|
|
int current = 0;
|
|
cudaGetDevice(¤t);
|
|
cudaSetDevice(state_->device);
|
|
state_.reset();
|
|
cudaSetDevice(current);
|
|
}
|
|
|
|
std::vector<double> MolecularReplacementGPU::RotationFunction(const SearchReflections &refl,
|
|
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
|
|
const std::vector<gemmi::Mat33> &rotations) {
|
|
const ScopedGpuPin pin(state_->device);
|
|
cudaStream_t stream = state_->stream;
|
|
const int n = static_cast<int>(refl.hkl.size());
|
|
const std::vector<gemmi::Vec3> s = CopyVectors(refl, cell, sg);
|
|
const int nops = static_cast<int>(s.size() / std::max(n, 1));
|
|
double w_mean = 0;
|
|
for (double e2 : refl.e2)
|
|
w_mean += e2;
|
|
w_mean /= n;
|
|
double w_ss = 0;
|
|
for (double e2 : refl.e2)
|
|
w_ss += (e2 - w_mean) * (e2 - w_mean);
|
|
const std::vector<float> e2(refl.e2.begin(), refl.e2.end());
|
|
auto d_s = upload(float3s(s), stream);
|
|
auto d_e2 = upload(e2, stream);
|
|
auto d_bin = upload(refl.bin, stream);
|
|
auto d_bin_n = upload(bin_counts(refl.bin, refl.bins), stream);
|
|
std::vector<Rot> rot(rotations.size());
|
|
for (size_t i = 0; i < rotations.size(); i++)
|
|
rot[i] = transposed(rotations[i]);
|
|
auto d_rot = upload(rot, stream);
|
|
const size_t batch = std::min(ROTATION_BATCH, std::max<size_t>(rotations.size(), 1));
|
|
CudaDevicePtr<float> scratch(batch * std::max(n, 1));
|
|
CudaDevicePtr<double> d_out(std::max<size_t>(rotations.size(), 1));
|
|
for (size_t first = 0; first < rotations.size(); first += batch) {
|
|
const int count = static_cast<int>(std::min(batch, rotations.size() - first));
|
|
rotation_function<<<count, THREADS, 0, stream>>>(state_->transform, d_rot.get() + first, d_s, d_e2, d_bin,
|
|
d_bin_n, n, nops, refl.bins, w_mean, w_ss, scratch,
|
|
d_out.get() + first);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
std::vector<double> out(rotations.size());
|
|
cuda_err(cudaMemcpyAsync(out.data(), d_out.get(), out.size() * sizeof(double), cudaMemcpyDeviceToHost, stream));
|
|
cuda_err(cudaStreamSynchronize(stream));
|
|
return out;
|
|
}
|
|
|
|
// What the translation function needs of one setup, worked out on the host: the gather's lists - for every
|
|
// point of the half grid in cuFFT's layout ((u nv + v) nwh + l), the (term, pair) entries landing on it, in
|
|
// entry order, and the same for the cross terms with a known partial structure - and the terms' vectors,
|
|
// phases and weights as the device takes them. Made once where a setup is used for several functions in a row;
|
|
// uploaded by each, so that the device holds it only while a function runs.
|
|
struct TranslationSetupGPU {
|
|
bool cross = false;
|
|
std::vector<int> offset, entry, cross_offset, cross_entry, bin_n;
|
|
std::vector<float3> s;
|
|
std::vector<float2> phase;
|
|
std::vector<float> weight;
|
|
};
|
|
|
|
namespace {
|
|
|
|
TranslationSetupGPU prepare_translation(const TranslationSetup &setup, bool cross) {
|
|
TranslationSetupGPU out;
|
|
out.cross = cross;
|
|
const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1;
|
|
const size_t cells = static_cast<size_t>(nu) * nv * nwh;
|
|
std::vector<int> &offset = out.offset;
|
|
offset.assign(cells + 1, 0);
|
|
for (size_t e = 0; e < setup.target_u.size(); e++)
|
|
if (setup.target_u[e] >= 0)
|
|
offset[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh + setup.target_l[e] + 1]++;
|
|
for (size_t c = 0; c < cells; c++)
|
|
offset[c + 1] += offset[c];
|
|
out.entry.resize(offset[cells]);
|
|
{
|
|
std::vector<int> fill(offset.begin(), offset.end() - 1);
|
|
for (size_t e = 0; e < setup.target_u.size(); e++)
|
|
if (setup.target_u[e] >= 0)
|
|
out.entry[fill[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh
|
|
+ setup.target_l[e]]++] = static_cast<int>(e);
|
|
}
|
|
if (cross) {
|
|
std::vector<int> &cross_offset = out.cross_offset;
|
|
cross_offset.assign(cells + 1, 0);
|
|
for (size_t e = 0; e < setup.cross_term.size(); e++)
|
|
cross_offset[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh + setup.cross_l[e] + 1]++;
|
|
for (size_t c = 0; c < cells; c++)
|
|
cross_offset[c + 1] += cross_offset[c];
|
|
out.cross_entry.resize(cross_offset[cells]);
|
|
std::vector<int> fill(cross_offset.begin(), cross_offset.end() - 1);
|
|
for (size_t e = 0; e < setup.cross_term.size(); e++)
|
|
out.cross_entry[fill[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh
|
|
+ setup.cross_l[e]]++] = static_cast<int>(e);
|
|
}
|
|
out.phase.resize(setup.phase.size());
|
|
for (size_t i = 0; i < out.phase.size(); i++)
|
|
out.phase[i] = make_float2(static_cast<float>(setup.phase[i].real()), static_cast<float>(setup.phase[i].imag()));
|
|
out.weight.assign(setup.weight.begin(), setup.weight.end());
|
|
out.s = float3s(setup.s);
|
|
out.bin_n = bin_counts(setup.bin, setup.bins);
|
|
return out;
|
|
}
|
|
|
|
// A prepared setup on the device, with the transform's plan, for one translation function.
|
|
struct TranslationDevice {
|
|
CudaDevicePtr<int> cross_offset, cross_entry, cross_term, cross_op, bin, bin_n, pair_j, pair_k, offset, entry;
|
|
CudaDevicePtr<char> cross_conj;
|
|
CudaDevicePtr<float3> s;
|
|
CudaDevicePtr<float2> phase;
|
|
CudaDevicePtr<float> weight;
|
|
CudaFFTPlan plan;
|
|
};
|
|
|
|
void upload_translation(const TranslationSetup &setup, const TranslationSetupGPU &p, cudaStream_t stream,
|
|
TranslationDevice &out) {
|
|
out.cross_offset = upload(p.cross_offset, stream);
|
|
out.cross_entry = upload(p.cross_entry, stream);
|
|
out.cross_term = upload(setup.cross_term, stream);
|
|
out.cross_op = upload(setup.cross_op, stream);
|
|
out.cross_conj = upload(setup.cross_conj, stream);
|
|
out.s = upload(p.s, stream);
|
|
out.phase = upload(p.phase, stream);
|
|
out.weight = upload(p.weight, stream);
|
|
out.bin = upload(setup.bin, stream);
|
|
out.bin_n = upload(p.bin_n, stream);
|
|
out.pair_j = upload(setup.pair_j, stream);
|
|
out.pair_k = upload(setup.pair_k, stream);
|
|
out.offset = upload(p.offset, stream);
|
|
out.entry = upload(p.entry, stream);
|
|
out.plan = CudaFFTPlan(CudaFFTShape{.n = {setup.nu, setup.nv, setup.nw}, .type = CUFFT_C2R}, stream);
|
|
}
|
|
|
|
} // namespace
|
|
|
|
std::unique_ptr<TranslationSetupGPU, void (*)(TranslationSetupGPU *)> MolecularReplacementGPU::Prepare(
|
|
const TranslationSetup &setup, bool known_part) {
|
|
return {new TranslationSetupGPU(prepare_translation(setup, known_part)), [](TranslationSetupGPU *p) { delete p; }};
|
|
}
|
|
|
|
namespace {
|
|
|
|
// What upload() of `count` values of `bytes` each takes.
|
|
size_t uploaded(size_t count, size_t bytes) {
|
|
return std::max<size_t>(count, 1) * bytes;
|
|
}
|
|
|
|
int stat_block_count(size_t points) {
|
|
return static_cast<int>(std::min<size_t>(1024, (points + THREADS - 1) / THREADS));
|
|
}
|
|
|
|
} // namespace
|
|
|
|
size_t MolecularReplacementGPU::RotationBytes(const SearchReflections &refl, const gemmi::SpaceGroup &sg,
|
|
size_t rotations) {
|
|
const size_t n = refl.hkl.size(), nops = sg.operations().sym_ops.size();
|
|
const size_t batch = std::min(ROTATION_BATCH, std::max<size_t>(rotations, 1));
|
|
return uploaded(n * nops, sizeof(float3)) + 2 * uploaded(n, sizeof(float)) + uploaded(refl.bins, sizeof(int))
|
|
+ uploaded(rotations, sizeof(Rot)) + batch * std::max<size_t>(n, 1) * sizeof(float)
|
|
+ uploaded(rotations, sizeof(double));
|
|
}
|
|
|
|
size_t MolecularReplacementGPU::TranslationBytes(const TranslationSetup &setup, size_t rotations, bool known) {
|
|
const size_t cells = static_cast<size_t>(setup.nu) * setup.nv * (setup.nw / 2 + 1);
|
|
const size_t points = static_cast<size_t>(setup.nu) * setup.nv * setup.nw;
|
|
const size_t terms = setup.Terms(), nops = setup.ops;
|
|
const size_t entries = std::count_if(setup.target_u.begin(), setup.target_u.end(), [](int u) { return u >= 0; });
|
|
const size_t nstat = static_cast<size_t>(stat_block_count(points)) * rotations;
|
|
// The setup, as upload_translation() puts it on the device
|
|
size_t bytes = uploaded(known ? cells + 1 : 0, sizeof(int)) + uploaded(known ? setup.cross_term.size() : 0, sizeof(int))
|
|
+ 2 * uploaded(setup.cross_term.size(), sizeof(int)) + uploaded(setup.cross_conj.size(), sizeof(char))
|
|
+ uploaded(terms * nops, sizeof(float3)) + uploaded(terms * nops, sizeof(float2))
|
|
+ 2 * uploaded(terms, sizeof(int)) + uploaded(setup.bins, sizeof(int))
|
|
+ 2 * uploaded(setup.pair_j.size(), sizeof(int)) + uploaded(cells + 1, sizeof(int))
|
|
+ uploaded(entries, sizeof(int))
|
|
+ std::max<size_t>(CudaFFTPlan::WorkBytes(CudaFFTShape{.n = {setup.nu, setup.nv, setup.nw},
|
|
.type = CUFFT_C2R}), 1);
|
|
// and the function's own buffers
|
|
bytes += uploaded(known ? terms : 0, sizeof(float2)) + terms * nops * sizeof(float2) + 2 * terms * sizeof(float)
|
|
+ cells * sizeof(float2) + points * sizeof(float)
|
|
+ nstat * (2 * sizeof(double) + sizeof(float) + sizeof(long long));
|
|
return bytes;
|
|
}
|
|
|
|
std::vector<TranslationPeak> MolecularReplacementGPU::TranslationFunction(
|
|
const TranslationSetup &setup, const std::vector<gemmi::Mat33> &rotations,
|
|
const std::vector<std::complex<double>> *known, TranslationSetupGPU *prepared) {
|
|
std::vector<TranslationPeak> out(rotations.size());
|
|
const int terms = static_cast<int>(setup.Terms()), nops = setup.ops;
|
|
const int npairs = static_cast<int>(setup.pair_j.size());
|
|
if (terms == 0 || rotations.empty() || (nops < 2 && known == nullptr))
|
|
return out;
|
|
const ScopedGpuPin pin(state_->device);
|
|
cudaStream_t stream = state_->stream;
|
|
const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1;
|
|
const size_t cells = static_cast<size_t>(nu) * nv * nwh, points = static_cast<size_t>(nu) * nv * nw;
|
|
TranslationSetupGPU made;
|
|
if (prepared == nullptr) {
|
|
made = prepare_translation(setup, known != nullptr);
|
|
prepared = &made;
|
|
}
|
|
TranslationDevice p;
|
|
upload_translation(setup, *prepared, stream, p);
|
|
std::vector<float2> known_f;
|
|
if (known != nullptr)
|
|
for (const std::complex<double> &k : *known)
|
|
known_f.push_back(make_float2(static_cast<float>(k.real()), static_cast<float>(k.imag())));
|
|
auto d_known = upload(known_f, stream);
|
|
CudaDevicePtr<float2> copies(static_cast<size_t>(terms) * nops);
|
|
CudaDevicePtr<float> term_sum(terms), scaled(terms);
|
|
CudaDevicePtr<float2> coef(cells);
|
|
CudaDevicePtr<float> map(points);
|
|
const int stat_blocks = stat_block_count(points);
|
|
const size_t nstat = static_cast<size_t>(stat_blocks) * rotations.size();
|
|
CudaDevicePtr<double> d_sum(nstat), d_sum2(nstat);
|
|
CudaDevicePtr<float> d_best(nstat);
|
|
CudaDevicePtr<long long> d_best_index(nstat);
|
|
|
|
const CudaFFTPlan &plan = p.plan; // on `stream`
|
|
{
|
|
for (size_t r = 0; r < rotations.size(); r++) {
|
|
const Rot rot = transposed(rotations[r]);
|
|
translation_copies<<<(terms + THREADS - 1) / THREADS, THREADS, 0, stream>>>(
|
|
state_->transform, rot, p.s, p.phase, terms, nops, copies, term_sum);
|
|
translation_weights<<<1, THREADS, 0, stream>>>(term_sum, p.weight, p.bin, p.bin_n, terms, setup.bins,
|
|
scaled);
|
|
translation_gather<<<static_cast<unsigned>((cells + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
|
|
p.offset, p.entry, copies, scaled, p.pair_j, p.pair_k, nops, npairs, cells, coef);
|
|
if (known != nullptr)
|
|
translation_gather_cross<<<static_cast<unsigned>((cells + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
|
|
p.cross_offset, p.cross_entry, copies, scaled, d_known, p.cross_term, p.cross_op, p.cross_conj,
|
|
nops, cells, coef);
|
|
cufft_err(plan.C2R(reinterpret_cast<cufftComplex *>(coef.get()), map.get()));
|
|
translation_stats<<<stat_blocks, THREADS, 0, stream>>>(
|
|
map, points, d_sum.get() + r * stat_blocks, d_sum2.get() + r * stat_blocks,
|
|
d_best.get() + r * stat_blocks, d_best_index.get() + r * stat_blocks);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
}
|
|
std::vector<double> sum(nstat), sum2(nstat);
|
|
std::vector<float> best(nstat);
|
|
std::vector<long long> best_index(nstat);
|
|
cuda_err(cudaMemcpyAsync(sum.data(), d_sum.get(), nstat * sizeof(double), cudaMemcpyDeviceToHost, stream));
|
|
cuda_err(cudaMemcpyAsync(sum2.data(), d_sum2.get(), nstat * sizeof(double), cudaMemcpyDeviceToHost, stream));
|
|
cuda_err(cudaMemcpyAsync(best.data(), d_best.get(), nstat * sizeof(float), cudaMemcpyDeviceToHost, stream));
|
|
cuda_err(cudaMemcpyAsync(best_index.data(), d_best_index.get(), nstat * sizeof(long long),
|
|
cudaMemcpyDeviceToHost, stream));
|
|
cuda_err(cudaStreamSynchronize(stream));
|
|
|
|
for (size_t r = 0; r < rotations.size(); r++) {
|
|
double s = 0, s2 = 0;
|
|
float b = -INFINITY;
|
|
long long bi = 0;
|
|
for (int k = 0; k < stat_blocks; k++) {
|
|
const size_t i = r * stat_blocks + k;
|
|
s += sum[i];
|
|
s2 += sum2[i];
|
|
if (best[i] > b) {
|
|
b = best[i];
|
|
bi = best_index[i];
|
|
}
|
|
}
|
|
const double count = static_cast<double>(points);
|
|
const double mean = s / count, sd = std::sqrt(std::max(0.0, s2 / count - mean * mean));
|
|
out[r].z = sd > 0 ? (b - mean) / sd : 0.0;
|
|
const long long w = bi % nw, v = (bi / nw) % nv, u = bi / (static_cast<long long>(nw) * nv);
|
|
out[r].t = gemmi::Fractional(static_cast<double>(u) / nu, static_cast<double>(v) / nv,
|
|
static_cast<double>(w) / nw);
|
|
}
|
|
return out;
|
|
}
|