leonarski_fandClaude Opus 5.5
51706641fa
The molecular transform's FFT went to the GPU only where its buffers took at most a quarter of the card (MapToFPhiGPUFits), the last quarter rule left in molecular replacement. On a large cell that is never: 4nwv's 768^3 box needs 5.45 GB, so on a 16 GB card it ran on the CPU in FFTW, single-threaded - 12.5 s of the search's 31 s. Now it is judged, as every other part of the search, against the search's budget (80% of the card less what is resident; the search holds nothing else yet when it runs), and the device writes only the cropped transform the search keeps (CropHalfL, moved to ModelFFT), so it needs 4.46 GB instead of 5.45. cuFFT and FFTW agree to float rounding, not bit for bit, so wherever the transform moves to the GPU the search's functions can differ in their last bits: on 4nwv on 16 GB and 8.55 GB cards it moves (it now does on any card above about 5.6 GB, before above 21.8 GB); the small boxes were on the GPU before and give the same values (the same kernel, cropped). Measured on the RTX 5080, timing benchmark: 4nwv's transform 14.1 -> 2.2 s, MR 31.6 -> 20.1 s, run 93.0 -> 89.2 s; MR's peak device memory unchanged at 8.7 GB (the FFT's own 5.2 GB comes before anything else is held). On this card the reflection files, maps, placed models, reports and the MR candidate lines of myob, cytc and 4nwv still came out byte-identical to rc175-metal-merge. [MolecularReplacement] passes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
635 lines
29 KiB
Plaintext
635 lines
29 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;
|
|
|
|
struct Transform {
|
|
const float2 *data;
|
|
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);
|
|
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);
|
|
}
|
|
|
|
// 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;
|
|
}
|
|
|
|
struct MolecularReplacementGPUState {
|
|
CudaStream stream;
|
|
CudaDevicePtr<float2> grid;
|
|
Transform transform{};
|
|
};
|
|
|
|
MolecularReplacementGPU::MolecularReplacementGPU(const MolecularTransform &ft)
|
|
: state_(std::make_unique<MolecularReplacementGPUState>()) {
|
|
const gemmi::FPhiGrid<float> &g = ft.Grid();
|
|
std::vector<float2> host(g.data.size());
|
|
for (size_t i = 0; i < host.size(); i++)
|
|
host[i] = make_float2(g.data[i].real(), g.data[i].imag());
|
|
state_->grid = upload(host, state_->stream);
|
|
state_->transform = {state_->grid.get(), g.nu, g.nv, g.nw, static_cast<float>(ft.Box()),
|
|
static_cast<float>(ft.Blur())};
|
|
cuda_err(cudaStreamSynchronize(state_->stream));
|
|
}
|
|
|
|
MolecularReplacementGPU::~MolecularReplacementGPU() = default;
|
|
|
|
std::vector<double> MolecularReplacementGPU::RotationFunction(const SearchReflections &refl,
|
|
const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
|
|
const std::vector<gemmi::Mat33> &rotations) {
|
|
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;
|
|
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;
|
|
}
|