Files
Jungfraujoch/image_analysis/structure_refinement/MolecularReplacementGPU.h
T
leonarski_fandClaude Opus 5.5 51706641fa 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>
2026-10-15 13:41:48 +02:00

60 lines
3.4 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
// The molecular-replacement search on a GPU (CUDA builds only; the header itself needs no CUDA): the
// rotation function and the translation function of MolecularReplacement.h, the same functions as the
// CPU computes - the same trilinear interpolation of the same molecular transform, the same bins, the
// same coefficients on the same grid - in float on the device, with cuFFT for FFTW. The two agree to
// rounding, not bit for bit. Every reduction runs in a fixed order and nothing is summed with an
// atomic, so the GPU is deterministic on its own.
#include <memory>
#include <vector>
#include "MolecularReplacement.h"
// 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;
struct TranslationSetupGPU;
class MolecularReplacementGPU {
public:
// Uploads the molecular transform, on the calling thread's current device. Throws JFJochException on
// a CUDA failure, as every call below does.
explicit MolecularReplacementGPU(const MolecularTransform &ft);
~MolecularReplacementGPU();
std::vector<double> RotationFunction(const SearchReflections &refl, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, const std::vector<gemmi::Mat33> &rotations);
// The translation function of one group at each of `rotations`; with `known`, beside a known partial
// structure (KnownStructureFactors). `prepared`: Prepare() of this setup, for a setup used again and again;
// made for this call otherwise.
std::vector<TranslationPeak> TranslationFunction(const TranslationSetup &setup,
const std::vector<gemmi::Mat33> &rotations,
const std::vector<std::complex<double>> *known = nullptr,
TranslationSetupGPU *prepared = nullptr);
// What the translation function needs of `setup`, worked out once on the host; with `known_part`, for the
// function beside a known partial structure.
std::unique_ptr<TranslationSetupGPU, void (*)(TranslationSetupGPU *)> Prepare(const TranslationSetup &setup,
bool known_part);
// The device memory a call takes for as long as it runs, and gives back when it returns: RotationFunction()
// of `rotations` against `refl`'s reflections in `sg`, and TranslationFunction() of `rotations` (with
// `known`, beside a known partial structure).
static size_t RotationBytes(const SearchReflections &refl, const gemmi::SpaceGroup &sg, size_t rotations);
static size_t TranslationBytes(const TranslationSetup &setup, size_t rotations, bool known);
private:
std::unique_ptr<MolecularReplacementGPUState> state_;
};