Model validation: scale fits on the GPU, change-of-basis probe gridded once per frame on all threads

The scale fit of every model-validation fit (fit_model: the real model, the indexing probe's relabellings
and every null replicate) was FitModelScale on the host - a long-double Levenberg-Marquardt per solvent grid
point, ~80% of the validation's CPU on a 400 A cell. It now runs on the GPU through ModelScaleGPU::FitSolvent,
the engine the rigid body and molecular replacement already use, from a small pool (ModelScaleGPUPool: up to
four engines with their own streams, sized on the card's total memory, at most an eighth of it). The grid,
the fits and the winner rule are FitModelScale's; the sums are taken in double on the device, so R-factors
agree with the host's to rounding. The CPU build, and a GPU build without a card, keep FitModelScale.

The change-of-basis probe gridded each candidate frame with gemmi's serial put_model_density_on_grid, once
for the data as merged and once more for each alternative indexing. It now grids each frame once, with
ModelGrid's parallel PutModelDensityOnGrid (the same grid, bit for bit), from threads outside the pool so the
gridding spreads over it, and scores every labelling on that F_calc. Exact.

4nwv (122k atoms, I222, 400 A cell): model validation 33.2 s -> 14.0 s; every logged R, sigma and decision
the same.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3

image_analysis/structure_refinement/ModelScaling.cpp

@@ -16,6 +16,13 @@
 
 #include "../../common/ParallelFor.h"
 
+#ifdef JFJOCH_USE_CUDA
+#include "../../common/CUDAWrapper.h"
+#include "../indexing/CUDAMemHelpers.h"
+#include "ModelScaleGPU.h"
+#include "RigidBodyGPUEngine.h"   // ModelScaleGPUTooFewReflections, RigidBodyGPUEngine::CurrentDevice
+#endif
+
 namespace {
 
 // gemmi::Scaling<float> as a grid point fits it: the same starting point, parameters, model values,
@@ -415,3 +422,102 @@ double SolventBMax(const gemmi::Model &model) {
     const double b_max = ModelScaleBox{}.b_hi;
     return n > 0 ? std::max(b_max, 2.0 * sum / static_cast<double>(n)) : b_max;
 }
+
+#ifdef JFJOCH_USE_CUDA
+struct ModelScaleGPUPool::Engine {
+    CudaStream stream;
+    ModelScaleGPU scale;
+    CudaDevicePtr<float2> fcalc, fmask;
+    explicit Engine(size_t max_points) : scale(stream, max_points), fcalc(max_points), fmask(max_points) {}
+};
+
+size_t ModelScaleGPUPool::DeviceBytes(size_t max_points, size_t engines) {
+    return engines * (ModelScaleGPU::DeviceBytes(max_points) + 2 * max_points * sizeof(float2));
+}
+
+ModelScaleGPUPool::ModelScaleGPUPool(int device, size_t max_points, size_t engines)
+    : device_(device), max_points_(max_points) {
+    const int current = RigidBodyGPUEngine::CurrentDevice();
+    set_gpu(device_);
+    for (size_t i = 0; i < engines; i++) {
+        engines_.push_back(std::make_unique<Engine>(max_points));
+        idle_.push_back(engines_.back().get());
+    }
+    set_gpu(current);
+}
+
+ModelScaleGPUPool::~ModelScaleGPUPool() {
+    const int current = RigidBodyGPUEngine::CurrentDevice();
+    set_gpu(device_);
+    engines_.clear();
+    set_gpu(current);
+}
+
+ModelScaleReport ModelScaleGPUPool::Fit(gemmi::Scaling<float> &scaling, double b_hi, size_t nthreads) {
+    const size_t n = scaling.points.size();
+    std::vector<std::array<int, 3>> hkl(n);
+    std::vector<double> stol2(n);
+    std::vector<float> fobs(n), sigma(n);
+    std::vector<float2> fc(n), fm(n);
+    for (size_t i = 0; i < n; i++) {
+        const auto &p = scaling.points[i];
+        hkl[i] = {p.hkl[0], p.hkl[1], p.hkl[2]};
+        stol2[i] = p.stol2;
+        fobs[i] = p.fobs;
+        sigma[i] = p.sigma;
+        fc[i] = {p.fcmol.real(), p.fcmol.imag()};
+        fm[i] = {p.fmask.real(), p.fmask.imag()};
+    }
+    std::vector<std::array<double, 6>> constraints;
+    for (const gemmi::Vec6 &c : scaling.constraint_matrix)
+        constraints.push_back(c);
+    double frac[9];
+    for (int i = 0; i < 3; i++)
+        for (int j = 0; j < 3; j++)
+            frac[3 * i + j] = scaling.cell.frac.mat[i][j];
+
+    Engine *e;
+    {
+        std::unique_lock lock(m_);
+        cv_.wait(lock, [&] { return !idle_.empty(); });
+        e = idle_.back();
+        idle_.pop_back();
+    }
+    auto release = [&] {
+        std::lock_guard lock(m_);
+        idle_.push_back(e);
+        cv_.notify_one();
+    };
+    const int current = RigidBodyGPUEngine::CurrentDevice();
+    ModelSolventFit fit;
+    try {
+        set_gpu(device_);
+        e->scale.SetPoints(hkl, stol2, fobs, sigma, constraints, frac);
+        if (cudaMemcpyAsync(e->fcalc, fc.data(), n * sizeof(float2), cudaMemcpyHostToDevice, e->stream) != cudaSuccess
+            || cudaMemcpyAsync(e->fmask, fm.data(), n * sizeof(float2), cudaMemcpyHostToDevice, e->stream) != cudaSuccess)
+            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "ModelScaleGPUPool: the upload failed");
+        fit = e->scale.FitSolvent(e->fcalc, e->fmask, b_hi);
+    } catch (const ModelScaleGPUTooFewReflections &) {
+        set_gpu(current);
+        release();
+        return FitModelScale(scaling, {.b_hi = b_hi}, nthreads);
+    } catch (...) {
+        set_gpu(current);
+        release();
+        throw;
+    }
+    set_gpu(current);
+    release();
+
+    scaling.k_sol = fit.k_sol;
+    scaling.b_sol = fit.b_sol;
+    scaling.k_overall = fit.scale.k_overall;
+    const double *b = fit.scale.b_star;
+    scaling.b_star = {b[0], b[1], b[2], b[3], b[4], b[5]};
+    ModelScaleReport report;
+    report.n_points = static_cast<int>(n);
+    report.n_grid = fit.n_grid;
+    report.r_work_fit = fit.r;
+    return report;
+}
+#endif