Convert quantiles' and reductions' scratch RMM vectors to CUDA buffers

cpp/src/quantiles/quantile.cu

@@ -10,6 +10,7 @@
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/sorting.hpp>
+#include <cudf/detail/utilities/buffer_factories.hpp>
 #include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/detail/valid_if.cuh>
 #include <cudf/dictionary/detail/iterator.cuh>
@@ -21,7 +22,6 @@
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
-#include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
 
 #include <cuda/functional>
@@ -76,7 +76,7 @@ struct quantile_functor {
     auto d_output = mutable_column_device_view::create(output->mutable_view(), stream);
 
     auto q_device =
-      cudf::detail::make_device_uvector(q, stream, cudf::get_current_device_resource_ref());
+      cudf::detail::make_device_buffer(q, stream, cudf::get_current_device_resource_ref());
 
     if (!cudf::is_dictionary(input.type())) {
       auto sorted_data =

cpp/src/quantiles/quantiles.cu

@@ -9,6 +9,7 @@
 #include <cudf/detail/gather.cuh>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/sorting.hpp>
+#include <cudf/detail/utilities/buffer_factories.hpp>
 #include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/quantiles.hpp>
 #include <cudf/table/table_view.hpp>
@@ -42,7 +43,7 @@ std::unique_ptr<table> quantiles(table_view const& input,
     });
 
   auto const q_device =
-    cudf::detail::make_device_uvector_async(q, stream, cudf::get_current_device_resource_ref());
+    cudf::detail::make_device_buffer_async(q, stream, cudf::get_current_device_resource_ref());
 
   auto quantile_idx_iter = cuda::transform_iterator(q_device.begin(), quantile_idx_lookup);
 

cpp/src/reductions/histogram.cu

@@ -21,6 +21,7 @@
 #include <cuco/operator.hpp>
 #include <cuco/static_set.cuh>
 #include <cuda/atomic>
+#include <cuda/buffer>
 #include <cuda/functional>
 #include <cuda/iterator>
 #include <cuda/std/tuple>
@@ -138,7 +139,7 @@ compute_row_frequencies(table_view const& input,
   std::size_t const num_rows = input.num_rows();
 
   // Construct a vector to store reduced counts and init to zero
-  rmm::device_uvector<histogram_count_type> reduction_results(num_rows, stream, mr);
+  cuda::device_buffer<histogram_count_type> reduction_results(stream, mr, num_rows, cuda::no_init);
   thrust::uninitialized_fill(rmm::exec_policy_nosync(stream, temp_mr),
                              reduction_results.begin(),
                              reduction_results.end(),

cpp/src/reductions/nth_element.cu

@@ -10,9 +10,9 @@
 #include <cudf/reduction/detail/reduction_functions.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
-#include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
 
+#include <cuda/buffer>
 #include <cuda/functional>
 #include <cuda/iterator>
 #include <cuda/stream>
@@ -38,7 +38,8 @@ std::unique_ptr<cudf::scalar> nth_element(column_view const& col,
       cuda::transform_iterator(cudf::detail::make_validity_iterator(*dcol),
                                cuda::proclaim_return_type<size_type>(
                                  [] __device__(auto b) { return static_cast<size_type>(b); }));
-    rmm::device_uvector<size_type> null_skipped_index(col.size(), stream);
+    cuda::device_buffer<size_type> null_skipped_index(
+      stream, cudf::get_current_device_resource_ref(), col.size(), cuda::no_init);
     // null skipped index for valids only.
     thrust::inclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
                            bitmask_iterator,

cpp/src/reductions/scan/ewm.cu

@@ -9,11 +9,13 @@
 #include <cudf/column/column_factories.hpp>
 #include <cudf/detail/iterator.cuh>
 #include <cudf/detail/null_mask.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
 #include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
 
+#include <cuda/buffer>
 #include <cuda/functional>
 #include <cuda/iterator>
 #include <cuda/std/utility>
@@ -130,9 +132,10 @@ struct ewma_noadjust_no_nulls_functor : public ewma_functor_base<T> {
 * Example: {1, NULL, 3, 4, NULL, NULL, 7}
         -> {0, 0     1, 0, 0,    1,    2}
 */
-rmm::device_uvector<cudf::size_type> null_roll_up(column_view const& input, cuda::stream_ref stream)
+cuda::device_buffer<cudf::size_type> null_roll_up(column_view const& input, cuda::stream_ref stream)
 {
-  rmm::device_uvector<cudf::size_type> output(input.size(), stream);
+  cuda::device_buffer<cudf::size_type> output(
+    stream, cudf::get_current_device_resource_ref(), input.size(), cuda::no_init);
 
   auto device_view = column_device_view::create(input, stream);
   auto invalid_it  = cuda::transform_iterator(
@@ -156,10 +159,11 @@ rmm::device_uvector<T> compute_ewma_adjust(column_view const& input,
                                            rmm::device_async_resource_ref mr)
 {
   rmm::device_uvector<T> output(input.size(), stream);
-  rmm::device_uvector<pair_type<T>> pairs(input.size(), stream);
+  cuda::device_buffer<pair_type<T>> pairs(
+    stream, cudf::get_current_device_resource_ref(), input.size(), cuda::no_init);
 
   if (input.has_nulls()) {
-    rmm::device_uvector<cudf::size_type> nullcnt = null_roll_up(input, stream);
+    cuda::device_buffer<cudf::size_type> nullcnt = null_roll_up(input, stream);
     auto device_view                             = column_device_view::create(input, stream);
     auto valid_it = cudf::detail::make_validity_iterator(*device_view);
     auto data =
@@ -228,13 +232,15 @@ rmm::device_uvector<T> compute_ewma_noadjust(column_view const& input,
                                              rmm::device_async_resource_ref mr)
 {
   rmm::device_uvector<T> output(input.size(), stream);
-  rmm::device_uvector<pair_type<T>> pairs(input.size(), stream);
-  rmm::device_uvector<cudf::size_type> nullcnt =
-    [&input, stream]() -> rmm::device_uvector<cudf::size_type> {
+  cuda::device_buffer<pair_type<T>> pairs(
+    stream, cudf::get_current_device_resource_ref(), input.size(), cuda::no_init);
+  cuda::device_buffer<cudf::size_type> nullcnt =
+    [&input, stream]() -> cuda::device_buffer<cudf::size_type> {
     if (input.has_nulls()) {
       return null_roll_up(input, stream);
     } else {
-      return rmm::device_uvector<cudf::size_type>(input.size(), stream);
+      return cuda::device_buffer<cudf::size_type>(
+        stream, cudf::get_current_device_resource_ref(), input.size(), cuda::no_init);
     }
   }();
   // denominators are all 1 and do not need to be computed

cpp/src/reductions/segmented/nunique.cu

@@ -15,6 +15,7 @@
 #include <cudf/types.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
+#include <cuda/buffer>
 #include <cuda/iterator>
 #include <cuda/stream>
 #include <thrust/transform.h>
@@ -23,13 +24,13 @@ namespace cudf {
 namespace reduction {
 namespace detail {
 namespace {
-template <typename ComparatorType>
+template <typename ComparatorType, typename OffsetIterator, typename LabelIterator>
 struct is_unique_fn {
   column_device_view const d_col;
   ComparatorType row_equal;
   null_policy null_handling;
-  size_type const* offsets;
-  size_type const* labels;
+  OffsetIterator offsets;
+  LabelIterator labels;
 
   __device__ size_type operator()(size_type idx) const
   {
@@ -58,13 +59,12 @@ std::unique_ptr<cudf::column> segmented_nunique(column_view const& col,
     auto const row_equal =
       comparator.equal_to<false>(cudf::nullate::DYNAMIC{col.has_nulls()}, null_equality::EQUAL);
 
-    auto labels = rmm::device_uvector<size_type>(col.size(), stream, temp_mr);
+    auto labels = cuda::device_buffer<size_type>(stream, temp_mr, col.size(), cuda::no_init);
     cudf::detail::label_segments(
       offsets.begin(), offsets.end(), labels.begin(), labels.end(), stream);
-    auto fn = is_unique_fn<decltype(row_equal)>{
-      *d_col, row_equal, null_handling, offsets.data(), labels.data()};
+    auto fn = is_unique_fn{*d_col, row_equal, null_handling, offsets.begin(), labels.begin()};
 
-    auto identifiers = rmm::device_uvector<size_type>(col.size(), stream, temp_mr);
+    auto identifiers = cuda::device_buffer<size_type>(stream, temp_mr, col.size(), cuda::no_init);
     thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
                       cuda::counting_iterator<size_type>{0},
                       cuda::counting_iterator<size_type>{col.size()},

cpp/src/reductions/unique_count.cu

@@ -9,9 +9,11 @@
 #include <cudf/reduction/unique_count.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/utilities/default_stream.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 
 #include <rmm/exec_policy.hpp>
 
+#include <cuda/buffer>
 #include <cuda/iterator>
 #include <cuda/stream>
 #include <thrust/count.h>
@@ -32,7 +34,7 @@ cudf::size_type unique_count(table_view const& keys,
     // Using a temporary buffer for intermediate transform results from the lambda containing
     // the comparator speeds up compile-time significantly without much degradation in
     // runtime performance over using the comparator directly in thrust::count_if.
-    auto d_results = rmm::device_uvector<bool>(keys.num_rows(), stream, temp_mr);
+    auto d_results = cuda::device_buffer<bool>(stream, temp_mr, keys.num_rows(), cuda::no_init);
     thrust::transform(rmm::exec_policy_nosync(stream, temp_mr),
                       cuda::counting_iterator<size_type>{0},
                       cuda::counting_iterator<size_type>{keys.num_rows()},