[ANN_BENCH] Remove Stream Pool Creation from Individual Wrappers

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
@@ -144,30 +144,6 @@ 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<non_blocking_stream> global_stream_pool(0);
-inline std::mutex gsp_mutex;
-}  // namespace detail
-#endif
-
 /**
  * Get a stream associated with the current benchmark thread.
  *

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 <raft/core/host_mdspan.hpp>
 #include <raft/core/logger.hpp>
 #include <raft/core/operators.hpp>
+#include <raft/core/resource/cuda_stream_pool.hpp>
 #include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/util/cudart_utils.hpp>
 
+#include <rmm/cuda_stream_pool.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
+#include <rmm/mr/cuda_async_managed_memory_resource.hpp>
 #include <rmm/mr/failure_callback_resource_adaptor.hpp>
-#include <rmm/mr/managed_memory_resource.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
 #include <rmm/resource_ref.hpp>
@@ -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<shared_raft_resources>()}
   {
+    raft::resource::set_cuda_stream_pool(*res_, std::make_shared<rmm::cuda_stream_pool>(1));
   }
 
   configured_raft_resources(configured_raft_resources&&);

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<T>, public algo_gpu {
 template <typename T, typename IdxT>
 void cuvs_ivf_flat<T, IdxT>::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<rmm::cuda_stream_pool>(n_streams));
   index_ = std::make_shared<cuvs::neighbors::ivf_flat::index<T, IdxT>>(
     std::move(cuvs::neighbors::ivf_flat::build(
       handle_,

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<T, IdxT>::load(const std::string& file)
 template <typename T, typename IdxT>
 void cuvs_ivf_pq<T, IdxT>::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<rmm::cuda_stream_pool>(n_streams));
   auto dataset_v = raft::make_device_matrix_view<const T, IdxT>(dataset, IdxT(nrow), dim_);
   std::make_shared<cuvs::neighbors::ivf_pq::index<IdxT>>(
     std::move(cuvs::neighbors::ivf_pq::build(handle_, index_params_, dataset_v)))

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<T, IdxT>::load(const std::string& file)
 template <typename T, typename IdxT>
 void cuvs_ivf_rabitq<T, IdxT>::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<rmm::cuda_stream_pool>(n_streams));
   auto dataset_v = raft::make_device_matrix_view<const T, IdxT>(dataset, IdxT(nrow), dim_);
   std::make_shared<cuvs::neighbors::ivf_rabitq::index<IdxT>>(
     std::move(cuvs::neighbors::ivf_rabitq::build(handle_, index_params_, dataset_v)))

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<T>, public algo_gpu {
 template <typename T>
 void cuvs_ivf_sq<T>::build(const T* dataset, size_t nrow)
 {
-  size_t n_streams = 1;
-  raft::resource::set_cuda_stream_pool(handle_, std::make_shared<rmm::cuda_stream_pool>(n_streams));
   index_ = std::make_shared<cuvs::neighbors::ivf_sq::index<uint8_t>>(
     std::move(cuvs::neighbors::ivf_sq::build(
       handle_,