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

Molecular replacement: one GPU plan per merge on model validation's budget, the transform cropped, output unchanged

Browse Source 
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>
This commit is contained in:
leonarski_f
2026-10-15 13:33:39 +02:00
co-authored by Claude Opus 5.5
parent 1ebb28e3ab
commit 011b98838e
12 changed files with 644 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/CPU_DATA_ANALYSIS_DECISIONS.md
+6 -1
View File
@@ -322,5 +322,10 @@ less: it may relabel the **enantiomorph** of the answer, past the anomalous chec
reindexes; another group with the answer's rotations, or the twin hypothesis' subgroup, is written
beside the answer as `<prefix>_mr.mtz`/`.cif` with the model as placed there, and the answer's own
files stay what the run decided. The rotation and translation functions, the rigid body and the
structure factors run on the GPU where there is one. The keys, thresholds and table are described in
structure factors run on the GPU where there is one, decided for each merge before its search starts on
the budget of §14.8 (80% of the card's total memory, less what a merge still at work holds there): the
molecular transform, one rigid-body engine shared by every fit of the merge and the largest score engine
take their place in that order where they still fit, more rigid-body engines are added only within half
of the budget, and the score engines and the functions' working buffers share what is left, waiting for
each other rather than spilling past it. The keys, thresholds and table are described in
[the results report](RUGNUX_REPORT.md#molecular-replacement).
image_analysis/structure_refinement/ModelValidation.cpp
+18 -9
View File
@@ -1883,10 +1883,14 @@ ModelValidationResult Validate(const std::vector<MergedReflection> &merged,
} // namespace
size_t ModelGPUBudget(size_t card, size_t resident) {
const size_t usable = card / 5 * 4;
return usable > resident ? usable - resident : 0;
}
ModelValidationGPUPlan PlanModelValidationGPU(const ModelValidationGPUNeed &need, size_t cards) {
ModelValidationGPUPlan plan;
const size_t usable = need.card / 5 * 4;
plan.budget = usable > need.resident ? usable - need.resident : 0;
plan.budget = ModelGPUBudget(need.card, need.resident);
auto take = [&](size_t bytes, size_t limit) {
if (bytes == 0 || plan.bytes + bytes > limit)
return false;
@@ -1999,7 +2003,8 @@ const gemmi::SpaceGroup *AdoptModelFrame(const ModelValidationResult &validation
ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection> &merged,
const gemmi::Structure &placed, const gemmi::SpaceGroup &data_space_group,
double k_sol, double b_sol, Logger &logger, size_t nthreads) {
double k_sol, double b_sol, Logger &logger, size_t nthreads,
size_t gpu_resident_bytes) {
ModelValidationResult result;
gemmi::Structure st = placed;
st.setup_cell_images();
@@ -2018,13 +2023,16 @@ ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection>
ModelStructureFactorsGPU *gpu = nullptr;
#ifdef JFJOCH_USE_CUDA
std::unique_ptr<ModelStructureFactorsGPU> engine;
// Read beside molecular replacement's engines, which keep to a quarter of the card
// (MolecularReplacement.cpp), so on the GPU where it fits a quarter as well.
if (get_gpu_count() > 0)
// After molecular replacement, whose engines are gone: alone on the card beside what is resident.
if (get_gpu_count() > 0) {
cuda_trim_allocation_pool();
engine = StructureFactorsEngine(RigidBodyGPUEngine::CurrentDevice(), st.models[0], st.cell, *sg, d_min, logger);
if (engine && engine->DeviceBytes() > engine->DeviceTotalMemory() / 4) {
logger.Info("Model validation: structure factors to {:.2f} A on the CPU - {} of GPU memory, over a quarter "
"of the card's {}", d_min, GBytes(engine->DeviceBytes()), GBytes(engine->DeviceTotalMemory()));
}
if (engine && engine->DeviceBytes() > ModelGPUBudget(engine->DeviceTotalMemory(), gpu_resident_bytes)) {
logger.Info("Model validation: structure factors to {:.2f} A on the CPU - {} of GPU memory, over the "
"budget of 80% of the card's {}{}", d_min, GBytes(engine->DeviceBytes()),
GBytes(engine->DeviceTotalMemory()),
gpu_resident_bytes > 0 ? fmt::format(" less {} held by the merge", GBytes(gpu_resident_bytes)) : "");
engine.reset();
}
if (engine)
@@ -2039,6 +2047,7 @@ ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection>
}
#else
(void) logger;
(void) gpu_resident_bytes;
#endif
if (gpu == nullptr) {
gemmi::DensityCalculator<Table, float> dc;
image_analysis/structure_refinement/ModelValidation.h
+9 -2
View File
@@ -308,6 +308,11 @@ struct ModelValidationGPUPlan {
bool may_run_beside = false;
};
// What model validation and molecular replacement may hold on a card: 80% of its TOTAL memory - never what
// is free at the moment, so the same input on the same card always runs the same way - less what is
// resident there already (a merge engine still at work beside them).
size_t ModelGPUBudget(size_t card, size_t resident);
// The one rule for what a validation puts on its GPU, decided before anything is reserved, from the
// card's TOTAL memory and never from what is free at the moment, so the same input on the same card
// always runs the same way. Everything the validation holds is alive at once - the engines are reserved
@@ -360,10 +365,12 @@ void KeepModelVerdict(ModelValidationResult &remade, const ModelValidationResult
// F(+) - F(-) on the phases of F_calc + solvent, and its reading at the model's atoms - anomalous_pairs,
// the scatterers' mean, the sites and anomalous_hands_disagree, by the same code and the same rule as the
// validation's. No placement, no scale fit, no null, no file. For a caller that has already placed the
// model and judged its fit (the --mr hand check). On the GPU where there is one.
// model and judged its fit (the --mr hand check). On the GPU where the structure factors fit the budget
// (ModelGPUBudget) beside `gpu_resident_bytes`.
ModelValidationResult ReadModelAnomalousMap(const std::vector<MergedReflection> &merged,
const gemmi::Structure &placed, const gemmi::SpaceGroup &data_space_group,
double k_sol, double b_sol, Logger &logger, size_t nthreads);
double k_sol, double b_sol, Logger &logger, size_t nthreads,
size_t gpu_resident_bytes = 0);
// Reference intensities computed from an atomic model: |F_model|^2, with a flat bulk-solvent
// contribution at the standard constants, for every reflection to `d_min`. This is the same
image_analysis/structure_refinement/MolecularReplacement.cpp
+318 -144
View File
@@ -38,14 +38,15 @@
#include "ModelScaling.h"
#include "ModelValidation.h"
#include "RigidBodyRefine.h"
#ifdef JFJOCH_USE_CUDA
#include "MolecularReplacementGPU.h"
#ifdef JFJOCH_USE_CUDA
#include "RigidBodyGPU.h"
#include "RigidBodyGPUEngine.h"
#include "ModelStructureFactorsGPU.h"
#include "ModelScaleGPU.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
#include "../indexing/CUDAMemHelpers.h" // cuda_trim_allocation_pool
#endif
namespace {
@@ -171,6 +172,24 @@ 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 {
@@ -282,16 +301,22 @@ PlacementScore r_factors(const gemmi::AsuData<std::complex<float>> &fmodel, cons
// Model validation's structure-factor engines on the GPU, one per group, copy count and resolution, made
// where the first placement of that kind is scored, each with the scale fit on the GPU beside it - or none,
// where there is no GPU, the cell is too small for its gridding or the two would take over a quarter of the
// card. Decided on the card's total memory, so the same input always takes the same path.
// where there is no GPU, the cell is too small for its gridding or the two do not fit the room the search's
// plan leaves them (PlanMolecularReplacementGPU). Decided on the engine's size, so the same input always takes
// the same path.
//
// Together they keep within that quarter too. A search asks for one per group, copy count and resolution,
// and kept to the end of the search they summed to well over the card on a large cell (seven of 1.4 GB
// beside the rigid-body engines). So an engine that does not fit beside the ones kept takes the place of
// the least recently used of those no score is using, and waits for a score to finish where every one is
// in use; one taken out is made again where it is next asked for. Every engine computes the same numbers
// whenever it is made, so this changes when the card holds them, not what a score gives.
// Together they keep within that room, which the rotation and translation functions' working buffers share
// (Hold). A search asks for one per group, copy count and resolution, and kept to the end of the search they
// summed to well over the card on a large cell (seven of 1.4 GB beside the rigid-body engines). So an engine,
// or a function's buffers, that do not fit beside what is held take the place of the least recently used
// engine no score is using, and wait for a score or a function to finish where every one is in use; an engine
// taken out is made again where it is next asked for. Every engine computes the same numbers whenever it is
// made, so this changes when the card holds them, not what a score gives.
#ifdef JFJOCH_USE_CUDA
std::string in_gb(size_t bytes) {
return fmt::format("{:.2f} GB", bytes / 1e9);
}
class ScoreEngines {
public:
struct Engine {
@@ -304,50 +329,63 @@ public:
CudaDevicePtr<std::complex<float>> fcalc, fmask;
};
// The room of a merge's plan, set before its placements are scored; `scores`: whether they go on the GPU.
void SetRoom(size_t room, bool scores) {
std::lock_guard lock(m_);
room_ = room;
scores_ = scores;
}
bool Fits(size_t bytes) {
std::lock_guard lock(m_);
return bytes <= room_;
}
// The engine for this kind of placement, kept on the card for as long as the returned pointer lives.
std::shared_ptr<Engine> get(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, double d_min) {
if (get_gpu_count() == 0)
return nullptr;
size_t atoms = 0;
for (const gemmi::Chain &ch : st.models[0].chains)
for (const gemmi::Residue &r : ch.residues)
atoms += r.atoms.size();
std::unique_lock lock(m_);
Entry &e = engines_[{sg.xhm(), atoms, d_min}];
if (!scores_)
return nullptr;
Entry &e = engines_[{sg.xhm(), gemmi::count_atom_sites(st.models[0]), d_min}];
e.last_use = ++clock_;
// Made by the first to ask for it, without holding up those asking for other ones.
cv_.wait(lock, [&] { return !e.making; });
if (e.engine || e.refused)
return Lease(e);
e.making = true;
const size_t room = room_;
lock.unlock();
std::unique_ptr<ModelStructureFactorsGPU> sf;
size_t points = 0, bytes = 0, quarter = 0;
size_t points = 0, bytes = 0;
try {
sf = std::make_unique<ModelStructureFactorsGPU>(RigidBodyGPUEngine::CurrentDevice(), st.models[0],
st.cell, sg, d_min);
points = sf->Reflections();
bytes = sf->DeviceBytes() + ModelScaleGPU::DeviceBytes(points) + 2 * points * sizeof(std::complex<float>);
quarter = sf->DeviceTotalMemory() / 4;
bytes = Bytes(*sf);
} catch (const std::exception &) {
cuda_clear_error();
sf.reset();
}
lock.lock();
if (!sf || !sf->Supported() || bytes > quarter) {
if (!sf || !sf->Supported() || bytes > room) {
e.refused = true;
e.making = false;
cv_.notify_all();
return nullptr;
}
const int card = sf->Device();
cv_.wait(lock, [&] { return MakeRoom(card, bytes, quarter); });
cv_.wait(lock, [&] { return MakeRoom(card, bytes); });
resident_[card] += bytes;
lock.unlock();
std::shared_ptr<Engine> engine;
try {
// The functions' buffers come from the stream-ordered pool, which keeps what they gave back; the
// structure factors are allocated outside it, and could not use it.
cuda_trim_allocation_pool();
sf->Reserve();
const int device = RigidBodyGPUEngine::CurrentDevice();
set_gpu(card);
@@ -376,6 +414,39 @@ public:
return Lease(e);
}
// What a score engine for `sf` takes: its structure factors, the scale fit and the two arrays between them.
static size_t Bytes(const ModelStructureFactorsGPU &sf) {
const size_t points = sf.Reflections();
return sf.DeviceBytes() + ModelScaleGPU::DeviceBytes(points) + 2 * points * sizeof(std::complex<float>);
}
// `bytes` of the room on the current card for a function's working buffers, held for as long as the result
// lives.
class Held {
public:
Held(ScoreEngines &owner, int card, size_t bytes) : owner_(owner), card_(card), bytes_(bytes) {}
~Held() {
{
std::lock_guard lock(owner_.m_);
owner_.resident_[card_] -= bytes_;
}
owner_.cv_.notify_all();
}
Held(const Held &) = delete;
Held &operator=(const Held &) = delete;
private:
ScoreEngines &owner_;
int card_;
size_t bytes_;
};
std::unique_ptr<Held> Hold(size_t bytes) {
const int card = RigidBodyGPUEngine::CurrentDevice();
std::unique_lock lock(m_);
cv_.wait(lock, [&] { return MakeRoom(card, bytes); });
resident_[card] += bytes;
return std::make_unique<Held>(*this, card, bytes);
}
private:
struct Entry {
std::shared_ptr<Engine> engine;
@@ -386,7 +457,7 @@ private:
};
// The entry's engine for a score: the copy it holds is what marks the engine in use, and handing it back
// is what may let an engine waiting for room in.
// is what may let an engine or a function waiting for room in.
std::shared_ptr<Engine> Lease(Entry &e) {
if (!e.engine)
return nullptr;
@@ -399,10 +470,10 @@ private:
});
}
// Whether `bytes` more fit the card's quarter beside the engines kept there, once the least recently
// used of those no score is using have been taken out. Under m_.
bool MakeRoom(int card, size_t bytes, size_t quarter) {
while (resident_[card] + bytes > quarter) {
// Whether `bytes` more fit the room on the card beside what is held there, once the least recently used
// engines no score is using have been taken out. Under m_.
bool MakeRoom(int card, size_t bytes) {
while (resident_[card] + bytes > room_) {
Entry *oldest = nullptr;
for (auto &[key, entry] : engines_)
if (entry.engine && entry.card == card && entry.engine.use_count() == 1
@@ -418,6 +489,8 @@ private:
std::map<std::tuple<std::string, size_t, double>, Entry> engines_;
std::map<int, size_t> resident_;
size_t room_ = 0;
bool scores_ = false;
uint64_t clock_ = 0;
std::mutex m_;
std::condition_variable cv_;
@@ -582,37 +655,6 @@ PlacementScore score_placement(const gemmi::Structure &st, const gemmi::SpaceGro
return out;
}
// The rigid-body engines for one model in one group, where there is a GPU.
struct RigidBodyEngines {
#ifdef JFJOCH_USE_CUDA
std::unique_ptr<RigidBodyGPUPool> pool;
#endif
RigidBodyGPUPool *get() {
#ifdef JFJOCH_USE_CUDA
return pool.get();
#else
return nullptr;
#endif
}
};
RigidBodyEngines make_engines(const gemmi::Structure &st, const gemmi::SpaceGroup &sg,
const std::vector<MergedReflection> &merged, double d_min, size_t engines,
Logger &logger) {
RigidBodyEngines out;
#ifdef JFJOCH_USE_CUDA
const double finest = RigidBodyLadder(d_min).back();
size_t observations = 0;
for (const MergedReflection &r : merged)
if (!std::isnan(r.F) && r.d >= finest)
++observations;
out.pool = RigidBodyGPUPool::Create(st.models[0], st.cell, sg, d_min, observations, engines, logger);
#else
(void) st; (void) sg; (void) merged; (void) d_min; (void) engines; (void) logger;
#endif
return out;
}
// The placement through the rigid body, and its R-factors where the rigid body left it. A search
// placement is a few degrees and a grid step from its answer, well inside the rigid body's reach, so
// the step is taken without the R-free check model validation makes before committing one.
@@ -642,21 +684,20 @@ PlacementScore refine_placement(gemmi::Structure &st, const gemmi::SpaceGroup &s
// The copies of `assembly` (its chains split evenly, in order) as a crystal of group `sg`, each refined beside
// the others, and scored.
PlacementScore carry_copies(const gemmi::Structure &assembly, int copies, const gemmi::SpaceGroup &sg,
const std::vector<MergedReflection> &merged, double d_min, size_t nthreads,
ScoreEngines *scorers, Logger &logger) {
const std::vector<MergedReflection> &merged, double d_min, RigidBodyGPUPool *pool,
size_t nthreads, ScoreEngines *scorers) {
std::vector<gemmi::Structure> body = SplitCopies(assembly, copies);
for (gemmi::Structure &b : body) {
b.spacegroup_hm = sg.xhm();
b.setup_cell_images();
}
RigidBodyEngines engines = make_engines(body[0], sg, merged, d_min, 1, logger);
for (int k = 0; k < copies; k++) {
std::vector<gemmi::Structure> rest;
for (int o = 0; o < copies; o++)
if (o != k)
rest.push_back(body[o]);
const gemmi::Structure fixed = AssembleCopies(rest);
refine_rigid_body(body[k], sg, merged, d_min, engines.get(), nthreads, nullptr, &fixed.models[0]);
refine_rigid_body(body[k], sg, merged, d_min, pool, nthreads, nullptr, &fixed.models[0]);
}
return score_placement(AssembleCopies(body), sg, merged, d_min, nthreads, scorers);
}
@@ -1154,12 +1195,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;
#ifdef JFJOCH_USE_CUDA
if (get_gpu_count() > 0 && MapToFPhiGPUFits(dc.grid))
grid_ = MapToFPhiGPU(dc.grid);
else
if (get_gpu_count() > 0 && MapToFPhiGPUFits(dc.grid)) {
full = MapToFPhiGPU(dc.grid);
cuda_trim_allocation_pool(); // its buffers, which the engines allocated outside the pool cannot use
} else
#endif
grid_ = MapToFPhi(dc.grid);
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);
d_min_ = d_min;
blur_ = dc.blur;
}
@@ -2177,9 +2224,64 @@ double ExpectedLogLikelihoodGain(const SearchModel &model, const MRMerge &merge,
return ellg;
}
#ifdef JFJOCH_USE_CUDA
// What one merge's placements need of the GPU, worked out on the CPU: the rigid body's pool - one for every fit
// of the merge, sized for the largest, the most copies the cell holds (Matthews) in any of its groups, planned
// but without engines; null where the GPU cannot grid them - and the largest score engine.
struct MergeGPUSize {
std::unique_ptr<RigidBodyGPUPool> pool;
size_t rigid_body = 0;
size_t rigid_body_engines = 0;
size_t structure_factors = 0;
};
MergeGPUSize size_merge_gpu(const SearchModel &model, const MRMerge &merge, double d_min, int device) {
set_gpu(device);
MergeGPUSize out;
const gemmi::UnitCell cell(merge.cell);
const std::vector<gemmi::SpaceGroup> groups = SameRotationGroups(merge.space_group);
const std::vector<double> pm = MatthewsProbabilities(cell.volume / static_cast<double>(merge.space_group.operations().order()),
model.mass_da);
int copies_max = 1;
for (size_t n = 0; n < pm.size(); n++)
if (pm[n] >= MATTHEWS_P_MIN)
copies_max = static_cast<int>(n) + 1;
// Every fit is given the reflections of its finest zone.
size_t observations = 0;
for (const MergedReflection &r : *merge.merged)
if (!std::isnan(r.F) && r.d >= RigidBodyLadder(d_min).back())
++observations;
// As many fits as run at once: a group's first copies side by side, the groups side by side, or a null's
// replicates.
out.rigid_body_engines = std::max<size_t>(groups.size() * RESCORED_PER_GROUP, NULL_REPLICATES);
Logger quiet = Logger::Buffered();
bool rigid_body_gpu = true; // every group's fits can be gridded
for (const gemmi::SpaceGroup &sg : groups) {
const gemmi::Structure one = PlaceSearchModel(model, gemmi::Mat33(), gemmi::Fractional(0, 0, 0), cell, sg);
const gemmi::Structure most = AssembleCopies(std::vector<gemmi::Structure>(copies_max, one));
for (const gemmi::Structure *st : {&one, &most}) {
if (!out.pool)
out.pool = RigidBodyGPUPool::Plan(st->models[0], cell, sg, d_min, observations, quiet);
rigid_body_gpu = rigid_body_gpu && out.pool && out.pool->Cover(st->models[0], sg, quiet);
}
try { // as ScoreEngines::get() makes it, which leaves the scores on the CPU where it cannot
const ModelStructureFactorsGPU sf(device, most.models[0], cell, sg, d_min);
if (sf.Supported())
out.structure_factors = std::max(out.structure_factors, ScoreEngines::Bytes(sf));
} catch (const std::exception &) {
cuda_clear_error();
}
}
if (rigid_body_gpu)
out.rigid_body = out.pool->EngineBytes();
return out;
}
#endif
// The search and the decision for one search model.
// `gpu_budget`: what the search may hold on its card (PlanMolecularReplacementGPU).
MolecularReplacementResult search_model(const SearchModel &model, const std::vector<MRMerge> &merges, Logger &logger,
size_t nthreads) {
size_t nthreads, size_t gpu_budget) {
MolecularReplacementResult result;
result.residues_read = model.residues_read;
result.residues_kept = model.residues_kept;
@@ -2215,6 +2317,13 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
return result;
}
#ifdef JFJOCH_USE_CUDA
// Each merge's GPU sizes; the answer merge's found on the CPU while the transform is made.
std::vector<std::future<MergeGPUSize>> sized(merges.size());
if (get_gpu_count() > 0)
sized[0] = std::async(std::launch::async, size_merge_gpu, std::cref(model), std::cref(merges[0]),
result.d_min_score, RigidBodyGPUEngine::CurrentDevice());
#endif
t0 = std::chrono::steady_clock::now();
std::optional<MolecularTransform> ft_made;
try {
@@ -2224,11 +2333,12 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
return result;
}
const MolecularTransform &ft = *ft_made;
[[maybe_unused]] const size_t transform_bytes = ft.Grid().data.size() * sizeof(std::complex<float>);
#ifdef JFJOCH_USE_CUDA
// The rotation and translation functions on the GPU where there is one; the rest of the search is
// the same either way.
// The rotation and translation functions on the GPU where the transform fits the budget (the plan's first
// part, PlanMolecularReplacementGPU); the rest of the search is the same either way.
std::unique_ptr<MolecularReplacementGPU> gpu;
if (get_gpu_count() > 0) {
if (get_gpu_count() > 0 && transform_bytes <= gpu_budget) {
try {
gpu = std::make_unique<MolecularReplacementGPU>(ft);
} catch (const JFJochException &e) {
@@ -2245,6 +2355,85 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
result.tncs.present ? " - tNCS" : "");
const std::vector<gemmi::Mat33> sphere = UniformRotations(ROTATION_STEP_DEG);
ScoreEngines scorers;
#ifdef JFJOCH_USE_CUDA
// Each merge's rigid-body engines, reserved by its plan and kept until its choice is made (choose_in).
std::vector<std::unique_ptr<RigidBodyGPUPool>> rigid_body(merges.size());
#endif
const auto rigid_body_pool = [&](size_t m) -> RigidBodyGPUPool * {
#ifdef JFJOCH_USE_CUDA
return rigid_body[m].get();
#else
(void) m;
return nullptr;
#endif
};
// What merge m's search and placements put on the GPU, decided at once before its search starts
// (PlanMolecularReplacementGPU), on the sizes size_merge_gpu() found - for the answer's merge while the
// transform was made.
const auto plan_gpu = [&](size_t m) {
#ifdef JFJOCH_USE_CUDA
if (get_gpu_count() == 0)
return;
for (auto &pool : rigid_body)
pool.reset(); // the other merge's choice is made
if (!sized[m].valid())
sized[m] = std::async(std::launch::deferred, size_merge_gpu, std::cref(model), std::cref(merges[m]),
result.d_min_score, RigidBodyGPUEngine::CurrentDevice());
MergeGPUSize size = sized[m].get();
MolecularReplacementGPUNeed need;
need.budget = gpu_budget;
need.transform = transform_bytes;
need.rigid_body = size.rigid_body;
need.rigid_body_engines = size.rigid_body_engines;
need.structure_factors = size.structure_factors;
const MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need);
scorers.SetRoom(plan.room, plan.structure_factors);
if (plan.rigid_body_engines > 0) {
size.pool->Reserve(plan.rigid_body_engines);
rigid_body[m] = std::move(size.pool);
}
logger.Info("Molecular replacement: {} on the GPU: transform {}, rigid body {}, scores {} - {} held of a {} "
"budget, {} left for the score engines and the functions' buffers",
merges[m].name, plan.search ? in_gb(need.transform) : "on the CPU",
plan.rigid_body_engines > 0 ? fmt::format("{} x {}", plan.rigid_body_engines, in_gb(need.rigid_body))
: "on the CPU",
plan.structure_factors ? "up to " + in_gb(need.structure_factors) + " each" : "on the CPU",
in_gb(plan.bytes), in_gb(gpu_budget), in_gb(plan.room));
#else
(void) m;
#endif
};
// A rotation or a translation function on the GPU where the search is there and the call's working buffers
// fit the room (held for the call), on the CPU otherwise - the same function, equal to rounding.
const auto rotation_function = [&](const SearchReflections &refl, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, const std::vector<gemmi::Mat33> &rotations) {
#ifdef JFJOCH_USE_CUDA
const size_t bytes = MolecularReplacementGPU::RotationBytes(refl, sg, rotations.size());
if (gpu && scorers.Fits(bytes)) {
const auto held = scorers.Hold(bytes);
return gpu->RotationFunction(refl, cell, sg, rotations);
}
#endif
return RotationFunction(ft, refl, cell, sg, rotations, nthreads);
};
const auto translation_function = [&](const TranslationSetup &setup, const std::vector<gemmi::Mat33> &rotations,
const std::vector<std::complex<double>> *known,
TranslationSetupGPU *prepared) {
#ifdef JFJOCH_USE_CUDA
const size_t bytes = MolecularReplacementGPU::TranslationBytes(setup, rotations.size(), known != nullptr);
if (gpu && scorers.Fits(bytes)) {
const auto held = scorers.Hold(bytes);
return gpu->TranslationFunction(setup, rotations, known, prepared);
}
#else
(void) prepared;
#endif
std::vector<TranslationPeak> tf(rotations.size());
ParallelFor(static_cast<int>(rotations.size()), nthreads, [&](int i) {
tf[i] = TranslationFunction(ft, setup, rotations[i], known);
});
return tf;
};
std::vector<gemmi::Mat33> orientations; // of each candidate's placement, before the rigid body
// One merge: the rotation function over its point group, then for each of its groups the
@@ -2271,12 +2460,8 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
d_llg = std::max(d_llg, std::min(result.d_min_search, ds[LLG_TERMS - 1]));
const SearchReflections lrefl = MakeSearchReflections(*merge.merged, cell, merge.space_group, d_llg,
result.d_max_search);
#ifdef JFJOCH_USE_CUDA
const std::vector<double> rf = gpu ? gpu->RotationFunction(refl, cell, merge.space_group, rotations)
: RotationFunction(ft, refl, cell, merge.space_group, rotations, nthreads);
#else
const std::vector<double> rf = RotationFunction(ft, refl, cell, merge.space_group, rotations, nthreads);
#endif
plan_gpu(m);
const std::vector<double> rf = rotation_function(refl, cell, merge.space_group, rotations);
double rf_mean = 0, rf_sd = 0;
for (double x : rf)
rf_mean += x;
@@ -2318,12 +2503,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
for (int axis = 0; axis < 3; axis++)
for (double sign : {1.0, -1.0})
trial.push_back(axis_rotation(axis, sign * step).multiply(peak_rotations[k]));
#ifdef JFJOCH_USE_CUDA
const std::vector<double> v = gpu ? gpu->RotationFunction(refl, cell, merge.space_group, trial)
: RotationFunction(ft, refl, cell, merge.space_group, trial, nthreads);
#else
const std::vector<double> v = RotationFunction(ft, refl, cell, merge.space_group, trial, nthreads);
#endif
const std::vector<double> v = rotation_function(refl, cell, merge.space_group, trial);
size_t next = 0;
for (size_t k = 0; k < peaks.size(); k++) {
if (!active[k])
@@ -2354,17 +2534,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
std::vector<std::future<std::vector<TranslationPeak>>> tf_running;
for (const gemmi::SpaceGroup &sg : groups)
tf_running.push_back(std::async(std::launch::async, [&, sg] {
const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg);
std::vector<TranslationPeak> tf(peaks.size());
#ifdef JFJOCH_USE_CUDA
if (gpu)
tf = gpu->TranslationFunction(setup, peak_rotations);
else
#endif
ParallelFor(static_cast<int>(peaks.size()), nthreads, [&](int i) {
tf[i] = TranslationFunction(ft, setup, peak_rotations[i]);
});
return tf;
return translation_function(MakeTranslationSetup(refl, cell, sg), peak_rotations, nullptr, nullptr);
}));
std::vector<TranslationPeak> tf_all;
for (auto &r : tf_running) {
@@ -2406,19 +2576,9 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
std::iota(by_tf.begin(), by_tf.end(), 0);
std::stable_sort(by_tf.begin(), by_tf.end(), [&](size_t a, size_t b) { return tfs[a].z > tfs[b].z; });
if (!seeds.empty()) {
const TranslationSetup setup = MakeTranslationSetup(refl, cell, sg);
std::vector<TranslationPeak> tf_seed(seeds.size());
#ifdef JFJOCH_USE_CUDA
const std::vector<double> rf_seed = gpu ? gpu->RotationFunction(refl, cell, merge.space_group, seeds)
: RotationFunction(ft, refl, cell, merge.space_group, seeds, nthreads);
if (gpu)
tf_seed = gpu->TranslationFunction(setup, seeds);
else
#else
const std::vector<double> rf_seed = RotationFunction(ft, refl, cell, merge.space_group, seeds, nthreads);
#endif
for (size_t k = 0; k < seeds.size(); k++)
tf_seed[k] = TranslationFunction(ft, setup, seeds[k]);
const std::vector<double> rf_seed = rotation_function(refl, cell, merge.space_group, seeds);
const std::vector<TranslationPeak> tf_seed = translation_function(MakeTranslationSetup(refl, cell, sg),
seeds, nullptr, nullptr);
std::vector<size_t> first;
for (size_t k = 0; k < seeds.size(); k++) {
first.push_back(rots.size());
@@ -2447,7 +2607,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
};
// --- the first copy: the best few through packing, the rigid body and R-free ---
RigidBodyEngines engines;
RigidBodyGPUPool *const engines = rigid_body_pool(m);
bool have = false;
gemmi::Mat33 orientation;
// Where not even one copy fits the cell, only the best translation-function placement is
@@ -2497,11 +2657,10 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
if (f.refine)
to_refine.push_back(&f);
if (!to_refine.empty()) {
engines = make_engines(*to_refine[0]->placed, sg, *merge.merged, d_min, to_refine.size(), logger);
std::vector<std::future<void>> running;
for (First *f : to_refine)
running.push_back(std::async(std::launch::async, [&, f] {
f->score = refine_placement(*f->placed, sg, *merge.merged, d_min, engines.get(), nthreads,
f->score = refine_placement(*f->placed, sg, *merge.merged, d_min, engines, nthreads,
&scorers, &f->moved);
f->clash = ClashFraction(*f->placed);
}));
@@ -2555,9 +2714,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
// --- further copies, one at a time, with the copies placed as a known partial structure ---
// Each copy is refined on its own and all of them scored together.
std::optional<TranslationSetup> setup2;
#ifdef JFJOCH_USE_CUDA
std::unique_ptr<TranslationSetupGPU, void (*)(TranslationSetupGPU *)> setup2_gpu{nullptr, nullptr};
#endif
// The first `count` copies refined in turn, each with all the others fixed, and the assembly scored.
const auto refine_beside_others = [&](int count) {
for (int c = 0; c < count; c++) {
@@ -2567,7 +2724,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
rest.push_back(copies[o]);
const gemmi::Structure fixed = AssembleCopies(rest);
gemmi::Mat33 moved;
refine_rigid_body(copies[c], sg, *merge.merged, d_min, engines.get(), nthreads, &moved,
refine_rigid_body(copies[c], sg, *merge.merged, d_min, engines, nthreads, &moved,
&fixed.models[0]);
copy_rotation[c] = moved.multiply(copy_rotation[c]);
copy_position[c] = centroid(copies[c]);
@@ -2610,17 +2767,11 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
}
rknown_llg = RefinedLogLikelihoodGain(rsetup(), rknown, model.error_A, &known_scale);
});
std::vector<TranslationPeak> tf2(rots2.size());
#ifdef JFJOCH_USE_CUDA
if (gpu) {
if (!setup2_gpu)
setup2_gpu = gpu->Prepare(*setup2, true);
tf2 = gpu->TranslationFunction(*setup2, rots2, &known, setup2_gpu.get());
} else
if (gpu && !setup2_gpu)
setup2_gpu = gpu->Prepare(*setup2, true);
#endif
ParallelFor(static_cast<int>(rots2.size()), nthreads, [&](int i) {
tf2[i] = TranslationFunction(ft, *setup2, rots2[i], &known);
});
const std::vector<TranslationPeak> tf2 = translation_function(*setup2, rots2, &known, setup2_gpu.get());
std::vector<size_t> by_tf2(rots2.size());
std::iota(by_tf2.begin(), by_tf2.end(), 0);
std::stable_sort(by_tf2.begin(), by_tf2.end(), [&](size_t a, size_t b) { return tf2[a].z > tf2[b].z; });
@@ -2691,7 +2842,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
for (Trial &trial : trials) {
std::vector<gemmi::Structure> all = copies;
gemmi::Mat33 moved;
refine_rigid_body(trial.copy, sg, *merge.merged, d_min, engines.get(), nthreads, &moved,
refine_rigid_body(trial.copy, sg, *merge.merged, d_min, engines, nthreads, &moved,
&others.models[0]);
all.push_back(trial.copy);
trial.assembly = AssembleCopies(all);
@@ -2837,8 +2988,6 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
// The null of candidate i.
const auto run_null = [&](int i) {
const MRCandidate &c = result.candidates[i];
RigidBodyEngines engines = make_engines(*c.placed, c.space_group, *merges[merge].merged,
result.d_min_score, NULL_REPLICATES, logger);
// One copy: its replicates placed by the translation function, as it was.
std::function<std::vector<gemmi::Fractional>(const std::vector<gemmi::Mat33> &)> translate;
if (c.copies == 1)
@@ -2847,24 +2996,16 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
gemmi::UnitCell(merges[merge].cell),
merges[merge].space_group,
result.d_min_search, result.d_max_search);
const TranslationSetup setup = MakeTranslationSetup(refl, gemmi::UnitCell(merges[merge].cell),
c.space_group);
std::vector<TranslationPeak> tf(turned.size());
#ifdef JFJOCH_USE_CUDA
if (gpu)
tf = gpu->TranslationFunction(setup, turned);
else
#endif
ParallelFor(static_cast<int>(turned.size()), nthreads, [&](int i) {
tf[i] = TranslationFunction(ft, setup, turned[i]);
});
const std::vector<TranslationPeak> tf = translation_function(
MakeTranslationSetup(refl, gemmi::UnitCell(merges[merge].cell), c.space_group), turned,
nullptr, nullptr);
std::vector<gemmi::Fractional> out;
for (const TranslationPeak &p : tf)
out.push_back(p.t);
return out;
};
return null_sigma(*c.placed, c.copies, orientations[i], c.space_group, *merges[merge].merged,
result.d_min_score, engines.get(), nthreads, &scorers, translate);
result.d_min_score, rigid_body_pool(merge), nthreads, &scorers, translate);
};
// The first candidate's null alone; where it fails, the nulls of all the other placed candidates side
// by side, each taken where the choice comes to it.
@@ -2889,7 +3030,6 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
std::vector<Option> options{{gemmi::Fractional(0, 0, 0), c.placed,
score_placement(*c.placed, c.space_group, merged, d_min, nthreads, &scorers,
nullptr, nullptr, &v), c.clash_fraction}};
RigidBodyEngines engines;
for (int o = 0; o < 8; o++) {
const gemmi::Fractional shift((v.x + 0.5 * (o & 1)) / 2, (v.y + 0.5 * ((o >> 1) & 1)) / 2,
(v.z + 0.5 * ((o >> 2) & 1)) / 2);
@@ -2907,9 +3047,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
SetModelPositions(moved->models[0], pos);
Option opt{shift, moved, {}, ClashFraction(*moved, c.copies)};
if (opt.clash <= 2 * CLASH_MAX) {
if (engines.get() == nullptr)
engines = make_engines(*moved, c.space_group, merged, d_min, 1, logger);
refine_rigid_body(*moved, c.space_group, merged, d_min, engines.get(), nthreads);
refine_rigid_body(*moved, c.space_group, merged, d_min, rigid_body_pool(merge), nthreads);
opt.clash = ClashFraction(*moved, c.copies);
opt.score = score_placement(*moved, c.space_group, merged, d_min, nthreads, &scorers, nullptr,
nullptr, &v);
@@ -3000,7 +3138,8 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
// they build the same one.
const MRCandidate &b = result.candidates[best];
const PlacementScore carried = carry_copies(*b.placed, b.copies, c.space_group, *merges[merge].merged,
result.d_min_score, nthreads, &scorers, logger);
result.d_min_score, rigid_body_pool(merge), nthreads,
&scorers);
logger.Info("Molecular replacement: {}'s {} copies carried into {}: R-free {:.3f} against {:.3f}",
b.space_group.xhm(), b.copies, c.space_group.xhm(), carried.r_free, b.r_free);
if (!(carried.r_free - b.r_free >= R_FREE_MARGIN) || !(b.r_free < CARRY_R_FREE_MAX)) {
@@ -3125,9 +3264,40 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
} // namespace
MolecularReplacementGPUPlan PlanMolecularReplacementGPU(const MolecularReplacementGPUNeed &need) {
MolecularReplacementGPUPlan plan;
auto take = [&](size_t bytes, size_t limit) {
if (bytes == 0 || plan.bytes + bytes > limit)
return false;
plan.bytes += bytes;
return true;
};
plan.search = take(need.transform, need.budget);
plan.rigid_body_engines = need.rigid_body_engines > 0 && take(need.rigid_body, need.budget) ? 1 : 0;
plan.structure_factors = take(need.structure_factors, need.budget);
// More engines only within half of the budget (see the declaration).
while (plan.rigid_body_engines > 0 && plan.rigid_body_engines < need.rigid_body_engines
&& take(need.rigid_body, need.budget / 2))
plan.rigid_body_engines++;
plan.room = need.budget - (plan.search ? need.transform : 0) - plan.rigid_body_engines * need.rigid_body;
return plan;
}
MolecularReplacementResult RunMolecularReplacement(const std::string &model_path, const std::vector<MRMerge> &merges,
Logger &logger, size_t nthreads) {
Logger &logger, size_t nthreads, size_t gpu_resident_bytes) {
const auto t0 = std::chrono::steady_clock::now();
// What the search may hold on the card: the budget model validation keeps to. What the stream-ordered pool
// holds with no buffer in it - freed by the steps before - goes back to the driver first, as before a
// validation: the engines are allocated outside the pool and cannot use it.
size_t gpu_budget = 0;
#ifdef JFJOCH_USE_CUDA
if (get_gpu_count() > 0) {
cuda_trim_allocation_pool();
gpu_budget = ModelGPUBudget(get_gpu_total_memory(RigidBodyGPUEngine::CurrentDevice()), gpu_resident_bytes);
}
#else
(void) gpu_resident_bytes;
#endif
SearchModel model = PrepareSearchModel(model_path);
if (!model.failure_reason.empty()) {
MolecularReplacementResult result;
@@ -3145,7 +3315,7 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
if (!pae_path.empty())
logger.Info("Molecular replacement: {} rigid domain(s) in the predicted aligned error {}", domains.size(), pae_path);
MolecularReplacementResult result = search_model(model, merges, logger, nthreads);
MolecularReplacementResult result = search_model(model, merges, logger, nthreads, gpu_budget);
result.model_path = model_path;
result.pae_domains = static_cast<int>(domains.size());
result.t_model += t_model;
@@ -3155,14 +3325,16 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
// own.
if (result.ok && result.chosen < 0 && result.decision.starts_with("NOT_PLACED") && domains.size() >= 2) {
const auto t_domains = std::chrono::steady_clock::now();
// Each domain on a thread of its own, its log kept and replayed in order.
// Each domain on a thread of its own, its log kept and replayed in order, and with an equal share of the
// budget.
std::vector<Logger> logs;
for (size_t i = 0; i < domains.size(); i++)
logs.push_back(Logger::Buffered());
std::vector<std::future<MolecularReplacementResult>> running;
for (size_t i = 0; i < domains.size(); i++)
running.push_back(std::async(std::launch::async, [&, i] {
return search_model(DomainModel(model, domains[i]), merges, logs[i], nthreads);
return search_model(DomainModel(model, domains[i]), merges, logs[i], nthreads,
gpu_budget / domains.size());
}));
std::vector<MolecularReplacementResult> found;
for (size_t i = 0; i < domains.size(); i++) {
@@ -3214,13 +3386,15 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
}
void CheckEnantiomorphHand(MolecularReplacementResult &result, const std::vector<MergedReflection> &merged,
const gemmi::SpaceGroup &data_space_group, Logger &logger, size_t nthreads) {
const gemmi::SpaceGroup &data_space_group, Logger &logger, size_t nthreads,
size_t gpu_resident_bytes) {
const MRCandidate &c = result.candidates[result.chosen];
const auto t0 = std::chrono::steady_clock::now();
Logger quiet = Logger::Buffered();
ModelValidationResult v;
try {
v = ReadModelAnomalousMap(merged, *c.placed, data_space_group, c.k_sol, c.b_sol, quiet, nthreads);
v = ReadModelAnomalousMap(merged, *c.placed, data_space_group, c.k_sol, c.b_sol, quiet, nthreads,
gpu_resident_bytes);
} catch (const std::runtime_error &e) { // a GPU failure: no reading, so no relabelling
logger.Warning("Molecular replacement: the anomalous map could not be read ({})", e.what());
}
image_analysis/structure_refinement/MolecularReplacement.h
+34 -2
View File
@@ -370,6 +370,36 @@ struct MolecularReplacementResult {
double t_domains = 0; // the PAE domains' own searches, wall clock (they run side by side)
};
// What a molecular-replacement search puts on its GPU, for one merge: the molecular transform the rotation and
// translation functions read, the rigid body's engines and the score engines, and what the functions' working
// buffers may take beside them. Bytes of one of each; 0 where that part cannot run on the GPU.
struct MolecularReplacementGPUNeed {
size_t budget = 0; // ModelGPUBudget (ModelValidation.h), or a PAE domain's share of it
size_t transform = 0;
size_t rigid_body = 0;
size_t rigid_body_engines = 0; // how many rigid-body fits can run at once
size_t structure_factors = 0; // the largest score engine of the merge's placements
};
struct MolecularReplacementGPUPlan {
bool search = false; // the transform held on the GPU, the functions run there
size_t rigid_body_engines = 0; // 0: the rigid body on the CPU
bool structure_factors = false; // the scores on the GPU
size_t room = 0; // for the score engines and the functions' working buffers
size_t bytes = 0; // the transform, the rigid-body engines and one score engine
};
// The one rule for what a search puts on its GPU, decided before anything is reserved, from the budget model
// validation keeps to (80% of the card's TOTAL memory less what is resident). The transform and the rigid
// body's engines are held for the whole merge, so it is their sum, with the largest score engine, that is held
// against the budget. Each part takes one where it still fits, in the order of what it saves: the transform
// (the rotation and translation functions), the rigid body, the scores; a part that does not fit runs on the
// CPU. More rigid-body engines only within half of the budget, as model validation adds them. What is left of
// the budget beside the transform and the engines is the room the score engines - kept for as long as they fit,
// the least recently used given up first - and the functions' working buffers share, each waiting for the other
// to be done rather than spilling past it.
MolecularReplacementGPUPlan PlanMolecularReplacementGPU(const MolecularReplacementGPUNeed &need);
// Search `merges` (the answer's first) for the predicted model at `model_path`, in every space group with
// the rotations of each merge's group - one copy, then further copies one at a time while each new one
// stands out, packs and lowers R-free - and decide. A candidate is PLACED where a copy's translation
@@ -377,8 +407,9 @@ struct MolecularReplacementResult {
// null. The model chooses among the answer merge's groups where one placed candidate leads
// every other by an R-free margin; it chooses the twin hypothesis' merge only where none of the
// answer's groups can hold the model at all (Matthews, packing) and a group of the subgroup can.
// `gpu_resident_bytes`: what is held on the card meanwhile by what came before (a merge engine still at work).
MolecularReplacementResult RunMolecularReplacement(const std::string &model_path, const std::vector<MRMerge> &merges,
Logger &logger, size_t nthreads);
Logger &logger, size_t nthreads, size_t gpu_resident_bytes = 0);
// The hand check of an ENANTIOMORPH decision, filled into result.hand_check: taken into the answer by
// --model's rule for a model's enantiomorph - the model has to fit against its own null, which the choice
@@ -386,4 +417,5 @@ MolecularReplacementResult RunMolecularReplacement(const std::string &model_path
// `data_space_group`) must not come out inverted, read by model validation's own code on the placement
// and the bulk solvent MR scored. Where it is not taken, the decision becomes ENANTIOMORPH_VETOED.
void CheckEnantiomorphHand(MolecularReplacementResult &result, const std::vector<MergedReflection> &merged,
const gemmi::SpaceGroup &data_space_group, Logger &logger, size_t nthreads);
const gemmi::SpaceGroup &data_space_group, Logger &logger, size_t nthreads,
size_t gpu_resident_bytes = 0);
image_analysis/structure_refinement/MolecularReplacementGPU.cu
+126 -64
View File
@@ -409,11 +409,65 @@ std::vector<double> MolecularReplacementGPU::RotationFunction(const SearchReflec
return out;
}
// What the translation function needs of one setup, on the device: the gather's lists - for every point of the
// half grid in cuFFT's layout ((u nv + v) nwh + l), the (term, pair) entries landing on it, in entry order, and
// the same for the cross terms with a known partial structure - the terms' vectors, phases, weights and bins,
// and the transform's plan. Made once where a setup is used for several functions in a row.
// What the translation function needs of one setup, worked out on the host: the gather's lists - for every
// point of the half grid in cuFFT's layout ((u nv + v) nwh + l), the (term, pair) entries landing on it, in
// entry order, and the same for the cross terms with a known partial structure - and the terms' vectors,
// phases and weights as the device takes them. Made once where a setup is used for several functions in a row;
// uploaded by each, so that the device holds it only while a function runs.
struct TranslationSetupGPU {
bool cross = false;
std::vector<int> offset, entry, cross_offset, cross_entry, bin_n;
std::vector<float3> s;
std::vector<float2> phase;
std::vector<float> weight;
};
namespace {
TranslationSetupGPU prepare_translation(const TranslationSetup &setup, bool cross) {
TranslationSetupGPU out;
out.cross = cross;
const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1;
const size_t cells = static_cast<size_t>(nu) * nv * nwh;
std::vector<int> &offset = out.offset;
offset.assign(cells + 1, 0);
for (size_t e = 0; e < setup.target_u.size(); e++)
if (setup.target_u[e] >= 0)
offset[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh + setup.target_l[e] + 1]++;
for (size_t c = 0; c < cells; c++)
offset[c + 1] += offset[c];
out.entry.resize(offset[cells]);
{
std::vector<int> fill(offset.begin(), offset.end() - 1);
for (size_t e = 0; e < setup.target_u.size(); e++)
if (setup.target_u[e] >= 0)
out.entry[fill[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh
+ setup.target_l[e]]++] = static_cast<int>(e);
}
if (cross) {
std::vector<int> &cross_offset = out.cross_offset;
cross_offset.assign(cells + 1, 0);
for (size_t e = 0; e < setup.cross_term.size(); e++)
cross_offset[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh + setup.cross_l[e] + 1]++;
for (size_t c = 0; c < cells; c++)
cross_offset[c + 1] += cross_offset[c];
out.cross_entry.resize(cross_offset[cells]);
std::vector<int> fill(cross_offset.begin(), cross_offset.end() - 1);
for (size_t e = 0; e < setup.cross_term.size(); e++)
out.cross_entry[fill[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh
+ setup.cross_l[e]]++] = static_cast<int>(e);
}
out.phase.resize(setup.phase.size());
for (size_t i = 0; i < out.phase.size(); i++)
out.phase[i] = make_float2(static_cast<float>(setup.phase[i].real()), static_cast<float>(setup.phase[i].imag()));
out.weight.assign(setup.weight.begin(), setup.weight.end());
out.s = float3s(setup.s);
out.bin_n = bin_counts(setup.bin, setup.bins);
return out;
}
// A prepared setup on the device, with the transform's plan, for one translation function.
struct TranslationDevice {
CudaDevicePtr<int> cross_offset, cross_entry, cross_term, cross_op, bin, bin_n, pair_j, pair_k, offset, entry;
CudaDevicePtr<char> cross_conj;
CudaDevicePtr<float3> s;
@@ -422,67 +476,74 @@ struct TranslationSetupGPU {
CudaFFTPlan plan;
};
namespace {
std::unique_ptr<TranslationSetupGPU> prepare_translation(const TranslationSetup &setup, bool cross, cudaStream_t stream) {
auto out = std::make_unique<TranslationSetupGPU>();
const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1;
const size_t cells = static_cast<size_t>(nu) * nv * nwh;
std::vector<int> offset(cells + 1, 0);
for (size_t e = 0; e < setup.target_u.size(); e++)
if (setup.target_u[e] >= 0)
offset[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh + setup.target_l[e] + 1]++;
for (size_t c = 0; c < cells; c++)
offset[c + 1] += offset[c];
std::vector<int> entry(offset[cells]);
{
std::vector<int> fill(offset.begin(), offset.end() - 1);
for (size_t e = 0; e < setup.target_u.size(); e++)
if (setup.target_u[e] >= 0)
entry[fill[(static_cast<size_t>(setup.target_u[e]) * nv + setup.target_v[e]) * nwh
+ setup.target_l[e]]++] = static_cast<int>(e);
}
std::vector<int> cross_offset(cells + 1, 0), cross_entry;
if (cross) {
for (size_t e = 0; e < setup.cross_term.size(); e++)
cross_offset[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh + setup.cross_l[e] + 1]++;
for (size_t c = 0; c < cells; c++)
cross_offset[c + 1] += cross_offset[c];
cross_entry.resize(cross_offset[cells]);
std::vector<int> fill(cross_offset.begin(), cross_offset.end() - 1);
for (size_t e = 0; e < setup.cross_term.size(); e++)
cross_entry[fill[(static_cast<size_t>(setup.cross_u[e]) * nv + setup.cross_v[e]) * nwh
+ setup.cross_l[e]]++] = static_cast<int>(e);
}
out->cross_offset = upload(cross_offset, stream);
out->cross_entry = upload(cross_entry, stream);
out->cross_term = upload(setup.cross_term, stream);
out->cross_op = upload(setup.cross_op, stream);
out->cross_conj = upload(setup.cross_conj, stream);
std::vector<float2> phase(setup.phase.size());
for (size_t i = 0; i < phase.size(); i++)
phase[i] = make_float2(static_cast<float>(setup.phase[i].real()), static_cast<float>(setup.phase[i].imag()));
const std::vector<float> weight(setup.weight.begin(), setup.weight.end());
out->s = upload(float3s(setup.s), stream);
out->phase = upload(phase, stream);
out->weight = upload(weight, stream);
out->bin = upload(setup.bin, stream);
out->bin_n = upload(bin_counts(setup.bin, setup.bins), stream);
out->pair_j = upload(setup.pair_j, stream);
out->pair_k = upload(setup.pair_k, stream);
out->offset = upload(offset, stream);
out->entry = upload(entry, stream);
out->plan = CudaFFTPlan(CudaFFTShape{.n = {nu, nv, nw}, .type = CUFFT_C2R}, stream);
cuda_err(cudaStreamSynchronize(stream)); // the host vectors go out of scope
return out;
void upload_translation(const TranslationSetup &setup, const TranslationSetupGPU &p, cudaStream_t stream,
TranslationDevice &out) {
out.cross_offset = upload(p.cross_offset, stream);
out.cross_entry = upload(p.cross_entry, stream);
out.cross_term = upload(setup.cross_term, stream);
out.cross_op = upload(setup.cross_op, stream);
out.cross_conj = upload(setup.cross_conj, stream);
out.s = upload(p.s, stream);
out.phase = upload(p.phase, stream);
out.weight = upload(p.weight, stream);
out.bin = upload(setup.bin, stream);
out.bin_n = upload(p.bin_n, stream);
out.pair_j = upload(setup.pair_j, stream);
out.pair_k = upload(setup.pair_k, stream);
out.offset = upload(p.offset, stream);
out.entry = upload(p.entry, stream);
out.plan = CudaFFTPlan(CudaFFTShape{.n = {setup.nu, setup.nv, setup.nw}, .type = CUFFT_C2R}, stream);
}
} // namespace
std::unique_ptr<TranslationSetupGPU, void (*)(TranslationSetupGPU *)> MolecularReplacementGPU::Prepare(
const TranslationSetup &setup, bool known_part) {
return {prepare_translation(setup, known_part, state_->stream).release(),
[](TranslationSetupGPU *p) { delete p; }};
return {new TranslationSetupGPU(prepare_translation(setup, known_part)), [](TranslationSetupGPU *p) { delete p; }};
}
namespace {
// What upload() of `count` values of `bytes` each takes.
size_t uploaded(size_t count, size_t bytes) {
return std::max<size_t>(count, 1) * bytes;
}
int stat_block_count(size_t points) {
return static_cast<int>(std::min<size_t>(1024, (points + THREADS - 1) / THREADS));
}
} // namespace
size_t MolecularReplacementGPU::RotationBytes(const SearchReflections &refl, const gemmi::SpaceGroup &sg,
size_t rotations) {
const size_t n = refl.hkl.size(), nops = sg.operations().sym_ops.size();
const size_t batch = std::min(ROTATION_BATCH, std::max<size_t>(rotations, 1));
return uploaded(n * nops, sizeof(float3)) + 2 * uploaded(n, sizeof(float)) + uploaded(refl.bins, sizeof(int))
+ uploaded(rotations, sizeof(Rot)) + batch * std::max<size_t>(n, 1) * sizeof(float)
+ uploaded(rotations, sizeof(double));
}
size_t MolecularReplacementGPU::TranslationBytes(const TranslationSetup &setup, size_t rotations, bool known) {
const size_t cells = static_cast<size_t>(setup.nu) * setup.nv * (setup.nw / 2 + 1);
const size_t points = static_cast<size_t>(setup.nu) * setup.nv * setup.nw;
const size_t terms = setup.Terms(), nops = setup.ops;
const size_t entries = std::count_if(setup.target_u.begin(), setup.target_u.end(), [](int u) { return u >= 0; });
const size_t nstat = static_cast<size_t>(stat_block_count(points)) * rotations;
// The setup, as upload_translation() puts it on the device
size_t bytes = uploaded(known ? cells + 1 : 0, sizeof(int)) + uploaded(known ? setup.cross_term.size() : 0, sizeof(int))
+ 2 * uploaded(setup.cross_term.size(), sizeof(int)) + uploaded(setup.cross_conj.size(), sizeof(char))
+ uploaded(terms * nops, sizeof(float3)) + uploaded(terms * nops, sizeof(float2))
+ 2 * uploaded(terms, sizeof(int)) + uploaded(setup.bins, sizeof(int))
+ 2 * uploaded(setup.pair_j.size(), sizeof(int)) + uploaded(cells + 1, sizeof(int))
+ uploaded(entries, sizeof(int))
+ std::max<size_t>(CudaFFTPlan::WorkBytes(CudaFFTShape{.n = {setup.nu, setup.nv, setup.nw},
.type = CUFFT_C2R}), 1);
// and the function's own buffers
bytes += uploaded(known ? terms : 0, sizeof(float2)) + terms * nops * sizeof(float2) + 2 * terms * sizeof(float)
+ cells * sizeof(float2) + points * sizeof(float)
+ nstat * (2 * sizeof(double) + sizeof(float) + sizeof(long long));
return bytes;
}
std::vector<TranslationPeak> MolecularReplacementGPU::TranslationFunction(
@@ -496,12 +557,13 @@ std::vector<TranslationPeak> MolecularReplacementGPU::TranslationFunction(
cudaStream_t stream = state_->stream;
const int nu = setup.nu, nv = setup.nv, nw = setup.nw, nwh = nw / 2 + 1;
const size_t cells = static_cast<size_t>(nu) * nv * nwh, points = static_cast<size_t>(nu) * nv * nw;
std::unique_ptr<TranslationSetupGPU> made;
TranslationSetupGPU made;
if (prepared == nullptr) {
made = prepare_translation(setup, known != nullptr, stream);
prepared = made.get();
made = prepare_translation(setup, known != nullptr);
prepared = &made;
}
TranslationSetupGPU &p = *prepared;
TranslationDevice p;
upload_translation(setup, *prepared, stream, p);
std::vector<float2> known_f;
if (known != nullptr)
for (const std::complex<double> &k : *known)
@@ -511,7 +573,7 @@ std::vector<TranslationPeak> MolecularReplacementGPU::TranslationFunction(
CudaDevicePtr<float> term_sum(terms), scaled(terms);
CudaDevicePtr<float2> coef(cells);
CudaDevicePtr<float> map(points);
const int stat_blocks = static_cast<int>(std::min<size_t>(1024, (points + THREADS - 1) / THREADS));
const int stat_blocks = stat_block_count(points);
const size_t nstat = static_cast<size_t>(stat_blocks) * rotations.size();
CudaDevicePtr<double> d_sum(nstat), d_sum2(nstat);
CudaDevicePtr<float> d_best(nstat);
image_analysis/structure_refinement/MolecularReplacementGPU.h
+7 -1
View File
@@ -44,11 +44,17 @@ public:
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
// 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_;
};
image_analysis/structure_refinement/RigidBodyGPU.cpp
+32 -38
View File
@@ -10,6 +10,7 @@
#include <ceres/rotation.h>
#include "gemmi/dencalc.hpp" // DensityCalculator
#include "gemmi/calculate.hpp" // count_atom_sites
#include "gemmi/it92.hpp" // IT92 x-ray form factors
#include "gemmi/scaling.hpp" // Scaling
@@ -132,17 +133,33 @@ std::unique_ptr<RigidBodyGPUPool> RigidBodyGPUPool::Plan(const gemmi::Model &mod
size_t max_observations, Logger &logger) {
if (get_gpu_count() == 0)
return nullptr;
std::unique_ptr<RigidBodyGPUPool> pool(new RigidBodyGPUPool);
pool->cell_ = cell;
pool->d_min_ = d_min;
pool->max_observations_ = max_observations;
pool->cap_ = std::make_unique<RigidBodyGPUCapacity>();
pool->cap_->observations = max_observations;
pool->cap_->rows = max_observations;
if (!pool->Cover(model, sg, logger))
return nullptr;
try {
std::unique_ptr<RigidBodyGPUPool> pool(new RigidBodyGPUPool);
// Sized for the largest zone of the ladder, which every fit in the validation walks a part of.
RigidBodyGPUCapacity cap;
const size_t ops = sg.operations().sym_ops.size();
for (double zone_d : RigidBodyLadder(d_min)) {
const RigidBodyGPUZone &zone = pool->Zone(model, cell, sg, zone_d);
pool->device_ = RigidBodyGPUEngine::CurrentDevice();
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
return pool;
}
bool RigidBodyGPUPool::Cover(const gemmi::Model &model, const gemmi::SpaceGroup &sg, Logger &logger) {
try {
// Sized for the largest zone of the ladder, which every fit walks a part of.
RigidBodyGPUCapacity &cap = *cap_;
for (double zone_d : RigidBodyLadder(d_min_)) {
const RigidBodyGPUZone &zone = Zone(model, cell_, sg, zone_d);
if (!RigidBodyGPUEngine::Supports(zone)) {
logger.Info("Model validation: the rigid body runs on the CPU - the cell is too small for the "
"GPU's gridding at {:.1f} A", zone_d);
return nullptr;
return false;
}
cap.atoms = std::max(cap.atoms, zone.atoms.size());
cap.grid_points = std::max(cap.grid_points, static_cast<size_t>(zone.nu) * zone.nv * zone.nw);
@@ -151,12 +168,8 @@ std::unique_ptr<RigidBodyGPUPool> RigidBodyGPUPool::Plan(const gemmi::Model &mod
cap.pairs = std::max(cap.pairs, RigidBodyGPUEngine::PairBound(zone));
cap.fft_work_bytes = std::max(cap.fft_work_bytes, RigidBodyGPUEngine::FFTWorkBytes(zone.nu, zone.nv, zone.nw));
}
cap.observations = max_observations;
cap.rows = max_observations;
cap.terms = max_observations * ops;
pool->cap_ = std::make_unique<RigidBodyGPUCapacity>(cap);
pool->device_ = RigidBodyGPUEngine::CurrentDevice();
return pool;
cap.terms = std::max(cap.terms, max_observations_ * sg.operations().sym_ops.size());
return true;
} catch (const JFJochException &e) {
throw RigidBodyGPUFailure(e.what());
}
@@ -181,36 +194,16 @@ void RigidBodyGPUPool::Reserve(size_t engines) {
idle_.push_back(e.get());
}
std::unique_ptr<RigidBodyGPUPool> RigidBodyGPUPool::Create(const gemmi::Model &model, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, double d_min,
size_t max_observations, size_t max_engines,
Logger &logger) {
std::unique_ptr<RigidBodyGPUPool> pool = Plan(model, cell, sg, d_min, max_observations, logger);
if (!pool)
return nullptr;
const size_t bytes = pool->EngineBytes();
const size_t total = get_gpu_total_memory(pool->device_);
const size_t fit = bytes > 0 ? total / 4 / bytes : 0;
if (fit == 0) {
logger.Info("Model validation: the rigid body runs on the CPU - one GPU engine needs {:.0f} MB, over a "
"quarter of the card's {:.0f} MB", bytes / 1e6, total / 1e6);
return nullptr;
}
pool->Reserve(std::min<size_t>(fit, std::max<size_t>(max_engines, 1)));
logger.Info("Model validation: the rigid body runs on the GPU, {} engine(s) of {:.0f} MB", pool->Engines(),
bytes / 1e6);
return pool;
}
RigidBodyGPUPool::~RigidBodyGPUPool() = default;
std::shared_ptr<const RigidBodyGPUZone> RigidBodyGPUPool::ObservedZone(
const gemmi::Model &model, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg,
const gemmi::AsuData<gemmi::ValueSigma<float>> &fobs, double d_min) {
const RigidBodyGPUZone &model_zone = Zone(model, cell, sg, d_min);
const size_t atoms = gemmi::count_atom_sites(model);
std::lock_guard lock(observed_m_);
for (const ObservedEntry &o : observed_)
if (o.d_min == d_min && SameObservations(o.fobs, fobs))
if (o.group == sg.xhm() && o.atoms == atoms && o.d_min == d_min && SameObservations(o.fobs, fobs))
return o.zone;
const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min);
std::vector<gemmi::Miller> hkl;
@@ -220,16 +213,17 @@ std::shared_ptr<const RigidBodyGPUZone> RigidBodyGPUPool::ObservedZone(
const SymmetryComposition composition(grid, d_min, hkl);
auto zone = std::make_shared<RigidBodyGPUZone>(model_zone);
AddObservations(*zone, cell, composition, fobs);
observed_.push_back({d_min, fobs, zone});
observed_.push_back({sg.xhm(), atoms, d_min, fobs, zone});
return zone;
}
const RigidBodyGPUZone &RigidBodyGPUPool::Zone(const gemmi::Model &model, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, double d_min) {
const std::tuple<std::string, size_t, double> key{sg.xhm(), gemmi::count_atom_sites(model), d_min};
std::lock_guard lock(zones_m_);
auto it = zones_.find(d_min);
auto it = zones_.find(key);
if (it == zones_.end())
it = zones_.emplace(d_min, std::make_unique<RigidBodyGPUZone>(ModelZone(model, cell, sg, d_min))).first;
it = zones_.emplace(key, std::make_unique<RigidBodyGPUZone>(ModelZone(model, cell, sg, d_min))).first;
return *it->second;
}
image_analysis/structure_refinement/RigidBodyGPU.h
+19 -12
View File
@@ -16,6 +16,7 @@
#include <optional>
#include <stdexcept>
#include <string>
#include <tuple>
#include <vector>
#include "RigidBodyRefine.h"
@@ -32,12 +33,13 @@ public:
explicit RigidBodyGPUFailure(const std::string &what) : std::runtime_error(what) {}
};
// A few engines, reserved once for a whole validation: each is a stream and the buffers for one fit at
// a time, sized for the finest zone of the ladder to d_min. The real fit and the null's replicates each
// take one for the length of their fit, and wait for one when all are taken. Engines are
// interchangeable and every kernel deterministic, so which replicate gets which engine does not change
// a number. A pool belongs to one model: what a zone needs of its atoms - their densities, radii and mask
// radii, which the placement does not change - is worked out once per zone and kept.
// A few engines, reserved once for a whole validation or search: each is a stream and the buffers for one
// fit at a time, sized for the finest zone of the ladder to d_min. Every fit takes one for its length, and
// waits for one when all are taken. Engines are interchangeable and every kernel deterministic, so which fit
// gets which engine does not change a number. A pool belongs to one cell and one d_min, and serves every
// model and group it was made for (Plan, Cover): what a zone needs of a model's atoms - their densities,
// radii and mask radii, which the placement does not change - is worked out once per group, atom count and
// zone and kept.
class RigidBodyGPUPool {
public:
// A pool without engines yet, sized for the finest zone of the ladder to d_min: EngineBytes() is
@@ -47,13 +49,12 @@ public:
static std::unique_ptr<RigidBodyGPUPool> Plan(const gemmi::Model &model, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, double d_min,
size_t max_observations, Logger &logger);
// Plan(), then as many engines as fit a quarter of the card's total memory, at most `max_engines`
// (molecular replacement's rule). Null, logged, where not even one fits.
static std::unique_ptr<RigidBodyGPUPool> Create(const gemmi::Model &model, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, double d_min,
size_t max_observations, size_t max_engines, Logger &logger);
~RigidBodyGPUPool();
// Sizes the engines, before Reserve(), for fits of `model` in `sg` as well - in the pool's cell, to its
// d_min, with at most its observations. False, logged, where the cell is too small for the gather.
bool Cover(const gemmi::Model &model, const gemmi::SpaceGroup &sg, Logger &logger);
size_t EngineBytes() const;
// Reserves `engines` engines, engine i on the i-th card from Device(), so a validation uses every
// card there is. Engines that cannot be set up throw RigidBodyGPUFailure.
@@ -77,14 +78,20 @@ public:
private:
RigidBodyGPUPool() = default;
int device_ = 0;
gemmi::UnitCell cell_;
double d_min_ = 0;
size_t max_observations_ = 0;
std::unique_ptr<RigidBodyGPUCapacity> cap_;
std::vector<std::unique_ptr<RigidBodyGPUEngine>> engines_;
std::vector<RigidBodyGPUEngine *> idle_;
std::mutex m_;
std::condition_variable cv_;
std::map<double, std::unique_ptr<RigidBodyGPUZone>> zones_;
// By group, atom count and zone.
std::map<std::tuple<std::string, size_t, double>, std::unique_ptr<RigidBodyGPUZone>> zones_;
std::mutex zones_m_;
struct ObservedEntry {
std::string group;
size_t atoms;
double d_min;
gemmi::AsuData<gemmi::ValueSigma<float>> fobs;
std::shared_ptr<const RigidBodyGPUZone> zone;
rugnux/RugnuxScaleMerge.cpp
+4 -2
View File
@@ -4302,17 +4302,18 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
// memory goes back before the search reserves its rigid-body engines on the card.
if (!ledger_ahead.valid())
rsm.reset();
const size_t merge_on_card = rsm && rsm->OnGPU() ? rsm->DeviceBytes() : 0;
std::vector<MRMerge> merges{{"MERGE", &sm.merged, *result.consensus_cell, experiment_.GetSpaceGroupOrP1()}};
if (model_twin)
merges.push_back({"TWIN_SUBGROUP_MERGE", &model_twin->merged, model_twin->cell, model_twin->space_group,
model_twin->law});
MolecularReplacementResult mr = RunMolecularReplacement(config_.model_path, merges, logger,
static_cast<size_t>(config_.nthreads));
static_cast<size_t>(config_.nthreads), merge_on_card);
if (mr.chosen >= 0) {
const MRCandidate &c = mr.candidates[mr.chosen];
if (mr.decision == "ENANTIOMORPH")
CheckEnantiomorphHand(mr, sm.merged, experiment_.GetSpaceGroupOrP1(), logger,
static_cast<size_t>(config_.nthreads));
static_cast<size_t>(config_.nthreads), merge_on_card);
if (mr.decision == "ENANTIOMORPH") {
// A change of label only (see ValidateAgainstModel): the two groups index identically.
experiment_.SetSpaceGroup(c.space_group);
@@ -4346,6 +4347,7 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
const UnitCell cell = own ? own->cell : *result.consensus_cell;
ModelValidationSchedule placed_elsewhere;
placed_elsewhere.fit_test = false;
placed_elsewhere.gpu_resident_bytes = merge_on_card;
phase("Validating against model (molecular replacement)");
ModelValidationResult v = ValidateAgainstModel(
own ? own->merged : sm.merged, cell, config_.model_path,
tests/ModelValidationTest.cpp
+3 -2
View File
@@ -1276,12 +1276,13 @@ TEST_CASE("ModelValidation_RigidBodyTranslationDerivativeIsExact", "[ModelValida
#ifdef JFJOCH_USE_CUDA
namespace {
// A pool of `engines`: decided on the card's total memory, so there is one wherever there is a GPU.
// A pool of `engines`, wherever there is a GPU.
std::unique_ptr<RigidBodyGPUPool> TestPool(const gemmi::Model &model, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, double d_min, size_t observations,
size_t engines, Logger &logger) {
auto pool = RigidBodyGPUPool::Create(model, cell, sg, d_min, observations, engines, logger);
auto pool = RigidBodyGPUPool::Plan(model, cell, sg, d_min, observations, logger);
REQUIRE(pool);
pool->Reserve(engines);
return pool;
}
}
tests/MolecularReplacementTest.cpp
+68 -1
View File
@@ -93,6 +93,72 @@ namespace {
}
}
namespace {
constexpr size_t GB = 1000000000;
// A search's parts, sized as a large cell measures: the transform 0.84 GB, rigid-body engines of 2.10 GB
// (nine fits at once), score engines of up to 1.86 GB; the budget 80% of the card.
MolecularReplacementGPUNeed LargeCellNeed(size_t card) {
MolecularReplacementGPUNeed need;
need.budget = card / 5 * 4;
need.transform = 835500000;
need.rigid_body = 2101000000;
need.rigid_body_engines = 9;
need.structure_factors = 1860000000;
return need;
}
}
TEST_CASE("MolecularReplacement_GPUPlanOnCardSizes", "[MolecularReplacement]") {
size_t last_rb = 0;
for (size_t card : {6 * GB, 8 * GB, 16 * GB, 24 * GB, 64 * GB}) {
const MolecularReplacementGPUNeed need = LargeCellNeed(card);
const MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need);
CAPTURE(card, plan.rigid_body_engines, plan.bytes, plan.room);
CHECK(plan.bytes <= need.budget);
// Every part on the GPU on every one of these cards; only the engine count grows with it.
CHECK(plan.search);
CHECK(plan.structure_factors);
CHECK(plan.rigid_body_engines >= std::max<size_t>(last_rb, 1));
// What the transform and the engines leave is the room, and it holds the largest score engine.
CHECK(plan.room == need.budget - need.transform - plan.rigid_body_engines * need.rigid_body);
CHECK(plan.room >= need.structure_factors);
last_rb = plan.rigid_body_engines;
}
// 16 GB, a 12.8 GB budget: one of each, 4.80 GB, and a second engine would pass half of it.
CHECK(PlanMolecularReplacementGPU(LargeCellNeed(16 * GB)).rigid_body_engines == 1);
CHECK(PlanMolecularReplacementGPU(LargeCellNeed(24 * GB)).rigid_body_engines == 3);
CHECK(PlanMolecularReplacementGPU(LargeCellNeed(64 * GB)).rigid_body_engines == 9);
}
TEST_CASE("MolecularReplacement_GPUPlanPartThatDoesNotFitGoesToTheCPU", "[MolecularReplacement]") {
// 6 GB, a 4.8 GB budget, which the three parts fill to 4.80 GB: with a larger score engine, the scores
// go to the CPU and the transform and the rigid body stay.
MolecularReplacementGPUNeed need = LargeCellNeed(6 * GB);
need.structure_factors = 2 * GB;
MolecularReplacementGPUPlan plan = PlanMolecularReplacementGPU(need);
CHECK(plan.search);
CHECK(plan.rigid_body_engines == 1);
CHECK(!plan.structure_factors);
CHECK(plan.bytes <= need.budget);
// A part the GPU cannot grid (0 bytes) is never placed; the room is then what the transform leaves.
need = LargeCellNeed(16 * GB);
need.rigid_body = 0;
plan = PlanMolecularReplacementGPU(need);
CHECK(plan.rigid_body_engines == 0);
CHECK(plan.room == need.budget - need.transform);
// No budget (no GPU, or all of it held beside): everything on the CPU.
need = LargeCellNeed(16 * GB);
need.budget = 0;
plan = PlanMolecularReplacementGPU(need);
CHECK(!plan.search);
CHECK(plan.rigid_body_engines == 0);
CHECK(!plan.structure_factors);
CHECK(plan.room == 0);
}
TEST_CASE("MolecularReplacement_UniformRotationsCoverSO3", "[MolecularReplacement]") {
const auto rotations = UniformRotations(6.0);
const double step = 6.0 * M_PI / 180.0;
@@ -443,8 +509,9 @@ TEST_CASE("MolecularReplacement_RigidBodyBesideFixedCopy", "[MolecularReplacemen
#ifdef JFJOCH_USE_CUDA
if (get_gpu_count() > 0) {
gemmi::Structure gpu = start;
auto pool = RigidBodyGPUPool::Create(gpu.models[0], cell, sg, 3.0, fobs.v.size(), 1, logger);
auto pool = RigidBodyGPUPool::Plan(gpu.models[0], cell, sg, 3.0, fobs.v.size(), logger);
REQUIRE(pool);
pool->Reserve(1);
RefineRigidBody(gpu.models[0], cell, sg, fobs, 3.0, logger, 4, pool.get(), &first.models[0]);
CHECK(worst(gpu) < 0.1);
}