From 3189f71e8fdf00eb5cfd41311b3820c0c35880bc Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 26 Oct 2026 22:30:31 +0100 Subject: [PATCH] Refinement: the engine polls for its evaluations and the pair count comes back pinned; the slots beyond one per card go to the cards the ladder's long trials leave free The GPU engine's waits on its stream (every evaluation, the refits' transforms) and the model density's wait for the (brick, atom) pair count were blocking event and stream synchronisations; where a weight trial has a card to itself, as on a box with a card per trial, the wake-up of the sleeping thread is on the critical path of every evaluation. Both now poll, and the pair count and overflow flag come back into pinned memory instead of a pageable int (a pageable copy back is staged and blocks the host for the whole stream). OwnSlotDevices places the slots beyond one per card on the cards the ladder's long trials (those refining B) leave free, so that on four cards the three B rungs have a card each and the three B-held rungs share the fourth, instead of two cards each carrying a B rung beside a B-held one. The outcome does not depend on the placement. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu --- .../structure_refinement/ModelDensityGPU.cu | 14 +++++++++----- .../structure_refinement/ModelDensityGPU.h | 1 + image_analysis/structure_refinement/Refinement.cpp | 14 ++++++++++---- .../structure_refinement/RefinementGPU.cu | 13 ++++++++++++- 4 files changed, 32 insertions(+), 10 deletions(-) diff --git a/image_analysis/structure_refinement/ModelDensityGPU.cu b/image_analysis/structure_refinement/ModelDensityGPU.cu index 5ada2e63b..7b26d7dfa 100644 --- a/image_analysis/structure_refinement/ModelDensityGPU.cu +++ b/image_analysis/structure_refinement/ModelDensityGPU.cu @@ -311,6 +311,7 @@ ModelDensityGPU::ModelDensityGPU(cudaStream_t stream, size_t max_atoms_, size_t value_sorted = CudaDevicePtr(np, sync); brick_start = CudaDevicePtr(nb, sync); brick_end = CudaDevicePtr(nb, sync); + host_count = CudaHostPtr(2); cub_bytes = CubBytes(na, np); cub_temp = CudaDevicePtr(std::max(cub_bytes, 1), sync); } @@ -351,11 +352,14 @@ void ModelDensityGPU::Compute(const float4 *d_pos, float *d_grid) { cuda_err(cudaGetLastError()); size_t temp = cub_bytes; cuda_err(cub::DeviceScan::ExclusiveSum(cub_temp.get(), temp, count.get(), offset.get(), n_atoms + 1, stream)); - int npairs = 0; - int over = 0; - cuda_err(cudaMemcpyAsync(&npairs, offset.get() + n_atoms, sizeof(int), cudaMemcpyDeviceToHost, stream)); - cuda_err(cudaMemcpyAsync(&over, overflow.get(), sizeof(int), cudaMemcpyDeviceToHost, stream)); - cuda_err(cudaStreamSynchronize(stream)); + // the count and the flag into pinned memory, the stream polled for them: this wait is on every evaluation + cuda_err(cudaMemcpyAsync(host_count.get(), offset.get() + n_atoms, sizeof(int), cudaMemcpyDeviceToHost, stream)); + cuda_err(cudaMemcpyAsync(host_count.get() + 1, overflow.get(), sizeof(int), cudaMemcpyDeviceToHost, stream)); + cudaError_t err; + while ((err = cudaStreamQuery(stream)) == cudaErrorNotReady) {} + cuda_err(err); + const int npairs = host_count[0]; + const int over = host_count[1]; if (over) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "model density: an atom's box over the bricks per axis"); if (static_cast(npairs) > max_pairs) diff --git a/image_analysis/structure_refinement/ModelDensityGPU.h b/image_analysis/structure_refinement/ModelDensityGPU.h index a40270627..3d0742b7e 100644 --- a/image_analysis/structure_refinement/ModelDensityGPU.h +++ b/image_analysis/structure_refinement/ModelDensityGPU.h @@ -87,6 +87,7 @@ private: CudaDevicePtr box; 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 CudaDevicePtr cub_temp; size_t cub_bytes = 0; }; diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index 7b11909a4..fc14a2ebd 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -1020,19 +1020,25 @@ std::unique_ptr MakeOwnSlot(const gemmi::Structure &before_topology, co } // The devices of the slots beside the main one, `wanted` of them at most: the GPUs in turn from the main -// engine's, each taking engines while they fit in its budget. -1 for each one on the CPU, where the main one and -// those beside it are no more than the threads. +// engine's, each taking engines while they fit in its budget. The slots beyond one per card go to the cards the +// ladder's long trials (those refining B) leave free: RunTrials gives the long trials to the first slots, one +// card each, so a short trial then shares a card with a short one, not with a long one. -1 for each one on the +// CPU, where the main one and those beside it are no more than the threads. std::vector OwnSlotDevices(const Context &main, int wanted, const RefinementSettings &settings) { #ifdef JFJOCH_USE_CUDA if (main.engine->OnGPU()) { const int n = get_gpu_count(); + const int longest = static_cast(settings.weight_ladder.size()) * + static_cast(std::count_if(settings.adp_ladder.begin(), settings.adp_ladder.end(), + [](double adp) { return adp > 0; })); + const int free = n - longest; // cards without a long trial, after the first slots' cards // a slot's engine and its scale pool, the robust fit's buffers included const size_t slot_bytes = main.engine_bytes + ModelScaleGPUPool::DeviceBytes(main.refl.Size(), 1, true); std::vector used(n, 0); used[main.device] = slot_bytes; std::vector devices; for (int k = 1; k <= wanted; k++) { - const int d = (main.device + k) % n; + const int d = (k < n || free <= 0) ? (main.device + k) % n : (main.device + longest + (k - n) % free) % n; const size_t usable = get_gpu_total_memory(d) / 5 * 4; const size_t budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0; if (used[d] + slot_bytes <= budget) { @@ -1149,7 +1155,7 @@ std::vector RunTrials(const std::vector &trial_slots, for (size_t s = 0; s < running; s++) trial_slots[s].ctx->nthreads = std::max(1, settings.nthreads / running); // The trials that refine B are the longest: they go to the slots first, so that two of them share a slot (or a - // card) only where the slots are fewer than they are. + // card) only where the slots are fewer than they are; the first slots are on cards of their own (OwnSlotDevices). std::vector order(trials.size()); std::iota(order.begin(), order.end(), 0); std::stable_sort(order.begin(), order.end(), diff --git a/image_analysis/structure_refinement/RefinementGPU.cu b/image_analysis/structure_refinement/RefinementGPU.cu index 56a5c00e1..ec7160f5b 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cu +++ b/image_analysis/structure_refinement/RefinementGPU.cu @@ -394,6 +394,17 @@ size_t ReduceBytes(size_t n) { return bytes; } +// The engine's wait on its stream, polling: an evaluation is a millisecond or a few, and where a slot has a +// card to itself the wake-up of a sleeping thread is on the critical path of every one of them. +class StreamPoller { +public: + cudaError_t Wait(cudaStream_t stream) const { + cudaError_t err; + while ((err = cudaStreamQuery(stream)) == cudaErrorNotReady) {} + return err; + } +}; + } // namespace struct RefinementGPUEngineImpl { @@ -412,7 +423,7 @@ struct RefinementGPUEngineImpl { size_t work_bytes; int device = 0; CudaStream stream; - StreamWaiter waiter; // the engine's waits on the stream, asleep + StreamPoller waiter; // the engine's waits on the stream, polled int n_atoms = 0; CudaDevicePtr pos;