mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 88 Wiki Activity
Files
rc181
Jungfraujoch/image_analysis/structure_refinement/RefinementGPU.cpp
T
leonarski_fandClaude Opus 5.5 67ad89de11 Refinement GPU: the transform buffers in lanes - a slot's own where the card holds one, shared otherwise 
The previous commit put every engine on a card on one set of working buffers, one evaluation at a
time. On a T4 two evaluations side by side cost only 9% more each than one, so serialising them
gives back the concurrency of two slots. The buffers (grid, spectrum, cuFFT work area, the mask's
own grid and spectrum) now come in RefinementGPULane: a slot takes a lane of its own where the card's
budget holds its engine and a lane, else only its engine on the card's least used lane (under that
lane's lock from launches to wait). The bulk-solvent mask's buffers stay one per card, under their
own lock. OwnSlotPlaces plans both per card.

4YBB (engine 2.02 GB, lane 1.48 GB, card 1.73 GB, scale pool 0.91 GB): RTX 5080 3 slots on 3 lanes,
512 s (4 slots on one lane 500 s, 3 slots unshared 516 s); a T4 (12.52 GB budget) 3 slots on 2 lanes,
an L4 5 on 4. Refined model CIF byte-identical. 51 cases / 22236 assertions pass.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37
2026-10-28 02:43:28 +01:00

482 lines
22 KiB
C++
Raw Permalink Blame History

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "RefinementGPU.h"
#include <algorithm>
#include <cmath>
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include "gemmi/dencalc.hpp"
#include "gemmi/it92.hpp"
#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent
#include "../../common/ParallelFor.h"
#include "ModelGrid.h" // IonFormFactor, AnisoRadiusB
#include "RefinementGPUEngine.h"
namespace {
using Table = gemmi::IT92<float>;
// One atom's density as the gather takes it, anisotropic with U (as PutModelDensityOnGrid() grids it), with `blur`.
ModelDensityAtom AnisoDensityAtom(const gemmi::Atom &atom, const gemmi::SMat33<float> &u,
const gemmi::DensityCalculator<Table, float> &dc) {
using CReal = Table::Coef::coef_type;
const auto &coef = IonFormFactor<float>(atom.element.elem, atom.charge);
const auto aniso_b = u.scaled(CReal(gemmi::u_to_b())).added_kI(CReal(dc.blur));
const CReal b_max = AnisoRadiusB(aniso_b);
const double radius = dc.estimate_radius(coef.precalculate_density_iso(b_max, 0.0f), b_max);
const auto precal = coef.precalculate_density_aniso_b(aniso_b, 0.0f);
ModelDensityAtom a{};
a.occ = atom.occ;
a.aniso = 1;
a.radius = static_cast<float>(radius);
a.max_r2 = static_cast<float>(radius * radius);
for (int k = 0; k < 5; k++) {
a.a[k] = precal.a[k];
const gemmi::SMat33<float> &m = precal.b[k];
const float e[6] = {m.u11, m.u22, m.u33, m.u12, m.u13, m.u23};
std::copy(e, e + 6, a.b[k]);
}
return a;
}
// One atom's density as the gather takes it, isotropic, with `blur`.
ModelDensityAtom DensityAtom(const gemmi::Atom &atom, double b_iso, const gemmi::DensityCalculator<Table, float> &dc) {
using CReal = Table::Coef::coef_type;
const auto &coef = IonFormFactor<float>(atom.element.elem, atom.charge);
const CReal b = static_cast<CReal>(b_iso + dc.blur);
const auto precal = coef.precalculate_density_iso(b, 0.0f);
const double radius = dc.estimate_radius(precal, b);
ModelDensityAtom a{};
a.occ = atom.occ;
a.aniso = 0;
a.radius = static_cast<float>(radius);
a.max_r2 = static_cast<float>(radius * radius);
for (int k = 0; k < 5; k++) {
a.a[k] = precal.a[k];
a.b[k][0] = precal.b[k];
}
return a;
}
std::vector<const gemmi::Atom *> Atoms(const gemmi::Model &model) {
std::vector<const gemmi::Atom *> atoms;
for (const gemmi::Chain &ch : model.chains)
for (const gemmi::Residue &r : ch.residues)
for (const gemmi::Atom &a : r.atoms)
atoms.push_back(&a);
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,
const RefinementReflections &refl, const gemmi::Model &model, double b_max, size_t nthreads,
size_t restraint_records, size_t restraint_slots)
: device_(device), cell_(cell), sg_(sg), d_min_(d_min), nthreads_(nthreads),
setup_(std::make_shared<RefinementGPUSetup>()), cards_(std::make_shared<Cards>()) {
const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min);
setup_->grid.nu = grid.nu;
setup_->grid.nv = grid.nv;
setup_->grid.nw = grid.nw;
for (int i = 0; i < 3; i++)
for (int j = 0; j < 3; j++) {
setup_->grid.orth[3 * i + j] = cell.orth.mat[i][j];
setup_->grid.frac[3 * i + j] = cell.frac.mat[i][j];
}
setup_->volume = cell.volume;
// the mask on a grid of its own, as BulkSolventMask has it
const gemmi::Grid<float> mask_grid = BulkSolventMask::MaskGrid(cell, sg, d_min);
setup_->mask_grid = setup_->grid;
setup_->mask_grid.nu = mask_grid.nu;
setup_->mask_grid.nv = mask_grid.nv;
setup_->mask_grid.nw = mask_grid.nw;
const gemmi::GroupOps gops = sg.operations();
setup_->den = gemmi::Op::DEN;
setup_->centring = static_cast<double>(gops.cen_ops.size());
for (const gemmi::Op &op : gops.sym_ops) {
std::array<int, 12> o{};
for (int i = 0; i < 3; i++) {
for (int j = 0; j < 3; j++)
o[3 * i + j] = op.rot[i][j];
o[9 + i] = op.tran[i];
}
setup_->sym_ops.push_back(o);
}
for (const gemmi::Op::Tran &cen : gops.cen_ops)
for (const gemmi::Op &op : gops.sym_ops) {
ModelMaskOp m{};
for (int i = 0; i < 3; i++) {
for (int j = 0; j < 3; j++)
m.rot[3 * i + j] = static_cast<double>(op.rot[i][j]) / gemmi::Op::DEN;
m.tran[i] = static_cast<double>(op.tran[i] + cen[i]) / gemmi::Op::DEN;
}
setup_->mask_ops.push_back(m);
}
for (size_t m = 0; m < refl.Size(); m++) {
setup_->rows.push_back({refl.hkl[m][0], refl.hkl[m][1], refl.hkl[m][2]});
setup_->inv_d2.push_back(refl.inv_d2[m]);
setup_->fo.push_back(refl.fo[m]);
setup_->work.push_back(!refl.free[m]);
setup_->centric.push_back(refl.centric[m]);
}
// a twin law's mates not observed: structure factors only
for (const gemmi::Miller &h : refl.mate_hkl) {
setup_->rows.push_back({h[0], h[1], h[2]});
setup_->inv_d2.push_back(cell.calculate_1_d2(h));
setup_->fo.push_back(0);
setup_->work.push_back(0);
setup_->centric.push_back(0);
}
setup_->mate = refl.mate;
// The adjoint of the composition on cuFFT's half-u spectrum: the term of row m and operator o, at
// k = hR, goes to k where k_u > 0, conjugated to -k where k_u < 0, and to both on the planes k_u = 0
// and k_u = nu/2, which hold their own Friedel mates.
const int nu = grid.nu, nv = grid.nv, nw = grid.nw;
auto point = [&](int ku, int kv, int kw) {
return ku + (nu / 2 + 1) * (gemmi::modulo(kv, nv) + nv * gemmi::modulo(kw, nw));
};
std::vector<std::pair<int, int>> entries;
const int n_ops = static_cast<int>(gops.sym_ops.size());
const std::vector<gemmi::Miller> row_hkl = refl.RowHkl();
for (size_t m = 0; m < row_hkl.size(); m++)
for (int o = 0; o < n_ops; o++) {
const gemmi::Miller k = gops.sym_ops[o].apply_to_hkl(row_hkl[m]);
const int term = (static_cast<int>(m) * n_ops + o) * 2;
if (k[0] > 0 && 2 * k[0] != nu)
entries.emplace_back(point(k[0], k[1], k[2]), term);
else if (k[0] < 0 && -2 * k[0] != nu)
entries.emplace_back(point(-k[0], -k[1], -k[2]), term + 1);
else {
const int ku = std::abs(k[0]);
entries.emplace_back(point(ku, k[1], k[2]), term);
entries.emplace_back(point(ku, -k[1], -k[2]), term + 1);
}
}
std::stable_sort(entries.begin(), entries.end(),
[](const auto &a, const auto &b) { return a.first < b.first; });
for (size_t e = 0; e < entries.size(); e++) {
if (e == 0 || entries[e].first != entries[e - 1].first) {
setup_->scatter_point.push_back(entries[e].first);
setup_->scatter_start.push_back(static_cast<int>(e));
}
setup_->scatter_term.push_back(entries[e].second);
}
setup_->scatter_start.push_back(static_cast<int>(entries.size()));
// The gather's reserve for every atom at b_max with the largest blur a refinement sets (B down to 0).
dc_ = Calculator(RefinementBlur(cell, sg, d_min, 0.0));
std::vector<ModelDensityAtom> widest;
for (const gemmi::Atom *a : Atoms(model))
widest.push_back(DensityAtom(*a, b_max, dc_));
setup_->max_atoms = widest.size();
setup_->max_pairs = ModelDensityGPU::PairBound(setup_->grid, widest);
supported_ = !setup_->rows.empty() && ModelDensityGPU::Supports(setup_->grid, widest) &&
setup_->mask_ops.size() <= ModelMaskGPU::MAX_OPS;
if (supported_) {
fft_work_ = RefinementGPUEngine::FFTWorkBytes(device, *setup_);
// the restraints, where they are not given: about thirty records and sixty slots per atom
if (restraint_records == 0 && restraint_slots == 0) {
restraint_records = 30 * widest.size();
restraint_slots = 60 * widest.size();
}
bytes_.engine = RefinementGPUEngine::DeviceBytes(*setup_, restraint_records, restraint_slots);
bytes_.card = RefinementGPUEngine::CardBytes(*setup_);
bytes_.lane = RefinementGPUEngine::LaneBytes(*setup_, fft_work_);
}
}
// One like `main` on `device`: the same setup, and the card's part shared with the engines of main's kind there.
RefinementEngineGPU(const RefinementEngineGPU &main, int device, size_t nthreads, bool own_lane)
: device_(device), cell_(main.cell_), sg_(main.sg_), d_min_(main.d_min_), nthreads_(nthreads), setup_(main.setup_),
dc_(Calculator(RefinementBlur(cell_, sg_, d_min_, 0.0))), supported_(main.supported_),
fft_work_(RefinementGPUEngine::FFTWorkBytes(device, *setup_)), bytes_(main.bytes_), own_lane_(own_lane), cards_(main.cards_) {}
bool Supported() const { return supported_; }
RefinementGPUBytes Bytes() const { return bytes_; }
void Reserve() {
engine_ = std::make_unique<RefinementGPUEngine>(device_, *setup_, cards_->Get(device_, *setup_), fft_work_, own_lane_);
}
void SetBlur(double blur) override {
dc_ = Calculator(blur);
last_b_.clear();
engine_->SetBlur(blur);
}
void SetRowModel(const XrayRowModel &rows) override {
std::vector<std::array<float, 2>> fmask(rows.fmask.size());
for (size_t m = 0; m < fmask.size(); m++)
fmask[m] = {rows.fmask[m].real(), rows.fmask[m].imag()};
// per row of the engine: the mates' rows have no observation, and nothing of theirs is read but fmask
std::vector<float> k_total = rows.k_total, d = rows.d, var = rows.var;
k_total.resize(setup_->rows.size(), 0.0f);
d.resize(setup_->rows.size(), 0.0f);
var.resize(setup_->rows.size(), 1.0f);
engine_->SetRowModel(k_total, fmask, d, var, rows.twin_fraction);
}
bool OnGPU() const override { return true; }
void Fcalc(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
SetAtoms(model);
std::vector<std::array<float, 2>> out;
engine_->Fcalc(out);
f.resize(out.size());
for (size_t m = 0; m < out.size(); m++)
f[m] = {out[m][0], out[m][1]};
}
void Fmask(const gemmi::Model &model, std::vector<std::complex<float>> &f) override {
const gemmi::SolventMasker masker(gemmi::AtomicRadiiSet::Refmac);
std::vector<ModelMaskAtom> mask_atoms;
for (const gemmi::Atom *a : Atoms(model)) {
if ((masker.ignore_hydrogen && a->is_hydrogen()) || (masker.ignore_zero_occupancy_atoms && a->occ <= 0))
continue;
const gemmi::Fractional f0 = cell_.fractionalize(a->pos);
double frac[3] = {f0.x, f0.y, f0.z};
for (double &x : frac) {
x -= std::floor(x);
if (x >= 1.0)
x = 0.0;
}
mask_atoms.push_back({frac[0], frac[1], frac[2],
masker.constant_r + masker.rprobe +
gemmi::refmac_radius_for_bulk_solvent(a->element.elem)});
}
std::vector<std::array<float, 2>> out;
engine_->Fmask(mask_atoms, out);
f.resize(out.size());
for (size_t m = 0; m < out.size(); m++)
f[m] = {out[m][0], out[m][1]};
}
double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) override {
SetAtoms(model);
double value = 0;
{
const RefinementStage stage("gpu.evaluate");
value = engine_->Evaluate(grad != nullptr);
}
if (grad != nullptr) {
Widen(engine_->HostGradient(), atom_ptrs_.size(), *grad, nthreads_);
CopyUGradient();
}
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);
double value = 0;
{
const RefinementStage stage("gpu.evaluate");
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, restraints);
}
if (grad != nullptr) {
Widen(engine_->HostGradient(), atom_ptrs_.size(), *grad, nthreads_);
CopyUGradient();
}
if (restraint_grad != nullptr)
Widen(engine_->HostRestraintGradient(), atom_ptrs_.size(), *restraint_grad, nthreads_);
return value;
}
void SetRestraints(int which, const RestraintSet &set) override {
engine_->SetRestraints(which, set.records, set.plane_atoms, set.slot_start, set.slot_index, set.n_slots);
}
double EvaluateRestraints(int which, const std::vector<std::array<double, 4>> &xyzb,
std::vector<std::array<double, 4>> *grad) override {
std::vector<std::array<float, 4>> x(xyzb.size()), g;
for (size_t i = 0; i < xyzb.size(); i++)
x[i] = {static_cast<float>(xyzb[i][0]), static_cast<float>(xyzb[i][1]), static_cast<float>(xyzb[i][2]),
static_cast<float>(xyzb[i][3])};
double value = 0;
{
const RefinementStage stage("gpu.restraints");
value = engine_->EvaluateRestraints(which, x, 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]};
}
return value;
}
private:
// The U derivatives of the last evaluation, where some atom is anisotropic.
void CopyUGradient() {
if (!AnyAniso()) {
grad_u_.clear();
return;
}
const float *g = engine_->HostUGradient();
grad_u_.resize(atom_ptrs_.size());
for (size_t i = 0; i < atom_ptrs_.size(); i++)
for (int c = 0; c < 6; c++)
grad_u_[i][c] = g[6 * i + c];
}
gemmi::DensityCalculator<Table, float> Calculator(double blur) const {
gemmi::DensityCalculator<Table, float> dc;
dc.d_min = d_min_;
dc.rate = 1.5;
dc.grid.unit_cell = cell_;
dc.grid.spacegroup = &sg_;
dc.blur = blur;
dc.cutoff = REFINEMENT_DENSITY_CUTOFF;
return dc;
}
// The atoms' grid positions to the device, and their densities where some B moved (made again only for the
// atoms whose B moved).
void SetAtoms(const gemmi::Model &model) {
std::optional<RefinementStage> stage("gpu.atoms_host");
// 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);
last_key_.assign(atoms.size(), AtomKey());
last_aniso_.assign(atoms.size(), AtomAniso());
atoms_.resize(atoms.size());
pos_.resize(atoms.size());
}
changed_.assign(atoms.size(), 0);
const int n3[3] = {setup_->grid.nu, setup_->grid.nv, setup_->grid.nw};
const std::vector<AtomAniso> *aniso = aniso_ != nullptr && aniso_->size() == atoms.size() ? aniso_ : nullptr;
ParallelChunks(n, ThreadsForWork(atoms.size(), nthreads_, 4096), [&](int lo, int hi) {
for (int i = lo; i < hi; i++) {
// made again where its B or its anisotropy moved, or where the solvent build put it in, took it out or
// relabelled it
const AtomKey key{atoms[i]->occ, static_cast<int>(atoms[i]->element.elem), atoms[i]->charge};
const AtomAniso none;
const AtomAniso &an = aniso != nullptr ? (*aniso)[i] : none;
if (!(atoms[i]->b_iso == last_b_[i]) || !(key == last_key_[i]) || an.on != last_aniso_[i].on ||
(an.on && an.dev != last_aniso_[i].dev)) {
atoms_[i] = an.on ? AnisoDensityAtom(*atoms[i], AnisoU(an, atoms[i]->b_iso), dc_)
: DensityAtom(*atoms[i], atoms[i]->b_iso, dc_);
last_b_[i] = atoms[i]->b_iso;
last_key_[i] = key;
last_aniso_[i] = an;
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};
}
});
stage.reset();
const RefinementStage upload("gpu.upload");
if (std::find(changed_.begin(), changed_.end(), 1) != changed_.end())
engine_->SetDensities(atoms_.data(), atoms_.size());
engine_->SetPositions(pos_.data(), pos_.size());
}
int device_;
gemmi::UnitCell cell_;
const gemmi::SpaceGroup &sg_;
double d_min_;
size_t nthreads_;
std::shared_ptr<RefinementGPUSetup> setup_; // read only once made: shared with the engines like this one
gemmi::DensityCalculator<Table, float> dc_;
bool supported_ = false;
size_t fft_work_ = 0;
RefinementGPUBytes bytes_;
bool own_lane_ = true;
// The cards' shared parts of the engines like this one, each kept while an engine on its card is.
struct Cards {
std::mutex m;
std::map<int, std::weak_ptr<RefinementGPUCard>> card;
std::shared_ptr<RefinementGPUCard> Get(int device, const RefinementGPUSetup &setup) {
const std::lock_guard lock(m);
std::shared_ptr<RefinementGPUCard> c = card[device].lock();
if (!c) {
c = RefinementGPUEngine::MakeCard(device, setup);
card[device] = c;
}
return c;
}
};
std::shared_ptr<Cards> cards_;
std::unique_ptr<RefinementGPUEngine> engine_;
std::vector<const gemmi::Atom *> atom_ptrs_; // the model's atoms at the last SetAtoms()
std::vector<float> last_b_;
struct AtomKey {
float occ = -1;
int element = -1;
signed char charge = 0;
bool operator==(const AtomKey &o) const { return occ == o.occ && element == o.element && charge == o.charge; }
};
std::vector<AtomKey> last_key_;
std::vector<AtomAniso> last_aniso_;
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
std::unique_ptr<RefinementEngine> MakeRefinementEngineGPU(int device, const gemmi::UnitCell &cell,
const gemmi::SpaceGroup &sg, double d_min,
const RefinementReflections &refl,
const gemmi::Model &model, double b_max, size_t budget,
RefinementGPUBytes *bytes, size_t nthreads,
size_t restraint_records, size_t restraint_slots) {
auto engine = std::make_unique<RefinementEngineGPU>(device, cell, sg, d_min, refl, model, b_max, nthreads,
restraint_records, restraint_slots);
if (bytes != nullptr)
*bytes = engine->Bytes();
if (!engine->Supported() || engine->Bytes().engine + engine->Bytes().card + engine->Bytes().lane > budget)
return nullptr;
engine->Reserve();
return engine;
}
std::unique_ptr<RefinementEngine> MakeRefinementEngineGPU(const RefinementEngine &like, int device, size_t nthreads,
bool own_lane) {
auto engine = std::make_unique<RefinementEngineGPU>(dynamic_cast<const RefinementEngineGPU &>(like), device, nthreads,
own_lane);
engine->Reserve();
return engine;
}
Reference in new issue View Git Blame Copy Permalink