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)

image_analysis/structure_refinement/RefinementGPU.cpp

@@ -209,6 +209,33 @@ public:
         return value;
     }
 
+    double EvaluateWithRestraints(const gemmi::Model &model, int which, std::vector<std::array<double, 4>> *grad,
+                                  std::vector<std::array<double, 4>> *restraint_grad, double &restraints) override {
+        SetAtoms(model);
+        std::vector<std::array<float, 4>> g, gr;
+        double value = 0;
+        {
+            const RefinementStage stage("gpu.evaluate");
+            value = engine_->EvaluateWithRestraints(grad != nullptr ? &g : nullptr, which, [&](std::array<float, 4> *x) {
+                for (size_t i = 0; i < atom_ptrs_.size(); i++) {
+                    const gemmi::Atom &a = *atom_ptrs_[i];
+                    x[i] = {static_cast<float>(a.pos.x), static_cast<float>(a.pos.y), static_cast<float>(a.pos.z), a.b_iso};
+                }
+            }, restraint_grad != nullptr ? &gr : nullptr, restraints);
+        }
+        if (grad != nullptr) {
+            grad->resize(g.size());
+            for (size_t i = 0; i < g.size(); i++)
+                (*grad)[i] = {g[i][0], g[i][1], g[i][2], g[i][3]};
+        }
+        if (restraint_grad != nullptr) {
+            restraint_grad->resize(gr.size());
+            for (size_t i = 0; i < gr.size(); i++)
+                (*restraint_grad)[i] = {gr[i][0], gr[i][1], gr[i][2], gr[i][3]};
+        }
+        return value;
+    }
+
     void SetRestraints(int which, const RestraintSet &set) override {
         engine_->SetRestraints(which, set.records, set.plane_atoms, set.slot_start, set.slot_index, set.n_slots);
     }
@@ -247,7 +274,8 @@ private:
     // atoms whose B moved).
     void SetAtoms(const gemmi::Model &model) {
         std::optional<RefinementStage> stage("gpu.atoms_host");
-        const std::vector<const gemmi::Atom *> atoms = Atoms(model);
+        atom_ptrs_ = Atoms(model);
+        const std::vector<const gemmi::Atom *> &atoms = atom_ptrs_;
         const int n = static_cast<int>(atoms.size());
         if (last_b_.size() != atoms.size()) {
             last_b_.assign(atoms.size(), NAN);
@@ -284,6 +312,7 @@ private:
     bool supported_ = false;
     size_t fft_work_ = 0, bytes_ = 0;
     std::unique_ptr<RefinementGPUEngine> engine_;
+    std::vector<const gemmi::Atom *> atom_ptrs_;   // the model's atoms at the last SetAtoms()
     std::vector<float> last_b_;
     std::vector<ModelDensityAtom> atoms_;
 };