From 257bf2fe2beab65f2f1726b04dc7efb746942902 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 26 Oct 2026 15:22:22 +0100 Subject: [PATCH] Refinement: trial slots counted with their scale pool and both restraint sets; on the CPU no more than -N GPU: a slot was packed onto a card by its engine's bytes alone. The per-slot estimate now adds the slot's ModelScaleGPUPool with the robust fit's buffers (ModelScaleGPUPool::DeviceBytes(n, 1, true)), and RefinementGPUEngine::DeviceBytes counts the two restraint sets SetRestraints() holds, not one. The 80%-of-card budget is unchanged, so where these do not fit fewer slots are packed instead of an allocation failing inside a running trial. CPU: OwnSlotDevices made a slot per wanted trial whatever the thread count, so -N 1 ran six trial threads with six models, topologies and engines in memory. The main slot and those beside it are now at most -N. Side by side equals sequential, so no result changes. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit 3c04c01158266c267f8e5ffaaaa98d6d8fd2a747) --- .../structure_refinement/Refinement.cpp | 17 ++++++++++------- .../structure_refinement/RefinementGPU.cu | 5 +++-- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index 0cd8605a0..ce6969a41 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -1013,29 +1013,32 @@ 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. +// 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. 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(); + // 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] = main.engine_bytes; + used[main.device] = slot_bytes; std::vector devices; for (int k = 1; k <= wanted; k++) { const int d = (main.device + k) % 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] + main.engine_bytes <= budget) { - used[d] += main.engine_bytes; + if (used[d] + slot_bytes <= budget) { + used[d] += slot_bytes; devices.push_back(d); } } return devices; } -#else - (void) settings; #endif - return std::vector(std::max(wanted, 0), -1); + // on the CPU, no more trials at once than threads: each slot is a model, a topology and an engine of its own + const int beside = std::min(wanted, static_cast(settings.nthreads) - 1); + return std::vector(std::max(beside, 0), -1); } // One weight trial: `cycles` macro-cycles from pos0 and b0, of the positions at X-ray weight w and, where adp_w is diff --git a/image_analysis/structure_refinement/RefinementGPU.cu b/image_analysis/structure_refinement/RefinementGPU.cu index 0108cd549..33fa3baff 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cu +++ b/image_analysis/structure_refinement/RefinementGPU.cu @@ -620,8 +620,9 @@ size_t RefinementGPUEngine::DeviceBytes(const RefinementGPUSetup &s, size_t fft_ bytes += ModelMaskGPU::DeviceBytes(GridPoints(s.mask_grid)); if (!SameGrid(s.mask_grid, s.grid)) bytes += GridPoints(s.mask_grid) * sizeof(float) + ComplexPoints(s.mask_grid) * sizeof(float2); - bytes += max_records * (sizeof(RestraintRecord) + sizeof(double)) + max_slots * (sizeof(float4) + sizeof(int)) + - na * (2 * sizeof(float4) + sizeof(int)); + // the two restraint sets of SetRestraints(), each up to these + bytes += 2 * (max_records * (sizeof(RestraintRecord) + sizeof(double)) + max_slots * (sizeof(float4) + sizeof(int)) + + na * (2 * sizeof(float4) + sizeof(int))); return bytes; }