diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index 93d5383c1..b6b0b6d78 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -259,18 +259,17 @@ public: ctx_.MoveRiders(); ctx_.evaluations++; std::vector> gx, gr; - double xray = 0; + double xray = 0, restraints = 0; { - const RefinementStage s1("eval.xray"); - xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr); + const RefinementStage s1("eval.engine"); + if (restraint_set_ >= 0) + xray = ctx_.engine->EvaluateWithRestraints(ctx_.Model(), restraint_set_, gradient != nullptr ? &gx : nullptr, + gradient != nullptr ? &gr : nullptr, restraints); + else + xray = ctx_.engine->Evaluate(ctx_.Model(), gradient != nullptr ? &gx : nullptr); } if (gradient != nullptr) ctx_.FoldRiders(gx); - double restraints = 0; - if (restraint_set_ >= 0) { - const RefinementStage s2("eval.restraints"); - restraints = ctx_.engine->EvaluateRestraints(restraint_set_, ctx_.XyzB(), gradient != nullptr ? &gr : nullptr); - } *cost = weight_ * xray + restraints; if (gradient != nullptr) for (size_t k = 0; k < param_atoms_.size(); k++) { diff --git a/image_analysis/structure_refinement/RefinementGPU.cpp b/image_analysis/structure_refinement/RefinementGPU.cpp index bcac87259..642fcf319 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cpp +++ b/image_analysis/structure_refinement/RefinementGPU.cpp @@ -209,6 +209,33 @@ public: 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) { + 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]}; + } + return value; + } + void SetRestraints(int which, const RestraintSet &set) override { engine_->SetRestraints(which, set.records, set.plane_atoms, set.slot_start, set.slot_index, set.n_slots); } @@ -247,7 +274,8 @@ private: // atoms whose B moved). void SetAtoms(const gemmi::Model &model) { std::optional stage("gpu.atoms_host"); - const std::vector atoms = Atoms(model); + 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); @@ -284,6 +312,7 @@ private: bool supported_ = false; size_t fft_work_ = 0, bytes_ = 0; std::unique_ptr engine_; + std::vector atom_ptrs_; // the model's atoms at the last SetAtoms() std::vector last_b_; std::vector atoms_; }; diff --git a/image_analysis/structure_refinement/RefinementGPU.cu b/image_analysis/structure_refinement/RefinementGPU.cu index 1c0c7407a..78ba751f0 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cu +++ b/image_analysis/structure_refinement/RefinementGPU.cu @@ -364,6 +364,10 @@ struct RefinementGPUEngineImpl { 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; + CudaHostPtr host_grad, host_xyzb, host_restraint_grad; CudaDevicePtr mask_atoms; CudaDevicePtr grid; CudaDevicePtr spectrum; @@ -389,6 +393,10 @@ struct RefinementGPUEngineImpl { 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); + host_restraint_grad = CudaHostPtr(na); mask_atoms = CudaDevicePtr(na, sync); grad = CudaDevicePtr(na, sync); grid = CudaDevicePtr(GridPoints(setup.grid), sync); @@ -478,8 +486,80 @@ struct RefinementGPUEngineImpl { nll, coef); cuda_err(cudaGetLastError()); } + + // The X-ray target into host_sums[0] and, with `grad`, its gradient into host_grad: queued, not waited for. + void QueueEvaluate(bool grad) { + const int rows = static_cast(setup.rows.size()); + Transform(); + Compose(true); + host_sums[0] = 0; + if (rows > 0) { + size_t temp = reduce_bytes; + cuda_err(cub::DeviceReduce::Sum(reduce_temp.get(), temp, nll.get(), nll_sum.get(), rows, stream)); + cuda_err(cudaMemcpyAsync(host_sums.get(), nll_sum, sizeof(double), cudaMemcpyDeviceToHost, stream)); + } + if (grad) { + const ModelDensityGrid &g = setup.grid; + cuda_err(cudaMemsetAsync(spectrum, 0, ComplexPoints(g) * sizeof(float2), stream)); + const int points = static_cast(setup.scatter_point.size()); + if (points > 0) { + scatter_kernel<<>>(coef, row_hkl, ops, + static_cast(setup.sym_ops.size()), setup.den, + scatter_point, scatter_start, scatter_term, + points, spectrum); + cuda_err(cudaGetLastError()); + } + cufft_err(inverse.C2R(reinterpret_cast(spectrum.get()), this->grid)); + if (n_atoms > 0) { + const size_t threads = static_cast(n_atoms) * WARP; + atom_gradient_kernel<<>>(density->DeviceAtoms(), pos, + density->DeviceBoxes(), n_atoms, geom, + this->grid, Norm(), this->grad); + cuda_err(cudaGetLastError()); + cuda_err(cudaMemcpyAsync(host_grad.get(), this->grad, n_atoms * sizeof(float4), cudaMemcpyDeviceToHost, + stream)); + } + } + } + + // Restraint set `which` at host_xyzb into host_sums[1] and, with `grad`, its gradient into + // host_restraint_grad: queued, not waited for. + void QueueRestraints(int which, bool grad) { + Restraints &r = restraints[which]; + if (r.atoms > 0) + cuda_err(cudaMemcpyAsync(r.xyzb, host_xyzb.get(), r.atoms * sizeof(float4), cudaMemcpyHostToDevice, stream)); + host_sums[1] = 0; + if (r.records > 0) { + restraint_kernel<<>>(r.rec, r.records, r.plane_atoms, r.xyzb, r.contrib, + r.energy); + cuda_err(cudaGetLastError()); + size_t temp = r.reduce_bytes; + cuda_err(cub::DeviceReduce::Sum(r.reduce_temp.get(), temp, r.energy.get(), r.sum.get(), r.records, stream)); + cuda_err(cudaMemcpyAsync(host_sums.get() + 1, r.sum, sizeof(double), cudaMemcpyDeviceToHost, stream)); + } + if (grad && r.atoms > 0) { + if (r.records > 0) { + slot_gather_kernel<<>>(r.contrib, r.slot_start, r.slot_index, + r.atoms, r.grad); + cuda_err(cudaGetLastError()); + } else + cuda_err(cudaMemsetAsync(r.grad, 0, r.atoms * sizeof(float4), stream)); + cuda_err(cudaMemcpyAsync(host_restraint_grad.get(), r.grad, r.atoms * sizeof(float4), cudaMemcpyDeviceToHost, + stream)); + } + } }; +namespace { + +void CopyOut(const float4 *from, int n, std::vector> &to) { + to.resize(n); + std::copy(reinterpret_cast *>(from), reinterpret_cast *>(from) + n, + to.begin()); +} + +} // namespace + size_t RefinementGPUEngine::FFTWorkBytes(int device, const ModelDensityGrid &grid) { DeviceGuard guard(device); return std::max(CudaFFTPlan::WorkBytes(ForwardShape(grid)), CudaFFTPlan::WorkBytes(InverseShape(grid))); @@ -590,40 +670,32 @@ void RefinementGPUEngine::Fmask(const std::vector &mask_atoms, st double RefinementGPUEngine::Evaluate(std::vector> *grad) { RefinementGPUEngineImpl &e = *impl_; DeviceGuard guard(e.device); - const int rows = static_cast(e.setup.rows.size()); - e.Transform(); - e.Compose(true); - double sum = 0; - if (rows > 0) { - size_t temp = e.reduce_bytes; - cuda_err(cub::DeviceReduce::Sum(e.reduce_temp.get(), temp, e.nll.get(), e.nll_sum.get(), rows, e.stream)); - cuda_err(cudaMemcpyAsync(&sum, e.nll_sum, sizeof(double), cudaMemcpyDeviceToHost, e.stream)); - } - if (grad != nullptr) { - const ModelDensityGrid &g = e.setup.grid; - cuda_err(cudaMemsetAsync(e.spectrum, 0, ComplexPoints(g) * sizeof(float2), e.stream)); - const int points = static_cast(e.setup.scatter_point.size()); - if (points > 0) { - scatter_kernel<<>>(e.coef, e.row_hkl, e.ops, - static_cast(e.setup.sym_ops.size()), e.setup.den, - e.scatter_point, e.scatter_start, e.scatter_term, - points, e.spectrum); - cuda_err(cudaGetLastError()); - } - 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.density->DeviceAtoms(), e.pos, - e.density->DeviceBoxes(), e.n_atoms, e.geom, - e.grid, e.Norm(), e.grad); - cuda_err(cudaGetLastError()); - } - grad->resize(e.n_atoms); - if (e.n_atoms > 0) - cuda_err(cudaMemcpyAsync(grad->data(), e.grad, e.n_atoms * sizeof(float4), cudaMemcpyDeviceToHost, e.stream)); - } + e.QueueEvaluate(grad != nullptr); cuda_err(cudaStreamSynchronize(e.stream)); - return sum; + 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::function *)> &xyzb, + std::vector> *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); + xyzb(reinterpret_cast *>(e.host_xyzb.get())); + e.QueueRestraints(which, restraint_grad != nullptr); + cuda_err(cudaStreamSynchronize(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]; } void RefinementGPUEngine::SetRestraints(int which, const std::vector &records, @@ -666,29 +738,12 @@ double RefinementGPUEngine::EvaluateRestraints(int which, const std::vector(xyzb.size()) != r.atoms) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: restraint atoms do not match"); - if (r.atoms > 0) - cuda_err(cudaMemcpyAsync(r.xyzb, xyzb.data(), xyzb.size() * sizeof(float4), cudaMemcpyHostToDevice, e.stream)); - double sum = 0; - if (r.records > 0) { - restraint_kernel<<>>(r.rec, r.records, r.plane_atoms, r.xyzb, r.contrib, - r.energy); - cuda_err(cudaGetLastError()); - size_t temp = r.reduce_bytes; - cuda_err(cub::DeviceReduce::Sum(r.reduce_temp.get(), temp, r.energy.get(), r.sum.get(), r.records, e.stream)); - cuda_err(cudaMemcpyAsync(&sum, r.sum, sizeof(double), cudaMemcpyDeviceToHost, e.stream)); - } - if (grad != nullptr) { - grad->resize(r.atoms); - if (r.atoms > 0) { - if (r.records > 0) { - slot_gather_kernel<<>>(r.contrib, r.slot_start, r.slot_index, - r.atoms, r.grad); - cuda_err(cudaGetLastError()); - } else - cuda_err(cudaMemsetAsync(r.grad, 0, r.atoms * sizeof(float4), e.stream)); - cuda_err(cudaMemcpyAsync(grad->data(), r.grad, r.atoms * sizeof(float4), cudaMemcpyDeviceToHost, e.stream)); - } - } + // the staging buffer free again cuda_err(cudaStreamSynchronize(e.stream)); - return sum; + std::copy(xyzb.begin(), xyzb.end(), reinterpret_cast *>(e.host_xyzb.get())); + e.QueueRestraints(which, grad != nullptr); + cuda_err(cudaStreamSynchronize(e.stream)); + if (grad != nullptr) + CopyOut(e.host_restraint_grad, r.atoms, *grad); + return e.host_sums[1]; } diff --git a/image_analysis/structure_refinement/RefinementGPUEngine.h b/image_analysis/structure_refinement/RefinementGPUEngine.h index cd67a2456..71eef17ff 100644 --- a/image_analysis/structure_refinement/RefinementGPUEngine.h +++ b/image_analysis/structure_refinement/RefinementGPUEngine.h @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -71,6 +72,11 @@ public: // The sum of E of set `which` at `xyzb` (x, y, z, B per atom) and dE/d(x, y, z, B) per atom. 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); // The most atoms of a plane restraint the kernel takes. static constexpr int MAX_PLANE_ATOMS = 32; diff --git a/image_analysis/structure_refinement/RefinementXray.cpp b/image_analysis/structure_refinement/RefinementXray.cpp index 61120304d..5ac7873f7 100644 --- a/image_analysis/structure_refinement/RefinementXray.cpp +++ b/image_analysis/structure_refinement/RefinementXray.cpp @@ -313,6 +313,15 @@ RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector> ModelXyzB(const gemmi::Model &model) { + std::vector> out; + for (const gemmi::Chain &ch : model.chains) + for (const gemmi::Residue &r : ch.residues) + for (const gemmi::Atom &a : r.atoms) + out.push_back({a.pos.x, a.pos.y, a.pos.z, a.b_iso}); + return out; +} + namespace { std::mutex stage_mutex; std::map> stage_times; diff --git a/image_analysis/structure_refinement/RefinementXray.h b/image_analysis/structure_refinement/RefinementXray.h index fc37a9e85..24fd13756 100644 --- a/image_analysis/structure_refinement/RefinementXray.h +++ b/image_analysis/structure_refinement/RefinementXray.h @@ -67,6 +67,9 @@ struct RFactors { }; RFactors ComputeRFactors(const RefinementReflections &refl, const std::vector> &fmodel); +// x, y, z and B of every atom of `model`, in model order. +std::vector> ModelXyzB(const gemmi::Model &model); + // What evaluates the targets for a set of atom positions and B: the CPU below, or a GPU (RefinementGPU.h). The // atoms are the model's, all of them, in model order; the grid and the reflections are fixed when the engine is // made, the row model each macro-cycle, the blur where SetBlur is called. @@ -95,6 +98,14 @@ public: grad->assign(xyzb.size(), {0, 0, 0, 0}); return ::EvaluateRestraints(restraints_[which], xyzb, grad); } + // Evaluate() and EvaluateRestraints(which) at the atoms of `model` as they are, in one go: returns the X-ray + // target, and the restraints' sum in `restraints`. + virtual double EvaluateWithRestraints(const gemmi::Model &model, int which, std::vector> *grad, + std::vector> *restraint_grad, double &restraints) { + const double xray = Evaluate(model, grad); + restraints = EvaluateRestraints(which, ModelXyzB(model), restraint_grad); + return xray; + } private: RestraintSet restraints_[2];