From d122118d7393dbccfc44b44eb12bb5a4be30608b Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 26 Oct 2026 22:28:10 +0100 Subject: [PATCH] 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 Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu --- .../structure_refinement/Refinement.cpp | 101 ++++++++++++------ .../structure_refinement/RefinementGPU.cpp | 72 ++++++++----- .../structure_refinement/RefinementGPU.cu | 40 ++++--- .../RefinementGPUEngine.h | 21 ++-- .../RefinementMinimiser.h | 39 +++++-- 5 files changed, 174 insertions(+), 99 deletions(-) diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index d2f5993b4..86a4d72a8 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/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 rider_parent; + std::vector rider_start, rider_index; // SetRiderParents() std::vector rider_offset; int evaluations = 0; const std::atomic *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> &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 &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 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(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(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> &g) const { + const int n = static_cast(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> 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(std::clamp(p[k], ctx_.b_floor, ctx_.b_ceiling)); - } + const int n = static_cast(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(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> gx, gr; + std::vector> &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(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 ¶m_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> 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 MakeOwnSlot(const gemmi::Structure &before_topology, co ctx.mask = std::make_unique(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 &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 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 &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 *from, size_t n, std::vector> &to, size_t nthreads) { + to.resize(n); + ParallelChunks(static_cast(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> *grad) override { SetAtoms(model); - std::vector> 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> *grad, std::vector> *restraint_grad, double &restraints) override { SetAtoms(model); - std::vector> g, gr; double value = 0; { const RefinementStage stage("gpu.evaluate"); - value = engine_->EvaluateWithRestraints(grad != nullptr ? &g : nullptr, which, [&](std::array *x) { + value = engine_->EvaluateWithRestraints(grad != nullptr, which, [&](std::array *x) { for (size_t i = 0; i < atom_ptrs_.size(); i++) { const gemmi::Atom &a = *atom_ptrs_[i]; x[i] = {static_cast(a.pos.x), static_cast(a.pos.y), static_cast(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 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 &atoms = atom_ptrs_; const int n = static_cast(atoms.size()); if (last_b_.size() != atoms.size()) { last_b_.assign(atoms.size(), NAN); atoms_.resize(atoms.size()); + pos_.resize(atoms.size()); } - std::vector> pos(atoms.size()); - std::vector 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(f.x * n3[0]), static_cast(f.y * n3[1]), static_cast(f.z * n3[2]), 0}; + pos_[i] = {static_cast(f.x * n3[0]), static_cast(f.y * n3[1]), static_cast(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 atom_ptrs_; // the model's atoms at the last SetAtoms() std::vector last_b_; PinnedVector atoms_; // pinned: a B rung sends them every evaluation (37 MB on a large model) + PinnedVector> pos_; + std::vector changed_; }; } // namespace diff --git a/image_analysis/structure_refinement/RefinementGPU.cu b/image_analysis/structure_refinement/RefinementGPU.cu index 14bcce975..c04fd5988 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cu +++ b/image_analysis/structure_refinement/RefinementGPU.cu @@ -427,7 +427,6 @@ struct RefinementGPUEngineImpl { int n_atoms = 0; CudaDevicePtr pos; - CudaHostPtr 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 host_sums; @@ -462,7 +461,6 @@ struct RefinementGPUEngineImpl { const auto sync = CudaAlloc::Synchronous; const size_t na = std::max(setup.max_atoms, 1), nr = std::max(setup.rows.size(), 1); pos = CudaDevicePtr(na, sync); - host_pos = CudaHostPtr(na); host_sums = CudaHostPtr(2); host_grad = CudaHostPtr(na); host_xyzb = CudaHostPtr(na); @@ -731,16 +729,13 @@ void RefinementGPUEngine::SetDensities(const ModelDensityAtom *atoms, size_t n) e.n_atoms = static_cast(n); } -void RefinementGPUEngine::SetPositions(const std::vector> &pos) { +void RefinementGPUEngine::SetPositions(const std::array *pos, size_t n) { RefinementGPUEngineImpl &e = *impl_; DeviceGuard guard(e.device); - if (static_cast(pos.size()) != e.n_atoms) + if (static_cast(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 *>(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> &f) { @@ -779,33 +774,34 @@ void RefinementGPUEngine::Fmask(const std::vector &mask_atoms, st cuda_err(e.waiter.Wait(e.stream)); } -double RefinementGPUEngine::Evaluate(std::vector> *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> *grad, int which, +const std::array *RefinementGPUEngine::HostGradient() const { + return reinterpret_cast *>(impl_->host_grad.get()); +} + +const std::array *RefinementGPUEngine::HostRestraintGradient() const { + return reinterpret_cast *>(impl_->host_restraint_grad.get()); +} + +double RefinementGPUEngine::EvaluateWithRestraints(bool grad, int which, const std::function *)> &xyzb, - std::vector> *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 *>(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]; } diff --git a/image_analysis/structure_refinement/RefinementGPUEngine.h b/image_analysis/structure_refinement/RefinementGPUEngine.h index a44d177f1..4b8c86d7c 100644 --- a/image_analysis/structure_refinement/RefinementGPUEngine.h +++ b/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> &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 *pos, size_t n); // F of the atoms at the rows. void Fcalc(std::vector> &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 &mask_atoms, std::vector> &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> *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 *HostGradient() const; + const std::array *HostRestraintGradient() const; // Restraint set `which` (0 or 1), as RestraintSet holds it. void SetRestraints(int which, const std::vector &records, const std::vector &plane_atoms, @@ -75,10 +80,10 @@ public: double EvaluateRestraints(int which, const std::vector> &xyzb, std::vector> *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> *grad, int which, - const std::function *)> &xyzb, - std::vector> *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 *)> &xyzb, + bool restraint_grad, double &restraints); // The most atoms of a plane restraint the kernel takes. static constexpr int MAX_PLANE_ATOMS = 32; diff --git a/image_analysis/structure_refinement/RefinementMinimiser.h b/image_analysis/structure_refinement/RefinementMinimiser.h index 84743d338..1a19b2677 100644 --- a/image_analysis/structure_refinement/RefinementMinimiser.h +++ b/image_analysis/structure_refinement/RefinementMinimiser.h @@ -9,6 +9,8 @@ #include #include +#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 &x, double &f, std::vector &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 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 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 + void Sweep(size_t n, Fn fn) const { + ParallelChunks(static_cast(n), ThreadsForWork(n, nthreads_), [&](int lo, int hi) { + for (int k = lo; k < hi; k++) + fn(static_cast(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 &g, std::vector &d) const { const size_t n = g.size(), m = s_.size(); - std::vector q = g, a(m), rho(m); + std::vector &q = q_; // kept from call to call: made anew, its pages fault in every time + q.assign(g.begin(), g.end()); + std::vector 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> s_, y_; std::deque sy_; // s.y of each pair, as taken when it was kept + std::vector spare_s_, spare_y_; // the pair the memory last let go, for the next one to reuse + mutable std::vector q_; // Direction()'s working vector std::vector h_; };