From 9f6af2fc7e10f1be18c405afd6c4d86b644519a8 Mon Sep 17 00:00:00 2001
From: Divye Gala <divyegala@gmail.com>
Date: Mon, 5 Feb 2024 20:38:53 -0500
Subject: [PATCH] Launch `neighborhood_recall` kernel on CUDA stream (#2156)

Authors:
  - Divye Gala (https://github.com/divyegala)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: https://github.com/rapidsai/raft/pull/2156
---
 cpp/include/raft/stats/detail/neighborhood_recall.cuh | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/cpp/include/raft/stats/detail/neighborhood_recall.cuh b/cpp/include/raft/stats/detail/neighborhood_recall.cuh
index 11d0448167..3f05225e74 100644
--- a/cpp/include/raft/stats/detail/neighborhood_recall.cuh
+++ b/cpp/include/raft/stats/detail/neighborhood_recall.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@
 #include <raft/core/math.hpp>
 #include <raft/core/mdspan_types.hpp>
 #include <raft/core/operators.hpp>
+#include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resources.hpp>
 
 #include <cub/cub.cuh>
@@ -108,7 +109,7 @@ void neighborhood_recall(
   auto constexpr kThreadsPerBlock = 32;
   auto const num_blocks           = indices.extent(0);
 
-  neighborhood_recall<<<num_blocks, kThreadsPerBlock>>>(
+  neighborhood_recall<<<num_blocks, kThreadsPerBlock, 0, raft::resource::get_cuda_stream(res)>>>(
     indices, ref_indices, distances, ref_distances, recall_score, eps);
 }