diff --git a/dali/operators/reader/gds_mem.cc b/dali/operators/reader/gds_mem.cc index bc8c80eb8a3..368efd3f56c 100644 --- a/dali/operators/reader/gds_mem.cc +++ b/dali/operators/reader/gds_mem.cc @@ -1,4 +1,4 @@ -// Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -239,9 +239,9 @@ void GDSStagingEngine::wait_buffers(std::unique_lock &lock) { lock.unlock(); CUDA_CALL(cudaEventSynchronize(ready_)); lock.lock(); - for (void *ptr : scheduled_) - ready_buffers_.insert(ptr); - scheduled_.clear(); + // Another thread may have submitted more copies and re-recorded ready_ while unlocked. + // Check the current event before reclaiming the current scheduled buffers. + return_ready(); } void GDSStagingEngine::return_unused(GDSStagingBuffer &&staging_buffer) { diff --git a/dali/operators/reader/gds_mem_test.cu b/dali/operators/reader/gds_mem_test.cu index 59bcaadb42c..06067033cd6 100644 --- a/dali/operators/reader/gds_mem_test.cu +++ b/dali/operators/reader/gds_mem_test.cu @@ -1,4 +1,4 @@ -// Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -14,6 +14,9 @@ #include #include +#include +#include +#include #include #include #include @@ -186,6 +189,79 @@ TEST(GDSMem, StagingEngine) { }); } +TEST(GDSMemTest, ConcurrentStagingBufferReuse) { + SkipIfIncompatible([&] { + CUFileDriverScope scope; + CUDAStream copy_stream = CUDAStream::Create(true); + // A small pool and immediate commits make multiple readers wait on the same event, + // which is repeatedly re-recorded as other readers submit more copies. + GDSStagingEngine engine(-1, 2, 0); + engine.set_stream(copy_stream); + constexpr int num_threads = 8; + constexpr int iterations = 32; + auto sample_value = [](int t, int i) -> uint8_t { + return (t * iterations + i) % 255 + 1; + }; + const size_t copy_size = std::min(engine.chunk_size(), 1 << 20); + const size_t output_size = num_threads * iterations * copy_size; + auto output = mm::alloc_raw_unique(output_size); + int device; + CUDA_CALL(cudaGetDevice(&device)); + std::atomic start{false}; + std::vector errors(num_threads); + std::vector threads; + for (int t = 0; t < num_threads; t++) { + threads.emplace_back([&, t] { + while (!start.load()) + std::this_thread::yield(); + try { + CUDA_CALL(cudaSetDevice(device)); + CUDAStream fill_stream = CUDAStream::Create(true); + for (int i = 0; i < iterations; i++) { + auto buffer = engine.get_staging_buffer(); + try { + // Complete the write before submitting the copy, just as a synchronous + // cuFileRead would. Different readers and iterations write different data. + uint8_t value = sample_value(t, i); + CUDA_CALL(cudaMemsetAsync(buffer.at(0), value, copy_size, fill_stream)); + CUDA_CALL(cudaStreamSynchronize(fill_stream)); + } catch (...) { + engine.return_unused(std::move(buffer)); + throw; + } + size_t offset = (t * iterations + i) * copy_size; + engine.copy_to_client(output.get() + offset, copy_size, std::move(buffer)); + } + } catch (...) { + errors[t] = std::current_exception(); + } + }); + } + start = true; + for (auto &thread : threads) + thread.join(); + engine.commit(); + CUDA_CALL(cudaStreamSynchronize(copy_stream)); + for (auto &error : errors) { + if (error) + std::rethrow_exception(error); + } + std::vector host(output_size); + CUDA_CALL(cudaMemcpy(host.data(), output.get(), output_size, cudaMemcpyDeviceToHost)); + for (int t = 0; t < num_threads; t++) { + for (int i = 0; i < iterations; i++) { + uint8_t expected = sample_value(t, i); + size_t offset = (t * iterations + i) * copy_size; + auto begin = host.begin() + offset; + auto end = begin + copy_size; + auto mismatch = std::find_if(begin, end, [=](uint8_t value) { return value != expected; }); + ASSERT_TRUE(mismatch == end) + << "Reader " << t << ", iteration " << i << ", byte " << mismatch - begin; + } + } + }); +} + namespace { __global__ void fill_kernel(int *target, int size, int start) { int idx = blockIdx.x * blockDim.x + threadIdx.x;