Merge mr-gpumem (8d17d0522) into rc175: RESULT-CHANGING (owner approved, float rounding): molecular replacement's CPU transform FFT on all threads, the result independent of the thread count

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PXc1zD1yxHqaR1DGdCTHL4

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<float> *in, float *
     fftwf_execute_dft_c2r(PlanFor(nu, nv, nw, true), reinterpret_cast<fftwf_complex *>(in), out);
 }
 
+gemmi::FPhiGrid<float> MapToFPhiInParallel(const gemmi::Grid<float> &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<size_t>(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<int>(plane), static_cast<int>(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<size_t>(v) * nu, half + static_cast<size_t>(v) * nu);
+    });
+    ParallelFor(reach + 1, nthreads, [&](int l) {
+        fftwf_execute_dft(planes, half + l * plane, half + l * plane);
+    });
+
+    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(nu, nk, reach + 1);
+    const float norm = static_cast<float>(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<size_t>(nu) * gemmi::modulo(k, nv);
+            std::complex<float> *out = &hkl.data[static_cast<size_t>(nu) * (gemmi::modulo(k, nk) + static_cast<size_t>(nk) * l)];
+            for (int u = 0; u < nu; u++)
+                out[u] = std::complex<float>(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<float> MapFromFPhi(const gemmi::FPhiGrid<float> &hkl) {
     if (hkl.axis_order == gemmi::AxisOrder::ZYX || !hkl.half_l)
         gemmi::fail("MapFromFPhi(): only an XYZ half-l grid is supported");

image_analysis/structure_refinement/ModelFFT.h

@@ -28,6 +28,13 @@ gemmi::Grid<float> MapFromFPhi(const gemmi::FPhiGrid<float> &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<float> *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<float> MapToFPhiInParallel(const gemmi::Grid<float> &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).

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