leonarski_f
8718065a7c
The score draws the mask on a grid of its own (min(d_min/3, 0.4 A), BulkSolventMask) on the host; on a GPU engine that was a CPU mask and transform every cycle, per trial slot on one thread. Now the engine masks on that grid itself: - RefinementGPUSetup::mask_grid (BulkSolventMask::MaskGrid, the grid the host mask uses); where it is not the structure factors' grid the engine reserves a grid, a spectrum and an R2C plan for it (counted in DeviceBytes and the FFT work area), and Fmask masks (ModelMaskGPU), transforms and reads the rows there. Refit takes the engine's Fmask on a GPU, BulkSolventMask on the CPU. - ModelMaskGPU's shrink: each edge point marks its stencil as solvent instead of every macromolecule point searching its stencil for an edge - the same idempotent stores to the same points, so the same mask to the bit (ModelMaskGPU_* tests unchanged), and a thin shell of points visited instead of half the grid. - BulkSolventMask finds its symmetry orbits at its first mask: a GPU refinement never draws one on the host (1.1 s of trial-slot setup on the 3.6 A set). The card's mask is gemmi's but at points exactly on an atom's radius (ModelMaskGPU.h): R-work/R-free of the 14 parity sets on the GPU are the CPU engine's to 4 digits. Refinement deterministic run to run on the GPU (7 timing sets, refined model and report identical over 5 runs of successive builds that did not change results). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 89b5d00907c9ad99e8b389f6272a3724f981cd10)
341 lines
15 KiB
C++
341 lines
15 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 "../../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;
|
|
// 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]);
|
|
}
|
|
|
|
// 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_);
|
|
// 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;
|
|
}
|
|
|
|
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) {
|
|
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]};
|
|
}
|
|
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");
|
|
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());
|
|
}
|
|
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<const gemmi::Atom *> atom_ptrs_; // the model's atoms at the last SetAtoms()
|
|
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;
|
|
}
|