Refinement GPU: the X-ray target and the restraints in one submission with one wait

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)
This commit is contained in:
leonarski_f committed 2026-10-26 01:38:23 +01:00
1 parent 001e61c1a4
commit 5569f3d1f2
6 files changed
+175 -66

No files matched your search

@@ -11,6 +11,7 @@
#include <array>
#include <cstddef>
#include <functional>
#include <memory>
#include <vector>
@@ -71,6 +72,11 @@ public:
// 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;