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)
88 lines
4.8 KiB
C++
88 lines
4.8 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
|
|
ModelDensityGrid mask_grid; // the bulk-solvent mask's (BulkSolventMask::MaskGrid), 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 RefinementGPUSetup &setup);
|
|
|
|
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), on the mask's own grid.
|
|
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_;
|
|
};
|