diff --git a/tests/CUDAMemHelpersTest.cpp b/tests/CUDAMemHelpersTest.cpp index e095333a7..2d0047319 100644 --- a/tests/CUDAMemHelpersTest.cpp +++ b/tests/CUDAMemHelpersTest.cpp @@ -141,55 +141,69 @@ TEST_CASE("IndexerThreadPool_FFTIndexerWaitsForGPUWorkBeside", "[CUDAMemHelpers] x.ImportIndexingSettings(settings); const std::vector recip{{0.1, 0, 0}, {0, 0.1, 0}, {0, 0, 0.1}}; - // The whole card, held by the work beside. - std::optional>> held; - const auto fill = [&held] { - held.emplace(); - for (size_t chunk = size_t(256) << 20; chunk >= (size_t(1) << 20); chunk /= 2) - while (true) { - try { - held->emplace_back(chunk, CudaAlloc::Synchronous); - } catch (const JFJochException &) { - break; + // Every card, held by the work beside: an indexer worker takes its card round-robin (pin_gpu), so on + // a machine with several the one it lands on is not the test thread's. Topped up right before each + // Run, so memory something else gave back in between does not let the indexer fit. + const int gpus = get_gpu_count(); + std::vector>> held(gpus); + const auto fill = [&] { + for (int d = 0; d < gpus; d++) { + ScopedGpuPin pin(d); + for (size_t chunk = size_t(256) << 20; chunk >= (size_t(1) << 20); chunk /= 2) + while (true) { + try { + held[d].emplace_back(chunk, CudaAlloc::Synchronous); + } catch (const JFJochException &) { + break; + } } - } + } + }; + const auto release = [&] { + for (int d = 0; d < gpus; d++) { + ScopedGpuPin pin(d); + held[d].clear(); + } }; // The pools' indexers go back to the allocation pool when they are destroyed; trim it after them, so - // the next test finds the card free. + // the next test finds the cards free. { IndexerThreadPool probe(settings, IndexerConstruction::OnFirstUse, /*beside=*/true); IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse); - fill(); { GPUWorkBeside work; + fill(); CHECK(probe.Run(x, recip).error != nullptr); } { std::optional image_pass(std::in_place); + GPUWorkBeside work; + fill(); std::thread ends([&] { std::this_thread::sleep_for(std::chrono::milliseconds(500)); - held.reset(); + release(); image_pass.reset(); }); - GPUWorkBeside work; CHECK(probe.Run(x, recip).error == nullptr); ends.join(); } - - fill(); { std::optional work(std::in_place); + fill(); std::thread ends([&] { std::this_thread::sleep_for(std::chrono::milliseconds(500)); - held.reset(); + release(); work.reset(); }); CHECK(pool.Run(x, recip).error == nullptr); ends.join(); } } - cuda_trim_allocation_pool(); + for (int d = 0; d < gpus; d++) { + ScopedGpuPin pin(d); + cuda_trim_allocation_pool(); + } } #endif