diff --git a/cpp/benchmarks/bitmask/bitmask_and.cpp b/cpp/benchmarks/bitmask/bitmask_and.cpp index 9baed4b12ffb..b6ddb293e79e 100644 --- a/cpp/benchmarks/bitmask/bitmask_and.cpp +++ b/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(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}) diff --git a/cpp/include/cudf/detail/null_mask.cuh b/cpp/include/cudf/detail/null_mask.cuh index 6d33cc35cb01..1e3a63b70873 100644 --- a/cpp/include/cudf/detail/null_mask.cuh +++ b/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 +template 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(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(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(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()); - - // 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. @@ -403,7 +402,8 @@ rmm::device_uvector inplace_segmented_bitmask_binop( CUDF_EXPECTS(segment_offsets.size() >= 2, "At least one segment needs to be passed for bitwise operations"); - rmm::device_uvector d_null_counts(segment_offsets.size() - 1, stream, mr); + auto const num_segments = static_cast(segment_offsets.size() - 1); + rmm::device_uvector 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); @@ -411,22 +411,29 @@ 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(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<<>>( - 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(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 + <<>>(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; }