diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt index 90d23d9aef..9e3b69fdbe 100644 --- a/cpp/bench/ann/CMakeLists.txt +++ b/cpp/bench/ann/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= @@ -181,6 +181,7 @@ function(ConfigureAnnBench) sqlite3 Threads::Threads $<$:CUDA::cudart_static> + $<$:rmm::rmm> $ $ ) @@ -426,8 +427,8 @@ if(CUVS_ANN_BENCH_SINGLE_EXE) target_include_directories(ANN_BENCH PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) target_link_libraries( - ANN_BENCH PRIVATE raft::raft nlohmann_json::nlohmann_json sqlite3 benchmark::benchmark dl - $<$:CUDA::nvtx3> + ANN_BENCH PRIVATE $ nlohmann_json::nlohmann_json sqlite3 + benchmark::benchmark dl $<$:CUDA::nvtx3> ) set_target_properties( ANN_BENCH 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..f699bdad1a 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,9 +15,11 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -127,6 +129,7 @@ class configured_raft_resources { { raft::resource::set_large_workspace_resource( *res_, raft::mr::device_resource{shared_res_->get_large_memory_resource()}); + raft::resource::set_cuda_stream_pool(*res_, std::make_shared(1)); } /** Default constructor creates all resources anew. */ 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..8ed5adee26 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 @@ -15,7 +15,6 @@ #include #include #include -#include #include #include @@ -88,9 +87,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..363d065e91 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 @@ -20,7 +20,6 @@ #include #include #include -#include #include @@ -118,9 +117,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..0c41ee1c96 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 @@ -18,7 +18,6 @@ #include #include #include -#include #include @@ -98,9 +97,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..5bf0098eaa 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 @@ -12,9 +12,7 @@ #include #include #include -#include #include -#include #include #include @@ -83,8 +81,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_,