From the nsys/perf profiles of the copy chain (GPU busy 39% between copies, the thread mostly in host work and syncs): - The known-structure translation function's lists, device arrays and cuFFT plan are made once per group (MolecularReplacementGPU::Prepare) instead of on every copy (~260 MB of uploads and a plan per call before). - The known copies' CopyFactors and likelihood, which rank the peaks, are taken while the GPU runs the translation function. - The structure-factor engines are made per entry, without one lock held across all groups' creation. - The transform's index wrap divides only for an index outside the grid (CPU and GPU). All exact: same arrays, same kernels, same order. 4mtk MR, alternating 3 pairs: 9.5 -> 9.0 s. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
55 lines
3.0 KiB
C++
55 lines
3.0 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` on the device, made once; with `known_part`, for the
|
|
// function beside a known partial structure.
|
|
std::unique_ptr<TranslationSetupGPU, void (*)(TranslationSetupGPU *)> Prepare(const TranslationSetup &setup,
|
|
bool known_part);
|
|
|
|
private:
|
|
std::unique_ptr<MolecularReplacementGPUState> state_;
|
|
};
|