RESULT-CHANGING: molecular replacement's transform FFT on the GPU wherever it fits the search's budget

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>

docs/GPU_DATA_ANALYSIS.md

@@ -168,7 +168,8 @@ space instead of by symmetrizing the grid).
 ## Molecular replacement
 
 - **Molecular transform** (`MolecularReplacementGPU.cu`, `MapToFPhiGPU`) - the Fourier transform of
-  the search model's density in a large box, by cuFFT, where it fits a quarter of the card.
+  the search model's density in a large box, by cuFFT, where its buffers fit the search's budget (80% of
+  the card, the search holding nothing else yet), kept only as far out as it is read.
 - **Rotation function** (`MolecularReplacementGPU.cu`) - for every trial orientation, the model's
   intensities summed over the copies, normalised by resolution bin and correlated with the data's
   normalised intensities.

image_analysis/structure_refinement/ModelFFT.cpp

@@ -110,3 +110,20 @@ gemmi::Grid<float> MapFromFPhi(const gemmi::FPhiGrid<float> &hkl) {
     fftwf_free(out);
     return map;
 }
+
+// The rows keep their length: cut to 2 reach + 1 as well, the rotation function read them 15% slower on the
+// GPU (measured on a large cell).
+gemmi::FPhiGrid<float> CropHalfL(const gemmi::FPhiGrid<float> &full, int reach) {
+    gemmi::FPhiGrid<float> out;
+    out.unit_cell = full.unit_cell;
+    out.spacegroup = full.spacegroup;
+    out.axis_order = full.axis_order;
+    out.half_l = true;
+    const int nu = full.nu, n = 2 * reach + 1;
+    out.set_size_without_checking(nu, n, reach + 1);
+    for (int l = 0; l <= reach; l++)
+        for (int k = -reach; k <= reach; k++)
+            std::copy_n(&full.data[static_cast<size_t>(nu) * (gemmi::modulo(k, full.nv) + static_cast<size_t>(full.nv) * l)],
+                        nu, &out.data[static_cast<size_t>(nu) * (gemmi::modulo(k, n) + static_cast<size_t>(n) * l)]);
+    return out;
+}

image_analysis/structure_refinement/ModelFFT.h

@@ -22,3 +22,8 @@ gemmi::FPhiGrid<float> MapToFPhi(const gemmi::Grid<float> &map);
 // exp(-2 pi i h.x) on an (nu, nv, 2*(nw-1)) grid, NaN coefficients read as zero), computed with FFTW
 // under the same planning rules.
 gemmi::Grid<float> MapFromFPhi(const gemmi::FPhiGrid<float> &hkl);
+
+// The half-l grid `full` with only the indices |k|, l <= reach kept, on a grid of full.nu x (2 reach + 1) x
+// (reach + 1) that an index wraps into as it wraps into the whole one (the molecular transform's, which is
+// read no further out).
+gemmi::FPhiGrid<float> CropHalfL(const gemmi::FPhiGrid<float> &full, int reach);

image_analysis/structure_refinement/MolecularReplacement.cpp

@@ -172,24 +172,6 @@ int good_even_size(int n) {
     return n;
 }
 
-// The half-l grid `full` with only the indices |k|, l <= reach kept, on a grid of full.nu x (2 reach + 1) x
-// (reach + 1) that an index wraps into as it wraps into the whole one. The rows keep their length: cut to
-// 2 reach + 1 as well, the rotation function read them 15% slower on the GPU (measured on a large cell).
-gemmi::FPhiGrid<float> crop_half_l(const gemmi::FPhiGrid<float> &full, int reach) {
-    gemmi::FPhiGrid<float> out;
-    out.unit_cell = full.unit_cell;
-    out.spacegroup = full.spacegroup;
-    out.axis_order = full.axis_order;
-    out.half_l = true;
-    const int nu = full.nu, n = 2 * reach + 1;
-    out.set_size_without_checking(nu, n, reach + 1);
-    for (int l = 0; l <= reach; l++)
-        for (int k = -reach; k <= reach; k++)
-            std::copy_n(&full.data[static_cast<size_t>(nu) * (gemmi::modulo(k, full.nv) + static_cast<size_t>(full.nv) * l)],
-                        nu, &out.data[static_cast<size_t>(nu) * (gemmi::modulo(k, n) + static_cast<size_t>(n) * l)]);
-    return out;
-}
-
 // --- placements scored against the data: R-work / R-free after the bulk solvent and the scale ---
 
 struct PlacementScore {
@@ -1179,7 +1161,8 @@ double RotationDistanceDeg(const gemmi::Mat33 &a, const gemmi::Mat33 &b,
 
 // --- the molecular transform ---
 
-MolecularTransform::MolecularTransform(const SearchModel &model, double d_min, size_t nthreads) {
+MolecularTransform::MolecularTransform(const SearchModel &model, double d_min, size_t nthreads,
+                                       [[maybe_unused]] size_t gpu_budget) {
     // On a P1 box six times the model's radius: the transform of a body of radius r turns its phase
     // by pi r / box per grid step, and trilinear interpolation reads it to a few percent at pi/6.
     // Interpolating the molecular transform of one model in place of a Fourier summation per
@@ -1195,18 +1178,18 @@ MolecularTransform::MolecularTransform(const SearchModel &model, double d_min, s
     dc.set_grid_cell_and_spacegroup(st);
     dc.set_refmac_compatible_blur(st.models[0]);
     PutModelDensityOnGrid(dc, st.models[0], {}, nthreads);
-    gemmi::FPhiGrid<float> full;
+    // The transform is read at |s| <= 1 / d_min only - the search's reflections, turned - and interpolation
+    // reads one step further out: of the grid the FFT makes, under half, which is all that is kept, here and
+    // on the GPU. The 2% are for a reflection's resolution taken in a cell slightly different from its search's.
+    const int reach = static_cast<int>(std::ceil(1.02 * box_ / d_min)) + 1;
 #ifdef JFJOCH_USE_CUDA
-    if (get_gpu_count() > 0 && MapToFPhiGPUFits(dc.grid)) {
-        full = MapToFPhiGPU(dc.grid);
+    // On the GPU where its buffers fit the budget - the search holds nothing else yet.
+    if (get_gpu_count() > 0 && MapToFPhiGPUBytes(dc.grid, reach) <= gpu_budget) {
+        grid_ = MapToFPhiGPU(dc.grid, reach);
         cuda_trim_allocation_pool();   // its buffers, which the engines allocated outside the pool cannot use
     } else
 #endif
-        full = MapToFPhi(dc.grid);
-    // The transform is read at |s| <= 1 / d_min only - the search's reflections, turned - and interpolation
-    // reads one step further out: of the grid the FFT made, under half, which is all that is kept, here and
-    // on the GPU. The 2% are for a reflection's resolution taken in a cell slightly different from its search's.
-    grid_ = crop_half_l(full, static_cast<int>(std::ceil(1.02 * box_ / d_min)) + 1);
+        grid_ = CropHalfL(MapToFPhi(dc.grid), reach);
     d_min_ = d_min;
     blur_ = dc.blur;
 }
@@ -2327,7 +2310,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
     t0 = std::chrono::steady_clock::now();
     std::optional<MolecularTransform> ft_made;
     try {
-        ft_made.emplace(model, result.d_min_search, nthreads);
+        ft_made.emplace(model, result.d_min_search, nthreads, gpu_budget);
     } catch (const std::exception &e) {   // the GPU's FFT
         result.failure_reason = std::string("the molecular transform: ") + e.what();
         return result;

image_analysis/structure_refinement/MolecularReplacement.h

@@ -90,7 +90,8 @@ double RotationDistanceDeg(const gemmi::Mat33 &a, const gemmi::Mat33 &b,
 // The structure factors of one copy of the model, at its centroid, at any reciprocal vector.
 class MolecularTransform {
 public:
-    MolecularTransform(const SearchModel &model, double d_min, size_t nthreads);
+    // `gpu_budget`: what its FFT may take on the GPU; on the CPU where that is too little.
+    MolecularTransform(const SearchModel &model, double d_min, size_t nthreads, size_t gpu_budget = 0);
     // F at the Cartesian reciprocal vector s (1/A), interpolated trilinearly.
     std::complex<double> At(const gemmi::Vec3 &s) const;
     // The grid At() interpolates (half-l, XYZ), its box edge in A and the blur it is unblurred by.

image_analysis/structure_refinement/MolecularReplacementGPU.cu

@@ -263,14 +263,15 @@ __global__ void translation_stats(const float *map, size_t n, double *sum, doubl
 }
 
 // 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;
+// 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), v = static_cast<int>((i / nu) % nv), l = static_cast<int>(i / (static_cast<size_t>(nu) * nv));
+    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)];
@@ -313,17 +314,15 @@ std::vector<int> bin_counts(const std::vector<int> &bin, int bins) {
 
 }  // namespace
 
-bool MapToFPhiGPUFits(const gemmi::Grid<float> &map) {
+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 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;
+    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) {
+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);
@@ -332,7 +331,7 @@ gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &map) {
     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);
+    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);
@@ -341,7 +340,7 @@ gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &map) {
     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, norm, d_out);
+            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),

image_analysis/structure_refinement/MolecularReplacementGPU.h

@@ -15,12 +15,11 @@
 
 #include "MolecularReplacement.h"
 
-// MapToFPhi (ModelFFT.h) with cuFFT on the calling thread's current device: the same transform, the same
-// half-l layout and conventions, to float rounding. For the molecular transform of a large model, whose box is
-// hundreds of millions of points. MapToFPhiGPUFits: whether its device buffers take at most a quarter of the
-// card's total memory - decided on the total, so the same input always takes the same path.
-bool MapToFPhiGPUFits(const gemmi::Grid<float> &map);
-gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &map);
+// CropHalfL(MapToFPhi(map), reach) (ModelFFT.h) with cuFFT on the calling thread's current device: the same
+// transform, the same half-l layout and conventions, to float rounding. For the molecular transform of a large
+// model, whose box is hundreds of millions of points. MapToFPhiGPUBytes: the device memory it takes while it runs.
+size_t MapToFPhiGPUBytes(const gemmi::Grid<float> &map, int reach);
+gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &map, int reach);
 
 struct MolecularReplacementGPUState;
 struct MolecularReplacementGPUGroup;

tests/MolecularReplacementTest.cpp

@@ -558,8 +558,8 @@ TEST_CASE("MolecularReplacement_GPUMatchesCPU", "[MolecularReplacement]") {
     gemmi::Structure st = PlaceSearchModel(model, truth, t, cell, gemmi::get_spacegroup_p1());
     dc.set_grid_cell_and_spacegroup(st);
     dc.put_model_density_on_grid(st.models[0]);
-    REQUIRE(MapToFPhiGPUFits(dc.grid));
-    const auto f_cpu = MapToFPhi(dc.grid), f_gpu = MapToFPhiGPU(dc.grid);
+    const int reach = dc.grid.nv / 3;   // inside the half grid, as the transform's own reach is
+    const auto f_cpu = CropHalfL(MapToFPhi(dc.grid), reach), f_gpu = MapToFPhiGPU(dc.grid, reach);
     REQUIRE(f_gpu.data.size() == f_cpu.data.size());
     double worst = 0, largest = 0;
     for (size_t i = 0; i < f_cpu.data.size(); i++) {