removing deprecated rmm policy usage

rapidsai · rapids-bot · May 27, 2021 · Apr 20, 2021 · Apr 20, 2021 · Apr 23, 2021
commit ec1cca59ab1076fdd083aac97b1b8a698570e984
diff --git a/cpp/include/raft/sparse/mst/detail/mst_solver_inl.cuh b/cpp/include/raft/sparse/mst/detail/mst_solver_inl.cuh
@@ -82,12 +82,12 @@ MST_solver<vertex_t, edge_t, weight_t>::MST_solver(
   //Initially, color holds the vertex id as color
   auto policy = rmm::exec_policy(stream);
   if (initialize_colors_) {
-    thrust::sequence(policy->on(stream), color.begin(), color.end(), 0);
-    thrust::sequence(policy->on(stream), color_index, color_index + v, 0);
+    thrust::sequence(policy, color.begin(), color.end(), 0);
+    thrust::sequence(policy, color_index, color_index + v, 0);
   } else {
     raft::copy(color.data().get(), color_index, v, stream);
   }
-  thrust::sequence(policy->on(stream), next_color.begin(), next_color.end(), 0);
+  thrust::sequence(policy, next_color.begin(), next_color.end(), 0);
 }
 
 template <typename vertex_t, typename edge_t, typename weight_t>
@@ -212,22 +212,22 @@ weight_t MST_solver<vertex_t, edge_t, weight_t>::alteration_max() {
   auto policy = rmm::exec_policy(stream);
   rmm::device_vector<weight_t> tmp(e);
   thrust::device_ptr<const weight_t> weights_ptr(weights);
-  thrust::copy(policy->on(stream), weights_ptr, weights_ptr + e, tmp.begin());
+  thrust::copy(policy, weights_ptr, weights_ptr + e, tmp.begin());
   //sort tmp weights
-  thrust::sort(policy->on(stream), tmp.begin(), tmp.end());
+  thrust::sort(policy, tmp.begin(), tmp.end());
 
   //remove duplicates
-  auto new_end = thrust::unique(policy->on(stream), tmp.begin(), tmp.end());
+  auto new_end = thrust::unique(policy, tmp.begin(), tmp.end());
 
   //min(a[i+1]-a[i])/2
   auto begin =
     thrust::make_zip_iterator(thrust::make_tuple(tmp.begin(), tmp.begin() + 1));
   auto end =
     thrust::make_zip_iterator(thrust::make_tuple(new_end - 1, new_end));
   auto init = tmp[1] - tmp[0];
-  auto max = thrust::transform_reduce(policy->on(stream), begin, end,
-                                      alteration_functor<weight_t>(), init,
-                                      thrust::minimum<weight_t>());
+  auto max =
+    thrust::transform_reduce(policy, begin, end, alteration_functor<weight_t>(),
+                             init, thrust::minimum<weight_t>());
   return max / static_cast<weight_t>(2);
 }
 
@@ -306,9 +306,9 @@ void MST_solver<vertex_t, edge_t, weight_t>::label_prop(vertex_t* mst_src,
 template <typename vertex_t, typename edge_t, typename weight_t>
 void MST_solver<vertex_t, edge_t, weight_t>::min_edge_per_vertex() {
   auto policy = rmm::exec_policy(stream);
-  thrust::fill(policy->on(stream), min_edge_color.begin(), min_edge_color.end(),
+  thrust::fill(policy, min_edge_color.begin(), min_edge_color.end(),
                std::numeric_limits<weight_t>::max());
-  thrust::fill(policy->on(stream), new_mst_edge.begin(), new_mst_edge.end(),
+  thrust::fill(policy, new_mst_edge.begin(), new_mst_edge.end(),
                std::numeric_limits<weight_t>::max());
 
   int n_threads = 32;
@@ -331,7 +331,7 @@ void MST_solver<vertex_t, edge_t, weight_t>::min_edge_per_supervertex() {
   auto nblocks = std::min((v + nthreads - 1) / nthreads, max_blocks);
 
   auto policy = rmm::exec_policy(stream);
-  thrust::fill(policy->on(stream), temp_src.begin(), temp_src.end(),
+  thrust::fill(policy, temp_src.begin(), temp_src.end(),
                std::numeric_limits<vertex_t>::max());
 
   vertex_t* color_ptr = color.data().get();
@@ -401,9 +401,8 @@ void MST_solver<vertex_t, edge_t, weight_t>::append_src_dst_pair(
     thrust::make_tuple(temp_src.end(), temp_dst.end(), temp_weights.end()));
 
   // copy new mst edges to final output
-  thrust::copy_if(policy->on(stream), temp_src_dst_zip_begin,
-                  temp_src_dst_zip_end, src_dst_zip_end,
-                  new_edges_functor<vertex_t, weight_t>());
+  thrust::copy_if(policy, temp_src_dst_zip_begin, temp_src_dst_zip_end,
+                  src_dst_zip_end, new_edges_functor<vertex_t, weight_t>());
 }
 
 }  // namespace mst

diff --git a/cpp/include/raft/sparse/mst/detail/utils.cuh b/cpp/include/raft/sparse/mst/detail/utils.cuh
@@ -17,8 +17,8 @@
 
 #pragma once
 
-#include <rmm/thrust_rmm_allocator.h>
 #include <iostream>
+#include <rmm/device_vector.hpp>
 #define MST_TIME
 
 namespace raft {

diff --git a/cpp/include/raft/sparse/mst/mst_solver.cuh b/cpp/include/raft/sparse/mst/mst_solver.cuh
@@ -17,9 +17,9 @@
 
 #pragma once
 
-#include <rmm/thrust_rmm_allocator.h>
 #include <raft/handle.hpp>
 #include <rmm/device_uvector.hpp>
+#include <rmm/device_vector.hpp>
 
 namespace raft {
 

diff --git a/cpp/test/mst.cu b/cpp/test/mst.cu
@@ -17,7 +17,6 @@
 #include <bits/stdc++.h>
 
 #include <gtest/gtest.h>
-#include <rmm/thrust_rmm_allocator.h>
 #include <iostream>
 #include <rmm/device_buffer.hpp>
 #include <vector>