mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 87 Wiki Activity
Files
e3782cabe4730cd2d4e73f2f47b87e0524a90c24
Jungfraujoch/image_analysis/structure_refinement/RefinementGPU.cpp
T
leonarski_f e3782cabe4 Refinement GPU: atom densities uploaded only when some B moved, one device copy, positions through a pinned buffer 
Every evaluation uploaded both copies of the atoms' densities (the gather's and the gradient's) and their
boxes, although in a positions-only minimisation no B changes between two Refits. The engine now keeps one
copy (the gradient kernel reads ModelDensityGPU's atoms and boxes), uploads it only where some atom's
density was made again, and sends the positions from a pinned buffer without waiting. The densities that
are made again (every atom in a B minimisation) and the grid positions are worked out on the refinement's
threads. Same bits.

9yl4: upload 3.29 -> 0.56 s, host side 1.29 -> 0.61 s, refinement 22.3 -> 17.7 s; 5reo 4.8 -> 4.4 s
(RTX 5080, -N 8, loaded box).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit 46c3b67c91aee0aec3473b0fb1a4945c00bda070)
(cherry picked from commit 0aebc98110c8155932af647d6204891aec662294)
2026-10-26 01:38:22 +01:00

306 lines
13 KiB
C++
Raw 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 <optional>
#include "gemmi/dencalc.hpp"
#include "gemmi/it92.hpp"
#include "gemmi/solmask.hpp" // SolventMasker, refmac_radius_for_bulk_solvent
#include "../../common/ParallelFor.h"
#include "RefinementGPUEngine.h"
namespace {
using Table = gemmi::IT92<float>;
// 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 = Table::get(atom.element, atom.charge, atom.serial);
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;
}
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)
: device_(device), cell_(cell), sg_(sg), d_min_(d_min), nthreads_(nthreads) {
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;
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]);
}
// 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());
for (size_t m = 0; m < refl.Size(); m++)
for (int o = 0; o < n_ops; o++) {
const gemmi::Miller k = gops.sym_ops[o].apply_to_hkl(refl.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_.grid);
// the restraints: about thirty records and slots per atom
bytes_ = RefinementGPUEngine::DeviceBytes(setup_, fft_work_, 30 * widest.size(), 60 * widest.size());
}
}
bool Supported() const { return supported_; }
size_t Bytes() const { return bytes_; }
void Reserve() { engine_ = std::make_unique<RefinementGPUEngine>(device_, setup_, fft_work_); }
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()};
engine_->SetRowModel(rows.k_total, fmask, rows.d, rows.var);
}
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);
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]};
}
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:
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;
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");
const std::vector<const gemmi::Atom *> atoms = Atoms(model);
const int n = static_cast<int>(atoms.size());
if (last_b_.size() != atoms.size()) {
last_b_.assign(atoms.size(), NAN);
atoms_.resize(atoms.size());
}
std::vector<std::array<float, 4>> pos(atoms.size());
std::vector<char> changed(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;
}
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_);
engine_->SetPositions(pos);
}
int device_;
gemmi::UnitCell cell_;
const gemmi::SpaceGroup &sg_;
double d_min_;
size_t nthreads_;
RefinementGPUSetup setup_;
gemmi::DensityCalculator<Table, float> dc_;
bool supported_ = false;
size_t fft_work_ = 0, bytes_ = 0;
std::unique_ptr<RefinementGPUEngine> engine_;
std::vector<float> last_b_;
std::vector<ModelDensityAtom> atoms_;
};
} // 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,
size_t *device_bytes, size_t nthreads) {
auto engine = std::make_unique<RefinementEngineGPU>(device, cell, sg, d_min, refl, model, b_max, nthreads);
if (device_bytes != nullptr)
*device_bytes = engine->Bytes();
if (!engine->Supported() || engine->Bytes() > budget)
return nullptr;
engine->Reserve();
return engine;
}
Reference in new issue View Git Blame Copy Permalink