diff --git a/cpp/src/neighbors/detail/cagra/compute_distance.hpp b/cpp/src/neighbors/detail/cagra/compute_distance.hpp index c1d0033a01..99bd6fc78d 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance.hpp @@ -253,10 +253,13 @@ struct dataset_descriptor_host { { if (!ready.load(std::memory_order_acquire)) { eval(stream); } // value is immutable at this point. - auto& [ptr, ready_stream] = std::get(value); - if (ready_stream != stream) { - RAFT_CUDA_TRY(cudaStreamWaitEvent(stream.get(), ready_event, 0)); - } + auto* ptr = std::get<0>(std::get(value)); + // Always order `stream` after the initialization kernel. Comparing the stream handles is not + // enough: the per-thread default stream (`cudaStreamPerThread`, which is also raft's default + // stream) has the same handle in every host thread but refers to a different stream in each. + // When several threads share a raft::resources (and thus this cached descriptor), skipping + // the wait lets a search kernel read the descriptor before it is initialized. + RAFT_CUDA_TRY(cudaStreamWaitEvent(stream.get(), ready_event, 0)); return ptr; } }; diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index ace28b7842..f35fc1ffb2 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -16,6 +16,8 @@ #include +#include + namespace cuvs::neighbors::cagra::detail { template , dataset_descriptor_host> value{kDefaultSize}; + /** Serializes lookup-or-insert, so that concurrent searches share a single descriptor. */ + std::mutex mutex; }; } // namespace descriptor_cache @@ -170,9 +174,17 @@ auto dataset_descriptor_init_with_cache(const raft::resources& res, -> dataset_descriptor_host { auto key = descriptor_cache::make_key(params, dataset, metric, dataset_norms); - auto& cache = - raft::resource::get_custom_resource>(res) - ->value; + auto* store = + raft::resource::get_custom_resource>(res); + auto& cache = store->value; + // Hold the lock across the lookup and the insertion. Otherwise threads that miss concurrently + // each create a descriptor and insert it under the same key (the LRU cache does not deduplicate + // keys), so concurrent searches can end up with different descriptors. When the last owner of + // one releases it, its device memory is freed with `cudaFreeAsync` on the stream that + // initialized it, which is not ordered after a kernel on another thread's stream that can still + // be reading it. Only host-side work happens under the lock (the device-side initialization is + // deferred to the first `dev_ptr()` call). + std::lock_guard lock(store->mutex); dataset_descriptor_host desc; if (!cache.get(key, &desc)) { desc =