Spread each segment of the segmented bitmask reduction across several blocks

cpp/benchmarks/bitmask/bitmask_and.cpp

@@ -84,6 +84,11 @@ void BM_segmented_bitmask_and(nvbench::state& state)
   set_throughputs(state);
 }
 
+void BM_segmented_bitmask_and_wide_masks(nvbench::state& state)
+{
+  BM_segmented_bitmask_and(state);
+}
+
 void BM_multi_segment_bitmask_and(nvbench::state& state)
 {
   auto const mask_size_bits = static_cast<size_t>(state.get_int64("mask_size_bits"));
@@ -134,6 +139,15 @@ NVBENCH_BENCH(BM_segmented_bitmask_and)
   .add_int64_axis("expected_masks_per_segment", {4, 8, 16})
   .add_int64_axis("mask_size_bits", {32, 64, 128});
 
+// The number of bits in a null mask is the row count of the table it belongs to, so masks of
+// hundreds of thousands of bits are the common case for the struct null mask reduction that this
+// kernel serves. The axes above are all narrow enough to fit in a handful of words.
+NVBENCH_BENCH(BM_segmented_bitmask_and_wide_masks)
+  .set_name("segmented_bitmask_and_wide_masks")
+  .add_int64_axis("num_segments", {8, 64, 512})
+  .add_int64_axis("expected_masks_per_segment", {4, 8})
+  .add_int64_axis("mask_size_bits", {100000, 1000000});
+
 NVBENCH_BENCH(BM_multi_segment_bitmask_and)
   .set_name("multi_segment_bitmask_and")
   .add_int64_axis("num_segments", {100, 1000, 10000})

cpp/include/cudf/detail/null_mask.cuh

@@ -117,32 +117,35 @@ CUDF_KERNEL void offset_bitmask_binop(Binop op,
 /**
  * @brief Performs a segmented binary operation on bitmasks with configurable bit offsets.
  *
- * For each segment in the input masks array, this kernel applies a binary reduction operation. Each
- * segment is processed by a separate warp, and the result is written directly to the destination
- * mask for that segment.
+ * For each segment in the input masks array, this kernel applies a binary reduction operation. The
+ * result is written directly to the destination mask for that segment.
  *
  * The kernel performs the following operations:
- * 1. Maps each warp to a segment defined by segment_offsets
+ * 1. Maps each block to a segment defined by segment_offsets and to a range of words within it
  * 2. For each segment, performs the binary operation on corresponding words of source bitmasks
  * 3. Counts the number of unset bits (nulls) in the resulting bitmask for each segment
- * 4. Writes the results to the destination mask and null counts array
+ * 4. Writes the results to the destination mask and accumulates into the null counts array
  *
+ * @tparam block_size      Number of threads per block
  * @tparam Binop           Type of binary operator
  *
  * @param op               The binary operator to apply to the bitmasks
  * @param num_segments     Number of segments to process
+ * @param blocks_per_segment Number of blocks cooperating on each segment
  * @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)
  * @param sources          Array of pointers to source bitmasks to be operated on
  * @param source_begin_bits Array of bit offsets from which each source mask is to be processed
  * @param source_size_bits The number of bits to process in each mask
  * @param segment_offsets  Array of indices defining the segments in the sources array
- * @param null_counts      Array where the count of unset bits for each segment will be written
+ * @param null_counts      Array the count of unset bits for each segment is accumulated into; must
+ * be zeroed before launch
  *
  */
-template <typename Binop>
+template <int block_size, typename Binop>
 CUDF_KERNEL void segmented_offset_bitmask_binop(Binop op,
                                                 size_type num_segments,
+                                                size_type blocks_per_segment,
                                                 bitmask_type** const destinations,
                                                 size_type destination_size,
                                                 bitmask_type const* const* const sources,
@@ -151,25 +154,17 @@ CUDF_KERNEL void segmented_offset_bitmask_binop(Binop op,
                                                 size_type const* const segment_offsets,
                                                 size_type* const null_counts)
 {
-  namespace cg = cooperative_groups;
-
-  // Create block level group
-  auto const block = cg::this_thread_block();
-
-  // Create warp-level group
-  auto const warp    = cg::tiled_partition<cudf::detail::warp_size>(block);
-  auto const warp_id = warp.meta_group_rank();
-  auto const lane    = warp.thread_rank();
+  // Each segment is processed by `blocks_per_segment` consecutive blocks, so that the words of a
+  // segment are spread across enough threads even when there are few segments.
+  auto const block_rank       = static_cast<size_type>(blockIdx.x);
+  auto const segment_id       = block_rank / blocks_per_segment;
+  auto const block_in_segment = block_rank % blocks_per_segment;
+  if (segment_id >= num_segments) { return; }
 
-  // Process one segment per warp.
-  auto const segment_id    = cudf::detail::grid_1d::global_thread_id() / warp.size();
   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);
@@ -178,9 +173,11 @@ CUDF_KERNEL void segmented_offset_bitmask_binop(Binop op,
   // Track null count (count of unset bits)
   size_type thread_null_count = 0;
 
-  // 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()) {
+  // Process the mask such that each thread of the segment's blocks handles different words
+  for (auto destination_word_index =
+         block_in_segment * block_size + static_cast<size_type>(threadIdx.x);
+       destination_word_index < destination_size;
+       destination_word_index += blocks_per_segment * block_size) {
     // Get the first mask word
     bitmask_type destination_word =
       detail::get_mask_offset_word(sources[segment_start],
@@ -216,11 +213,13 @@ CUDF_KERNEL void segmented_offset_bitmask_binop(Binop op,
     destination[destination_word_index] = destination_word;
   }
 
-  // Reduce the null counts across the warp
-  size_type warp_count = cg::reduce(warp, thread_null_count, cg::plus<size_type>());
-
-  // Only the first lane in the warp writes the result
-  if (lane == 0) { null_counts[segment_id] = warp_count; }
+  // Reduce the null counts across the block, then across the blocks of the segment
+  using BlockReduce = cub::BlockReduce<size_type, block_size>;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  auto const block_null_count = BlockReduce(temp_storage).Sum(thread_null_count);
+  if (threadIdx.x == 0 && block_null_count > 0) {
+    atomicAdd(&null_counts[segment_id], block_null_count);
+  }
 }
 
 // Forward declarations; defined later in this header but called from the templates below.
@@ -403,30 +402,38 @@ rmm::device_uvector<size_type> inplace_segmented_bitmask_binop(
   CUDF_EXPECTS(segment_offsets.size() >= 2,
                "At least one segment needs to be passed for bitwise operations");
 
-  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);
   auto d_segment_offsets =
     cudf::detail::make_device_uvector_async(segment_offsets, stream, temp_mr);
 
   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);
-  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,
-    dest_masks.data(),
-    dest_mask_size,
-    d_masks.data(),
-    d_begin_bits.data(),
-    mask_size_bits,
-    d_segment_offsets.data(),
-    d_null_counts.data());
+  auto constexpr max_blocks = 4096;
+
+  // A batch of a few wide segments needs the words of each segment split across several blocks to
+  // fill the device, while a batch of many segments already has enough parallelism with one block
+  // each.
+  auto const blocks_per_segment =
+    std::max(1,
+             std::min(util::div_rounding_up_safe<int>(dest_mask_size, block_size),
+                      max_blocks / num_segments));
+
+  CUDF_CUDA_TRY(cudaMemsetAsync(
+    d_null_counts.data(), 0, d_null_counts.size() * sizeof(size_type), stream.get()));
+  segmented_offset_bitmask_binop<block_size>
+    <<<num_segments * blocks_per_segment, block_size, 0, stream.get()>>>(op,
+                                                                         num_segments,
+                                                                         blocks_per_segment,
+                                                                         dest_masks.data(),
+                                                                         dest_mask_size,
+                                                                         d_masks.data(),
+                                                                         d_begin_bits.data(),
+                                                                         mask_size_bits,
+                                                                         d_segment_offsets.data(),
+                                                                         d_null_counts.data());
   CUDF_CHECK_CUDA(stream.get());
   return d_null_counts;
 }