leonarski_fandClaude Opus 5.5
e1f49dd729
From the nsys/perf profiles of the copy chain (GPU busy 39% between copies, the thread mostly in host work and syncs): - The known-structure translation function's lists, device arrays and cuFFT plan are made once per group (MolecularReplacementGPU::Prepare) instead of on every copy (~260 MB of uploads and a plan per call before). - The known copies' CopyFactors and likelihood, which rank the peaks, are taken while the GPU runs the translation function. - The structure-factor engines are made per entry, without one lock held across all groups' creation. - The transform's index wrap divides only for an index outside the grid (CPU and GPU). All exact: same arrays, same kernels, same order. 4mtk MR, alternating 3 pairs: 9.5 -> 9.0 s. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
586 lines
26 KiB
Plaintext
586 lines
26 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.
|
|
__global__ void half_u_to_half_w(const cufftComplex *in, int nu, int nv, int nw, float norm, float2 *out) {
|
|
const int nwh = nw / 2 + 1, nuh = nu / 2 + 1;
|
|
const size_t n = static_cast<size_t>(nu) * nv * nwh;
|
|
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), v = static_cast<int>((i / nu) % nv), l = static_cast<int>(i / (static_cast<size_t>(nu) * 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
|
|
|
|
bool MapToFPhiGPUFits(const gemmi::Grid<float> &map) {
|
|
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 half_w = static_cast<size_t>(map.nu) * map.nv * (map.nw / 2 + 1);
|
|
const size_t bytes = points * sizeof(float) + (half_u + half_w) * sizeof(cufftComplex);
|
|
size_t free_bytes = 0, total = 0;
|
|
cuda_err(cudaMemGetInfo(&free_bytes, &total));
|
|
return bytes <= total / 4;
|
|
}
|
|
|
|
gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &map) {
|
|
// 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, map.nv, map.nw / 2 + 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());
|
|
cufftHandle plan;
|
|
cufft_err(cufftPlan3d(&plan, map.nw, map.nv, map.nu, CUFFT_R2C));
|
|
try {
|
|
cufft_err(cufftSetStream(plan, 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, 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));
|
|
} catch (...) {
|
|
cufftDestroy(plan);
|
|
throw;
|
|
}
|
|
cufftDestroy(plan);
|
|
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, on the device: 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 - the terms' vectors, phases, weights and bins,
|
|
// and the transform's plan. Made once where a setup is used for several functions in a row.
|
|
struct TranslationSetupGPU {
|
|
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;
|
|
cufftHandle plan = 0;
|
|
~TranslationSetupGPU() {
|
|
if (plan != 0)
|
|
cufftDestroy(plan);
|
|
}
|
|
};
|
|
|
|
namespace {
|
|
|
|
std::unique_ptr<TranslationSetupGPU> prepare_translation(const TranslationSetup &setup, bool cross, cudaStream_t stream) {
|
|
auto out = std::make_unique<TranslationSetupGPU>();
|
|
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(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];
|
|
std::vector<int> entry(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)
|
|
entry[fill[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh
|
|
+ setup.target_l[e]]++] = static_cast<int>(e);
|
|
}
|
|
std::vector<int> cross_offset(cells + 1, 0), cross_entry;
|
|
if (cross) {
|
|
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];
|
|
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++)
|
|
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->cross_offset = upload(cross_offset, stream);
|
|
out->cross_entry = upload(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);
|
|
std::vector<float2> phase(setup.phase.size());
|
|
for (size_t i = 0; i < phase.size(); i++)
|
|
phase[i] = make_float2(static_cast<float>(setup.phase[i].real()), static_cast<float>(setup.phase[i].imag()));
|
|
const std::vector<float> weight(setup.weight.begin(), setup.weight.end());
|
|
out->s = upload(float3s(setup.s), stream);
|
|
out->phase = upload(phase, stream);
|
|
out->weight = upload(weight, stream);
|
|
out->bin = upload(setup.bin, stream);
|
|
out->bin_n = upload(bin_counts(setup.bin, setup.bins), stream);
|
|
out->pair_j = upload(setup.pair_j, stream);
|
|
out->pair_k = upload(setup.pair_k, stream);
|
|
out->offset = upload(offset, stream);
|
|
out->entry = upload(entry, stream);
|
|
cufft_err(cufftPlan3d(&out->plan, nu, nv, nw, CUFFT_C2R));
|
|
cuda_err(cudaStreamSynchronize(stream)); // the host vectors go out of scope
|
|
return out;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
std::unique_ptr<TranslationSetupGPU, void (*)(TranslationSetupGPU *)> MolecularReplacementGPU::Prepare(
|
|
const TranslationSetup &setup, bool known_part) {
|
|
return {prepare_translation(setup, known_part, state_->stream).release(),
|
|
[](TranslationSetupGPU *p) { delete p; }};
|
|
}
|
|
|
|
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;
|
|
std::unique_ptr<TranslationSetupGPU> made;
|
|
if (prepared == nullptr) {
|
|
made = prepare_translation(setup, known != nullptr, stream);
|
|
prepared = made.get();
|
|
}
|
|
TranslationSetupGPU &p = *prepared;
|
|
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 = static_cast<int>(std::min<size_t>(1024, (points + THREADS - 1) / THREADS));
|
|
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;
|
|
}
|