diff --git a/cpp/benchmarks/bitmask/bitmask_and.cpp b/cpp/benchmarks/bitmask/bitmask_and.cpp index 8058084de70..ff133706408 100644 --- a/cpp/benchmarks/bitmask/bitmask_and.cpp +++ b/cpp/benchmarks/bitmask/bitmask_and.cpp @@ -87,6 +87,8 @@ 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(state.get_int64("mask_size_bits")); @@ -137,6 +139,12 @@ 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}); +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}) diff --git a/cpp/include/cudf/detail/null_mask.cuh b/cpp/include/cudf/detail/null_mask.cuh index 6ca24b4462f..e3882fe4f43 100644 --- a/cpp/include/cudf/detail/null_mask.cuh +++ b/cpp/include/cudf/detail/null_mask.cuh @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -17,8 +18,6 @@ #include -#include -#include #include #include #include @@ -116,22 +115,23 @@ 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 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 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 @@ -140,13 +140,15 @@ CUDF_KERNEL void offset_bitmask_binop(Binop op, * @param segment_offsets Array of `num_segments + 1` indices defining the segments in the sources * array, segment `i` covering `[segment_offsets[i], segment_offsets[i + 1])`; behavior is undefined * unless the indices are non-decreasing and none exceeds the number of masks in `sources` - * @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 +template CUDF_KERNEL void segmented_offset_bitmask_binop(Binop op, bitmask_type identity, size_type num_segments, + size_type blocks_per_segment, bitmask_type** const destinations, size_type destination_size, bitmask_type const* const* const sources, @@ -155,20 +157,11 @@ 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(block); - auto const warp_id = warp.meta_group_rank(); - auto const lane = warp.thread_rank(); - - // Process one segment per warp. - auto const segment_id = cudf::detail::grid_1d::global_thread_id() / warp.size(); - - // Exit early if this warp doesn't have a valid segment + // 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(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; } auto const segment_start = segment_offsets[segment_id]; @@ -183,9 +176,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(threadIdx.x); + destination_word_index < destination_size; + destination_word_index += blocks_per_segment * block_size) { bitmask_type destination_word = identity; // Apply the binary operation with each source mask in the segment @@ -216,11 +211,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()); - - // 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; + __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. @@ -411,7 +408,9 @@ rmm::device_uvector inplace_segmented_bitmask_binop( CUDF_EXPECTS(segment_offsets.size() >= 2, "At least one segment needs to be passed for bitwise operations"); auto const num_segments = static_cast(segment_offsets.size() - 1); - rmm::device_uvector d_null_counts(num_segments, stream, mr); + // The kernel accumulates into the null counts, so they have to start at zero + auto d_null_counts = + cudf::detail::make_zeroed_device_uvector_async(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); @@ -419,22 +418,38 @@ rmm::device_uvector inplace_segmented_bitmask_binop( 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(block_size, cudf::detail::warp_size); - auto const num_blocks = util::div_rounding_up_safe(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<<>>( - op, - identity, - num_segments, - dest_masks.data(), - dest_mask_size, - d_masks.data(), - d_begin_bits.data(), - mask_size_bits, - d_segment_offsets.data(), - d_null_counts.data()); + + // Any block past this one would find no words left to process + auto const blocks_to_cover_segment = util::div_rounding_up_safe(dest_mask_size, block_size); + + // Splitting the segments beyond a few waves of resident blocks only adds blocks that queue behind + // the ones already running, and each block pays for its own reduction and atomic however few + // words it ends up owning. The number of waves is empirical: on an A100 both a single wave and an + // uncapped grid are measurably slower on the shapes the benchmarks cover. + auto constexpr target_waves = 8; + auto blocks_per_multiprocessor = int{}; + CUDF_CUDA_TRY(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks_per_multiprocessor, segmented_offset_bitmask_binop, block_size, 0)); + auto const block_budget_per_segment = + target_waves * blocks_per_multiprocessor * cudf::detail::num_multiprocessors() / num_segments; + + // A few wide segments need their words spread over blocks to fill the device, while many segments + // have enough parallelism with one block each and can exhaust the budget + auto const blocks_per_segment = + std::max(1, std::min(blocks_to_cover_segment, block_budget_per_segment)); + + segmented_offset_bitmask_binop + <<>>(op, + identity, + 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; } diff --git a/cpp/tests/bitmask/bitmask_tests.cpp b/cpp/tests/bitmask/bitmask_tests.cpp index 2d56a774c14..e33502ddcb2 100644 --- a/cpp/tests/bitmask/bitmask_tests.cpp +++ b/cpp/tests/bitmask/bitmask_tests.cpp @@ -731,6 +731,45 @@ TEST_F(MergeBitmaskTest, TestSegmentedBitmaskAndEmptySegments) } } +TEST_F(MergeBitmaskTest, TestSegmentedBitmaskAndMultipleBlocksPerSegment) +{ + // Wide enough that a segment spans more than one block of the reduction kernel + auto const num_rows = 100'003; + cudf::test::fixed_width_column_wrapper const col1( + cuda::make_counting_iterator(0), + cuda::make_counting_iterator(num_rows), + cudf::test::iterators::nulls_at_multiples_of(3)); + cudf::test::fixed_width_column_wrapper const col2( + cuda::make_counting_iterator(0), + cuda::make_counting_iterator(num_rows), + cudf::test::iterators::nulls_at_multiples_of(5)); + cudf::test::fixed_width_column_wrapper const col3( + cuda::make_counting_iterator(0), + cuda::make_counting_iterator(num_rows), + cudf::test::iterators::nulls_at_multiples_of(7)); + + auto const multiples_of = [&](int n) { return (num_rows - 1) / n + 1; }; + std::vector const expected_null_counts{ + multiples_of(3) + multiples_of(5) - multiples_of(15), multiples_of(7)}; + + std::vector const colviews{col1, col2, col3}; + std::vector const segment_offsets{0, 2, 3}; + 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, expected_null_counts); + + auto const [expected_mask, expected_null_count] = + cudf::bitmask_and(cudf::table_view({col1, col2})); + EXPECT_EQ(expected_null_count, expected_null_counts[0]); + CUDF_TEST_EXPECT_EQUAL_BUFFERS( + result_masks[0]->data(), expected_mask.data(), cudf::num_bitmask_words(num_rows)); + CUDF_TEST_EXPECT_EQUAL_BUFFERS(result_masks[1]->data(), + static_cast(col3).null_mask(), + cudf::num_bitmask_words(num_rows)); +} + TEST_F(MergeBitmaskTest, TestSegmentedBitmaskAndNoColumns) { std::vector const colviews{};