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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu

image_analysis/structure_refinement/ModelDensityGPU.cu

@@ -311,6 +311,7 @@ ModelDensityGPU::ModelDensityGPU(cudaStream_t stream, size_t max_atoms_, size_t
     value_sorted = CudaDevicePtr<int>(np, sync);
     brick_start = CudaDevicePtr<int>(nb, sync);
     brick_end = CudaDevicePtr<int>(nb, sync);
+    host_count = CudaHostPtr<int>(2);
     cub_bytes = CubBytes(na, np);
     cub_temp = CudaDevicePtr<char>(std::max<size_t>(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<size_t>(npairs) > max_pairs)

image_analysis/structure_refinement/ModelDensityGPU.h

@@ -87,6 +87,7 @@ private:
     CudaDevicePtr<int> box;
     CudaDevicePtr<int> overflow;          // an atom's box over MAX_BRICKS_PER_AXIS
     CudaDevicePtr<int> count, offset, key, value, key_sorted, value_sorted, brick_start, brick_end;
+    CudaHostPtr<int> host_count;          // the pair count and the overflow flag Compute() reads back, pinned
     CudaDevicePtr<char> cub_temp;
     size_t cub_bytes = 0;
 };

image_analysis/structure_refinement/Refinement.cpp

@@ -1020,19 +1020,25 @@ std::unique_ptr<OwnSlot> 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<int> 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<int>(settings.weight_ladder.size()) *
+                            static_cast<int>(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<size_t> used(n, 0);
         used[main.device] = slot_bytes;
         std::vector<int> 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<TrialOutcome> RunTrials(const std::vector<TrialSlot> &trial_slots,
     for (size_t s = 0; s < running; s++)
         trial_slots[s].ctx->nthreads = std::max<size_t>(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<size_t> order(trials.size());
     std::iota(order.begin(), order.end(), 0);
     std::stable_sort(order.begin(), order.end(),

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<float4> pos;