Refinement GPU: atom densities uploaded only when some B moved, one device copy, positions through a pinned buffer

Every evaluation uploaded both copies of the atoms' densities (the gather's and the gradient's) and their
boxes, although in a positions-only minimisation no B changes between two Refits. The engine now keeps one
copy (the gradient kernel reads ModelDensityGPU's atoms and boxes), uploads it only where some atom's
density was made again, and sends the positions from a pinned buffer without waiting. The densities that
are made again (every atom in a B minimisation) and the grid positions are worked out on the refinement's
threads. Same bits.

9yl4: upload 3.29 -> 0.56 s, host side 1.29 -> 0.61 s, refinement 22.3 -> 17.7 s; 5reo 4.8 -> 4.4 s
(RTX 5080, -N 8, loaded box).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit 46c3b67c91aee0aec3473b0fb1a4945c00bda070)
(cherry picked from commit 0aebc98110c8155932af647d6204891aec662294)

image_analysis/structure_refinement/RefinementGPU.cpp

@@ -11,6 +11,7 @@
 #include "gemmi/it92.hpp"
 #include "gemmi/solmask.hpp"   // SolventMasker, refmac_radius_for_bulk_solvent
 
+#include "../../common/ParallelFor.h"
 #include "RefinementGPUEngine.h"
 
 namespace {
@@ -48,8 +49,8 @@ std::vector<const gemmi::Atom *> Atoms(const gemmi::Model &model) {
 class RefinementEngineGPU : public RefinementEngine {
 public:
     RefinementEngineGPU(int device, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min,
-                        const RefinementReflections &refl, const gemmi::Model &model, double b_max)
-        : device_(device), cell_(cell), sg_(sg), d_min_(d_min) {
+                        const RefinementReflections &refl, const gemmi::Model &model, double b_max, size_t nthreads)
+        : device_(device), cell_(cell), sg_(sg), d_min_(d_min), nthreads_(nthreads) {
         const gemmi::Grid<float> grid = RigidBodyZoneGrid(cell, sg, d_min);
         setup_.grid.nu = grid.nu;
         setup_.grid.nv = grid.nv;
@@ -242,33 +243,42 @@ private:
         return dc;
     }
 
-    // The atoms' densities (made again only for the atoms whose B moved) and grid positions, to the device.
+    // The atoms' grid positions to the device, and their densities where some B moved (made again only for the
+    // 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);
+        const int n = static_cast<int>(atoms.size());
         if (last_b_.size() != atoms.size()) {
             last_b_.assign(atoms.size(), NAN);
             atoms_.resize(atoms.size());
         }
         std::vector<std::array<float, 4>> pos(atoms.size());
+        std::vector<char> changed(atoms.size(), 0);
         const int n3[3] = {setup_.grid.nu, setup_.grid.nv, setup_.grid.nw};
-        for (size_t i = 0; i < atoms.size(); 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;
+        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;
+                }
+                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};
             }
-            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};
-        }
+        });
         stage.reset();
         const RefinementStage upload("gpu.upload");
-        engine_->SetAtoms(atoms_, pos);
+        if (std::find(changed.begin(), changed.end(), 1) != changed.end())
+            engine_->SetDensities(atoms_);
+        engine_->SetPositions(pos);
     }
 
     int device_;
     gemmi::UnitCell cell_;
     const gemmi::SpaceGroup &sg_;
     double d_min_;
+    size_t nthreads_;
     RefinementGPUSetup setup_;
     gemmi::DensityCalculator<Table, float> dc_;
     bool supported_ = false;
@@ -284,8 +294,8 @@ std::unique_ptr<RefinementEngine> MakeRefinementEngineGPU(int device, const gemm
                                                           const gemmi::SpaceGroup &sg, double d_min,
                                                           const RefinementReflections &refl,
                                                           const gemmi::Model &model, double b_max, size_t budget,
-                                                          size_t *device_bytes) {
-    auto engine = std::make_unique<RefinementEngineGPU>(device, cell, sg, d_min, refl, model, b_max);
+                                                          size_t *device_bytes, size_t nthreads) {
+    auto engine = std::make_unique<RefinementEngineGPU>(device, cell, sg, d_min, refl, model, b_max, nthreads);
     if (device_bytes != nullptr)
         *device_bytes = engine->Bytes();
     if (!engine->Supported() || engine->Bytes() > budget)