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