Use raft::handle_t::get_thrust_policy to create thrust policy

cpp/include/raft/sparse/distance/coo_spmv.cuh

@@ -34,8 +34,6 @@
 
 #include <cusparse_v2.h>
 
-#include <rmm/exec_policy.hpp>
-
 namespace raft {
 namespace sparse {
 namespace distance {

cpp/include/raft/sparse/distance/coo_spmv_strategies/base_strategy.cuh

@@ -23,7 +23,6 @@
 
 #include <rmm/device_uvector.hpp>
 #include <rmm/device_vector.hpp>
-#include <rmm/exec_policy.hpp>
 
 namespace raft {
 namespace sparse {

cpp/include/raft/sparse/distance/coo_spmv_strategies/coo_mask_row_iterators.cuh

@@ -20,7 +20,6 @@
 #include "../utils.cuh"
 
 #include <rmm/device_uvector.hpp>
-#include <rmm/exec_policy.hpp>
 
 namespace raft {
 namespace sparse {

cpp/include/raft/sparse/distance/coo_spmv_strategies/hash_strategy.cuh

@@ -55,7 +55,7 @@ class hash_strategy : public coo_spmv_strategy<value_idx, value_t, tpb> {
                        rmm::device_uvector<value_idx> &mask_indptr,
                        std::tuple<value_idx, value_idx> &n_rows_divided,
                        cudaStream_t stream) {
-    auto policy = rmm::exec_policy(stream);
+    auto policy = this->config.handle.get_thrust_policy();
 
     auto less = thrust::copy_if(
       policy, thrust::make_counting_iterator(value_idx(0)),

cpp/include/raft/sparse/distance/l2_distance.cuh

@@ -245,9 +245,8 @@ class hellinger_expanded_distances_t : public distances_t<value_t> {
     : config_(&config), workspace(0, config.handle.get_stream()) {}
 
   void compute(value_t *out_dists) {
-    raft::mr::device::buffer<value_idx> coo_rows(
-      config_->handle.get_device_allocator(), config_->handle.get_stream(),
-      max(config_->b_nnz, config_->a_nnz));
+    rmm::device_uvector<value_idx> coo_rows(max(config_->b_nnz, config_->a_nnz),
+                                            config_->handle.get_stream());
 
     raft::sparse::convert::csr_to_coo(config_->b_indptr, config_->b_nrows,
                                       coo_rows.data(), config_->b_nnz,

cpp/include/raft/sparse/hierarchy/detail/agglomerative.cuh

@@ -21,7 +21,6 @@
 #include <raft/handle.hpp>
 #include <raft/mr/device/buffer.hpp>
 #include <rmm/device_uvector.hpp>
-#include <rmm/exec_policy.hpp>
 
 #include <thrust/device_ptr.h>
 #include <thrust/execution_policy.h>
@@ -224,7 +223,7 @@ void extract_flattened_clusters(const raft::handle_t &handle, value_idx *labels,
                                 const value_idx *children, size_t n_clusters,
                                 size_t n_leaves) {
   auto stream = handle.get_stream();
-  auto thrust_policy = rmm::exec_policy(rmm::cuda_stream_view{stream});
+  auto thrust_policy = handle.get_thrust_policy();
 
   // Handle special case where n_clusters == 1
   if (n_clusters == 1) {

cpp/include/raft/sparse/hierarchy/detail/connectivities.cuh

@@ -22,7 +22,6 @@
 
 #include <raft/linalg/unary_op.cuh>
 #include <rmm/device_uvector.hpp>
-#include <rmm/exec_policy.hpp>
 
 #include <raft/linalg/distance_type.h>
 #include <raft/sparse/hierarchy/common.h>
@@ -61,7 +60,7 @@ struct distance_graph_impl<raft::hierarchy::LinkageDistance::KNN_GRAPH,
            rmm::device_uvector<value_idx> &indices,
            rmm::device_uvector<value_t> &data, int c) {
     auto stream = handle.get_stream();
-    auto exec_policy = rmm::exec_policy(rmm::cuda_stream_view{stream});
+    auto thrust_policy = handle.get_thrust_policy();
 
     // Need to symmetrize knn into undirected graph
     raft::sparse::COO<value_t, value_idx> knn_graph_coo(stream);
@@ -77,7 +76,7 @@ struct distance_graph_impl<raft::hierarchy::LinkageDistance::KNN_GRAPH,
       knn_graph_coo.rows(), knn_graph_coo.cols(), knn_graph_coo.vals()));
 
     thrust::transform(
-      exec_policy, transform_in, transform_in + knn_graph_coo.nnz,
+      thrust_policy, transform_in, transform_in + knn_graph_coo.nnz,
       knn_graph_coo.vals(),
       [=] __device__(const thrust::tuple<value_idx, value_idx, value_t> &tup) {
         bool self_loop = thrust::get<0>(tup) == thrust::get<1>(tup);

cpp/include/raft/sparse/hierarchy/detail/mst.cuh

@@ -25,12 +25,9 @@
 #include <raft/sparse/selection/connect_components.cuh>
 #include <rmm/device_uvector.hpp>
 
-#include <rmm/exec_policy.hpp>
-
 #include <thrust/device_ptr.h>
 #include <thrust/execution_policy.h>
 #include <thrust/sort.h>
-#include <rmm/exec_policy.hpp>
 
 namespace raft {
 namespace hierarchy {

cpp/include/raft/sparse/linalg/spectral.cuh

@@ -63,7 +63,7 @@ void fit_embedding(const raft::handle_t &handle, int *rows, int *cols, T *vals,
   index_type maxiter = 4000;  //default reset value (when set to 0);
   value_type tol = 0.01;
   index_type restart_iter = 15 + neigvs;  //what cugraph is using
-  auto t_exe_p = thrust::cuda::par.on(stream);
+  auto t_exe_p = handle.get_thrust_policy();
   using thrust_exe_policy_t = decltype(t_exe_p);
 
   raft::eigen_solver_config_t<index_type, value_type> cfg{neigvs, maxiter,
@@ -83,8 +83,7 @@ void fit_embedding(const raft::handle_t &handle, int *rows, int *cols, T *vals,
     using value_type_t = value_type;
 
     std::pair<value_type_t, index_type_t> solve(
-      handle_t const &handle, thrust_exe_policy_t t_exe_policy,
-      size_type_t n_obs_vecs, size_type_t dim,
+      handle_t const &handle, size_type_t n_obs_vecs, size_type_t dim,
       value_type_t const *__restrict__ obs,
       index_type_t *__restrict__ codes) const {
       return std::make_pair<value_type_t, index_type_t>(0, 0);

cpp/include/raft/sparse/linalg/symmetrize.cuh

@@ -30,7 +30,6 @@
 
 #include <cuda_runtime.h>
 #include <stdio.h>
-#include <rmm/exec_policy.hpp>
 
 #include <algorithm>
 #include <iostream>

cpp/include/raft/sparse/mst/detail/mst_solver_inl.cuh

@@ -24,7 +24,6 @@
 
 #include <raft/cudart_utils.h>
 #include <rmm/device_buffer.hpp>
-#include <rmm/exec_policy.hpp>
 
 #include <thrust/device_ptr.h>
 #include <thrust/device_vector.h>
@@ -38,7 +37,6 @@
 #include <raft/cudart_utils.h>
 
 #include <rmm/device_buffer.hpp>
-#include <rmm/exec_policy.hpp>
 
 namespace raft {
 namespace mst {
@@ -88,7 +86,7 @@ MST_solver<vertex_t, edge_t, weight_t, alteration_t>::MST_solver(
   sm_count = handle_.get_device_properties().multiProcessorCount;
 
   //Initially, color holds the vertex id as color
-  auto policy = rmm::exec_policy(rmm::cuda_stream_view{stream});
+  auto policy = handle.get_thrust_policy();
   if (initialize_colors_) {
     thrust::sequence(policy, color.begin(), color.end(), 0);
     thrust::sequence(policy, color_index, color_index + v, 0);
@@ -227,7 +225,7 @@ template <typename vertex_t, typename edge_t, typename weight_t,
           typename alteration_t>
 alteration_t
 MST_solver<vertex_t, edge_t, weight_t, alteration_t>::alteration_max() {
-  auto policy = rmm::exec_policy(rmm::cuda_stream_view{stream});
+  auto policy = handle.get_thrust_policy();
   rmm::device_vector<weight_t> tmp(e);
   thrust::device_ptr<const weight_t> weights_ptr(weights);
   thrust::copy(policy, weights_ptr, weights_ptr + e, tmp.begin());
@@ -327,7 +325,7 @@ template <typename vertex_t, typename edge_t, typename weight_t,
           typename alteration_t>
 void MST_solver<vertex_t, edge_t, weight_t,
                 alteration_t>::min_edge_per_vertex() {
-  auto policy = rmm::exec_policy(rmm::cuda_stream_view{stream});
+  auto policy = handle.get_thrust_policy();
   thrust::fill(policy, min_edge_color.begin(), min_edge_color.end(),
                std::numeric_limits<alteration_t>::max());
   thrust::fill(policy, new_mst_edge.begin(), new_mst_edge.end(),
@@ -354,7 +352,7 @@ void MST_solver<vertex_t, edge_t, weight_t,
   auto nthreads = std::min(v, max_threads);
   auto nblocks = std::min((v + nthreads - 1) / nthreads, max_blocks);
 
-  auto policy = rmm::exec_policy(rmm::cuda_stream_view{stream});
+  auto policy = handle.get_thrust_policy();
   thrust::fill(policy, temp_src.begin(), temp_src.end(),
                std::numeric_limits<vertex_t>::max());
 
@@ -411,7 +409,7 @@ template <typename vertex_t, typename edge_t, typename weight_t,
           typename alteration_t>
 void MST_solver<vertex_t, edge_t, weight_t, alteration_t>::append_src_dst_pair(
   vertex_t* mst_src, vertex_t* mst_dst, weight_t* mst_weights) {
-  auto policy = rmm::exec_policy(rmm::cuda_stream_view{stream});
+  auto policy = handle.get_thrust_policy();
 
   auto curr_mst_edge_count = prev_mst_edge_count[0];
 

cpp/include/raft/sparse/op/reduce.cuh

@@ -31,7 +31,6 @@
 #include <cuda_runtime.h>
 #include <stdio.h>
 #include <rmm/device_uvector.hpp>
-#include <rmm/exec_policy.hpp>
 
 #include <algorithm>
 #include <iostream>
@@ -126,16 +125,15 @@ void max_duplicates(const raft::handle_t &handle,
                     const value_idx *rows, const value_idx *cols,
                     const value_t *vals, size_t nnz, size_t m, size_t n) {
   auto stream = handle.get_stream();
-
-  auto exec_policy = rmm::exec_policy(rmm::cuda_stream_view{stream});
+  auto thrust_policy = handle.get_thrust_policy();
 
   // compute diffs & take exclusive scan
   rmm::device_uvector<value_idx> diff(nnz + 1, stream);
 
   compute_duplicates_mask(diff.data(), rows, cols, nnz, stream);
 
-  thrust::exclusive_scan(thrust::cuda::par.on(stream), diff.data(),
-                         diff.data() + diff.size(), diff.data());
+  thrust::exclusive_scan(thrust_policy, diff.data(), diff.data() + diff.size(),
+                         diff.data());
 
   // compute final size
   value_idx size = 0;

cpp/include/raft/sparse/selection/connect_components.cuh

@@ -30,7 +30,6 @@
 #include <thrust/device_ptr.h>
 #include <thrust/sort.h>
 #include <rmm/device_uvector.hpp>
-#include <rmm/exec_policy.hpp>
 
 #include <cub/cub.cuh>
 
@@ -361,7 +360,7 @@ void connect_components(const raft::handle_t &handle,
   raft::sparse::op::compute_duplicates_mask(out_index.data(), colors.data(),
                                             nn_colors.data(), n_rows, stream);
 
-  thrust::exclusive_scan(thrust::cuda::par.on(stream), out_index.data(),
+  thrust::exclusive_scan(handle.get_thrust_policy(), out_index.data(),
                          out_index.data() + out_index.size(), out_index.data());
 
   // compute final size

cpp/include/raft/spectral/cluster_solvers.hpp

@@ -42,19 +42,16 @@ struct kmeans_solver_t {
                                                    size_type_t> const& config)
     : config_(config) {}
 
-  template <typename thrust_exe_policy_t>
   std::pair<value_type_t, index_type_t> solve(
-    handle_t const& handle, thrust_exe_policy_t t_exe_policy,
-    size_type_t n_obs_vecs, size_type_t dim,
+    handle_t const& handle, size_type_t n_obs_vecs, size_type_t dim,
     value_type_t const* __restrict__ obs,
     index_type_t* __restrict__ codes) const {
     RAFT_EXPECTS(obs != nullptr, "Null obs buffer.");
     RAFT_EXPECTS(codes != nullptr, "Null codes buffer.");
     value_type_t residual{};
     index_type_t iters{};
-    kmeans(handle, t_exe_policy, n_obs_vecs, dim, config_.n_clusters,
-           config_.tol, config_.maxIter, obs, codes, residual, iters,
-           config_.seed);
+    kmeans(handle, n_obs_vecs, dim, config_.n_clusters, config_.tol,
+           config_.maxIter, obs, codes, residual, iters, config_.seed);
     return std::make_pair(residual, iters);
   }