leonarski_f
5569f3d1f2
RefinementEngine::EvaluateWithRestraints evaluates both targets at the model's atoms. On the GPU the X-ray chain is queued first, the restraints' coordinates are written into a pinned buffer while the device works on it, the restraint kernels follow on the same stream, and the host waits once; the sums and gradients come back through pinned buffers. Same kernels on the same inputs, same bits. The CPU engine keeps calling the two in turn. 9yl4: evaluation 5.1 -> 4.7 ms, refinement 14.2 -> 13.4 s; 5reo 4.0 -> 3.9 s (RTX 5080, -N 8). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 51434b98b62a9e6322ffe811bf272a45ff6ce003) (cherry picked from commit 0bcfa3f8041cec4bb4bc376b8559f30c36d3ac80)
87 lines
4.7 KiB
C++
87 lines
4.7 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 <functional>
|
|
#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);
|
|
// Evaluate() and EvaluateRestraints() in one submission with one wait: the X-ray target is queued first, and
|
|
// `xyzb` writes the restraints' x, y, z, B of every atom while the device works on it.
|
|
double EvaluateWithRestraints(std::vector<std::array<float, 4>> *grad, int which,
|
|
const std::function<void(std::array<float, 4> *)> &xyzb,
|
|
std::vector<std::array<float, 4>> *restraint_grad, double &restraints);
|
|
|
|
// The most atoms of a plane restraint the kernel takes.
|
|
static constexpr int MAX_PLANE_ATOMS = 32;
|
|
|
|
private:
|
|
std::unique_ptr<RefinementGPUEngineImpl> impl_;
|
|
};
|