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;