Use cuda::device_buffer<std::byte> for null masks

cpp/benchmarks/bitmask/bitmask_and.cpp

@@ -44,12 +44,12 @@ auto setup_masks(nvbench::state& state)
   std::exclusive_scan(segments.begin(), segments.end(), segments.begin(), 0);
 
   // Create masks
-  std::vector<rmm::device_buffer> masks;
+  std::vector<cuda::device_buffer<uint8_t>> masks;
   std::vector<cudf::bitmask_type*> mask_pointers;
   masks.reserve(num_masks);
   std::generate_n(std::back_inserter(masks), num_masks, [mask_size_bits, seed, &mask_pointers]() {
     auto mask_pair = create_random_null_mask(mask_size_bits, null_probability, seed);
-    mask_pointers.push_back(static_cast<cudf::bitmask_type*>(mask_pair.first.data()));
+    mask_pointers.push_back(reinterpret_cast<cudf::bitmask_type*>(mask_pair.first.data()));
     return std::move(mask_pair.first);
   });
 

cpp/benchmarks/bitmask/set_null_mask.cpp

@@ -39,11 +39,11 @@ auto generate_test_data(cudf::size_type num_masks,
 
   auto valids = thrust::host_vector<bool>(num_masks, true);
 
-  std::vector<rmm::device_buffer> masks(num_masks);
+  std::vector<cuda::device_buffer<uint8_t>> masks(num_masks);
   std::vector<cudf::bitmask_type*> masks_ptr(num_masks);
   for (cudf::size_type i = 0; i < num_masks; ++i) {
     masks[i]     = cudf::create_null_mask(mask_size, cudf::mask_state::UNINITIALIZED);
-    masks_ptr[i] = static_cast<cudf::bitmask_type*>(masks[i].data());
+    masks_ptr[i] = reinterpret_cast<cudf::bitmask_type*>(masks[i].data());
   }
 
   return std::make_tuple(std::move(begin_bits),
@@ -57,19 +57,20 @@ auto generate_test_data(cudf::size_type num_masks,
 
 void BM_setnullmask(nvbench::state& state)
 {
-  auto const mask_size    = static_cast<cudf::size_type>(state.get_int64("mask_size"));
-  rmm::device_buffer mask = cudf::create_null_mask(mask_size, cudf::mask_state::UNINITIALIZED);
+  auto const mask_size = static_cast<cudf::size_type>(state.get_int64("mask_size"));
+  cuda::device_buffer<uint8_t> mask =
+    cudf::create_null_mask(mask_size, cudf::mask_state::UNINITIALIZED);
   auto begin = 0, end = mask_size;
 
   state.set_cuda_stream(nvbench::make_cuda_stream_view(cudf::get_default_stream().get()));
   auto const mem_stats_logger = cudf::memory_stats_logger();
 
-  state.exec(nvbench::exec_tag::sync | nvbench::exec_tag::timer,
-             [&](nvbench::launch& launch, auto& timer) {
-               timer.start();
-               cudf::set_null_mask(static_cast<cudf::bitmask_type*>(mask.data()), begin, end, true);
-               timer.stop();
-             });
+  state.exec(
+    nvbench::exec_tag::sync | nvbench::exec_tag::timer, [&](nvbench::launch& launch, auto& timer) {
+      timer.start();
+      cudf::set_null_mask(reinterpret_cast<cudf::bitmask_type*>(mask.data()), begin, end, true);
+      timer.stop();
+    });
 
   state.add_buffer_size(
     mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage");

cpp/benchmarks/common/generate_input.cu

@@ -785,11 +785,12 @@ std::unique_ptr<cudf::column> create_random_column<cudf::list_view>(data_profile
     thrust::device_pointer_cast(offsets.end())[-1] =
       current_child_column->size();  // Always include all elements
 
-    auto offsets_column = std::make_unique<cudf::column>(cudf::data_type{cudf::type_id::INT32},
-                                                         current_num_rows + 1,
-                                                         offsets.release(),
-                                                         rmm::device_buffer{},
-                                                         0);
+    auto offsets_column =
+      std::make_unique<cudf::column>(cudf::data_type{cudf::type_id::INT32},
+                                     current_num_rows + 1,
+                                     offsets.release(),
+                                     cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED),
+                                     0);
 
     auto [null_mask, null_count] = profile.get_null_probability().has_value()
                                      ? cudf::bools_to_mask(cudf::device_span<bool const>(valids))
@@ -1065,10 +1066,12 @@ std::unique_ptr<cudf::column> create_string_column(cudf::size_type num_rows,
   return std::move(table->release().front());
 }
 
-std::pair<rmm::device_buffer, cudf::size_type> create_random_null_mask(
+std::pair<cuda::device_buffer<uint8_t>, cudf::size_type> create_random_null_mask(
   cudf::size_type size, std::optional<double> null_probability, unsigned seed)
 {
-  if (not null_probability.has_value()) { return {rmm::device_buffer{}, 0}; }
+  if (not null_probability.has_value()) {
+    return {cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED), 0};
+  }
   CUDF_EXPECTS(*null_probability >= 0.0 and *null_probability <= 1.0,
                "Null probability must be within the range [0.0, 1.0]");
   if (*null_probability == 0.0f) {

cpp/benchmarks/common/generate_input.hpp

@@ -666,5 +666,5 @@ std::vector<cudf::type_id> mix_dtypes(std::pair<cudf::type_id, cudf::type_id> co
  * @param seed Optional, seed for the pseudo-random engine
  * @return null mask device buffer with random null mask data and null count
  */
-std::pair<rmm::device_buffer, cudf::size_type> create_random_null_mask(
+std::pair<cuda::device_buffer<uint8_t>, cudf::size_type> create_random_null_mask(
   cudf::size_type size, std::optional<double> null_probability = std::nullopt, unsigned seed = 1);

cpp/benchmarks/copying/copy_if_else.cpp

@@ -24,9 +24,9 @@ static void bench_copy_if_else(nvbench::state& state, nvbench::type_list<DataTyp
   auto const input = create_random_table({input_type, input_type, bool_type}, row_count{num_rows});
 
   if (!nulls) {
-    input->get_column(0).set_null_mask(rmm::device_buffer{}, 0);
-    input->get_column(1).set_null_mask(rmm::device_buffer{}, 0);
-    input->get_column(2).set_null_mask(rmm::device_buffer{}, 0);
+    input->get_column(0).set_null_mask(cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED), 0);
+    input->get_column(1).set_null_mask(cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED), 0);
+    input->get_column(2).set_null_mask(cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED), 0);
   }
 
   cudf::column_view lhs(input->view().column(0));

cpp/benchmarks/io/parquet/experimental/variant/extract.cpp

@@ -288,8 +288,12 @@ std::unique_ptr<cudf::column> build_variant_column(std::span<std::span<uint8_t c
       rmm::device_buffer{offsets.data(), offsets.size() * sizeof(int32_t), stream, mr};
     auto d_data = rmm::device_buffer{flat.data(), flat.size() * sizeof(uint8_t), stream, mr};
 
-    auto off_col = std::make_unique<cudf::column>(
-      cudf::data_type{cudf::type_id::INT32}, n + 1, std::move(d_offsets), rmm::device_buffer{}, 0);
+    auto off_col =
+      std::make_unique<cudf::column>(cudf::data_type{cudf::type_id::INT32},
+                                     n + 1,
+                                     std::move(d_offsets),
+                                     cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED),
+                                     0);
     auto data_col = std::make_unique<cudf::column>(cudf::data_type{cudf::type_id::UINT8},
                                                    static_cast<cudf::size_type>(flat.size()),
                                                    std::move(d_data),

cpp/benchmarks/replace/clamp.cpp

@@ -25,7 +25,8 @@ void bench_clamp(nvbench::state& state, nvbench::type_list<ClampType>)
 
   auto const dtype = cudf::type_to_id<ClampType>();
   auto const input = create_random_column(dtype, row_count{n_rows});
-  if (!include_nulls) input->set_null_mask(rmm::device_buffer{}, 0);
+  if (!include_nulls)
+    input->set_null_mask(cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED), 0);
 
   auto [low_scalar, high_scalar] = cudf::minmax(*input);
 

cpp/benchmarks/replace/nans.cpp

@@ -25,7 +25,8 @@ void bench_replace_nans(nvbench::state& state, nvbench::type_list<FloatingType>)
 
   auto const dtype = cudf::type_to_id<FloatingType>();
   auto const input = create_random_column(dtype, row_count{n_rows});
-  if (!include_nulls) input->set_null_mask(rmm::device_buffer{}, 0);
+  if (!include_nulls)
+    input->set_null_mask(cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED), 0);
 
   auto zero = cudf::make_fixed_width_scalar<FloatingType>(0);
 

cpp/examples/strings/custom_optimized.cu

@@ -143,7 +143,8 @@ std::unique_ptr<cudf::column> redact_strings(cudf::column_view const& names,
     *d_names, *d_visibilities, offsets.data(), chars.data());
 
   // create column from offsets vector (move only)
-  auto offsets_column = std::make_unique<cudf::column>(std::move(offsets), rmm::device_buffer{}, 0);
+  auto offsets_column = std::make_unique<cudf::column>(
+    std::move(offsets), cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED), 0);
 
   // create column for chars vector (no copy is performed)
   auto result = cudf::make_strings_column(

cpp/include/cudf/binaryop.hpp

@@ -317,7 +317,7 @@ bool is_supported_operation(data_type out, data_type lhs, data_type rhs, binary_
  * @param mr      Device memory resource used to allocate the returned valid mask
  * @return        Computed validity mask
  */
-std::pair<rmm::device_buffer, size_type> scalar_col_valid_mask_and(
+std::pair<cuda::device_buffer<uint8_t>, size_type> scalar_col_valid_mask_and(
   column_view const& col,
   scalar const& s,
   cuda::stream_ref stream           = cudf::get_default_stream(),

cpp/include/cudf/column/column.hpp

@@ -73,7 +73,9 @@ class column {
    * @param null_count The count of null elements.
    */
   template <typename T, CUDF_ENABLE_IF(cudf::is_numeric<T>() or cudf::is_chrono<T>())>
-  column(rmm::device_uvector<T>&& other, rmm::device_buffer&& null_mask, size_type null_count)
+  column(rmm::device_uvector<T>&& other,
+         cuda::device_buffer<uint8_t>&& null_mask,
+         size_type null_count)
     : _type{cudf::data_type{cudf::type_to_id<T>()}},
       _size{[&]() {
         CUDF_EXPECTS(
@@ -103,7 +105,7 @@ class column {
    * @param null_count Optional, the count of null elements.
    * @param children Optional, vector of child columns
    */
-  template <typename B1, typename B2 = rmm::device_buffer>
+  template <typename B1, typename B2 = cuda::device_buffer<uint8_t>>
   column(data_type dtype,
          size_type size,
          B1&& data,
@@ -166,7 +168,7 @@ class column {
    * `new_null_count` is 0.
    * @param new_null_count The count of null elements.
    */
-  void set_null_mask(rmm::device_buffer&& new_null_mask, size_type new_null_count);
+  void set_null_mask(cuda::device_buffer<uint8_t>&& new_null_mask, size_type new_null_count);
 
   /**
    * @brief Sets the column's null value indicator bitmask to `new_null_mask`.
@@ -180,7 +182,7 @@ class column {
    * @param stream The stream on which to perform the allocation and copy. Uses the default CUDF
    * stream if none is specified.
    */
-  void set_null_mask(rmm::device_buffer const& new_null_mask,
+  void set_null_mask(cuda::device_buffer<uint8_t> const& new_null_mask,
                      size_type new_null_count,
                      cuda::stream_ref stream = cudf::get_default_stream());
 
@@ -247,9 +249,9 @@ class column {
    * Returned by `column::release()`.
    */
   struct contents {
-    std::unique_ptr<rmm::device_buffer> data;       ///< data device memory buffer
-    std::unique_ptr<rmm::device_buffer> null_mask;  ///< null mask device memory buffer
-    std::vector<std::unique_ptr<column>> children;  ///< child columns
+    std::unique_ptr<rmm::device_buffer> data;                 ///< data device memory buffer
+    std::unique_ptr<cuda::device_buffer<uint8_t>> null_mask;  ///< null mask device memory buffer
+    std::vector<std::unique_ptr<column>> children;            ///< child columns
   };
 
   /**
@@ -324,9 +326,10 @@ class column {
   cudf::size_type _size{};                ///< The number of elements in the column
   rmm::device_buffer _data{};             ///< Dense, contiguous, type erased device memory
                                           ///< buffer containing the column elements
-  rmm::device_buffer _null_mask{};        ///< Bitmask used to represent null values.
-                                          ///< May be empty if `null_count() == 0`
-  mutable cudf::size_type _null_count{};  ///< The number of null elements
+  cuda::device_buffer<uint8_t> _null_mask = cudf::create_null_mask(
+    0, cudf::mask_state::UNALLOCATED);               ///< Bitmask used to represent null values.
+                                                     ///< May be empty if `null_count() == 0`
+  mutable cudf::size_type _null_count{};             ///< The number of null elements
   std::vector<std::unique_ptr<column>> _children{};  ///< Depending on element type, child
                                                      ///< columns may contain additional data
 };

cpp/include/cudf/column/column_factories.hpp

@@ -442,7 +442,7 @@ std::unique_ptr<column> make_strings_column(size_type num_strings,
                                             std::unique_ptr<column> offsets_column,
                                             rmm::device_buffer&& chars_buffer,
                                             size_type null_count,
-                                            rmm::device_buffer&& null_mask);
+                                            cuda::device_buffer<uint8_t>&& null_mask);
 
 /**
  * @brief Construct a LIST type column given offsets column, child column, null mask and null
@@ -501,7 +501,7 @@ std::unique_ptr<cudf::column> make_lists_column(size_type num_rows,
                                                 std::unique_ptr<column> offsets_column,
                                                 std::unique_ptr<column> child_column,
                                                 size_type null_count,
-                                                rmm::device_buffer&& null_mask);
+                                                cuda::device_buffer<uint8_t>&& null_mask);
 
 /**
  * @brief Create an empty LIST column
@@ -540,7 +540,7 @@ std::unique_ptr<cudf::column> make_structs_column(
   size_type num_rows,
   std::vector<std::unique_ptr<column>>&& child_columns,
   size_type null_count,
-  rmm::device_buffer&& null_mask,
+  cuda::device_buffer<uint8_t>&& null_mask,
   cuda::stream_ref stream           = cudf::get_default_stream(),
   rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
@@ -574,7 +574,7 @@ std::unique_ptr<cudf::column> create_structs_hierarchy(
   size_type num_rows,
   std::vector<std::unique_ptr<column>>&& child_columns,
   size_type null_count,
-  rmm::device_buffer&& null_mask,
+  cuda::device_buffer<uint8_t>&& null_mask,
   cuda::stream_ref stream           = cudf::get_default_stream(),
   rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 

cpp/include/cudf/concatenate.hpp

@@ -11,6 +11,7 @@
 #include <cudf/utilities/memory_resource.hpp>
 #include <cudf/utilities/span.hpp>
 
+#include <cuda/buffer>
 #include <cuda/stream>
 
 #include <memory>
@@ -39,7 +40,7 @@ namespace CUDF_EXPORT cudf {
  * @param stream CUDA stream used for device memory operations and kernel launches
  * @return Bitmasks of all the column views in the views vector
  */
-rmm::device_buffer concatenate_masks(
+cuda::device_buffer<uint8_t> concatenate_masks(
   std::span<column_view const> views,
   cuda::stream_ref stream           = cudf::get_default_stream(),
   rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());

cpp/include/cudf/detail/concatenate_masks.hpp

@@ -11,6 +11,7 @@
 
 #include <rmm/device_buffer.hpp>
 
+#include <cuda/buffer>
 #include <cuda/stream>
 
 #include <span>
@@ -55,9 +56,9 @@ size_type concatenate_masks(host_span<column_view const> views,
  *
  * @param stream CUDA stream used for device memory operations and kernel launches.
  */
-rmm::device_buffer concatenate_masks(std::span<column_view const> views,
-                                     cuda::stream_ref stream,
-                                     rmm::device_async_resource_ref mr);
+cuda::device_buffer<uint8_t> concatenate_masks(std::span<column_view const> views,
+                                               cuda::stream_ref stream,
+                                               rmm::device_async_resource_ref mr);
 
 }  // namespace detail
 }  // namespace cudf

cpp/include/cudf/detail/gather.cuh

@@ -357,23 +357,25 @@ struct column_gatherer_impl<list_view> {
         lists::detail::gather_list_nested(list.get_sliced_child(stream), gd, stream, output_mr);
 
       // return the final column
-      return make_lists_column(gather_map_size,
-                               std::move(gd.offsets),
-                               std::move(child),
-                               0,
-                               rmm::device_buffer{0, stream, output_mr});
+      return make_lists_column(
+        gather_map_size,
+        std::move(gd.offsets),
+        std::move(child),
+        0,
+        cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED, stream, output_mr));
     }
 
     // it's a leaf.  do a regular gather
     auto child =
       lists::detail::gather_list_leaf(list.get_sliced_child(stream), gd, stream, output_mr);
 
     // assemble final column
-    return make_lists_column(gather_map_size,
-                             std::move(gd.offsets),
-                             std::move(child),
-                             0,
-                             rmm::device_buffer{0, stream, output_mr});
+    return make_lists_column(
+      gather_map_size,
+      std::move(gd.offsets),
+      std::move(child),
+      0,
+      cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED, stream, output_mr));
   }
 };
 
@@ -494,7 +496,10 @@ struct column_gatherer_impl<struct_view> {
       gather_map_size,
       std::move(output_struct_members),
       0,
-      rmm::device_buffer{0, stream, output_mr},  // Null mask will be fixed up in cudf::gather().
+      cudf::create_null_mask(0,
+                             cudf::mask_state::UNALLOCATED,
+                             stream,
+                             output_mr),  // Null mask will be fixed up in cudf::gather().
       stream,
       output_mr);
   }