leonarski_f
e3782cabe4
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)
81 lines
4.2 KiB
C++
81 lines
4.2 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#pragma once
|
|
|
|
// The device half of restrained refinement on a GPU (RefinementGPU.h, CUDA builds only). The host half works
|
|
// out with gemmi everything that depends on the cell, the group and the reflections; this half grids the atoms,
|
|
// transforms, composes the symmetry, evaluates the likelihood target, takes its derivative back to every atom
|
|
// by the FFT method, and evaluates the geometry restraints. Every kernel is deterministic - gathers and fixed
|
|
// reduction trees, no floating-point atomics - so the same input gives the same bits on every run.
|
|
|
|
#include <array>
|
|
#include <cstddef>
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
#include "ModelDensityGPU.h"
|
|
#include "ModelMaskGPU.h"
|
|
#include "RefinementMath.h"
|
|
|
|
// Everything an engine needs that does not depend on the atoms.
|
|
struct RefinementGPUSetup {
|
|
ModelDensityGrid grid; // the structure-factor grid, u fastest
|
|
double volume = 0;
|
|
std::vector<std::array<int, 12>> sym_ops; // gemmi's Op: rot row-major and tran, in units of 1 / den
|
|
int den = 24;
|
|
double centring = 1; // the number of centring vectors
|
|
std::vector<ModelMaskOp> mask_ops; // every operator with every centring vector
|
|
std::vector<std::array<int, 3>> rows; // the reflections, in the order their numbers come back
|
|
std::vector<double> inv_d2;
|
|
std::vector<float> fo;
|
|
std::vector<char> work, centric;
|
|
// The adjoint of the composition, as a gather: the half-u spectrum points the gradient's coefficients
|
|
// land on (ascending), and for each the terms (row * ops + op) * 2 + conj adding to it, in order.
|
|
std::vector<int> scatter_point, scatter_start, scatter_term;
|
|
size_t max_atoms = 0, max_pairs = 0;
|
|
};
|
|
|
|
struct RefinementGPUEngineImpl;
|
|
|
|
class RefinementGPUEngine {
|
|
public:
|
|
static size_t DeviceBytes(const RefinementGPUSetup &setup, size_t fft_work_bytes, size_t max_records,
|
|
size_t max_slots);
|
|
static size_t FFTWorkBytes(int device, const ModelDensityGrid &grid);
|
|
|
|
RefinementGPUEngine(int device, const RefinementGPUSetup &setup, size_t fft_work_bytes);
|
|
~RefinementGPUEngine();
|
|
|
|
// n_cen exp(blur s^2 / 4) per reflection: prepare_asu_data()'s unblur.
|
|
void SetBlur(double blur);
|
|
// Per macro-cycle: F_model = k_total (F + fmask), and the likelihood's D and variance.
|
|
void SetRowModel(const std::vector<float> &k_total, const std::vector<std::array<float, 2>> &fmask,
|
|
const std::vector<float> &d, const std::vector<float> &var);
|
|
|
|
// The atoms' densities, as ModelStructureFactorsGPU takes them (the boxes of ModelDensityGPU::Boxes() made
|
|
// from them): called again only where some atom's density changed.
|
|
void SetDensities(const std::vector<ModelDensityAtom> &atoms);
|
|
// The atoms' grid positions (fractional times the grid size, not wrapped), in the order of SetDensities().
|
|
void SetPositions(const std::vector<std::array<float, 4>> &pos);
|
|
// F of the atoms at the rows.
|
|
void Fcalc(std::vector<std::array<float, 2>> &f);
|
|
// The bulk-solvent mask's F at the rows (the mask atoms fractional, wrapped).
|
|
void Fmask(const std::vector<ModelMaskAtom> &mask_atoms, std::vector<std::array<float, 2>> &f);
|
|
// -ln L over the working rows for the atoms last set and, with `grad`, dT/d(x, y, z, B) per atom.
|
|
double Evaluate(std::vector<std::array<float, 4>> *grad);
|
|
|
|
// Restraint set `which` (0 or 1), as RestraintSet holds it.
|
|
void SetRestraints(int which, const std::vector<RestraintRecord> &records, const std::vector<int> &plane_atoms,
|
|
const std::vector<int> &slot_start, const std::vector<int> &slot_index, int n_slots);
|
|
// The sum of E of set `which` at `xyzb` (x, y, z, B per atom) and dE/d(x, y, z, B) per atom.
|
|
double EvaluateRestraints(int which, const std::vector<std::array<float, 4>> &xyzb,
|
|
std::vector<std::array<float, 4>> *grad);
|
|
|
|
// The most atoms of a plane restraint the kernel takes.
|
|
static constexpr int MAX_PLANE_ATOMS = 32;
|
|
|
|
private:
|
|
std::unique_ptr<RefinementGPUEngineImpl> impl_;
|
|
};
|