leonarski_f
24b1711d64
RefinementStage sums the time of each stage over a refinement (Refit and its parts, the evaluations and their X-ray, restraint, host and upload shares, the repulsion build, the riding hydrogens, the rungs) and logs it at debug level at the end; rugnux_refine_test -v shows it. No change to any result. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 9491f54e0beccf36c65ebcec9afff90d08485f86) (cherry picked from commit 0ce206a88df960754476ce2f75ffaed90fbf3123)
296 lines
13 KiB
C++
296 lines
13 KiB
C++
// 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 "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)
|
|
: device_(device), cell_(cell), sg_(sg), d_min_(d_min) {
|
|
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' densities (made again only for the atoms whose B moved) and grid positions, to the device.
|
|
void SetAtoms(const gemmi::Model &model) {
|
|
std::optional<RefinementStage> stage("gpu.atoms_host");
|
|
const std::vector<const gemmi::Atom *> atoms = Atoms(model);
|
|
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());
|
|
const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw};
|
|
for (size_t i = 0; i < atoms.size(); 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;
|
|
}
|
|
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");
|
|
engine_->SetAtoms(atoms_, pos);
|
|
}
|
|
|
|
int device_;
|
|
gemmi::UnitCell cell_;
|
|
const gemmi::SpaceGroup &sg_;
|
|
double d_min_;
|
|
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) {
|
|
auto engine = std::make_unique<RefinementEngineGPU>(device, cell, sg, d_min, refl, model, b_max);
|
|
if (device_bytes != nullptr)
|
|
*device_bytes = engine->Bytes();
|
|
if (!engine->Supported() || engine->Bytes() > budget)
|
|
return nullptr;
|
|
engine->Reserve();
|
|
return engine;
|
|
}
|