diff --git a/cpp/bench/ann/src/common/util.hpp b/cpp/bench/ann/src/common/util.hpp index 6e4c57b35e..069d257717 100644 --- a/cpp/bench/ann/src/common/util.hpp +++ b/cpp/bench/ann/src/common/util.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -7,6 +7,10 @@ #include "ann_types.hpp" #include "cuda_stub.hpp" // cuda-related utils +#ifndef BUILD_CPU_ONLY +#include +#endif + #if __has_include() #define ANN_BENCH_NVTX3_HEADERS_FOUND #include @@ -145,25 +149,8 @@ struct cuda_timer { }; #ifndef BUILD_CPU_ONLY -// ATM, rmm::stream does not support passing in flags; hence this helper type. -struct non_blocking_stream { - non_blocking_stream() { cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking); } - ~non_blocking_stream() noexcept - { - if (stream_ != nullptr) { cudaStreamDestroy(stream_); } - } - non_blocking_stream(non_blocking_stream const&) = delete; - non_blocking_stream(non_blocking_stream&& other) noexcept { std::swap(stream_, other.stream_); } - auto operator=(non_blocking_stream const&) -> non_blocking_stream& = delete; - auto operator=(non_blocking_stream&&) -> non_blocking_stream& = delete; - [[nodiscard]] auto view() const noexcept -> cudaStream_t { return stream_; } - - private: - cudaStream_t stream_{nullptr}; -}; - namespace detail { -inline std::vector global_stream_pool(0); +inline std::vector global_stream_pool(0); inline std::mutex gsp_mutex; } // namespace detail #endif @@ -180,7 +167,9 @@ inline auto get_stream_from_global_pool() -> cudaStream_t #ifndef BUILD_CPU_ONLY std::lock_guard guard(detail::gsp_mutex); if (static_cast(detail::global_stream_pool.size()) < benchmark_n_threads) { - detail::global_stream_pool.resize(benchmark_n_threads); + while (static_cast(detail::global_stream_pool.size()) < benchmark_n_threads) { + detail::global_stream_pool.emplace_back(rmm::cuda_stream::flags::non_blocking); + } } return detail::global_stream_pool[benchmark_thread_id].view(); #else diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h index 1a276e8cc8..0ca96cb3c0 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -15,13 +15,15 @@ #include #include #include +#include #include #include +#include #include #include +#include #include -#include #include #include #include @@ -67,7 +69,7 @@ inline auto rmm_oom_callback(std::size_t bytes, void*) -> bool */ class shared_raft_resources { public: - using large_mr_type = rmm::mr::managed_memory_resource; + using large_mr_type = rmm::mr::cuda_async_managed_memory_resource; shared_raft_resources() try : large_mr_() { @@ -132,6 +134,7 @@ class configured_raft_resources { /** Default constructor creates all resources anew. */ configured_raft_resources() : configured_raft_resources{std::make_shared()} { + raft::resource::set_cuda_stream_pool(*res_, std::make_shared(1)); } configured_raft_resources(configured_raft_resources&&); diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h index ae334f70d7..9d73990986 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_flat_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -88,9 +88,6 @@ class cuvs_ivf_flat : public algo, public algo_gpu { template void cuvs_ivf_flat::build(const T* dataset, size_t nrow) { - // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping. - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); index_ = std::make_shared>( std::move(cuvs::neighbors::ivf_flat::build( handle_, diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h index 161563d8fe..0ddb41d19e 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_pq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -118,9 +118,6 @@ void cuvs_ivf_pq::load(const std::string& file) template void cuvs_ivf_pq::build(const T* dataset, size_t nrow) { - // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping. - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); auto dataset_v = raft::make_device_matrix_view(dataset, IdxT(nrow), dim_); std::make_shared>( std::move(cuvs::neighbors::ivf_pq::build(handle_, index_params_, dataset_v))) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h index ca8f77b808..d9b12da607 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_rabitq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -98,9 +98,6 @@ void cuvs_ivf_rabitq::load(const std::string& file) template void cuvs_ivf_rabitq::build(const T* dataset, size_t nrow) { - // Create a CUDA stream pool with 1 stream (besides main stream) for kernel/copy overlapping. - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); auto dataset_v = raft::make_device_matrix_view(dataset, IdxT(nrow), dim_); std::make_shared>( std::move(cuvs::neighbors::ivf_rabitq::build(handle_, index_params_, dataset_v))) diff --git a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h index 1503e6bb84..8a96951803 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ivf_sq_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -83,8 +83,6 @@ class cuvs_ivf_sq : public algo, public algo_gpu { template void cuvs_ivf_sq::build(const T* dataset, size_t nrow) { - size_t n_streams = 1; - raft::resource::set_cuda_stream_pool(handle_, std::make_shared(n_streams)); index_ = std::make_shared>( std::move(cuvs::neighbors::ivf_sq::build( handle_,