leonarski_fandClaude Opus 5.5
fbdf036e6f
Every RigidBodyGPUEngine call did cudaSetDevice(engine's card) and never set the caller back. Molecular replacement's group threads refine copies on the pool's engines - on any card where every engine on this one is taken - and then run the next translation function on the same thread. MolecularReplacementGPU ran its calls on whatever card the thread was on, so that function's buffers and cuFFT plan were made on the engine's card, beside the transform and the stream on the search's card: an illegal memory access, 9s02 with four cards visible, every run (the further-copies loop: refine a trial on another card's engine, then translation_function on the group thread). - Each engine call (and the constructor) now holds a DeviceGuard: on the engine's card for the call, back on the caller's after it. - MolecularReplacementGPU remembers the card its transform was uploaded to and runs every call (and its destructor) there, whichever card the calling thread is on. - RigidBodyGPU_CallsLeaveTheCallerOnItsCard: one engine per card, all taken from one thread; the thread's card is unchanged after each call (meaningful with two or more cards). No number changes: where every call runs is the same card as before for each object. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4
715 lines
33 KiB
Plaintext
715 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(cufftExecR2C(plan, 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 cufftHandle plan = p.plan;
|
|
{
|
|
cufft_err(cufftSetStream(plan, 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(cufftExecC2R(plan, 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;
|
|
}
|