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

@@ -259,18 +259,17 @@ public:
ctx_.MoveRiders();
ctx_.evaluations++;
std::vector<std::array<double, 4>> gx, gr;
double xray = 0;
double xray = 0, restraints = 0;
{
const RefinementStage s1("eval.xray");
xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr);
const RefinementStage s1("eval.engine");
if (restraint_set_ >= 0)
xray = ctx_.engine->EvaluateWithRestraints(ctx_.Model(), restraint_set_, gradient != nullptr ? &gx : nullptr,
gradient != nullptr ? &gr : nullptr, restraints);
else
xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr);
}
if (gradient != nullptr)
ctx_.FoldRiders(gx);
double restraints = 0;
if (restraint_set_ >= 0) {
const RefinementStage s2("eval.restraints");
restraints = ctx_.engine->EvaluateRestraints(restraint_set_, ctx_.XyzB(), gradient != nullptr ? &gr : nullptr);
}
*cost = weight_ * xray + restraints;
if (gradient != nullptr)
for (size_t k = 0; k < param_atoms_.size(); k++) {