Jungfraujoch/image_analysis/structure_refinement/MolecularReplacementGPU.cu

// 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;
}