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:
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;
|
||||
|
||||
Reference in new issue
Block a user