add support multi-stream for CompositeIndex::search

Author: rhdong

cpp/src/neighbors/composite/index.cu

@@ -18,6 +18,7 @@
 #include <cuvs/neighbors/composite/index.hpp>
 #include <cuvs/selection/select_k.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/cuda_stream_pool.hpp>
 #include <raft/linalg/add.cuh>
 #include <rmm/device_uvector.hpp>
 
@@ -47,29 +48,59 @@ void CompositeIndex<T, IdxT, OutputIdxT>::search(
   size_t num_indices = children_.size();
   size_t buffer_size = num_queries * K * num_indices;
 
-  auto stream  = raft::resource::get_cuda_stream(handle);
-  auto tmp_res = raft::resource::get_workspace_resource(handle);
+  auto main_stream = raft::resource::get_cuda_stream(handle);
+  auto tmp_res     = raft::resource::get_workspace_resource(handle);
 
-  rmm::device_uvector<out_index_type> neighbors_buffer(buffer_size, stream, tmp_res);
-  rmm::device_uvector<float> distances_buffer(buffer_size, stream, tmp_res);
+  rmm::device_uvector<out_index_type> neighbors_buffer(buffer_size, main_stream, tmp_res);
+  rmm::device_uvector<float> distances_buffer(buffer_size, main_stream, tmp_res);
+
+  std::vector<rmm::device_uvector<out_index_type>> temp_neighbors;
+  std::vector<rmm::device_uvector<float>> temp_distances;
+
+  for (size_t i = 0; i < num_indices; i++) {
+    temp_neighbors.emplace_back(num_queries * K, main_stream, tmp_res);
+    temp_distances.emplace_back(num_queries * K, main_stream, tmp_res);
+  }
+
+  raft::resource::wait_stream_pool_on_stream(handle);
 
   out_index_type offset = 0;
   out_index_type stride = K * num_indices;
 
   for (size_t i = 0; i < num_indices; i++) {
     const auto& sub_index = children_[i];
-    sub_index->search(handle, params, queries, neighbors, distances, filter);
+
+    auto stream = raft::resource::get_next_usable_stream(handle, i);
+
+    raft::resources stream_pool_handle(handle);
+    raft::resource::set_cuda_stream(stream_pool_handle, stream);
+
+    auto temp_neighbors_view =
+      raft::make_device_matrix_view<out_index_type, matrix_index_type, raft::row_major>(
+        temp_neighbors[i].data(), num_queries, K);
+    auto temp_distances_view =
+      raft::make_device_matrix_view<float, matrix_index_type, raft::row_major>(
+        temp_distances[i].data(), num_queries, K);
+
+    sub_index->search(
+      stream_pool_handle, params, queries, temp_neighbors_view, temp_distances_view, filter);
+
     if (offset != 0) {
-      raft::linalg::addScalar(
-        neighbors.data_handle(), neighbors.data_handle(), offset, neighbors.size(), stream);
+      raft::linalg::addScalar(temp_neighbors[i].data(),
+                              temp_neighbors[i].data(),
+                              offset,
+                              temp_neighbors[i].size(),
+                              stream);
     }
 
     raft::copy_matrix(
-      neighbors_buffer.data() + i * K, stride, neighbors.data_handle(), K, K, num_queries, stream);
+      neighbors_buffer.data() + i * K, stride, temp_neighbors[i].data(), K, K, num_queries, stream);
     raft::copy_matrix(
-      distances_buffer.data() + i * K, stride, distances.data_handle(), K, K, num_queries, stream);
+      distances_buffer.data() + i * K, stride, temp_distances[i].data(), K, K, num_queries, stream);
+
     offset += sub_index->size();
   }
+  raft::resource::sync_stream_pool(handle);
 
   auto distances_view = raft::make_device_matrix_view<const float, matrix_index_type>(
     distances_buffer.data(), num_queries, K * num_indices);