From e3782cabe4730cd2d4e73f2f47b87e0524a90c24 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 25 Oct 2026 13:30:59 +0100 Subject: [PATCH] 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 Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 46c3b67c91aee0aec3473b0fb1a4945c00bda070) (cherry picked from commit 0aebc98110c8155932af647d6204891aec662294) --- .../structure_refinement/ModelDensityGPU.h | 4 +++ .../structure_refinement/Refinement.cpp | 2 +- .../structure_refinement/RefinementGPU.cpp | 36 ++++++++++++------- .../structure_refinement/RefinementGPU.cu | 31 ++++++++-------- .../structure_refinement/RefinementGPU.h | 4 +-- .../RefinementGPUEngine.h | 8 +++-- 6 files changed, 52 insertions(+), 33 deletions(-) diff --git a/image_analysis/structure_refinement/ModelDensityGPU.h b/image_analysis/structure_refinement/ModelDensityGPU.h index d46941ab7..a40270627 100644 --- a/image_analysis/structure_refinement/ModelDensityGPU.h +++ b/image_analysis/structure_refinement/ModelDensityGPU.h @@ -66,6 +66,10 @@ public: // for the number of (brick, atom) pairs. void Compute(const float4 *d_pos, float *d_grid); + // The densities and the boxes of the atoms last set, on the device, for a kernel that walks the same boxes. + const ModelDensityAtom *DeviceAtoms() const { return atoms; } + const int *DeviceBoxes() const { return box; } + // What the gather kernel needs of the grid. struct Geometry { int nu, nv, nw; diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index 560d3ff93..070c38d85 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -611,7 +611,7 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg const size_t budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0; size_t bytes = 0; ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(), - ctx.b_ceiling, budget, &bytes); + ctx.b_ceiling, budget, &bytes, settings.nthreads); if (ctx.engine) ctx.scale_gpu = std::make_unique(device, ctx.refl.Size(), 1); else diff --git a/image_analysis/structure_refinement/RefinementGPU.cpp b/image_analysis/structure_refinement/RefinementGPU.cpp index 25afc251c..bcac87259 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cpp +++ b/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 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 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 stage("gpu.atoms_host"); const std::vector atoms = Atoms(model); + const int n = static_cast(atoms.size()); if (last_b_.size() != atoms.size()) { last_b_.assign(atoms.size(), NAN); atoms_.resize(atoms.size()); } std::vector> pos(atoms.size()); + std::vector 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(f.x * n3[0]), static_cast(f.y * n3[1]), static_cast(f.z * n3[2]), 0}; } - 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}; - } + }); 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 dc_; bool supported_ = false; @@ -284,8 +294,8 @@ std::unique_ptr 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(device, cell, sg, d_min, refl, model, b_max); + size_t *device_bytes, size_t nthreads) { + auto engine = std::make_unique(device, cell, sg, d_min, refl, model, b_max, nthreads); if (device_bytes != nullptr) *device_bytes = engine->Bytes(); if (!engine->Supported() || engine->Bytes() > budget) diff --git a/image_analysis/structure_refinement/RefinementGPU.cu b/image_analysis/structure_refinement/RefinementGPU.cu index 78d8b1e92..1c0c7407a 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cu +++ b/image_analysis/structure_refinement/RefinementGPU.cu @@ -363,7 +363,7 @@ struct RefinementGPUEngineImpl { int n_atoms = 0; CudaDevicePtr pos; - CudaDevicePtr box; + CudaHostPtr host_pos; // the positions' staging buffer, pinned CudaDevicePtr mask_atoms; CudaDevicePtr grid; CudaDevicePtr spectrum; @@ -378,7 +378,6 @@ struct RefinementGPUEngineImpl { size_t reduce_bytes = 0; CudaDevicePtr scatter_point, scatter_start, scatter_term; CudaDevicePtr grad; - CudaDevicePtr atoms_copy; // the density gather keeps its own; this one is the gradient's std::unique_ptr density; std::unique_ptr mask; CudaFFTPlan forward, inverse; @@ -389,9 +388,8 @@ 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); - box = CudaDevicePtr(3 * na, sync); + host_pos = CudaHostPtr(na); mask_atoms = CudaDevicePtr(na, sync); - atoms_copy = CudaDevicePtr(na, sync); grad = CudaDevicePtr(na, sync); grid = CudaDevicePtr(GridPoints(setup.grid), sync); spectrum = CudaDevicePtr(ComplexPoints(setup.grid), sync); @@ -536,21 +534,25 @@ void RefinementGPUEngine::SetRowModel(const std::vector &k_total, const s cuda_err(cudaStreamSynchronize(e.stream)); } -void RefinementGPUEngine::SetAtoms(const std::vector &atoms, const std::vector> &pos) { +void RefinementGPUEngine::SetDensities(const std::vector &atoms) { RefinementGPUEngineImpl &e = *impl_; DeviceGuard guard(e.device); - if (atoms.size() > e.setup.max_atoms || pos.size() != atoms.size()) + if (atoms.size() > e.setup.max_atoms) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: atoms over the reserve"); e.density->SetAtoms(e.setup.grid, atoms); e.n_atoms = static_cast(atoms.size()); - const std::vector b = ModelDensityGPU::Boxes(e.setup.grid, atoms); - if (!atoms.empty()) { - cuda_err(cudaMemcpyAsync(e.pos, pos.data(), pos.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream)); - cuda_err(cudaMemcpyAsync(e.atoms_copy, atoms.data(), atoms.size() * sizeof(ModelDensityAtom), - cudaMemcpyHostToDevice, e.stream)); - Upload(e.box, b, e.stream); - } +} + +void RefinementGPUEngine::SetPositions(const std::vector> &pos) { + RefinementGPUEngineImpl &e = *impl_; + DeviceGuard guard(e.device); + if (static_cast(pos.size()) != 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(cudaStreamSynchronize(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)); } void RefinementGPUEngine::Fcalc(std::vector> &f) { @@ -611,7 +613,8 @@ double RefinementGPUEngine::Evaluate(std::vector> *grad) { cufft_err(e.inverse.C2R(reinterpret_cast(e.spectrum.get()), e.grid)); if (e.n_atoms > 0) { const size_t threads = static_cast(e.n_atoms) * WARP; - atom_gradient_kernel<<>>(e.atoms_copy, e.pos, e.box, e.n_atoms, e.geom, + atom_gradient_kernel<<>>(e.density->DeviceAtoms(), e.pos, + e.density->DeviceBoxes(), e.n_atoms, e.geom, e.grid, e.Norm(), e.grad); cuda_err(cudaGetLastError()); } diff --git a/image_analysis/structure_refinement/RefinementGPU.h b/image_analysis/structure_refinement/RefinementGPU.h index 965fab1ad..764454679 100644 --- a/image_analysis/structure_refinement/RefinementGPU.h +++ b/image_analysis/structure_refinement/RefinementGPU.h @@ -16,9 +16,9 @@ // The engine for refining `model` against `refl` on `device`, or null where the GPU cannot take this case (an // atom's box wider than the cell, too many symmetry operators) or its memory is over `budget` bytes. // `b_max`: the largest B an atom may reach, which the gather's reserve is sized for. Throws JFJochException on a -// CUDA failure. +// CUDA failure. `nthreads`: the host threads that work out the atoms' densities. std::unique_ptr MakeRefinementEngineGPU(int device, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double d_min, const RefinementReflections &refl, const gemmi::Model &model, double b_max, size_t budget, - size_t *device_bytes = nullptr); + size_t *device_bytes = nullptr, size_t nthreads = 1); diff --git a/image_analysis/structure_refinement/RefinementGPUEngine.h b/image_analysis/structure_refinement/RefinementGPUEngine.h index c3281301b..cd67a2456 100644 --- a/image_analysis/structure_refinement/RefinementGPUEngine.h +++ b/image_analysis/structure_refinement/RefinementGPUEngine.h @@ -53,9 +53,11 @@ public: void SetRowModel(const std::vector &k_total, const std::vector> &fmask, const std::vector &d, const std::vector &var); - // The atoms, as ModelStructureFactorsGPU takes them: densities, grid positions (fractional times the grid - // size, not wrapped) and the boxes of ModelDensityGPU::Boxes(). - void SetAtoms(const std::vector &atoms, const std::vector> &pos); + // The atoms' densities, as ModelStructureFactorsGPU takes them (the boxes of ModelDensityGPU::Boxes() made + // from them): called again only where some atom's density changed. + void SetDensities(const std::vector &atoms); + // The atoms' grid positions (fractional times the grid size, not wrapped), in the order of SetDensities(). + void SetPositions(const std::vector> &pos); // 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).