Refinement: the per-evaluation host chain on the threads and in place

Byte-identical: every pass is element by element in the serial loop's order, or a copy that is no longer made.

- The engine's gradients are read in place from its pinned buffers and widened to double on the threads
  (RefinementGPUEngine::HostGradient/HostRestraintGradient); the positions go up from a pinned vector of the
  engine's host side; the model's atom pointers are kept between calls (the model is the context's own).
- Context: the riders moved and folded on the threads - a parent takes its riders in index order through a
  CSR list (SetRiderParents), as the serial walk added them.
- AtomTarget: the parameters set and the gradient assembled on the threads; the per-atom gradient vectors kept
  between evaluations (7 MB each on a large model: made anew, their pages fault in every time).
- RefinementMinimiser: the sweeps over the parameters (x + alpha d, s, y) on the threads it is given, the dot
  products the serial chains they were; Direction()'s working vector and a let-go pair's vectors reused.

A/B against the frozen combined2 binary, one lock hold each, outputs identical but provenance lines:
cytc refinement 1.600 -> 1.548 s (live), 1.724 -> 1.668 (CLI); thau 1.906 -> 1.885, 2.031 -> 2.027;
4nwv (with the commit before) 39.9 -> 35.5 s, the ladder 14.6 -> 12.7, each step 8.3 -> 7.2.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu

image_analysis/structure_refinement/Refinement.cpp

@@ -28,6 +28,7 @@
 #include "../../common/JFJochException.h"
 #include "../../common/JFJochMath.h"
 #include "../../common/Logger.h"
+#include "../../common/ParallelFor.h"
 #include "ModelFFT.h"
 #include "ModelScaling.h"
 #include "RefinementMath.h"
@@ -137,6 +138,7 @@ struct Context {
     double b_ceiling = B_CEILING_MAX;
     // Riding hydrogens: per atom the parent it rides on (-1: none), and where it sits from the parent.
     std::vector<int> rider_parent;
+    std::vector<int> rider_start, rider_index;   // SetRiderParents()
     std::vector<gemmi::Position> rider_offset;
     int evaluations = 0;
     const std::atomic<bool> *cancel = nullptr;   // RefinementSettings::cancel
@@ -145,20 +147,42 @@ struct Context {
 
     gemmi::Model &Model() { return st.models[0]; }
 
-    // The hydrogens moved with their parents, at the offsets they were last placed at, and given their B.
-    void MoveRiders() {
-        for (size_t i = 0; i < rider_parent.size(); i++)
-            if (rider_parent[i] >= 0) {
-                atoms[i]->pos = atoms[rider_parent[i]]->pos + rider_offset[i];
-                atoms[i]->b_iso = atoms[rider_parent[i]]->b_iso;
-            }
-    }
-    // What the target says of a hydrogen is said of its parent.
-    void FoldRiders(std::vector<std::array<double, 4>> &g) const {
+    // The riders' parents, and per parent its riders in index order (rider_start, rider_index: CSR over the atoms).
+    void SetRiderParents(const std::vector<int> &parents) {
+        rider_parent = parents;
+        rider_start.assign(atoms.size() + 1, 0);
+        for (int p : rider_parent)
+            if (p >= 0)
+                rider_start[p + 1]++;
+        for (size_t i = 0; i < atoms.size(); i++)
+            rider_start[i + 1] += rider_start[i];
+        rider_index.resize(rider_start.back());
+        std::vector<int> fill(rider_start.begin(), rider_start.end() - 1);
         for (size_t i = 0; i < rider_parent.size(); i++)
             if (rider_parent[i] >= 0)
-                for (int c = 0; c < 4; c++)
-                    g[rider_parent[i]][c] += g[i][c];
+                rider_index[fill[rider_parent[i]]++] = static_cast<int>(i);
+    }
+    // The hydrogens moved with their parents, at the offsets they were last placed at, and given their B.
+    void MoveRiders() {
+        const int n = static_cast<int>(rider_parent.size());
+        ParallelChunks(n, ThreadsForWork(n, nthreads, 16384), [&](int lo, int hi) {
+            for (int i = lo; i < hi; i++)
+                if (rider_parent[i] >= 0) {
+                    atoms[i]->pos = atoms[rider_parent[i]]->pos + rider_offset[i];
+                    atoms[i]->b_iso = atoms[rider_parent[i]]->b_iso;
+                }
+        });
+    }
+    // What the target says of a hydrogen is said of its parent: each parent takes its riders in index order, as
+    // the serial walk over the atoms added them (a rider has no riders of its own).
+    void FoldRiders(std::vector<std::array<double, 4>> &g) const {
+        const int n = static_cast<int>(rider_parent.size());
+        ParallelChunks(n, ThreadsForWork(n, nthreads, 16384), [&](int lo, int hi) {
+            for (int i = lo; i < hi; i++)
+                for (int r = rider_start[i]; r < rider_start[i + 1]; r++)
+                    for (int c = 0; c < 4; c++)
+                        g[i][c] += g[rider_index[r]][c];
+        });
     }
 
     std::vector<std::array<double, 4>> XyzB() const {
@@ -509,13 +533,16 @@ public:
     }
     // The atoms set to the parameters, their riders with them.
     void Set(const double *p) const {
-        for (size_t k = 0; k < param_atoms_.size(); k++) {
-            gemmi::Atom &a = *ctx_.atoms[param_atoms_[k]];
-            if (mode_ == Mode::Xyz)
-                a.pos = gemmi::Position(p[3 * k], p[3 * k + 1], p[3 * k + 2]);
-            else
-                a.b_iso = static_cast<float>(std::clamp(p[k], ctx_.b_floor, ctx_.b_ceiling));
-        }
+        const int n = static_cast<int>(param_atoms_.size());
+        ParallelChunks(n, ThreadsForWork(n, ctx_.nthreads, 16384), [&](int lo, int hi) {
+            for (int k = lo; k < hi; k++) {
+                gemmi::Atom &a = *ctx_.atoms[param_atoms_[k]];
+                if (mode_ == Mode::Xyz)
+                    a.pos = gemmi::Position(p[3 * k], p[3 * k + 1], p[3 * k + 2]);
+                else
+                    a.b_iso = static_cast<float>(std::clamp(p[k], ctx_.b_floor, ctx_.b_ceiling));
+            }
+        });
         ctx_.MoveRiders();
     }
     // the most a parameter moves in the first trial of a line search
@@ -525,7 +552,7 @@ public:
         const RefinementStage stage(mode_ == Mode::Xyz ? "eval.xyz" : "eval.b");
         Set(p);
         ctx_.evaluations++;
-        std::vector<std::array<double, 4>> gx, gr;
+        std::vector<std::array<double, 4>> &gx = gx_, &gr = gr_;
         double xray = 0, restraints = 0;
         {
             const RefinementStage s1("eval.engine");
@@ -536,18 +563,21 @@ public:
         if (gradient == nullptr)
             return std::isfinite(*cost);
         ctx_.FoldRiders(gx);
-        for (size_t k = 0; k < param_atoms_.size(); k++) {
-            const int i = param_atoms_[k];
-            if (mode_ == Mode::Xyz)
-                for (int c = 0; c < 3; c++)
-                    gradient[3 * k + c] = weight_ * gx[i][c] + gr[i][c];
-            else {
-                // held at a bound, B does not move further out: by the whole target, restraints included
-                const double g = weight_ * gx[i][3] + gr[i][3];
-                const bool held = (p[k] <= ctx_.b_floor && g > 0) || (p[k] >= ctx_.b_ceiling && g < 0);
-                gradient[k] = held ? 0.0 : g;
+        const int n = static_cast<int>(param_atoms_.size());
+        ParallelChunks(n, ThreadsForWork(n, ctx_.nthreads, 16384), [&](int lo, int hi) {
+            for (int k = lo; k < hi; k++) {
+                const int i = param_atoms_[k];
+                if (mode_ == Mode::Xyz)
+                    for (int c = 0; c < 3; c++)
+                        gradient[3 * k + c] = weight_ * gx[i][c] + gr[i][c];
+                else {
+                    // held at a bound, B does not move further out: by the whole target, restraints included
+                    const double g = weight_ * gx[i][3] + gr[i][3];
+                    const bool held = (p[k] <= ctx_.b_floor && g > 0) || (p[k] >= ctx_.b_ceiling && g < 0);
+                    gradient[k] = held ? 0.0 : g;
+                }
             }
-        }
+        });
         return std::isfinite(*cost);
     }
 
@@ -557,6 +587,9 @@ private:
     const std::vector<int> &param_atoms_;
     double weight_;
     int restraint_set_;
+    // the gradients per atom, kept from evaluation to evaluation (7 MB each on a large model: made anew, their
+    // pages fault in every time)
+    mutable std::vector<std::array<double, 4>> gx_, gr_;
 };
 
 // The minimiser's diagonal for the B factors of `param_atoms` (RefinementMinimiser::SetDiagonal): weight times the
@@ -998,7 +1031,7 @@ std::unique_ptr<OwnSlot> MakeOwnSlot(const gemmi::Structure &before_topology, co
     ctx.mask = std::make_unique<BulkSolventMask>(ctx.cell, *ctx.sg, ctx.d_min, nthreads);
     slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
     ctx.atoms = AtomList(ctx.Model());
-    ctx.rider_parent = slot->topo.rider_parent;
+    ctx.SetRiderParents(slot->topo.rider_parent);
 #ifdef JFJOCH_USE_CUDA
     if (main.engine->OnGPU()) {
         try {
@@ -1077,7 +1110,7 @@ TrialOutcome RunTrial(TrialSlot slot, const std::vector<int> &refinable, const s
     ctx.SetB(b0);
     // the positions and the B factors in turn, each with a memory of its own kept from macro-cycle to macro-cycle
     const AtomTarget xyz(ctx, AtomTarget::Mode::Xyz, refinable, w, 0), b(ctx, AtomTarget::Mode::B, refinable, adp_w, 1);
-    RefinementMinimiser xyz_minimiser, b_minimiser;
+    RefinementMinimiser xyz_minimiser(20, ctx.nthreads), b_minimiser(20, ctx.nthreads);
     double xray_scale = 0;
     for (int cycle = 0; cycle < cycles && !ctx.Cancelled(); cycle++) {
         PlaceRiders(ctx, topo);
@@ -1241,7 +1274,7 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector<M
         topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens);
         topo_stage.reset();
         ctx.atoms = AtomList(ctx.Model());
-        ctx.rider_parent = topo.rider_parent;
+        ctx.SetRiderParents(topo.rider_parent);
         PlaceRiders(ctx, topo);
     }
 

image_analysis/structure_refinement/RefinementGPU.cpp

@@ -46,6 +46,31 @@ std::vector<const gemmi::Atom *> Atoms(const gemmi::Model &model) {
     return atoms;
 }
 
+// Whether `atoms` is Atoms(model) still: as many atoms, the first and the last the same ones.
+bool SameAtoms(const gemmi::Model &model, const std::vector<const gemmi::Atom *> &atoms) {
+    size_t n = 0;
+    const gemmi::Atom *first = nullptr, *last = nullptr;
+    for (const gemmi::Chain &ch : model.chains)
+        for (const gemmi::Residue &r : ch.residues)
+            if (!r.atoms.empty()) {
+                if (first == nullptr)
+                    first = &r.atoms.front();
+                last = &r.atoms.back();
+                n += r.atoms.size();
+            }
+    return n > 0 && n == atoms.size() && first == atoms.front() && last == atoms.back();
+}
+
+// The engine's single-precision gradient into the caller's doubles, element by element in parallel: bit for bit
+// the serial loop's.
+void Widen(const std::array<float, 4> *from, size_t n, std::vector<std::array<double, 4>> &to, size_t nthreads) {
+    to.resize(n);
+    ParallelChunks(static_cast<int>(n), ThreadsForWork(n, nthreads, 16384), [&](int lo, int hi) {
+        for (int i = lo; i < hi; i++)
+            to[i] = {from[i][0], from[i][1], from[i][2], from[i][3]};
+    });
+}
+
 class RefinementEngineGPU : public RefinementEngine {
 public:
     RefinementEngineGPU(int device, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
@@ -201,44 +226,33 @@ public:
 
     double Evaluate(const gemmi::Model &model, std::vector<std::array<double, 4>> *grad) override {
         SetAtoms(model);
-        std::vector<std::array<float, 4>> g;
         double value = 0;
         {
             const RefinementStage stage("gpu.evaluate");
-            value = engine_->Evaluate(grad != nullptr ? &g : nullptr);
-        }
-        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]};
+            value = engine_->Evaluate(grad != nullptr);
         }
+        if (grad != nullptr)
+            Widen(engine_->HostGradient(), atom_ptrs_.size(), *grad, nthreads_);
         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) {
+            value = engine_->EvaluateWithRestraints(grad != 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]};
+            }, restraint_grad != nullptr, restraints);
         }
+        if (grad != nullptr)
+            Widen(engine_->HostGradient(), atom_ptrs_.size(), *grad, nthreads_);
+        if (restraint_grad != nullptr)
+            Widen(engine_->HostRestraintGradient(), atom_ptrs_.size(), *restraint_grad, nthreads_);
         return value;
     }
 
@@ -280,32 +294,34 @@ private:
     // atoms whose B moved).
     void SetAtoms(const gemmi::Model &model) {
         std::optional<RefinementStage> stage("gpu.atoms_host");
-        atom_ptrs_ = Atoms(model);
+        // the pointers kept from the last call where the model is the same one still (its context's own)
+        if (!SameAtoms(model, atom_ptrs_))
+            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);
             atoms_.resize(atoms.size());
+            pos_.resize(atoms.size());
         }
-        std::vector<std::array<float, 4>> pos(atoms.size());
-        std::vector<char> changed(atoms.size(), 0);
+        changed_.assign(atoms.size(), 0);
         const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw};
         ParallelChunks(n, ThreadsForWork(atoms.size(), nthreads_, 4096), [&](int lo, int hi) {
             for (int i = lo; i < hi; i++) {
                 if (!(atoms[i]->b_iso == last_b_[i])) {
                     atoms_[i] = DensityAtom(*atoms[i], atoms[i]->b_iso, dc_);
                     last_b_[i] = atoms[i]->b_iso;
-                    changed[i] = 1;
+                    changed_[i] = 1;
                 }
                 const gemmi::Fractional f = cell_.fractionalize(atoms[i]->pos);
-                pos[i] = {static_cast<float>(f.x * n3[0]), static_cast<float>(f.y * n3[1]), static_cast<float>(f.z * n3[2]), 0};
+                pos_[i] = {static_cast<float>(f.x * n3[0]), static_cast<float>(f.y * n3[1]), static_cast<float>(f.z * n3[2]), 0};
             }
         });
         stage.reset();
         const RefinementStage upload("gpu.upload");
-        if (std::find(changed.begin(), changed.end(), 1) != changed.end())
+        if (std::find(changed_.begin(), changed_.end(), 1) != changed_.end())
             engine_->SetDensities(atoms_.data(), atoms_.size());
-        engine_->SetPositions(pos);
+        engine_->SetPositions(pos_.data(), pos_.size());
     }
 
     int device_;
@@ -321,6 +337,8 @@ private:
     std::vector<const gemmi::Atom *> atom_ptrs_;   // the model's atoms at the last SetAtoms()
     std::vector<float> last_b_;
     PinnedVector<ModelDensityAtom> atoms_;   // pinned: a B rung sends them every evaluation (37 MB on a large model)
+    PinnedVector<std::array<float, 4>> pos_;
+    std::vector<char> changed_;
 };
 
 }  // namespace

image_analysis/structure_refinement/RefinementGPU.cu

@@ -427,7 +427,6 @@ struct RefinementGPUEngineImpl {
     int n_atoms = 0;
 
     CudaDevicePtr<float4> pos;
-    CudaHostPtr<float4> host_pos;   // the positions' staging buffer, pinned
     // What comes back, pinned, so that a copy back queues without waiting: the X-ray target and the restraints'
     // sums, their gradients, and the restraints' positions going up.
     CudaHostPtr<double> host_sums;
@@ -462,7 +461,6 @@ struct RefinementGPUEngineImpl {
         const auto sync = CudaAlloc::Synchronous;
         const size_t na = std::max<size_t>(setup.max_atoms, 1), nr = std::max<size_t>(setup.rows.size(), 1);
         pos = CudaDevicePtr<float4>(na, sync);
-        host_pos = CudaHostPtr<float4>(na);
         host_sums = CudaHostPtr<double>(2);
         host_grad = CudaHostPtr<float4>(na);
         host_xyzb = CudaHostPtr<float4>(na);
@@ -731,16 +729,13 @@ void RefinementGPUEngine::SetDensities(const ModelDensityAtom *atoms, size_t n)
     e.n_atoms = static_cast<int>(n);
 }
 
-void RefinementGPUEngine::SetPositions(const std::vector<std::array<float, 4>> &pos) {
+void RefinementGPUEngine::SetPositions(const std::array<float, 4> *pos, size_t n) {
     RefinementGPUEngineImpl &e = *impl_;
     DeviceGuard guard(e.device);
-    if (static_cast<int>(pos.size()) != e.n_atoms)
+    if (static_cast<int>(n) != e.n_atoms)
         throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: positions do not match the atoms");
-    // the staging buffer free again: the last copy out of it done
-    cuda_err(e.waiter.Wait(e.stream));
-    std::copy(pos.begin(), pos.end(), reinterpret_cast<std::array<float, 4> *>(e.host_pos.get()));
-    if (!pos.empty())
-        cuda_err(cudaMemcpyAsync(e.pos, e.host_pos, pos.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream));
+    if (n > 0)
+        cuda_err(cudaMemcpyAsync(e.pos, pos, n * sizeof(float4), cudaMemcpyHostToDevice, e.stream));
 }
 
 void RefinementGPUEngine::Fcalc(std::vector<std::array<float, 2>> &f) {
@@ -779,33 +774,34 @@ void RefinementGPUEngine::Fmask(const std::vector<ModelMaskAtom> &mask_atoms, st
     cuda_err(e.waiter.Wait(e.stream));
 }
 
-double RefinementGPUEngine::Evaluate(std::vector<std::array<float, 4>> *grad) {
+double RefinementGPUEngine::Evaluate(bool grad) {
     RefinementGPUEngineImpl &e = *impl_;
     DeviceGuard guard(e.device);
-    e.QueueEvaluate(grad != nullptr);
+    e.QueueEvaluate(grad);
     cuda_err(e.waiter.Wait(e.stream));
-    if (grad != nullptr)
-        CopyOut(e.host_grad, e.n_atoms, *grad);
     return e.host_sums[0];
 }
 
-double RefinementGPUEngine::EvaluateWithRestraints(std::vector<std::array<float, 4>> *grad, int which,
+const std::array<float, 4> *RefinementGPUEngine::HostGradient() const {
+    return reinterpret_cast<const std::array<float, 4> *>(impl_->host_grad.get());
+}
+
+const std::array<float, 4> *RefinementGPUEngine::HostRestraintGradient() const {
+    return reinterpret_cast<const std::array<float, 4> *>(impl_->host_restraint_grad.get());
+}
+
+double RefinementGPUEngine::EvaluateWithRestraints(bool grad, int which,
                                                    const std::function<void(std::array<float, 4> *)> &xyzb,
-                                                   std::vector<std::array<float, 4>> *restraint_grad,
-                                                   double &restraints) {
+                                                   bool restraint_grad, double &restraints) {
     RefinementGPUEngineImpl &e = *impl_;
     DeviceGuard guard(e.device);
     RefinementGPUEngineImpl::Restraints &r = e.restraints[which];
     if (r.atoms != e.n_atoms)
         throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: restraint atoms do not match");
-    e.QueueEvaluate(grad != nullptr);
+    e.QueueEvaluate(grad);
     xyzb(reinterpret_cast<std::array<float, 4> *>(e.host_xyzb.get()));
-    e.QueueRestraints(which, restraint_grad != nullptr);
+    e.QueueRestraints(which, restraint_grad);
     cuda_err(e.waiter.Wait(e.stream));
-    if (grad != nullptr)
-        CopyOut(e.host_grad, e.n_atoms, *grad);
-    if (restraint_grad != nullptr)
-        CopyOut(e.host_restraint_grad, r.atoms, *restraint_grad);
     restraints = e.host_sums[1];
     return e.host_sums[0];
 }

image_analysis/structure_refinement/RefinementGPUEngine.h

@@ -59,14 +59,19 @@ public:
     // from them): called again only where some atom's density changed. Read until the call returns; from pinned
     // memory they go up by DMA.
     void SetDensities(const ModelDensityAtom *atoms, size_t n);
-    // 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);
+    // The atoms' grid positions (fractional times the grid size, not wrapped), in the order of SetDensities():
+    // read until the next evaluation returns, so from a pinned buffer the copy goes up by DMA.
+    void SetPositions(const std::array<float, 4> *pos, size_t n);
     // 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);
+    // -ln L over the working rows for the atoms last set and, with `grad`, dT/d(x, y, z, B) per atom into
+    // HostGradient().
+    double Evaluate(bool grad);
+    // The gradients of the last evaluation, in the engine's pinned buffers: valid until the next one.
+    const std::array<float, 4> *HostGradient() const;
+    const std::array<float, 4> *HostRestraintGradient() const;
 
     // 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,
@@ -75,10 +80,10 @@ public:
     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);
+    // `xyzb` writes the restraints' x, y, z, B of every atom while the device works on it. The gradients into
+    // HostGradient() and HostRestraintGradient().
+    double EvaluateWithRestraints(bool grad, int which, const std::function<void(std::array<float, 4> *)> &xyzb,
+                                  bool restraint_grad, double &restraints);
 
     // The most atoms of a plane restraint the kernel takes.
     static constexpr int MAX_PLANE_ATOMS = 32;

image_analysis/structure_refinement/RefinementMinimiser.h

@@ -9,6 +9,8 @@
 #include <functional>
 #include <vector>
 
+#include "../../common/ParallelFor.h"
+
 // Limited-memory BFGS whose memory outlives one call, so that a refinement's macro-cycles - each a few dozen
 // iterations on a target whose scale and likelihood parameters were refitted in between - go on from what the
 // earlier ones learnt of the curvature instead of starting each from a steepest-descent step. The initial inverse
@@ -21,7 +23,8 @@ public:
     // f(x, g): the target at x, and its gradient in g (sized as x); false where it cannot be evaluated there.
     using Target = std::function<bool(const std::vector<double> &x, double &f, std::vector<double> &g)>;
 
-    explicit RefinementMinimiser(size_t memory = 20) : memory_(memory) {}
+    // `nthreads`: the threads the sweeps over the parameters run on (the dot products stay serial chains).
+    explicit RefinementMinimiser(size_t memory = 20, size_t nthreads = 1) : memory_(memory), nthreads_(nthreads) {}
 
     // Forget the pairs (the target changed beyond what the memory still describes, e.g. another weight).
     void Reset() {
@@ -72,8 +75,7 @@ public:
             double f_new = 0;
             bool accepted = false;
             for (int ls = 0; ls < 10; ls++) {
-                for (size_t k = 0; k < n; k++)
-                    x_new[k] = x[k] + alpha * d[k];
+                Sweep(n, [&](size_t k) { x_new[k] = x[k] + alpha * d[k]; });
                 const bool ok = target(x_new, f_new, g_new);
                 out.evaluations++;
                 if (ok && f_new <= f + 1e-4 * alpha * gd) {
@@ -86,21 +88,29 @@ public:
             }
             if (!accepted)
                 break;
-            std::vector<double> s(n), y(n);
-            for (size_t k = 0; k < n; k++) {
+            // the pair's vectors are the ones the memory last let go, where it did: made anew, their pages fault in
+            std::vector<double> s = std::move(spare_s_), y = std::move(spare_y_);
+            s.resize(n);
+            y.resize(n);
+            Sweep(n, [&](size_t k) {
                 s[k] = x_new[k] - x[k];
                 y[k] = g_new[k] - g[k];
-            }
+            });
             const double sy = Dot(s, y);
             if (sy > 1e-10 * std::sqrt(Dot(s, s) * Dot(y, y))) {
                 s_.push_back(std::move(s));
                 y_.push_back(std::move(y));
                 sy_.push_back(sy);
                 if (s_.size() > memory_) {
+                    spare_s_ = std::move(s_.front());
+                    spare_y_ = std::move(y_.front());
                     s_.pop_front();
                     y_.pop_front();
                     sy_.pop_front();
                 }
+            } else {
+                spare_s_ = std::move(s);
+                spare_y_ = std::move(y);
             }
             const double decrease = f - f_new;
             x.swap(x_new);
@@ -122,13 +132,24 @@ private:
         return s;
     }
 
+    // fn(k) for every k below n on the threads, each element on its own: the serial loop's bits.
+    template <typename Fn>
+    void Sweep(size_t n, Fn fn) const {
+        ParallelChunks(static_cast<int>(n), ThreadsForWork(n, nthreads_), [&](int lo, int hi) {
+            for (int k = lo; k < hi; k++)
+                fn(static_cast<size_t>(k));
+        });
+    }
+
     // d = -H g by the two-loop recursion, H0 = gamma diag(1/h). Every dot product is a serial chain of n
     // additions, so each sweep over the parameters takes the next pair's dot product while it applies the
     // previous pair's update to q - the same operations on each element in the same order as a sweep for
     // each would do, with the update's memory traffic hidden under the chain.
     void Direction(const std::vector<double> &g, std::vector<double> &d) const {
         const size_t n = g.size(), m = s_.size();
-        std::vector<double> q = g, a(m), rho(m);
+        std::vector<double> &q = q_;   // kept from call to call: made anew, its pages fault in every time
+        q.assign(g.begin(), g.end());
+        std::vector<double> a(m), rho(m);
         for (size_t i = 0; i < m; i++)
             rho[i] = 1.0 / sy_[i];
         double dot = m > 0 ? Dot(s_[m - 1], q) : 0.0, yhy = 0.0;
@@ -181,8 +202,10 @@ private:
             d[k] = -q[k];
     }
 
-    size_t memory_;
+    size_t memory_, nthreads_;
     std::deque<std::vector<double>> s_, y_;
     std::deque<double> sy_;   // s.y of each pair, as taken when it was kept
+    std::vector<double> spare_s_, spare_y_;   // the pair the memory last let go, for the next one to reuse
+    mutable std::vector<double> q_;           // Direction()'s working vector
     std::vector<double> h_;
 };