Define empty segments for segmented_offset_bitmask_binop via an identity seed

Seeding each segment's reduction with the identity of the binary operator, instead of
loading the segment's first mask, gives an empty segment a well defined result (an
all-valid destination mask, null count zero) and removes the read of sources[segment_start]
that a segment with no masks would otherwise perform out of bounds.

Also validate that segment offsets are non-decreasing and within the mask array, and stop
the bitmask benchmark's segment size generator from drawing negative sizes, which produced
non-monotonic offsets and an illegal memory access for a small expected_masks_per_segment.

Author: vuule

cpp/benchmarks/bitmask/bitmask_and.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -12,6 +12,8 @@
 
 #include <nvbench/nvbench.cuh>
 
+#include <algorithm>
+#include <numeric>
 #include <random>
 
 namespace {
@@ -34,7 +36,9 @@ auto setup_masks(nvbench::state& state)
   std::vector<cudf::size_type> segments(num_segments + 1);
   auto num_masks = 0;
   std::generate_n(segments.begin(), num_segments, [&normal_dist, &generator, &num_masks]() {
-    cudf::size_type segment_size = normal_dist(generator);
+    // The distribution is unbounded, so a small mean draws negative sizes; those would make the
+    // offsets below non-monotonic. Empty segments are valid input and are kept.
+    auto const segment_size = std::max(cudf::size_type{0}, cudf::size_type(normal_dist(generator)));
     num_masks += segment_size;
     return segment_size;
   });

cpp/include/cudf/detail/null_mask.cuh

@@ -130,6 +130,8 @@ CUDF_KERNEL void offset_bitmask_binop(Binop op,
  * @tparam Binop           Type of binary operator
  *
  * @param op               The binary operator to apply to the bitmasks
+ * @param identity         Identity element of `op`, used to seed each segment's reduction; a
+ * segment containing no masks therefore produces a destination filled with `identity`
  * @param num_segments     Number of segments to process
  * @param destinations     Array of pointers to destination bitmasks where results will be written
  * @param destination_size Size of each destination mask in bitmask words (not bits)
@@ -142,6 +144,7 @@ CUDF_KERNEL void offset_bitmask_binop(Binop op,
  */
 template <typename Binop>
 CUDF_KERNEL void segmented_offset_bitmask_binop(Binop op,
+                                                bitmask_type identity,
                                                 size_type num_segments,
                                                 bitmask_type** const destinations,
                                                 size_type destination_size,
@@ -162,14 +165,15 @@ CUDF_KERNEL void segmented_offset_bitmask_binop(Binop op,
   auto const lane    = warp.thread_rank();
 
   // Process one segment per warp.
-  auto const segment_id    = cudf::detail::grid_1d::global_thread_id() / warp.size();
+  auto const segment_id = cudf::detail::grid_1d::global_thread_id() / warp.size();
+
+  // Exit early if this warp doesn't have a valid segment, before indexing anything by `segment_id`
+  if (segment_id >= num_segments) { return; }
+
   auto const segment_start = segment_offsets[segment_id];
   auto const segment_end   = segment_offsets[segment_id + 1];
   auto const destination   = destinations[segment_id];
 
-  // Exit early if this warp doesn't have a valid segment
-  if (segment_id >= num_segments) { return; }
-
   // Calculate bit range information
   auto const last_bit_index    = source_size_bits - 1;
   auto const last_word_index   = cudf::word_index(last_bit_index);
@@ -181,15 +185,12 @@ CUDF_KERNEL void segmented_offset_bitmask_binop(Binop op,
   // Process the mask such that each thread in warp handles different words
   for (size_type destination_word_index = lane; destination_word_index < destination_size;
        destination_word_index += warp.size()) {
-    // Get the first mask word
-    bitmask_type destination_word =
-      detail::get_mask_offset_word(sources[segment_start],
-                                   destination_word_index,
-                                   source_begin_bits[segment_start],
-                                   source_begin_bits[segment_start] + source_size_bits);
+    // Seeding with the identity keeps an empty segment well defined and avoids indexing `sources`
+    // with a segment's start offset before knowing that the segment contains anything.
+    bitmask_type destination_word = identity;
 
     // Apply the binary operation with each source mask in the segment
-    for (size_type mask_pos = segment_start + 1; mask_pos < segment_end; mask_pos++) {
+    for (size_type mask_pos = segment_start; mask_pos < segment_end; mask_pos++) {
       destination_word =
         op(destination_word,
            detail::get_mask_offset_word(sources[mask_pos],
@@ -242,7 +243,8 @@ rmm::device_uvector<size_type> inplace_segmented_bitmask_binop(
   size_type mask_size_bits,
   host_span<size_type const> segment_offsets,
   cuda::stream_ref stream,
-  rmm::device_async_resource_ref mr);
+  rmm::device_async_resource_ref mr,
+  bitmask_type identity = ~bitmask_type{0});
 
 /**
  * @copydoc bitmask_binop(Binop op, host_span<bitmask_type const* const>, host_span<size_type>
@@ -280,7 +282,8 @@ segmented_bitmask_binop(Binop op,
                         size_type mask_size_bits,
                         host_span<size_type const> segment_offsets,
                         cuda::stream_ref stream,
-                        rmm::device_async_resource_ref mr)
+                        rmm::device_async_resource_ref mr,
+                        bitmask_type identity = ~bitmask_type{0})
 {
   auto const num_bytes = bitmask_allocation_size_bytes(mask_size_bits);
   CUDF_EXPECTS(
@@ -314,7 +317,8 @@ segmented_bitmask_binop(Binop op,
                                                      mask_size_bits,
                                                      segment_offsets,
                                                      stream,
-                                                     cudf::get_current_device_resource_ref());
+                                                     cudf::get_current_device_resource_ref(),
+                                                     identity);
 
   return std::pair(std::move(h_destination_masks),
                    cudf::detail::make_std_vector<size_type>(null_counts, stream));
@@ -376,9 +380,12 @@ size_type inplace_bitmask_binop(Binop op,
  * @param[in] masks Host span of pointers to source bitmasks to be operated on
  * @param[in] masks_begin_bits The bit offsets from which each source mask is to be operated on
  * @param[in] mask_size_bits The number of bits to be operated on in each mask
- * @param[in] segment_offsets Host span of offsets defining the segments for the operation
+ * @param[in] segment_offsets Host span of offsets defining the segments for the operation; must be
+ * non-decreasing. An empty segment writes `identity` to its destination mask
  * @param[in] stream CUDA stream used for device memory operations and kernel launches
  * @param[in] mr Device memory resource used to allocate output device vector of null counts
+ * @param[in] identity Identity element of `op`, defaulted to the all-set mask that is the identity
+ * of bitwise AND
  * @return A device vector containing the counts of unset bits in the destination mask corresponding
  * to each segment after the AND operation
  */
@@ -392,7 +399,8 @@ rmm::device_uvector<size_type> inplace_segmented_bitmask_binop(
   size_type mask_size_bits,
   host_span<size_type const> segment_offsets,
   cuda::stream_ref stream,
-  rmm::device_async_resource_ref mr)
+  rmm::device_async_resource_ref mr,
+  bitmask_type identity)
 {
   CUDF_EXPECTS(
     std::all_of(masks_begin_bits.begin(), masks_begin_bits.end(), [](auto b) { return b >= 0; }),
@@ -402,8 +410,14 @@ rmm::device_uvector<size_type> inplace_segmented_bitmask_binop(
                "Mask pointer cannot be null");
   CUDF_EXPECTS(segment_offsets.size() >= 2,
                "At least one segment needs to be passed for bitwise operations");
+  CUDF_EXPECTS(std::is_sorted(segment_offsets.begin(), segment_offsets.end()),
+               "Segment offsets must be non-decreasing");
+  CUDF_EXPECTS(
+    segment_offsets.front() >= 0 && segment_offsets.back() <= static_cast<size_type>(masks.size()),
+    "Segment offsets are out of the bounds of the mask array");
 
-  rmm::device_uvector<size_type> d_null_counts(segment_offsets.size() - 1, stream, mr);
+  auto const num_segments = static_cast<size_type>(segment_offsets.size() - 1);
+  rmm::device_uvector<size_type> d_null_counts(num_segments, stream, mr);
   auto temp_mr      = cudf::get_current_device_resource_ref();
   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);
@@ -413,13 +427,13 @@ rmm::device_uvector<size_type> inplace_segmented_bitmask_binop(
   auto constexpr block_size = 256;
   auto constexpr warps_per_block =
     util::div_rounding_up_safe<int>(block_size, cudf::detail::warp_size);
-  auto const num_blocks =
-    util::div_rounding_up_safe<int>(segment_offsets.size() - 1, warps_per_block);
+  auto const num_blocks = util::div_rounding_up_safe<int>(num_segments, warps_per_block);
   static_assert(block_size % cudf::detail::warp_size == 0,
                 "For segmented bitmask operations, block size must be a multiple of warp size");
   segmented_offset_bitmask_binop<<<num_blocks, block_size, 0, stream.get()>>>(
     op,
-    segment_offsets.size() - 1,
+    identity,
+    num_segments,
     dest_masks.data(),
     dest_mask_size,
     d_masks.data(),

cpp/tests/bitmask/bitmask_tests.cpp

@@ -675,6 +675,81 @@ TEST_F(MergeBitmaskTest, TestSegmentedBitmaskAndMultipleSegments)
   }
 }
 
+TEST_F(MergeBitmaskTest, TestSegmentedBitmaskAndEmptySegments)
+{
+  cudf::test::fixed_width_column_wrapper<bool> const bools_col1({0, 1, 0, 1, 1}, {0, 1, 1, 1, 0});
+  cudf::test::fixed_width_column_wrapper<bool> const bools_col2({0, 2, 1, 0, 255}, {1, 1, 0, 1, 0});
+
+  auto const num_rows = 5;
+  std::vector<cudf::column_view> const colviews{bools_col1, bools_col2};
+
+  // An empty segment reduces no masks, so it yields the identity of bitwise AND: every row valid.
+  auto const expect_all_valid = [&](void const* mask) {
+    EXPECT_EQ(cudf::null_count(static_cast<cudf::bitmask_type const*>(mask), 0, num_rows), 0);
+  };
+
+  // The reduction of both masks leaves rows 1 and 3 valid.
+  auto const expect_anded = [&](void const* mask) {
+    auto odd_indices = cudf::test::iterators::nulls_at_multiples_of(2);
+    auto const odd =
+      std::get<0>(cudf::test::detail::make_null_mask(odd_indices, odd_indices + num_rows));
+    CUDF_TEST_EXPECT_EQUAL_BUFFERS(mask, odd.data(), cudf::num_bitmask_words(num_rows));
+  };
+
+  // Leading, then trailing, then interior empty segment. The trailing case is the one that would
+  // index the mask array past its end if a segment's first mask were read unconditionally.
+  {
+    std::vector<cudf::size_type> const segment_offsets{0, 0, 2};
+    auto const [result_masks, result_null_count] =
+      cudf::segmented_bitmask_and(colviews, segment_offsets);
+    ASSERT_EQ(result_masks.size(), 2);
+    EXPECT_EQ(result_null_count[0], 0);
+    EXPECT_EQ(result_null_count[1], 3);
+    expect_all_valid(result_masks[0]->data());
+    expect_anded(result_masks[1]->data());
+  }
+
+  {
+    std::vector<cudf::size_type> const segment_offsets{0, 2, 2};
+    auto const [result_masks, result_null_count] =
+      cudf::segmented_bitmask_and(colviews, segment_offsets);
+    ASSERT_EQ(result_masks.size(), 2);
+    EXPECT_EQ(result_null_count[0], 3);
+    EXPECT_EQ(result_null_count[1], 0);
+    expect_anded(result_masks[0]->data());
+    expect_all_valid(result_masks[1]->data());
+  }
+
+  {
+    std::vector<cudf::size_type> const segment_offsets{0, 1, 1, 2};
+    auto const [result_masks, result_null_count] =
+      cudf::segmented_bitmask_and(colviews, segment_offsets);
+    ASSERT_EQ(result_masks.size(), 3);
+    EXPECT_EQ(result_null_count[0], 2);
+    EXPECT_EQ(result_null_count[1], 0);
+    EXPECT_EQ(result_null_count[2], 2);
+    CUDF_TEST_EXPECT_EQUAL_BUFFERS(result_masks[0]->data(),
+                                   static_cast<cudf::column_view>(bools_col1).null_mask(),
+                                   cudf::num_bitmask_words(num_rows));
+    expect_all_valid(result_masks[1]->data());
+    CUDF_TEST_EXPECT_EQUAL_BUFFERS(result_masks[2]->data(),
+                                   static_cast<cudf::column_view>(bools_col2).null_mask(),
+                                   cudf::num_bitmask_words(num_rows));
+  }
+}
+
+TEST_F(MergeBitmaskTest, TestSegmentedBitmaskAndRejectsDecreasingOffsets)
+{
+  cudf::test::fixed_width_column_wrapper<bool> const bools_col1({0, 1, 0, 1, 1}, {0, 1, 1, 1, 0});
+  cudf::test::fixed_width_column_wrapper<bool> const bools_col2({0, 2, 1, 0, 255}, {1, 1, 0, 1, 0});
+  std::vector<cudf::column_view> const colviews{bools_col1, bools_col2};
+
+  EXPECT_THROW(cudf::segmented_bitmask_and(colviews, std::vector<cudf::size_type>{0, 2, 1}),
+               cudf::logic_error);
+  EXPECT_THROW(cudf::segmented_bitmask_and(colviews, std::vector<cudf::size_type>{0, 3}),
+               cudf::logic_error);
+}
+
 TEST_F(MergeBitmaskTest, TestBitmaskOr)
 {
   cudf::test::fixed_width_column_wrapper<bool> const bools_col1({0, 1, 0, 1, 1}, {1, 1, 0, 0, 1});