leonarski_fandClaude Opus 5.5
011b98838e
What molecular replacement put on its GPU was decided part by part on a quarter of the card each: every group's rigid-body pool (up to three engines), every null's (up to nine), the score engines' cache, the --mr hand check's structure factors. The groups and the nulls run side by side, so the quarters added up: on 4nwv on a 16 GB card the search held 4 rigid-body engines of 1.88 GB, two score engines of 1.64 GB, the 1.82 GB molecular transform, two translation functions' 1.2 GB and the 2.4 GB the allocation pool kept from the validation before - 13.8 GB at its peak, 83% of the card. On an 8.55 GB card that pages. Now each merge's search is planned once, before it starts, on the budget model validation uses (ModelGPUBudget: 80% of the card's total, less what a merge engine still at work holds): the transform, one rigid-body engine and the largest score engine take their place in that order where they still fit, more rigid-body engines only within half of the budget, and what is left is a room the score engines (least recently used given up first, as before) and the rotation and translation functions' working buffers share, each waiting for the other to finish rather than spilling past it (PlanMolecularReplacementGPU). A part that does not fit runs on the CPU, decided on its size alone. - One rigid-body pool per merge serves every fit of the merge - the groups' first and further copies, the nulls, the tNCS shifts, the carried copies - sized for the largest (the most copies Matthews allows, in any group) and reserved once, instead of a pool made and freed for each. RigidBodyGPUPool::Cover widens a planned pool to another model or group; its zones are kept by group, atom count and resolution. RigidBodyGPUPool::Create (the quarter rule) is gone. - The sizes are worked out on the CPU while the transform is computed. - The molecular transform keeps only the indices it is read at (|k|, l <= box / d_min + 1): 0.84 GB instead of 1.82 on 4nwv, host and device. Its rows keep their full length: cut short too, the rotation function read them 15% slower. - The translation function's prepared setup is kept on the host and uploaded by each call, so the device holds it only while a function runs; the cross-term offsets are uploaded only where there are cross terms. - The allocation pool is trimmed before the search, after the transform's GPU FFT, before a score engine is made and before the hand check, whose structure factors are now judged against the budget. - rugnux passes the merge engine's bytes, where the finalist ledger keeps it, to the search, the hand check and the validation of the placed model. Measured on the RTX 5080 (16.6 GB), timing benchmark: peak device memory during MR 13.8 -> 8.7 GB on 4nwv (start 3.0 -> 0.66 GB after the trim), 1.67 -> 1.14 GB on cytc, 6.3 -> 0.76 GB on myob. On 4nwv the plan holds the transform 0.84 GB, one rigid-body engine of 2.10 GB and leaves 10.35 GB of a 13.29 GB budget to the score engines and the functions; a second engine would pass half of the budget (one engine against two: +1.2 s on the rescoring and the null). MR 31.7 -> 31.6 s on 4nwv. Reflection files, maps, placed models, reports and the MR candidate lines byte-identical to rc175-metal-merge on myob, cytc and 4nwv; [MolecularReplacement], [ModelValidation], RigidBody*, ModelScale*, [CUDAMemHelpers] pass. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
61 lines
3.4 KiB
C++
61 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"
|
|
|
|
// 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);
|
|
|
|
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_;
|
|
};
|