[ANN_BENCH] Remove Stream Pool Creation from Individual Wrappers

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,8 @@ 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, rmm::cuda_stream::flags::non_blocking));
   }
 
   configured_raft_resources(configured_raft_resources&&);

cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h

@@ -679,7 +679,8 @@ void cuvs_cagra<T, IdxT>::search_base(
         raft::resources composite_handle(handle_);
         size_t n_streams = cagra_indices.size();
         raft::resource::set_cuda_stream_pool(composite_handle,
-                                             std::make_shared<rmm::cuda_stream_pool>(n_streams));
+                                             std::make_shared<rmm::cuda_stream_pool>(
+                                               n_streams, rmm::cuda_stream::flags::non_blocking));
 
         cuvs::neighbors::composite::composite_index<T, IdxT, algo_base::index_type> composite(
           cagra_indices);

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_,