Revert "add more syncs, use thrust_policy"

This reverts commit 8d4d1a2.
rapidsai · Oct 3, 2024 · 366af06 · 366af06
1 parent 0409d12
commit 366af06
Showing 1 changed file with 5 additions and 14 deletions.
diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh
@@ -24,7 +24,6 @@
 #include <raft/core/error.hpp>
 #include <raft/core/host_mdarray.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
-#include <raft/core/resource/thrust_policy.hpp>
 #include <raft/core/resources.hpp>
 
 #include <raft/util/arch.cuh>  // raft::util::arch::SM_*
@@ -1163,19 +1162,15 @@ GNND<Data_t, Index_t>::GNND(raft::resources const& res, const BuildConfig& build
 {
   static_assert(NUM_SAMPLES <= 32);
 
-  thrust::fill(raft::resource::get_thrust_policy(res),
+  thrust::fill(thrust::device,
                dists_buffer_.data_handle(),
                dists_buffer_.data_handle() + dists_buffer_.size(),
                std::numeric_limits<float>::max());
-  thrust::fill(raft::resource::get_thrust_policy(res),
+  thrust::fill(thrust::device,
                reinterpret_cast<Index_t*>(graph_buffer_.data_handle()),
                reinterpret_cast<Index_t*>(graph_buffer_.data_handle()) + graph_buffer_.size(),
                std::numeric_limits<Index_t>::max());
-  thrust::fill(raft::resource::get_thrust_policy(res),
-               d_locks_.data_handle(),
-               d_locks_.data_handle() + d_locks_.size(),
-               0);
-  raft::resource::sync_stream(res);
+  thrust::fill(thrust::device, d_locks_.data_handle(), d_locks_.data_handle() + d_locks_.size(), 0);
 };
 
 template <typename Data_t, typename Index_t>
@@ -1194,7 +1189,7 @@ void GNND<Data_t, Index_t>::add_reverse_edges(Index_t* graph_ptr,
 template <typename Data_t, typename Index_t>
 void GNND<Data_t, Index_t>::local_join(cudaStream_t stream)
 {
-  thrust::fill(raft::resource::get_thrust_policy(res),
+  thrust::fill(thrust::device.on(stream),
                dists_buffer_.data_handle(),
                dists_buffer_.data_handle() + dists_buffer_.size(),
                std::numeric_limits<float>::max());
@@ -1213,7 +1208,6 @@ void GNND<Data_t, Index_t>::local_join(cudaStream_t stream)
     DEGREE_ON_DEVICE,
     d_locks_.data_handle(),
     l2_norms_.data_handle());
-  raft::resource::sync_stream(res);
 }
 
 template <typename Data_t, typename Index_t>
@@ -1245,11 +1239,10 @@ void GNND<Data_t, Index_t>::build(Data_t* data, const Index_t nrow, Index_t* out
                                        batch.offset());
   }
 
-  thrust::fill(raft::resource::get_thrust_policy(res),
+  thrust::fill(thrust::device.on(stream),
                (Index_t*)graph_buffer_.data_handle(),
                (Index_t*)graph_buffer_.data_handle() + graph_buffer_.size(),
                std::numeric_limits<Index_t>::max());
-  raft::resource::sync_stream(res);
 
   graph_.clear();
   graph_.init_random_graph();
@@ -1336,7 +1329,6 @@ void GNND<Data_t, Index_t>::build(Data_t* data, const Index_t nrow, Index_t* out
     graph_.sample_graph_new(thrust::raw_pointer_cast(graph_host_buffer_.data()), DEGREE_ON_DEVICE);
   }
 
-  raft::resource::sync_stream(res);
   graph_.update_graph(thrust::raw_pointer_cast(graph_host_buffer_.data()),
                       thrust::raw_pointer_cast(dists_host_buffer_.data()),
                       DEGREE_ON_DEVICE,
@@ -1422,7 +1414,6 @@ void build(raft::resources const& res,
 
   GNND<const T, int> nnd(res, build_config);
   nnd.build(dataset.data_handle(), dataset.extent(0), int_graph.data_handle());
-  raft::resource::sync_stream(res);
 
 #pragma omp parallel for
   for (size_t i = 0; i < static_cast<size_t>(dataset.extent(0)); i++) {