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

Refinement: the per-evaluation host chain on the threads and in place

Browse Source 
Byte-identical: every pass is element by element in the serial loop's order, or a copy that is no longer made.

- The engine's gradients are read in place from its pinned buffers and widened to double on the threads
  (RefinementGPUEngine::HostGradient/HostRestraintGradient); the positions go up from a pinned vector of the
  engine's host side; the model's atom pointers are kept between calls (the model is the context's own).
- Context: the riders moved and folded on the threads - a parent takes its riders in index order through a
  CSR list (SetRiderParents), as the serial walk added them.
- AtomTarget: the parameters set and the gradient assembled on the threads; the per-atom gradient vectors kept
  between evaluations (7 MB each on a large model: made anew, their pages fault in every time).
- RefinementMinimiser: the sweeps over the parameters (x + alpha d, s, y) on the threads it is given, the dot
  products the serial chains they were; Direction()'s working vector and a let-go pair's vectors reused.

A/B against the frozen combined2 binary, one lock hold each, outputs identical but provenance lines:
cytc refinement 1.600 -> 1.548 s (live), 1.724 -> 1.668 (CLI); thau 1.906 -> 1.885, 2.031 -> 2.027;
4nwv (with the commit before) 39.9 -> 35.5 s, the ladder 14.6 -> 12.7, each step 8.3 -> 7.2.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
This commit is contained in:
leonarski_fandClaude Fable 5.1 committed 2026-10-26 22:48:28 +01:00
1 parent b16873036c
commit d122118d73
5 files changed
+174 -99
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/structure_refinement/RefinementGPU.cpp
+45 -27
View File
@@ -46,6 +46,31 @@ std::vector<const gemmi::Atom *> Atoms(const gemmi::Model &model) {
return atoms;
}
// Whether `atoms` is Atoms(model) still: as many atoms, the first and the last the same ones.
bool SameAtoms(const gemmi::Model &model, const std::vector<const gemmi::Atom *> &atoms) {
size_t n = 0;
const gemmi::Atom *first = nullptr, *last = nullptr;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues)
if (!r.atoms.empty()) {
if (first == nullptr)
first = &r.atoms.front();
last = &r.atoms.back();
n += r.atoms.size();
}
return n > 0 && n == atoms.size() && first == atoms.front() && last == atoms.back();
}
// The engine's single-precision gradient into the caller's doubles, element by element in parallel: bit for bit
// the serial loop's.
void Widen(const std::array<float, 4> *from, size_t n, std::vector<std::array<double, 4>> &to, size_t nthreads) {
to.resize(n);
ParallelChunks(static_cast<int>(n), ThreadsForWork(n, nthreads, 16384), [&](int lo, int hi) {
for (int i = lo; i < hi; i++)
to[i] = {from[i][0], from[i][1], from[i][2], from[i][3]};
});
}
class RefinementEngineGPU : public RefinementEngine {
public:
RefinementEngineGPU(int device, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
@@ -201,44 +226,33 @@ public:
double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) override {
SetAtoms(model);
std::vector<std::array<float, 4>> g;
double value = 0;
{
const RefinementStage stage("gpu.evaluate");
value = engine_->Evaluate(grad != nullptr ? &g : nullptr);
}
if (grad != nullptr) {
grad->resize(g.size());
for (size_t i = 0; i < g.size(); i++)
(*grad)[i] = {g[i][0], g[i][1], g[i][2], g[i][3]};
value = engine_->Evaluate(grad != nullptr);
}
if (grad != nullptr)
Widen(engine_->HostGradient(), atom_ptrs_.size(), *grad, nthreads_);
return value;
}
double EvaluateWithRestraints(const gemmi::Model &model, int which, std::vector<std::array<double, 4>> *grad,
std::vector<std::array<double, 4>> *restraint_grad, double &restraints) override {
SetAtoms(model);
std::vector<std::array<float, 4>> g, gr;
double value = 0;
{
const RefinementStage stage("gpu.evaluate");
value = engine_->EvaluateWithRestraints(grad != nullptr ? &g : nullptr, which, [&](std::array<float, 4> *x) {
value = engine_->EvaluateWithRestraints(grad != nullptr, which, [&](std::array<float, 4> *x) {
for (size_t i = 0; i < atom_ptrs_.size(); i++) {
const gemmi::Atom &a = *atom_ptrs_[i];
x[i] = {static_cast<float>(a.pos.x), static_cast<float>(a.pos.y), static_cast<float>(a.pos.z), a.b_iso};
}
}, restraint_grad != nullptr ? &gr : nullptr, restraints);
}
if (grad != nullptr) {
grad->resize(g.size());
for (size_t i = 0; i < g.size(); i++)
(*grad)[i] = {g[i][0], g[i][1], g[i][2], g[i][3]};
}
if (restraint_grad != nullptr) {
restraint_grad->resize(gr.size());
for (size_t i = 0; i < gr.size(); i++)
(*restraint_grad)[i] = {gr[i][0], gr[i][1], gr[i][2], gr[i][3]};
}, restraint_grad != nullptr, restraints);
}
if (grad != nullptr)
Widen(engine_->HostGradient(), atom_ptrs_.size(), *grad, nthreads_);
if (restraint_grad != nullptr)
Widen(engine_->HostRestraintGradient(), atom_ptrs_.size(), *restraint_grad, nthreads_);
return value;
}
@@ -280,32 +294,34 @@ private:
// atoms whose B moved).
void SetAtoms(const gemmi::Model &model) {
std::optional<RefinementStage> stage("gpu.atoms_host");
atom_ptrs_ = Atoms(model);
// the pointers kept from the last call where the model is the same one still (its context's own)
if (!SameAtoms(model, atom_ptrs_))
atom_ptrs_ = Atoms(model);
const std::vector<const gemmi::Atom *> &atoms = atom_ptrs_;
const int n = static_cast<int>(atoms.size());
if (last_b_.size() != atoms.size()) {
last_b_.assign(atoms.size(), NAN);
atoms_.resize(atoms.size());
pos_.resize(atoms.size());
}
std::vector<std::array<float, 4>> pos(atoms.size());
std::vector<char> changed(atoms.size(), 0);
changed_.assign(atoms.size(), 0);
const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw};
ParallelChunks(n, ThreadsForWork(atoms.size(), nthreads_, 4096), [&](int lo, int hi) {
for (int i = lo; i < hi; i++) {
if (!(atoms[i]->b_iso == last_b_[i])) {
atoms_[i] = DensityAtom(*atoms[i], atoms[i]->b_iso, dc_);
last_b_[i] = atoms[i]->b_iso;
changed[i] = 1;
changed_[i] = 1;
}
const gemmi::Fractional f = cell_.fractionalize(atoms[i]->pos);
pos[i] = {static_cast<float>(f.x * n3[0]), static_cast<float>(f.y * n3[1]), static_cast<float>(f.z * n3[2]), 0};
pos_[i] = {static_cast<float>(f.x * n3[0]), static_cast<float>(f.y * n3[1]), static_cast<float>(f.z * n3[2]), 0};
}
});
stage.reset();
const RefinementStage upload("gpu.upload");
if (std::find(changed.begin(), changed.end(), 1) != changed.end())
if (std::find(changed_.begin(), changed_.end(), 1) != changed_.end())
engine_->SetDensities(atoms_.data(), atoms_.size());
engine_->SetPositions(pos);
engine_->SetPositions(pos_.data(), pos_.size());
}
int device_;
@@ -321,6 +337,8 @@ private:
std::vector<const gemmi::Atom *> atom_ptrs_; // the model's atoms at the last SetAtoms()
std::vector<float> last_b_;
PinnedVector<ModelDensityAtom> atoms_; // pinned: a B rung sends them every evaluation (37 MB on a large model)
PinnedVector<std::array<float, 4>> pos_;
std::vector<char> changed_;
};
} // namespace