From 8d17d052273ccf3bf2b1f4501b00759c2d980713 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 15:48:47 +0200 Subject: [PATCH] RESULT-CHANGING: molecular replacement's CPU transform FFT on all threads Where the molecular transform's FFT runs on the CPU - a CPU-only build, --no-gpu, a card too small for it - it was one single-threaded FFTW 3D r2c: 12.5 s of 13 s of the transform on 4nwv's 768^3 box, one core busy while fifteen waited. Now MapToFPhiInParallel does it as two passes over the threads: the r2c along w of every (u, v) column, then the 2D transform of each plane l the search keeps (l <= reach), cropping as CropHalfL does. Each column and each plane is done by one plan, so the result does not depend on the thread count; it agrees with FFTW's 3D plan to float rounding, not bit for bit. The GPU path is unchanged. CPU-only (--no-gpu), timing benchmark: 4nwv's transform 13.1 -> 2.4 s (provisional, load average 13-25). What the rounding moved: 4nwv and myob came out byte-identical; on cytc (three copies, tNCS) the null's replicates, placed by the translation function, read the fit at +22.8 sd instead of +18.7 (MR_HAND_FIT_SIGMA), and the placement - same R-work, R-free and every MR candidate line - came out at an origin-equivalent position 43 A from the previous one, so its model file and maps differ while the answer (ENANTIOMORPH, P 32 2 1, placement CC 0.744) does not. GPU runs byte-identical on all three. [MolecularReplacement] passes. Co-Authored-By: Claude Opus 5.5 --- .../structure_refinement/ModelFFT.cpp | 61 +++++++++++++++++++ .../structure_refinement/ModelFFT.h | 7 +++ .../MolecularReplacement.cpp | 2 +- 3 files changed, 69 insertions(+), 1 deletion(-) diff --git a/image_analysis/structure_refinement/ModelFFT.cpp b/image_analysis/structure_refinement/ModelFFT.cpp index aee3be8eb..57b482481 100644 --- a/image_analysis/structure_refinement/ModelFFT.cpp +++ b/image_analysis/structure_refinement/ModelFFT.cpp @@ -15,6 +15,7 @@ #include "gemmi/fail.hpp" #include "../../common/FFTWPlannerLock.h" +#include "../../common/ParallelFor.h" namespace { @@ -83,6 +84,66 @@ void MapFromFPhiExecute(int nu, int nv, int nw, std::complex *in, float * fftwf_execute_dft_c2r(PlanFor(nu, nv, nw, true), reinterpret_cast(in), out); } +gemmi::FPhiGrid MapToFPhiInParallel(const gemmi::Grid &map, int reach, size_t nthreads) { + if (map.axis_order == gemmi::AxisOrder::ZYX) + gemmi::fail("MapToFPhiInParallel(): ZYX order is not supported"); + const int nu = map.nu, nv = map.nv, nw = map.nw, nwh = nw / 2 + 1, nk = 2 * reach + 1; + const size_t plane = static_cast(nu) * nv; + float *in = fftwf_alloc_real(map.data.size()); + fftwf_complex *half = fftwf_alloc_complex(plane * nwh); + if (in == nullptr || half == nullptr) { + fftwf_free(in); + fftwf_free(half); + gemmi::fail("MapToFPhiInParallel(): cannot allocate the FFT buffers"); + } + // Made once here; executing a plan on new arrays is safe from several threads at once. Unaligned, as a + // column or a plane starts wherever it does. + fftwf_plan columns = nullptr, planes = nullptr; + { + std::lock_guard lock(FFTWPlannerMutex()); + // The nu columns of one v, w along the stride of a plane. + fftwf_iodim col_dim{nw, static_cast(plane), static_cast(plane)}; + fftwf_iodim col_many{nu, 1, 1}; + columns = fftwf_plan_guru_dft_r2c(1, &col_dim, 1, &col_many, in, half, FFTW_ESTIMATE | FFTW_UNALIGNED); + fftwf_iodim plane_dims[2] = {{nv, nu, nu}, {nu, 1, 1}}; + planes = fftwf_plan_guru_dft(2, plane_dims, 0, nullptr, half, half, FFTW_FORWARD, FFTW_ESTIMATE | FFTW_UNALIGNED); + } + if (columns == nullptr || planes == nullptr) + gemmi::fail("MapToFPhiInParallel(): no FFTW plan for the grid"); + std::copy(map.data.begin(), map.data.end(), in); + ParallelFor(nv, nthreads, [&](int v) { + fftwf_execute_dft_r2c(columns, in + static_cast(v) * nu, half + static_cast(v) * nu); + }); + ParallelFor(reach + 1, nthreads, [&](int l) { + fftwf_execute_dft(planes, half + l * plane, half + l * plane); + }); + + gemmi::FPhiGrid 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(nu, nk, reach + 1); + const float norm = static_cast(map.unit_cell.volume / map.point_count()); + // Conjugated and scaled as MapToFPhi does it, and cropped as CropHalfL crops it. + ParallelFor(reach + 1, nthreads, [&](int l) { + for (int k = -reach; k <= reach; k++) { + const fftwf_complex *row = half + l * plane + static_cast(nu) * gemmi::modulo(k, nv); + std::complex *out = &hkl.data[static_cast(nu) * (gemmi::modulo(k, nk) + static_cast(nk) * l)]; + for (int u = 0; u < nu; u++) + out[u] = std::complex(row[u][0] * norm, -row[u][1] * norm); + } + }); + { + std::lock_guard lock(FFTWPlannerMutex()); + fftwf_destroy_plan(columns); + fftwf_destroy_plan(planes); + } + fftwf_free(in); + fftwf_free(half); + return hkl; +} + gemmi::Grid MapFromFPhi(const gemmi::FPhiGrid &hkl) { if (hkl.axis_order == gemmi::AxisOrder::ZYX || !hkl.half_l) gemmi::fail("MapFromFPhi(): only an XYZ half-l grid is supported"); diff --git a/image_analysis/structure_refinement/ModelFFT.h b/image_analysis/structure_refinement/ModelFFT.h index 3ebb0754f..d59d64984 100644 --- a/image_analysis/structure_refinement/ModelFFT.h +++ b/image_analysis/structure_refinement/ModelFFT.h @@ -28,6 +28,13 @@ gemmi::Grid MapFromFPhi(const gemmi::FPhiGrid &hkl); // (nu, nv, nw) map before MapFromFPhi's 1/V. Both from fftwf_malloc, as a plan requires. void MapFromFPhiExecute(int nu, int nv, int nw, std::complex *in, float *out); +// CropHalfL(MapToFPhi(map), reach) on `nthreads` threads, for a map of hundreds of millions of points (the +// molecular transform's box), whose single-threaded transform is minutes of CPU: the r2c along w of every +// (u, v) column, then the 2D transform of every plane l <= reach - each column and each plane by one plan, so +// the result does not depend on the thread count. The same transform as MapToFPhi to float rounding, not bit +// for bit: FFTW's 3D plan orders its arithmetic its own way. +gemmi::FPhiGrid MapToFPhiInParallel(const gemmi::Grid &map, int reach, size_t nthreads); + // 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). diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index dee1aa224..47d165261 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -1190,7 +1190,7 @@ MolecularTransform::MolecularTransform(const SearchModel &model, double d_min, s cuda_trim_allocation_pool(); // its buffers, which the engines allocated outside the pool cannot use } else #endif - grid_ = CropHalfL(MapToFPhi(dc.grid), reach); + grid_ = MapToFPhiInParallel(dc.grid, reach, nthreads); d_min_ = d_min; blur_ = dc.blur; }