Use simplified rmm::exec_policy (#6939)

Updates libcudf to use the new, simplified `rmm::exec_policy` and include the new refactored headers `rmm/exec_policy.hpp` and `rmm/device_vector.hpp`

The new `exec_policy` can be passed directly to Thrust, no longer any need to call `rmm::exec_policy(stream)->on(stream)`.

Depends on rapidsai/rmm#647

CHANGELOG.md

@@ -12,6 +12,7 @@
 - PR #6275 Update to official libcu++ on Github
 - PR #6838 Fix `columns` & `index` handling in dataframe constructor
 - PR #6750 Remove **kwargs from string/categorical methods
+- PR #6939 Use simplified `rmm::exec_policy`
 
 ## Bug Fixes
 

cpp/docs/TRANSITIONGUIDE.md

@@ -155,7 +155,7 @@ namespace detail{
         RMM_ALLOC(...,stream);
         CUDA_TRY(cudaMemcpyAsync(...,stream.value()));
         kernel<<<..., stream>>>(...);
-        thrust::algorithm(rmm::exec_policy(stream)->on(stream), ...);
+        thrust::algorithm(rmm::exec_policy(stream), ...);
         stream.synchronize();
         RMM_FREE(...,stream);
     }

cpp/include/cudf/column/column_factories.hpp

@@ -19,8 +19,8 @@
 #include <cudf/types.hpp>
 #include <cudf/utilities/traits.hpp>
 
-#include <rmm/thrust_rmm_allocator.h>
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_vector.hpp>
 
 namespace cudf {
 /**

cpp/include/cudf/detail/aggregation/aggregation.cuh

@@ -24,6 +24,7 @@
 #include <cudf/table/table_device_view.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/exec_policy.hpp>
 
 namespace cudf {
 namespace detail {
@@ -423,10 +424,7 @@ struct identity_initializer {
   std::enable_if_t<is_supported<T, k>(), void> operator()(mutable_column_view const& col,
                                                           rmm::cuda_stream_view stream)
   {
-    thrust::fill(rmm::exec_policy(stream)->on(stream.value()),
-                 col.begin<T>(),
-                 col.end<T>(),
-                 get_identity<T, k>());
+    thrust::fill(rmm::exec_policy(stream), col.begin<T>(), col.end<T>(), get_identity<T, k>());
   }
 
   template <typename T, aggregation::Kind k>

cpp/include/cudf/detail/copy_if.cuh

@@ -283,7 +283,7 @@ struct scatter_gather_functor {
   {
     rmm::device_uvector<cudf::size_type> indices(output_size, stream);
 
-    thrust::copy_if(rmm::exec_policy(stream)->on(stream.value()),
+    thrust::copy_if(rmm::exec_policy(stream),
                     thrust::counting_iterator<cudf::size_type>(0),
                     thrust::counting_iterator<cudf::size_type>(input.size()),
                     indices.begin(),

cpp/include/cudf/detail/gather.cuh

@@ -34,8 +34,9 @@
 #include <cudf/utilities/traits.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
-#include <rmm/thrust_rmm_allocator.h>
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_vector.hpp>
+#include <rmm/exec_policy.hpp>
 
 #include <algorithm>
 
@@ -125,19 +126,16 @@ void gather_helper(InputItr source_itr,
 {
   using map_type = typename std::iterator_traits<MapIterator>::value_type;
   if (nullify_out_of_bounds) {
-    thrust::gather_if(rmm::exec_policy(stream)->on(stream.value()),
+    thrust::gather_if(rmm::exec_policy(stream),
                       gather_map_begin,
                       gather_map_end,
                       gather_map_begin,
                       source_itr,
                       target_itr,
                       bounds_checker<map_type>{0, source_size});
   } else {
-    thrust::gather(rmm::exec_policy(stream)->on(stream.value()),
-                   gather_map_begin,
-                   gather_map_end,
-                   source_itr,
-                   target_itr);
+    thrust::gather(
+      rmm::exec_policy(stream), gather_map_begin, gather_map_end, source_itr, target_itr);
   }
 }
 

cpp/include/cudf/detail/groupby/sort_helper.hpp

@@ -21,8 +21,8 @@
 #include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
 
-#include <rmm/thrust_rmm_allocator.h>
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_vector.hpp>
 
 namespace cudf {
 namespace groupby {

cpp/include/cudf/detail/indexalator.cuh

@@ -314,7 +314,7 @@ struct input_indexalator : base_indexalator<input_indexalator> {
  * Example output iterator usage.
  * @code
  *  auto result_itr = indexalator_factory::create_output_iterator(indices->mutable_view());
- *  thrust::lower_bound(rmm::exec_policy(stream)->on(stream),
+ *  thrust::lower_bound(rmm::exec_policy(stream),
  *                      input->begin<Element>(),
  *                      input->end<Element>(),
  *                      values->begin<Element>(),

cpp/include/cudf/detail/reduction.cuh

@@ -20,10 +20,10 @@
 
 #include <cudf/utilities/type_dispatcher.hpp>
 
-#include <rmm/thrust_rmm_allocator.h>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/device_scalar.hpp>
+#include <rmm/exec_policy.hpp>
 
 #include <cub/device/device_reduce.cuh>
 
@@ -209,7 +209,7 @@ std::unique_ptr<scalar> reduce(InputIterator d_in,
   // compute the result value from intermediate value in device
   using ScalarType = cudf::scalar_type_t<OutputType>;
   auto result      = new ScalarType(OutputType{0}, true, stream, mr);
-  thrust::for_each_n(rmm::exec_policy(stream)->on(stream.value()),
+  thrust::for_each_n(rmm::exec_policy(stream),
                      intermediate_result.data(),
                      1,
                      [dres = result->data(), cop, valid_count, ddof] __device__(auto i) {

cpp/include/cudf/detail/scatter.cuh

@@ -30,6 +30,7 @@
 #include <cudf/utilities/traits.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/exec_policy.hpp>
 
 namespace cudf {
 namespace detail {
@@ -69,7 +70,7 @@ auto scatter_to_gather(MapIterator scatter_map_begin,
 
   // Convert scatter map to a gather map
   thrust::scatter(
-    rmm::exec_policy(stream)->on(stream.value()),
+    rmm::exec_policy(stream),
     thrust::make_counting_iterator<MapValueType>(0),
     thrust::make_counting_iterator<MapValueType>(std::distance(scatter_map_begin, scatter_map_end)),
     scatter_map_begin,
@@ -94,7 +95,7 @@ struct column_scatterer_impl {
 
     // NOTE use source.begin + scatter rows rather than source.end in case the
     // scatter map is smaller than the number of source rows
-    thrust::scatter(rmm::exec_policy(stream)->on(stream.value()),
+    thrust::scatter(rmm::exec_policy(stream),
                     source.begin<Type>(),
                     source.begin<Type>() + cudf::distance(scatter_map_begin, scatter_map_end),
                     scatter_map_begin,
@@ -180,7 +181,7 @@ struct column_scatterer_impl<dictionary32, MapIterator> {
     auto source_itr  = indexalator_factory::make_input_iterator(source_view.indices());
     auto new_indices = std::make_unique<column>(target_view.get_indices_annotated(), stream, mr);
     auto target_itr  = indexalator_factory::make_output_iterator(new_indices->mutable_view());
-    thrust::scatter(rmm::exec_policy(stream)->on(stream.value()),
+    thrust::scatter(rmm::exec_policy(stream),
                     source_itr,
                     source_itr + std::distance(scatter_map_begin, scatter_map_end),
                     scatter_map_begin,
@@ -262,8 +263,7 @@ std::unique_ptr<table> scatter(
     auto bounds      = bounds_checker<MapType>{begin, end};
     CUDF_EXPECTS(
       std::distance(scatter_map_begin, scatter_map_end) ==
-        thrust::count_if(
-          rmm::exec_policy(stream)->on(stream.value()), scatter_map_begin, scatter_map_end, bounds),
+        thrust::count_if(rmm::exec_policy(stream), scatter_map_begin, scatter_map_end, bounds),
       "Scatter map index out of bounds");
   }
 

cpp/include/cudf/detail/unary.hpp

@@ -20,6 +20,7 @@
 #include <cudf/unary.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/exec_policy.hpp>
 
 namespace cudf {
 namespace detail {
@@ -54,7 +55,7 @@ std::unique_ptr<column> true_if(
   auto output_mutable_view = output->mutable_view();
   auto output_data         = output_mutable_view.data<bool>();
 
-  thrust::transform(rmm::exec_policy(stream)->on(stream.value()), begin, end, output_data, p);
+  thrust::transform(rmm::exec_policy(stream), begin, end, output_data, p);
 
   return output;
 }

cpp/include/cudf/lists/detail/gather.cuh

@@ -21,6 +21,7 @@
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
+#include <rmm/exec_policy.hpp>
 
 #include <thrust/transform_scan.h>
 
@@ -82,7 +83,7 @@ gather_data make_gather_data(cudf::lists_column_view const& source_column,
   // generate the compacted outgoing offsets.
   auto count_iter = thrust::make_counting_iterator<int32_t>(0);
   thrust::transform_exclusive_scan(
-    rmm::exec_policy(stream)->on(stream.value()),
+    rmm::exec_policy(stream),
     count_iter,
     count_iter + offset_count,
     dst_offsets_v.begin<int32_t>(),
@@ -106,7 +107,7 @@ gather_data make_gather_data(cudf::lists_column_view const& source_column,
 
   // generate the base offsets
   rmm::device_uvector<int32_t> base_offsets = rmm::device_uvector<int32_t>(output_count, stream);
-  thrust::transform(rmm::exec_policy(stream)->on(stream.value()),
+  thrust::transform(rmm::exec_policy(stream),
                     gather_map,
                     gather_map + output_count,
                     base_offsets.data(),

cpp/include/cudf/lists/detail/scatter.cuh

@@ -28,6 +28,7 @@
 #include <cudf/types.hpp>
 
 #include <rmm/device_uvector.hpp>
+#include <rmm/exec_policy.hpp>
 
 #include <thrust/binary_search.h>
 
@@ -148,7 +149,7 @@ rmm::device_uvector<unbound_list_view> list_vector_from_column(
 
   auto vector = rmm::device_uvector<unbound_list_view>(n_rows, stream, mr);
 
-  thrust::transform(rmm::exec_policy(stream)->on(stream.value()),
+  thrust::transform(rmm::exec_policy(stream),
                     thrust::make_counting_iterator<size_type>(0),
                     thrust::make_counting_iterator<size_type>(n_rows),
                     vector.begin(),
@@ -233,7 +234,7 @@ void print(std::string const& msg, column_view const& col, rmm::cuda_stream_view
 
   std::cout << msg << " = [";
   thrust::for_each_n(
-    rmm::exec_policy(stream)->on(stream.value()),
+    rmm::exec_policy(stream),
     thrust::make_counting_iterator<size_type>(0),
     col.size(),
     [c = col.template data<int32_t>()] __device__(auto const& i) { printf("%d,", c[i]); });
@@ -246,7 +247,7 @@ void print(std::string const& msg,
 {
   std::cout << msg << " == [";
 
-  thrust::for_each_n(rmm::exec_policy(stream)->on(stream.value()),
+  thrust::for_each_n(rmm::exec_policy(stream),
                      thrust::make_counting_iterator<size_type>(0),
                      scatter.size(),
                      [s = scatter.begin()] __device__(auto const& i) {
@@ -420,7 +421,7 @@ struct list_child_constructor {
     };
 
     // For each list-row, copy underlying elements to the child column.
-    thrust::for_each_n(rmm::exec_policy(stream)->on(stream.value()),
+    thrust::for_each_n(rmm::exec_policy(stream),
                        thrust::make_counting_iterator<size_type>(0),
                        list_vector.size(),
                        copy_child_values_for_list_index);
@@ -490,7 +491,7 @@ struct list_child_constructor {
         });
     };
 
-    thrust::for_each_n(rmm::exec_policy(stream)->on(stream.value()),
+    thrust::for_each_n(rmm::exec_policy(stream),
                        thrust::make_counting_iterator<size_type>(0),
                        list_vector.size(),
                        populate_string_views);
@@ -581,7 +582,7 @@ struct list_child_constructor {
         });
     };
 
-    thrust::for_each_n(rmm::exec_policy(stream)->on(stream.value()),
+    thrust::for_each_n(rmm::exec_policy(stream),
                        thrust::make_counting_iterator<size_type>(0),
                        list_vector.size(),
                        populate_child_list_views);
@@ -780,7 +781,7 @@ std::unique_ptr<column> scatter(
                                                mr);
 
   // Scatter.
-  thrust::scatter(rmm::exec_policy(stream)->on(stream.value()),
+  thrust::scatter(rmm::exec_policy(stream),
                   source_vector.begin(),
                   source_vector.end(),
                   scatter_map_begin,

cpp/include/cudf/scalar/scalar.hpp

@@ -21,7 +21,6 @@
 
 #include <cudf/fixed_point/fixed_point.hpp>
 
-#include <rmm/thrust_rmm_allocator.h>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/device_scalar.hpp>

cpp/include/cudf/strings/detail/copy_if_else.cuh

@@ -23,6 +23,7 @@
 #include <cudf/strings/strings_column_view.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/exec_policy.hpp>
 
 namespace cudf {
 namespace strings {
@@ -62,7 +63,6 @@ std::unique_ptr<cudf::column> copy_if_else(
   auto strings_count = std::distance(lhs_begin, lhs_end);
   if (strings_count == 0) return make_empty_strings_column(stream, mr);
 
-  auto execpol = rmm::exec_policy(stream);
   // create null mask
   auto valid_mask = cudf::detail::valid_if(
     thrust::make_counting_iterator<size_type>(0),
@@ -97,7 +97,7 @@ std::unique_ptr<cudf::column> copy_if_else(
   auto d_chars      = chars_column->mutable_view().template data<char>();
   // fill in chars
   thrust::for_each_n(
-    execpol->on(stream.value()),
+    rmm::exec_policy(stream),
     thrust::make_counting_iterator<size_type>(0),
     strings_count,
     [lhs_begin, rhs_begin, filter_fn, d_offsets, d_chars] __device__(size_type idx) {

cpp/include/cudf/strings/detail/copy_range.cuh

@@ -24,6 +24,7 @@
 #include <cudf/types.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/exec_policy.hpp>
 
 #include <thrust/for_each.h>
 #include <thrust/iterator/counting_iterator.h>
@@ -187,7 +188,7 @@ std::unique_ptr<column> copy_range(
     // copy to the chars column
 
     auto p_chars = (p_chars_column->mutable_view()).template data<char>();
-    thrust::for_each(rmm::exec_policy(stream)->on(stream.value()),
+    thrust::for_each(rmm::exec_policy(stream),
                      thrust::make_counting_iterator(0),
                      thrust::make_counting_iterator(target.size()),
                      [source_value_begin,

cpp/include/cudf/strings/detail/gather.cuh

@@ -24,6 +24,7 @@
 #include <cudf/strings/strings_column_view.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/exec_policy.hpp>
 
 namespace cudf {
 
@@ -70,7 +71,6 @@ std::unique_ptr<cudf::column> gather(
   auto strings_count = strings.size();
   if (output_count == 0) return make_empty_strings_column(stream, mr);
 
-  auto execpol        = rmm::exec_policy(stream);
   auto strings_column = column_device_view::create(strings.parent(), stream);
   auto d_strings      = *strings_column;
 
@@ -104,7 +104,7 @@ std::unique_ptr<cudf::column> gather(
       string_view d_str = d_strings.element<string_view>(index);
       memcpy(d_chars + d_offsets[idx], d_str.data(), d_str.size_bytes());
     };
-  thrust::for_each_n(execpol->on(stream.value()),
+  thrust::for_each_n(rmm::exec_policy(stream),
                      thrust::make_counting_iterator<size_type>(0),
                      output_count,
                      gather_chars);

cpp/include/cudf/strings/detail/merge.cuh

@@ -25,6 +25,7 @@
 #include <strings/utilities.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/exec_policy.hpp>
 
 namespace cudf {
 namespace strings {
@@ -54,7 +55,7 @@ std::unique_ptr<column> merge(strings_column_view const& lhs,
   using cudf::detail::side;
   size_type strings_count = static_cast<size_type>(std::distance(begin, end));
   if (strings_count == 0) return make_empty_strings_column(stream, mr);
-  auto execpol    = rmm::exec_policy(stream);
+
   auto lhs_column = column_device_view::create(lhs.parent(), stream);
   auto d_lhs      = *lhs_column;
   auto rhs_column = column_device_view::create(rhs.parent(), stream);
@@ -86,7 +87,7 @@ std::unique_ptr<column> merge(strings_column_view const& lhs,
     strings::detail::create_chars_child_column(strings_count, null_count, bytes, stream, mr);
   // merge the strings
   auto d_chars = chars_column->mutable_view().template data<char>();
-  thrust::for_each_n(execpol->on(stream.value()),
+  thrust::for_each_n(rmm::exec_policy(stream),
                      thrust::make_counting_iterator<size_type>(0),
                      strings_count,
                      [d_lhs, d_rhs, begin, d_offsets, d_chars] __device__(size_type idx) {