review feedback

cpp/include/raft/handle.hpp

@@ -113,7 +113,7 @@ class handle_t {
     std::lock_guard<std::mutex> _(mutex_);
     if (!cublas_initialized_) {
       CUBLAS_CHECK(cublasCreate(&cublas_handle_));
-      CUBLAS_CHECK(cublasSetStream(cublas_handle, stream_view_));
+      CUBLAS_CHECK(cublasSetStream(cublas_handle_, stream_view_));
       cublas_initialized_ = true;
     }
     return cublas_handle_;
@@ -163,6 +163,17 @@ class handle_t {
     }
   }
 
+  /**
+   * @brief synchronize subset of stream pool
+   * 
+   * @param[in] stream_indices the indices of the streams in the stream pool to synchronize
+   */
+  void sync_stream_pool(const std::vector<std::size_t> stream_indices) const {
+    for (const auto& stream_index : stream_indices) {
+      stream_pool_.get_stream(stream_index).synchronize();
+    }
+  }
+
   /**
    * @brief ask stream pool to wait on last event in main stream
    */

cpp/include/raft/spatial/knn/detail/knn_brute_force_faiss.cuh

@@ -270,8 +270,9 @@ void brute_force_knn_impl(std::vector<float *> &input, std::vector<int> &sizes,
     float *out_d_ptr = out_D + (i * k * n);
     int64_t *out_i_ptr = out_I + (i * k * n);
 
-    cudaStream_t stream =
-      n_internal_streams > 0 ? internalStreams.get_stream() : userStream;
+    cudaStream_t stream = n_internal_streams > 0
+                            ? internalStreams.get_stream().value()
+                            : userStream;
 
     switch (metric) {
       case raft::distance::DistanceType::Haversine:

cpp/test/handle.cpp

@@ -35,9 +35,10 @@ TEST(Raft, HandleDefault) {
 
 TEST(Raft, Handle) {
   // test stream pool creation
-  rmm::cuda_stream_pool stream_pool{4};
+  constexpr std::size_t n_streams = 4;
+  rmm::cuda_stream_pool stream_pool{n_streams};
   handle_t h(rmm::cuda_stream_default, stream_pool);
-  ASSERT_EQ(4, h.get_stream_pool().get_pool_size());
+  ASSERT_EQ(n_streams, h.get_stream_pool().get_pool_size());
 
   // test non default stream handle
   cudaStream_t stream;