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 <noreply@anthropic.com>
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-10-15 13:41:48 +02:00
1 parent 011b98838e
commit 51706641fa
8 files changed
+57 -52

No files matched your search

@@ -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<float> crop_half_l(const gemmi::FPhiGrid<float> &full, int reach) {
gemmi::FPhiGrid<float> 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<size_t>(nu) * (gemmi::modulo(k, full.nv) + static_cast<size_t>(full.nv) * l)],
nu, &out.data[static_cast<size_t>(nu) * (gemmi::modulo(k, n) + static_cast<size_t>(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<float> 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<int>(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<int>(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<MolecularTransform> 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;