mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 82 Wiki Activity

Merge mr-gpumem (51706641f) into rc175: RESULT-CHANGING (owner approved, float rounding): molecular replacement's transform FFT on the GPU wherever it fits the search's budget

Browse Source 
4nwv molecular replacement 31.7 -> 20.1 s; outputs identical on myob, cytochrome C and 4nwv on the RTX 5080.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-10-15 13:53:02 +02:00
co-authored by Claude Opus 5.5
parent 323df1561e 51706641fa
commit d91fc7ca3d
8 changed files with 57 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/GPU_DATA_ANALYSIS.md
+2 -1
View File
@@ -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.
image_analysis/structure_refinement/ModelFFT.cpp
+17
View File
@@ -110,3 +110,20 @@ gemmi::Grid<float> MapFromFPhi(const gemmi::FPhiGrid<float> &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<float> CropHalfL(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;
}
image_analysis/structure_refinement/ModelFFT.h
+5
View File
@@ -22,3 +22,8 @@ gemmi::FPhiGrid<float> MapToFPhi(const gemmi::Grid<float> &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<float> MapFromFPhi(const gemmi::FPhiGrid<float> &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<float> CropHalfL(const gemmi::FPhiGrid<float> &full, int reach);
image_analysis/structure_refinement/MolecularReplacement.cpp
+11 -28
View File
@@ -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;
image_analysis/structure_refinement/MolecularReplacement.h
+2 -1
View File
@@ -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<double> 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.
image_analysis/structure_refinement/MolecularReplacementGPU.cu
+13 -14
View File
@@ -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<size_t>(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<size_t>(nu) * nk * (reach + 1);
const size_t i = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
if (i >= n)
return;
const int u = static_cast<int>(i % nu), v = static_cast<int>((i / nu) % nv), l = static_cast<int>(i / (static_cast<size_t>(nu) * nv));
const int u = static_cast<int>(i % nu), k = static_cast<int>((i / nu) % nk), l = static_cast<int>(i / (static_cast<size_t>(nu) * nk));
const int v = k <= reach ? k : k - nk + nv;
cufftComplex x;
if (u < nuh) {
x = in[u + static_cast<size_t>(nuh) * (v + static_cast<size_t>(nv) * l)];
@@ -313,17 +314,15 @@ std::vector<int> bin_counts(const std::vector<int> &bin, int bins) {
} // namespace
bool MapToFPhiGPUFits(const gemmi::Grid<float> &map) {
size_t MapToFPhiGPUBytes(const gemmi::Grid<float> &map, int reach) {
const size_t points = static_cast<size_t>(map.nu) * map.nv * map.nw;
const size_t half_u = static_cast<size_t>(map.nu / 2 + 1) * map.nv * map.nw;
const size_t half_w = static_cast<size_t>(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<size_t>(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<float> MapToFPhiGPU(const gemmi::Grid<float> &map) {
gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &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<float> MapToFPhiGPU(const gemmi::Grid<float> &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<float>(map.unit_cell.volume / map.point_count());
CudaStream stream;
auto d_map = upload(map.data, stream);
@@ -341,7 +340,7 @@ gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &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<<<static_cast<unsigned>((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<float>) == sizeof(float2));
cuda_err(cudaMemcpyAsync(hkl.data.data(), d_out.get(), hkl.data.size() * sizeof(float2),
image_analysis/structure_refinement/MolecularReplacementGPU.h
+5 -6
View File
@@ -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<float> &map);
gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &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<float> &map, int reach);
gemmi::FPhiGrid<float> MapToFPhiGPU(const gemmi::Grid<float> &map, int reach);
struct MolecularReplacementGPUState;
struct MolecularReplacementGPUGroup;
tests/MolecularReplacementTest.cpp
+2 -2
View File
@@ -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++) {