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++) {