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).