Convert headers' scratch rmm::device_uvector<T> to cuda::device_buffer<T>

Use CCCL buffers instead of RMM vectors where possible in inline
implementations in the headers.

Issue: rapidsai/build-planning#321

Author: KyleFromNVIDIA

cpp/include/cudf/detail/device_scalar.hpp

@@ -9,9 +9,9 @@
 #include <cudf/detail/utilities/host_vector.hpp>
 #include <cudf/detail/utilities/vector_factories.hpp>
 
-#include <rmm/device_uvector.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cuda/buffer>
 #include <cuda/stream>
 
 #include <type_traits>
@@ -43,23 +43,25 @@ class device_scalar {
   explicit device_scalar(
     cuda::stream_ref stream,
     rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref())
-    : _storage{1, stream, std::move(mr)}, bounce_buffer{make_pinned_vector<T>(1, stream)}
+    : _storage{stream, std::move(mr), 1, cuda::no_init},
+      bounce_buffer{make_pinned_vector<T>(1, stream)}
   {
   }
 
   explicit device_scalar(
     T const& initial_value,
     cuda::stream_ref stream,
     rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref())
-    : _storage{1, stream, std::move(mr)}, bounce_buffer{make_pinned_vector<T>(1, stream)}
+    : _storage{stream, std::move(mr), 1, cuda::no_init},
+      bounce_buffer{make_pinned_vector<T>(1, stream)}
   {
     set_value_async(initial_value, stream);
   }
 
   device_scalar(device_scalar const& other,
                 cuda::stream_ref stream,
                 rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref())
-    : _storage{other._storage, stream, mr}, bounce_buffer{make_pinned_vector<T>(1, stream)}
+    : _storage{stream, mr, other._storage}, bounce_buffer{make_pinned_vector<T>(1, stream)}
   {
   }
 
@@ -88,7 +90,7 @@ class device_scalar {
   [[nodiscard]] T const* data() const noexcept { return _storage.data(); }
 
  private:
-  rmm::device_uvector<T> _storage;
+  cuda::device_buffer<T> _storage;
   mutable cudf::detail::host_vector<T> bounce_buffer;
 };
 

cpp/include/cudf/detail/gather.cuh

@@ -8,6 +8,7 @@
 #include <cudf/detail/indexalator.cuh>
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/utilities/assert.cuh>
+#include <cudf/detail/utilities/buffer_factories.hpp>
 #include <cudf/detail/utilities/cuda.cuh>
 #include <cudf/detail/utilities/grid_1d.cuh>
 #include <cudf/detail/utilities/vector_factories.hpp>
@@ -583,10 +584,10 @@ void gather_bitmask(table_view const& source,
   std::transform(target.begin(), target.end(), target_masks.begin(), [](auto const& col) {
     return col->mutable_view().null_mask();
   });
-  auto d_target_masks = make_device_uvector_async(target_masks, stream, temp_mr);
+  auto d_target_masks = make_device_buffer_async(target_masks, stream, temp_mr);
 
   auto const device_source = table_device_view::create(source, stream, temp_mr);
-  auto d_valid_counts = make_zeroed_device_uvector_async<size_type>(target.size(), stream, temp_mr);
+  auto d_valid_counts = make_zeroed_device_buffer_async<size_type>(target.size(), stream, temp_mr);
 
   // Dispatch operation enum to get implementation
   auto const impl = [op]() {

cpp/include/cudf/detail/labeling/label_segments.cuh

@@ -9,9 +9,9 @@
 #include <cudf/types.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
-#include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
 
+#include <cuda/buffer>
 #include <cuda/iterator>
 #include <cuda/stream>
 #include <thrust/for_each.h>
@@ -166,11 +166,13 @@ void labels_to_offsets(InputIterator labels_begin,
 
   // This stores the unique label values.
   // Given the example above, we will have this array containing [0, 1, 4].
-  auto list_indices = rmm::device_uvector<OutputType>(num_segments, stream);
+  auto list_indices = cuda::device_buffer<OutputType>(
+    stream, cudf::get_current_device_resource_ref(), num_segments, cuda::no_init);
 
   // Stores the non-zero segment sizes.
   // Given the example above, we will have this array containing [4, 2, 4].
-  auto list_sizes = rmm::device_uvector<OutputType>(num_segments, stream);
+  auto list_sizes = cuda::device_buffer<OutputType>(
+    stream, cudf::get_current_device_resource_ref(), num_segments, cuda::no_init);
 
   // Count the numbers of labels in the each segment.
   auto const end = cudf::detail::reduce_by_key(labels_begin,  // keys

cpp/include/cudf/detail/null_mask.cuh

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <cudf/detail/device_scalar.hpp>
+#include <cudf/detail/utilities/buffer_factories.hpp>
 #include <cudf/detail/utilities/cuda.cuh>
 #include <cudf/detail/utilities/cuda.hpp>
 #include <cudf/detail/utilities/grid_1d.cuh>
@@ -313,7 +314,7 @@ segmented_bitmask_binop(Binop op,
     h_destination_masks_ptrs.push_back(h_destination_masks.back()->data());
   }
   auto destination_masks =
-    cudf::detail::make_device_uvector_async(h_destination_masks_ptrs, stream, temp_mr);
+    cudf::detail::make_device_buffer_async(h_destination_masks_ptrs, stream, temp_mr);
 
   // for destination size, pass number of words in each destination buffer instead of number of bits
   auto null_counts = inplace_segmented_bitmask_binop(op,
@@ -363,8 +364,8 @@ size_type inplace_bitmask_binop(Binop op,
   auto const temp_mr = mr.get_temporary_mr();
   cudf::detail::device_scalar<size_type> d_counter{0, stream, temp_mr};
 
-  auto d_masks      = cudf::detail::make_device_uvector_async(masks, stream, temp_mr);
-  auto d_begin_bits = cudf::detail::make_device_uvector_async(masks_begin_bits, stream, temp_mr);
+  auto d_masks      = cudf::detail::make_device_buffer_async(masks, stream, temp_mr);
+  auto d_begin_bits = cudf::detail::make_device_buffer_async(masks_begin_bits, stream, temp_mr);
 
   auto constexpr block_size = 256;
   cudf::detail::grid_1d config(dest_mask.size(), block_size);
@@ -427,10 +428,9 @@ rmm::device_uvector<size_type> inplace_segmented_bitmask_binop(
   auto d_null_counts = cudf::detail::make_zeroed_device_uvector_async<size_type>(
     num_segments, stream, mr.get_output_mr());
   auto const temp_mr = mr.get_temporary_mr();
-  auto d_masks       = cudf::detail::make_device_uvector_async(masks, stream, temp_mr);
-  auto d_begin_bits  = cudf::detail::make_device_uvector_async(masks_begin_bits, stream, temp_mr);
-  auto d_segment_offsets =
-    cudf::detail::make_device_uvector_async(segment_offsets, stream, temp_mr);
+  auto d_masks       = cudf::detail::make_device_buffer_async(masks, stream, temp_mr);
+  auto d_begin_bits  = cudf::detail::make_device_buffer_async(masks_begin_bits, stream, temp_mr);
+  auto d_segment_offsets = cudf::detail::make_device_buffer_async(segment_offsets, stream, temp_mr);
 
   auto constexpr block_size = 256;
 
@@ -701,7 +701,7 @@ std::vector<size_type> segmented_count_bits(bitmask_type const* bitmask,
     std::distance(indices_begin, indices_end), stream);
   std::copy(indices_begin, indices_end, std::back_inserter(h_indices));
   auto const temp_mr   = cudf::get_current_device_resource_ref();
-  auto const d_indices = make_device_uvector_async(h_indices, stream, temp_mr);
+  auto const d_indices = make_device_buffer_async(h_indices, stream, temp_mr);
 
   // Compute the bit counts over each segment.
   auto first_bit_indices_begin = cuda::transform_iterator(

cpp/include/cudf/detail/row_operator/primitive_lexicographic.cuh

@@ -11,8 +11,8 @@
 #include <cudf/detail/row_operator/lexicographic_common.cuh>
 #include <cudf/detail/row_operator/primitive_common.cuh>
 #include <cudf/detail/utilities/assert.cuh>
+#include <cudf/detail/utilities/buffer_factories.hpp>
 #include <cudf/detail/utilities/cuda.hpp>
-#include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/table/table_device_view.cuh>
 #include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
@@ -22,9 +22,9 @@
 #include <cudf/utilities/traits.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
-#include <rmm/device_uvector.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cuda/buffer>
 #include <cuda/stream>
 
 #include <type_traits>
@@ -150,9 +150,9 @@ class lexicographic_comparator {
                            host_span<null_order const> null_precedence = {},
                            cuda::stream_ref stream                     = cudf::get_default_stream())
     : _table{table_device_view::create(table, stream)},
-      _column_order{cudf::detail::make_device_uvector_async(
+      _column_order{cudf::detail::make_device_buffer_async(
         column_order, stream, cudf::get_current_device_resource_ref())},
-      _null_precedence{cudf::detail::make_device_uvector_async(
+      _null_precedence{cudf::detail::make_device_buffer_async(
         null_precedence, stream, cudf::get_current_device_resource_ref())}
   {
     if (not column_order.empty() or not null_precedence.empty()) {
@@ -180,8 +180,8 @@ class lexicographic_comparator {
                                                        rmm::device_async_resource_ref>;
 
   table_device_view_owner const _table;
-  rmm::device_uvector<order> const _column_order;
-  rmm::device_uvector<null_order> const _null_precedence;
+  cuda::device_buffer<order> const _column_order;
+  cuda::device_buffer<null_order> const _null_precedence;
 };
 
 }  // namespace cudf::detail::row::primitive

cpp/include/cudf/detail/utilities/batched_memset.hpp

@@ -7,7 +7,7 @@
 
 #include <cudf/detail/iterator.cuh>
 #include <cudf/detail/nvtx/ranges.hpp>
-#include <cudf/detail/utilities/vector_factories.hpp>
+#include <cudf/detail/utilities/buffer_factories.hpp>
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
@@ -39,7 +39,7 @@ void batched_memset(cudf::host_span<cudf::device_span<T> const> host_buffers,
   CUDF_FUNC_RANGE();
 
   // Copy buffer spans into device memory and then get sizes
-  auto buffers = cudf::detail::make_device_uvector_async(
+  auto buffers = cudf::detail::make_device_buffer_async(
     host_buffers, stream, cudf::get_current_device_resource_ref());
 
   // Vector of sizes of all buffer spans

cpp/include/cudf/reduction/detail/approx_distinct_count.cuh

@@ -11,9 +11,8 @@
 #include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
-#include <rmm/device_uvector.hpp>
-
 #include <cuco/hyperloglog_ref.cuh>
+#include <cuda/buffer>
 #include <cuda/functional>
 #include <cuda/std/span>
 #include <cuda/stream>
@@ -34,7 +33,7 @@ namespace detail {
  * non-owning sketch storage modes:
  *
  * - **Owning mode**: Allocates and manages its own sketch storage as an
- *   `rmm::device_uvector<int32_t>`. Used when constructing from a table or when copying
+ *   `cuda::device_buffer<int32_t>`. Used when constructing from a table or when copying
  *   from a span.
  *
  * - **Non-owning mode**: Operates on user-provided `cuda::std::span<int32_t>` sketch
@@ -260,18 +259,17 @@ class approx_distinct_count {
   /**
    * @brief Storage type supporting both owning and non-owning modes
    *
-   * - Owning: `device_uvector<register_type>` - allocates and manages storage
+   * - Owning: `device_buffer<register_type>` - allocates and manages storage
    * - Non-owning: `span<byte>` - operates on user-provided storage
    *
    * The public API always exposes storage as `span<byte>` via `sketch()`.
    */
   using storage_type =
-    std::variant<rmm::device_uvector<register_type>,  ///< Owning storage (allocated internally)
+    std::variant<cuda::device_buffer<register_type>,  ///< Owning storage (allocated internally)
                  cuda::std::span<cuda::std::byte>     ///< Non-owning storage (user-provided)
                  >;
 
-  // Declared before `_storage` so it outlives the owning `device_uvector`, which only holds a
-  // reference to this resource. Unused in non-owning (span) mode.
+  // Declared before `_storage` so it outlives the owning buffer. Unused in non-owning (span) mode.
   cuda::mr::any_resource<cuda::mr::device_accessible> _mr;  ///< Owns the sketch storage resource
   storage_type _storage;       ///< Sketch register storage (owning or non-owning)
   std::int32_t _precision;     ///< HLL precision parameter (determines 2^p registers)

cpp/include/cudf/reduction/detail/segmented_reduction.cuh

@@ -16,6 +16,7 @@
 #include <rmm/exec_policy.hpp>
 
 #include <cub/device/device_segmented_reduce.cuh>
+#include <cuda/buffer>
 #include <cuda/execution>
 #include <cuda/iterator>
 #include <cuda/std/execution>
@@ -130,8 +131,10 @@ void segmented_reduce(InputIterator d_in,
   auto const initial_value = op.template get_identity<IntermediateType>();
   auto const binary_op     = cudf::detail::cast_functor<IntermediateType>(op.get_binary_op());
 
-  rmm::device_uvector<IntermediateType> intermediate_result{static_cast<std::size_t>(num_segments),
-                                                            stream};
+  cuda::device_buffer<IntermediateType> intermediate_result{stream,
+                                                            cudf::get_current_device_resource_ref(),
+                                                            static_cast<std::size_t>(num_segments),
+                                                            cuda::no_init};
 
   auto const env = [&] {
     if constexpr (std::is_same_v<Op, op::variance> || std::is_same_v<Op, op::standard_deviation>) {

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

@@ -9,9 +9,9 @@
 #include <cudf/strings/detail/strings_column_factories.cuh>
 #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/std/optional>
@@ -66,7 +66,8 @@ std::unique_ptr<cudf::column> copy_if_else(StringIterLeft lhs_begin,
   if (null_count == 0) { null_mask = cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED); }
 
   // build vector of strings
-  rmm::device_uvector<string_index_pair> indices(strings_count, stream);
+  cuda::device_buffer<string_index_pair> indices(
+    stream, cudf::get_current_device_resource_ref(), strings_count, cuda::no_init);
   thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
                     cuda::counting_iterator<size_type>{0},
                     cuda::counting_iterator{static_cast<size_type>(strings_count)},

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

@@ -22,6 +22,7 @@
 #include <rmm/exec_policy.hpp>
 
 #include <cub/device/device_memcpy.cuh>
+#include <cuda/buffer>
 #include <cuda/functional>
 #include <cuda/iterator>
 #include <cuda/std/execution>
@@ -248,7 +249,8 @@ auto make_strings_children(SizeAndExecuteFunction size_and_exec_fn,
   };
 
   // Compute the output sizes
-  auto output_sizes        = rmm::device_uvector<size_type>(strings_count, stream);
+  auto output_sizes = cuda::device_buffer<size_type>(
+    stream, cudf::get_current_device_resource_ref(), strings_count, cuda::no_init);
   size_and_exec_fn.d_sizes = output_sizes.data();
   size_and_exec_fn.d_chars = nullptr;
   for_each_fn(size_and_exec_fn);