From 51706641fae540d90f09c26dcff667dc05ffe060 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 13:41:48 +0200 Subject: [PATCH] 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 --- docs/GPU_DATA_ANALYSIS.md | 3 +- .../structure_refinement/ModelFFT.cpp | 17 ++++++++ .../structure_refinement/ModelFFT.h | 5 +++ .../MolecularReplacement.cpp | 39 ++++++------------- .../MolecularReplacement.h | 3 +- .../MolecularReplacementGPU.cu | 27 +++++++------ .../MolecularReplacementGPU.h | 11 +++--- tests/MolecularReplacementTest.cpp | 4 +- 8 files changed, 57 insertions(+), 52 deletions(-) diff --git a/docs/GPU_DATA_ANALYSIS.md b/docs/GPU_DATA_ANALYSIS.md index 4573c4a69..3e99cd7f8 100644 --- a/docs/GPU_DATA_ANALYSIS.md +++ b/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. diff --git a/image_analysis/structure_refinement/ModelFFT.cpp b/image_analysis/structure_refinement/ModelFFT.cpp index ff5c43ef6..54830fa86 100644 --- a/image_analysis/structure_refinement/ModelFFT.cpp +++ b/image_analysis/structure_refinement/ModelFFT.cpp @@ -110,3 +110,20 @@ gemmi::Grid MapFromFPhi(const gemmi::FPhiGrid &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 CropHalfL(const gemmi::FPhiGrid &full, int reach) { + gemmi::FPhiGrid 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(nu) * (gemmi::modulo(k, full.nv) + static_cast(full.nv) * l)], + nu, &out.data[static_cast(nu) * (gemmi::modulo(k, n) + static_cast(n) * l)]); + return out; +} diff --git a/image_analysis/structure_refinement/ModelFFT.h b/image_analysis/structure_refinement/ModelFFT.h index 793d7153b..88a83fdb5 100644 --- a/image_analysis/structure_refinement/ModelFFT.h +++ b/image_analysis/structure_refinement/ModelFFT.h @@ -22,3 +22,8 @@ gemmi::FPhiGrid MapToFPhi(const gemmi::Grid &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 MapFromFPhi(const gemmi::FPhiGrid &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 CropHalfL(const gemmi::FPhiGrid &full, int reach); diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 395d4c232..c04f07bbf 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/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 crop_half_l(const gemmi::FPhiGrid &full, int reach) { - gemmi::FPhiGrid 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(nu) * (gemmi::modulo(k, full.nv) + static_cast(full.nv) * l)], - nu, &out.data[static_cast(nu) * (gemmi::modulo(k, n) + static_cast(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 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(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(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 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; diff --git a/image_analysis/structure_refinement/MolecularReplacement.h b/image_analysis/structure_refinement/MolecularReplacement.h index fb344a89b..f887e3058 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.h +++ b/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 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. diff --git a/image_analysis/structure_refinement/MolecularReplacementGPU.cu b/image_analysis/structure_refinement/MolecularReplacementGPU.cu index 5b4766dd9..23cf234f0 100644 --- a/image_analysis/structure_refinement/MolecularReplacementGPU.cu +++ b/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(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(nu) * nk * (reach + 1); const size_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (i >= n) return; - const int u = static_cast(i % nu), v = static_cast((i / nu) % nv), l = static_cast(i / (static_cast(nu) * nv)); + const int u = static_cast(i % nu), k = static_cast((i / nu) % nk), l = static_cast(i / (static_cast(nu) * nk)); + const int v = k <= reach ? k : k - nk + nv; cufftComplex x; if (u < nuh) { x = in[u + static_cast(nuh) * (v + static_cast(nv) * l)]; @@ -313,17 +314,15 @@ std::vector bin_counts(const std::vector &bin, int bins) { } // namespace -bool MapToFPhiGPUFits(const gemmi::Grid &map) { +size_t MapToFPhiGPUBytes(const gemmi::Grid &map, int reach) { const size_t points = static_cast(map.nu) * map.nv * map.nw; const size_t half_u = static_cast(map.nu / 2 + 1) * map.nv * map.nw; - const size_t half_w = static_cast(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(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 MapToFPhiGPU(const gemmi::Grid &map) { +gemmi::FPhiGrid MapToFPhiGPU(const gemmi::Grid &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 MapToFPhiGPU(const gemmi::Grid &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(map.unit_cell.volume / map.point_count()); CudaStream stream; auto d_map = upload(map.data, stream); @@ -341,7 +340,7 @@ gemmi::FPhiGrid MapToFPhiGPU(const gemmi::Grid &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<<((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) == sizeof(float2)); cuda_err(cudaMemcpyAsync(hkl.data.data(), d_out.get(), hkl.data.size() * sizeof(float2), diff --git a/image_analysis/structure_refinement/MolecularReplacementGPU.h b/image_analysis/structure_refinement/MolecularReplacementGPU.h index 7f32d7d76..1d68e31e0 100644 --- a/image_analysis/structure_refinement/MolecularReplacementGPU.h +++ b/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 &map); -gemmi::FPhiGrid MapToFPhiGPU(const gemmi::Grid &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 &map, int reach); +gemmi::FPhiGrid MapToFPhiGPU(const gemmi::Grid &map, int reach); struct MolecularReplacementGPUState; struct MolecularReplacementGPUGroup; diff --git a/tests/MolecularReplacementTest.cpp b/tests/MolecularReplacementTest.cpp index aea364215..e03353a9f 100644 --- a/tests/MolecularReplacementTest.cpp +++ b/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++) {