From 4a8e13ddfaf4fd801d142c4d893d3b4df784fcb8 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 26 Oct 2026 08:31:43 +0100 Subject: [PATCH] Refinement GPU engine: its waits on the stream asleep (blocking-sync event) rather than spinning Bitwise. cudaStreamSynchronize spins a core for the whole wait; the weight trials' engines wait on a shared card for most of a run, so about half of the user cycles of a refinement were the spin (PROFILE.md: 38-50%). A StreamWaiter (CUDAMemHelpers.h: an event made with cudaEventBlockingSync, recorded and waited for) now takes the refinement engine's waits - evaluation, Fcalc, Fmask, the uploads of positions, row model, blur and restraints. The short waits of the host-driven loops stay spinning: the scale fit's Levenberg-Marquardt (one launch of 66-120 us per step) and the gather's pair count; making those block too cost 2-3% of wall (wake-up latency), here the wall is unchanged. No process-wide cudaDeviceScheduleBlockingSync: that would change every GPU wait of the broker. CPU use (-N 8, time -v): 5reo 301-315% -> 190%, 5f6m 326% -> 181%, 9yl4 333% -> 222-228%, 7d98 307% -> 181%, 5f6m MD start 309-320% -> 151%. Wall, two interleaved runs each: 5reo 1.43-1.52 -> 1.43, 5f6m 2.19-2.20 -> 2.19-2.20, 9yl4 6.83-6.84 -> 6.88-6.89, 7d98 3.50-3.51 -> 3.58, 5f6m MD start 4.16-4.34 -> 4.16 s. Refined CIF md5 unchanged on the 5 GPU cases; side by side = one at a time. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit d23d2cf7bb99b30617a3f57c531b03db8bc03e05) --- image_analysis/indexing/CUDAMemHelpers.h | 24 +++++++++++++++++++ .../structure_refinement/RefinementGPU.cu | 21 ++++++++-------- 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/image_analysis/indexing/CUDAMemHelpers.h b/image_analysis/indexing/CUDAMemHelpers.h index fcbe55d72..37de3fe94 100644 --- a/image_analysis/indexing/CUDAMemHelpers.h +++ b/image_analysis/indexing/CUDAMemHelpers.h @@ -108,6 +108,30 @@ public: cudaEvent_t get() const { return event_; } }; +// cudaStreamSynchronize with the waiting thread asleep: an event made with cudaEventBlockingSync, recorded on the +// stream and waited for. cudaStreamSynchronize spins a core for the whole of the wait (the process's default +// schedule), which on a card shared by several engines - refinement's weight trials - is a core per engine for +// most of the run, taken from whatever else the host is doing. Made on the device of the streams it waits for; +// one thread at a time. +class StreamWaiter { + cudaEvent_t event_ = nullptr; +public: + StreamWaiter() { + if (cudaEventCreateWithFlags(&event_, cudaEventBlockingSync | cudaEventDisableTiming) != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Failed to create CUDA event"); + } + ~StreamWaiter() { + if (event_) cudaEventDestroy(event_); + } + StreamWaiter(const StreamWaiter&) = delete; + StreamWaiter& operator=(const StreamWaiter&) = delete; + + cudaError_t Wait(cudaStream_t stream) const { + const cudaError_t err = cudaEventRecord(event_, stream); + return err != cudaSuccess ? err : cudaEventSynchronize(event_); + } +}; + // How much freed memory the pool keeps rather than returning it to the driver. Returning it puts the // next allocation straight back on the path this exists to avoid, so hold the per-worker engine // buffers; but the card also has to fit the merge, which asks for several gigabytes of its own after diff --git a/image_analysis/structure_refinement/RefinementGPU.cu b/image_analysis/structure_refinement/RefinementGPU.cu index 4ba649599..ec50cac99 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cu +++ b/image_analysis/structure_refinement/RefinementGPU.cu @@ -364,6 +364,7 @@ struct RefinementGPUEngineImpl { size_t work_bytes; int device = 0; CudaStream stream; + StreamWaiter waiter; // the engine's waits on the stream, asleep int n_atoms = 0; CudaDevicePtr pos; @@ -622,7 +623,7 @@ void RefinementGPUEngine::SetBlur(double blur) { for (size_t m = 0; m < scale.size(); m++) scale[m] = e.setup.centring * std::exp(blur * 0.25 * e.setup.inv_d2[m]); Upload(e.row_scale, scale, e.stream); - cuda_err(cudaStreamSynchronize(e.stream)); + cuda_err(e.waiter.Wait(e.stream)); } void RefinementGPUEngine::SetRowModel(const std::vector &k_total, const std::vector> &fmask, @@ -635,7 +636,7 @@ void RefinementGPUEngine::SetRowModel(const std::vector &k_total, const s e.stream)); Upload(e.ml_d, d, e.stream); Upload(e.ml_var, var, e.stream); - cuda_err(cudaStreamSynchronize(e.stream)); + cuda_err(e.waiter.Wait(e.stream)); } void RefinementGPUEngine::SetDensities(const std::vector &atoms) { @@ -653,7 +654,7 @@ void RefinementGPUEngine::SetPositions(const std::vector> & 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)); + cuda_err(e.waiter.Wait(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)); @@ -667,7 +668,7 @@ void RefinementGPUEngine::Fcalc(std::vector> &f) { f.resize(e.setup.rows.size()); if (!f.empty()) cuda_err(cudaMemcpyAsync(f.data(), e.f, f.size() * sizeof(float2), cudaMemcpyDeviceToHost, e.stream)); - cuda_err(cudaStreamSynchronize(e.stream)); + cuda_err(e.waiter.Wait(e.stream)); } void RefinementGPUEngine::Fmask(const std::vector &mask_atoms, std::vector> &f) { @@ -692,14 +693,14 @@ void RefinementGPUEngine::Fmask(const std::vector &mask_atoms, st cuda_err(cudaGetLastError()); cuda_err(cudaMemcpyAsync(f.data(), e.fm, rows * sizeof(float2), cudaMemcpyDeviceToHost, e.stream)); } - cuda_err(cudaStreamSynchronize(e.stream)); + cuda_err(e.waiter.Wait(e.stream)); } double RefinementGPUEngine::Evaluate(std::vector> *grad) { RefinementGPUEngineImpl &e = *impl_; DeviceGuard guard(e.device); e.QueueEvaluate(grad != nullptr); - cuda_err(cudaStreamSynchronize(e.stream)); + cuda_err(e.waiter.Wait(e.stream)); if (grad != nullptr) CopyOut(e.host_grad, e.n_atoms, *grad); return e.host_sums[0]; @@ -717,7 +718,7 @@ double RefinementGPUEngine::EvaluateWithRestraints(std::vector *>(e.host_xyzb.get())); e.QueueRestraints(which, restraint_grad != nullptr); - cuda_err(cudaStreamSynchronize(e.stream)); + cuda_err(e.waiter.Wait(e.stream)); if (grad != nullptr) CopyOut(e.host_grad, e.n_atoms, *grad); if (restraint_grad != nullptr) @@ -756,7 +757,7 @@ void RefinementGPUEngine::SetRestraints(int which, const std::vector> &xyzb, @@ -767,10 +768,10 @@ double RefinementGPUEngine::EvaluateRestraints(int which, const std::vector(xyzb.size()) != r.atoms) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "refinement: restraint atoms do not match"); // the staging buffer free again - cuda_err(cudaStreamSynchronize(e.stream)); + cuda_err(e.waiter.Wait(e.stream)); std::copy(xyzb.begin(), xyzb.end(), reinterpret_cast *>(e.host_xyzb.get())); e.QueueRestraints(which, grad != nullptr); - cuda_err(cudaStreamSynchronize(e.stream)); + cuda_err(e.waiter.Wait(e.stream)); if (grad != nullptr) CopyOut(e.host_restraint_grad, r.atoms, *grad); return e.host_sums[1];