diff --git a/image_analysis/structure_refinement/ModelDensityGPU.cu b/image_analysis/structure_refinement/ModelDensityGPU.cu index 7b26d7dfa..a85f86bd6 100644 --- a/image_analysis/structure_refinement/ModelDensityGPU.cu +++ b/image_analysis/structure_refinement/ModelDensityGPU.cu @@ -254,16 +254,20 @@ size_t ModelDensityGPU::Bricks(int nu, int nv, int nw) { return static_cast((nu + BRICK - 1) / BRICK) * ((nv + BRICK - 1) / BRICK) * ((nw + BRICK - 1) / BRICK); } -std::vector ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const std::vector &atoms) { +void ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const ModelDensityAtom *atoms, size_t n, int *box) { const int n3[3] = {grid.nu, grid.nv, grid.nw}; double spacing[3]; for (int k = 0; k < 3; k++) spacing[k] = 1.0 / (n3[k] * std::sqrt(grid.frac[3 * k] * grid.frac[3 * k] + grid.frac[3 * k + 1] * grid.frac[3 * k + 1] + grid.frac[3 * k + 2] * grid.frac[3 * k + 2])); - std::vector box(3 * atoms.size()); - for (size_t i = 0; i < atoms.size(); i++) + for (size_t i = 0; i < n; i++) for (int k = 0; k < 3; k++) box[3 * i + k] = static_cast(std::ceil(atoms[i].radius / spacing[k])); +} + +std::vector ModelDensityGPU::Boxes(const ModelDensityGrid &grid, const std::vector &atoms) { + std::vector box(3 * atoms.size()); + Boxes(grid, atoms.data(), atoms.size(), box.data()); return box; } @@ -280,15 +284,19 @@ size_t ModelDensityGPU::PairBound(const ModelDensityGrid &grid, const std::vecto return pairs; } -bool ModelDensityGPU::Supports(const ModelDensityGrid &grid, const std::vector &atoms) { - const std::vector box = Boxes(grid, atoms); +bool ModelDensityGPU::Supports(const ModelDensityGrid &grid, const int *box, size_t n) { const int n3[3] = {grid.nu, grid.nv, grid.nw}; - for (size_t i = 0; i < box.size(); i++) + for (size_t i = 0; i < 3 * n; i++) if (2 * box[i] + 1 > n3[i % 3] || AxisBrickBound(box[i], n3[i % 3]) > MAX_BRICKS_PER_AXIS) return false; return true; } +bool ModelDensityGPU::Supports(const ModelDensityGrid &grid, const std::vector &atoms) { + const std::vector box = Boxes(grid, atoms); + return Supports(grid, box.data(), atoms.size()); +} + size_t ModelDensityGPU::DeviceBytes(size_t max_atoms, size_t max_pairs, size_t max_bricks) { const size_t na = max_atoms + 1; return na * (sizeof(ModelDensityAtom) + 3 * sizeof(int) + 2 * sizeof(int)) + 4 * max_pairs * sizeof(int) + @@ -314,14 +322,22 @@ ModelDensityGPU::ModelDensityGPU(cudaStream_t stream, size_t max_atoms_, size_t host_count = CudaHostPtr(2); cub_bytes = CubBytes(na, np); cub_temp = CudaDevicePtr(std::max(cub_bytes, 1), sync); + host_box = CudaHostPtr(3 * na); } void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const std::vector &atom_list) { - if (atom_list.size() > max_atoms || Bricks(grid.nu, grid.nv, grid.nw) > max_bricks) + SetAtoms(grid, atom_list.data(), atom_list.size()); +} + +void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const ModelDensityAtom *atom_list, size_t n) { + if (n > max_atoms || Bricks(grid.nu, grid.nv, grid.nw) > max_bricks) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: grid or atoms over the reserve"); + // The boxes into the pinned staging buffer, so that their copy and the atoms' (pinned where the caller keeps + // them so) go up by DMA instead of through the driver's staging of pageable memory. + Boxes(grid, atom_list, n, host_box); // An atom whose box is wider than the cell would reach a point through more than one image, which // gemmi's box walk adds and the gather's box test does not; such a cell is left to the CPU. - if (!Supports(grid, atom_list)) + if (!Supports(grid, host_box, n)) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: an atom wider than the cell"); geom.nu = grid.nu; geom.nv = grid.nv; @@ -334,14 +350,12 @@ void ModelDensityGPU::SetAtoms(const ModelDensityGrid &grid, const std::vector(grid.orth[4] * inv_n[1]); geom.m12 = static_cast(grid.orth[5] * inv_n[2]); geom.m22 = static_cast(grid.orth[8] * inv_n[2]); - const std::vector b = Boxes(grid, atom_list); - n_atoms = static_cast(atom_list.size()); + n_atoms = static_cast(n); if (n_atoms > 0) { - cuda_err(cudaMemcpyAsync(atoms, atom_list.data(), atom_list.size() * sizeof(ModelDensityAtom), - cudaMemcpyHostToDevice, stream)); - cuda_err(cudaMemcpyAsync(box, b.data(), b.size() * sizeof(int), cudaMemcpyHostToDevice, stream)); + cuda_err(cudaMemcpyAsync(atoms, atom_list, n * sizeof(ModelDensityAtom), cudaMemcpyHostToDevice, stream)); + cuda_err(cudaMemcpyAsync(box, host_box, 3 * n * sizeof(int), cudaMemcpyHostToDevice, stream)); } - // The uploads read host vectors that end with this call. + // The uploads read the caller's atoms, which it may change once this returns. cuda_err(cudaStreamSynchronize(stream)); } diff --git a/image_analysis/structure_refinement/ModelDensityGPU.h b/image_analysis/structure_refinement/ModelDensityGPU.h index 3d0742b7e..1a930d355 100644 --- a/image_analysis/structure_refinement/ModelDensityGPU.h +++ b/image_analysis/structure_refinement/ModelDensityGPU.h @@ -49,16 +49,20 @@ public: // gemmi's box around each atom (MakeAtomBox, ModelGrid.cpp): the points within ceil(radius / spacing) // of the nearest one along each axis, three per atom. static std::vector Boxes(const ModelDensityGrid &grid, const std::vector &atoms); + static void Boxes(const ModelDensityGrid &grid, const ModelDensityAtom *atoms, size_t n, int *box); // (brick, atom) pairs of the gather over at most. static size_t PairBound(const ModelDensityGrid &grid, const std::vector &atoms); // Whether the gather reproduces gemmi's box walk on this grid: every atom's box narrower than the cell, // so that no point is reached by two images of one atom. static bool Supports(const ModelDensityGrid &grid, const std::vector &atoms); + static bool Supports(const ModelDensityGrid &grid, const int *box, size_t n); ModelDensityGPU(cudaStream_t stream, size_t max_atoms, size_t max_pairs, size_t max_bricks); - // The grid and the atoms' densities. Throws if they are over the capacity or if !Supports(). + // The grid and the atoms' densities. Throws if they are over the capacity or if !Supports(). The atoms are + // read until the call returns; from pinned memory (PinnedVector) they go up by DMA. void SetAtoms(const ModelDensityGrid &grid, const std::vector &atoms); + void SetAtoms(const ModelDensityGrid &grid, const ModelDensityAtom *atoms, size_t n); // The density into d_grid (nu * nv * nw floats, every point written). d_pos: each atom's centre in // grid units - its fractional position times the grid size, as the CPU's DensityBox holds it, not @@ -85,6 +89,7 @@ private: CudaDevicePtr atoms; CudaDevicePtr box; + CudaHostPtr host_box; // the boxes' staging buffer, pinned CudaDevicePtr overflow; // an atom's box over MAX_BRICKS_PER_AXIS CudaDevicePtr count, offset, key, value, key_sorted, value_sorted, brick_start, brick_end; CudaHostPtr host_count; // the pair count and the overflow flag Compute() reads back, pinned diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index fc14a2ebd..d2f5993b4 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -916,9 +916,11 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg ctx.b_ceiling, budget, &bytes, settings.nthreads); ctx.device = device; ctx.engine_bytes = bytes; - if (ctx.engine) + if (ctx.engine) { ctx.scale_gpu = std::make_unique(device, ctx.refl.Size(), 1); - else + logger.Debug("Refinement: GPU engine {:.2f} GB of a {:.2f} GB budget (scale pool {:.2f} GB)", bytes / 1e9, + budget / 1e9, ModelScaleGPUPool::DeviceBytes(ctx.refl.Size(), 1, true) / 1e9); + } else logger.Info("Refinement: on the CPU - the GPU cannot take this case or {:.1f} GB is over its budget", bytes / 1e9); } catch (const JFJochException &e) { diff --git a/image_analysis/structure_refinement/RefinementGPU.cpp b/image_analysis/structure_refinement/RefinementGPU.cpp index f08b1df17..b9cf9d40c 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cpp +++ b/image_analysis/structure_refinement/RefinementGPU.cpp @@ -304,7 +304,7 @@ private: stage.reset(); const RefinementStage upload("gpu.upload"); if (std::find(changed.begin(), changed.end(), 1) != changed.end()) - engine_->SetDensities(atoms_); + engine_->SetDensities(atoms_.data(), atoms_.size()); engine_->SetPositions(pos); } @@ -320,7 +320,7 @@ private: std::unique_ptr engine_; std::vector atom_ptrs_; // the model's atoms at the last SetAtoms() std::vector last_b_; - std::vector atoms_; + PinnedVector atoms_; // pinned: a B rung sends them every evaluation (37 MB on a large model) }; } // namespace diff --git a/image_analysis/structure_refinement/RefinementGPU.cu b/image_analysis/structure_refinement/RefinementGPU.cu index ec7160f5b..14bcce975 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cu +++ b/image_analysis/structure_refinement/RefinementGPU.cu @@ -722,13 +722,13 @@ void RefinementGPUEngine::SetRowModel(const std::vector &k_total, const s cuda_err(e.waiter.Wait(e.stream)); } -void RefinementGPUEngine::SetDensities(const std::vector &atoms) { +void RefinementGPUEngine::SetDensities(const ModelDensityAtom *atoms, size_t n) { RefinementGPUEngineImpl &e = *impl_; DeviceGuard guard(e.device); - if (atoms.size() > e.setup.max_atoms) + if (n > 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()); + e.density->SetAtoms(e.setup.grid, atoms, n); + e.n_atoms = static_cast(n); } void RefinementGPUEngine::SetPositions(const std::vector> &pos) { diff --git a/image_analysis/structure_refinement/RefinementGPUEngine.h b/image_analysis/structure_refinement/RefinementGPUEngine.h index c64677c78..a44d177f1 100644 --- a/image_analysis/structure_refinement/RefinementGPUEngine.h +++ b/image_analysis/structure_refinement/RefinementGPUEngine.h @@ -56,8 +56,9 @@ public: const std::vector &d, const std::vector &var); // 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); + // from them): called again only where some atom's density changed. Read until the call returns; from pinned + // memory they go up by DMA. + void SetDensities(const ModelDensityAtom *atoms, size_t n); // 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. diff --git a/image_analysis/structure_refinement/RefinementMinimiser.h b/image_analysis/structure_refinement/RefinementMinimiser.h index 4fdff40c7..84743d338 100644 --- a/image_analysis/structure_refinement/RefinementMinimiser.h +++ b/image_analysis/structure_refinement/RefinementMinimiser.h @@ -122,29 +122,60 @@ private: return s; } - // d = -H g by the two-loop recursion, H0 = gamma diag(1/h). + // d = -H g by the two-loop recursion, H0 = gamma diag(1/h). Every dot product is a serial chain of n + // additions, so each sweep over the parameters takes the next pair's dot product while it applies the + // previous pair's update to q - the same operations on each element in the same order as a sweep for + // each would do, with the update's memory traffic hidden under the chain. void Direction(const std::vector &g, std::vector &d) const { const size_t n = g.size(), m = s_.size(); std::vector q = g, a(m), rho(m); - for (size_t i = m; i-- > 0;) { + for (size_t i = 0; i < m; i++) rho[i] = 1.0 / sy_[i]; - a[i] = rho[i] * Dot(s_[i], q); - for (size_t k = 0; k < n; k++) - q[k] -= a[i] * y_[i][k]; + double dot = m > 0 ? Dot(s_[m - 1], q) : 0.0, yhy = 0.0; + for (size_t i = m; i-- > 0;) { + a[i] = rho[i] * dot; + const double ai = a[i]; + const double *y = y_[i].data(); + dot = 0.0; + if (i > 0) { + const double *s = s_[i - 1].data(); + for (size_t k = 0; k < n; k++) { + q[k] -= ai * y[k]; + dot += s[k] * q[k]; + } + } else { + const double *yl = y_[m - 1].data(); + for (size_t k = 0; k < n; k++) { + q[k] -= ai * y[k]; + yhy += yl[k] * yl[k] / h_[k]; + } + } } - double gamma = 1.0; + const double gamma = m > 0 ? sy_[m - 1] / yhy : 1.0; if (m > 0) { - double yhy = 0; + const double *y = y_[0].data(); + for (size_t k = 0; k < n; k++) { + q[k] *= gamma / h_[k]; + dot += y[k] * q[k]; + } + } else for (size_t k = 0; k < n; k++) - yhy += y_[m - 1][k] * y_[m - 1][k] / h_[k]; - gamma = sy_[m - 1] / yhy; - } - for (size_t k = 0; k < n; k++) - q[k] *= gamma / h_[k]; + q[k] *= gamma / h_[k]; for (size_t i = 0; i < m; i++) { - const double b = rho[i] * Dot(y_[i], q); - for (size_t k = 0; k < n; k++) - q[k] += s_[i][k] * (a[i] - b); + // b on its own and (a[i] - b) in the loop, as before: written as a[i] - rho[i] * dot the compiler + // contracts it into a fused multiply-add, and the bits change + const double b = rho[i] * dot; + const double *s = s_[i].data(); + dot = 0.0; + if (i + 1 < m) { + const double *y = y_[i + 1].data(); + for (size_t k = 0; k < n; k++) { + q[k] += s[k] * (a[i] - b); + dot += y[k] * q[k]; + } + } else + for (size_t k = 0; k < n; k++) + q[k] += s[k] * (a[i] - b); } for (size_t k = 0; k < n; k++) d[k] = -q[k];