diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 664b13dd971c..37e4b4688972 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -735,18 +735,17 @@ add_library( src/groupby/common/m2_var_std.cu src/groupby/common/utils.cpp src/groupby/groupby.cu - src/groupby/hash/compute_global_memory_aggs.cu - src/groupby/hash/compute_global_memory_aggs_null.cu src/groupby/hash/compute_groupby.cu - src/groupby/hash/compute_mapping_indices.cu - src/groupby/hash/compute_mapping_indices_null.cu - src/groupby/hash/compute_shared_memory_aggs.cu src/groupby/hash/compute_single_pass_aggs.cu - src/groupby/hash/compute_single_pass_aggs_null.cu src/groupby/hash/extract_single_pass_aggs.cpp src/groupby/hash/groupby.cu src/groupby/hash/hash_compound_agg_finalizer.cu src/groupby/hash/output_utils.cu + src/groupby/hash/single_pass_argminmax.cu + src/groupby/hash/single_pass_minmax.cu + src/groupby/hash/single_pass_product.cu + src/groupby/hash/single_pass_sum_overflow.cu + src/groupby/hash/single_pass_sums.cu src/groupby/sort/aggregate.cpp src/groupby/sort/group_argmax.cu src/groupby/sort/group_argmin.cu diff --git a/cpp/src/groupby/hash/compute_global_memory_aggs.cu b/cpp/src/groupby/hash/compute_global_memory_aggs.cu deleted file mode 100644 index bd979be593cc..000000000000 --- a/cpp/src/groupby/hash/compute_global_memory_aggs.cu +++ /dev/null @@ -1,22 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "compute_global_memory_aggs.cuh" - -#include - -namespace cudf::groupby::detail::hash { - -template std::pair, rmm::device_uvector> -compute_global_memory_aggs(bitmask_type const* row_bitmask, - table_view const& values, - global_set_t const& key_set, - host_span h_agg_kinds, - device_span d_agg_kinds, - std::span is_agg_intermediate, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_global_memory_aggs.cuh b/cpp/src/groupby/hash/compute_global_memory_aggs.cuh deleted file mode 100644 index 1fdce3089bac..000000000000 --- a/cpp/src/groupby/hash/compute_global_memory_aggs.cuh +++ /dev/null @@ -1,189 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include "compute_global_memory_aggs.hpp" -#include "output_utils.hpp" -#include "single_pass_functors.cuh" - -#include -#include - -#include - -#include -#include -#include -#include - -#include -#include -#include -#include - -namespace cudf::groupby::detail::hash { - -/** - * @brief Compute and return an array mapping each input row to its corresponding key index in - * the input keys table. - * - * @tparam SetType Type of the key hash set - * @param row_bitmask Bitmask indicating which rows in the input keys table are valid - * @param set_ref Key hash set - * @param num_rows Number of rows in the input keys table - * @param stream CUDA stream used for device memory operations and kernel launches - * @return A device vector mapping each input row to its key index - */ -template -rmm::device_uvector compute_matching_keys(bitmask_type const* row_bitmask, - SetRef set_ref, - size_type num_rows, - cuda::stream_ref stream) -{ - // Mapping from each row in the input key/value into the indices of the key. - rmm::device_uvector key_indices(num_rows, stream); - - // Need to set to sentinel value for rows that are null (if any). - // The sentinel value will then be used to identify null rows instead of using the bitmask. - thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator(0), - cuda::counting_iterator(num_rows), - key_indices.begin(), - [set_ref, row_bitmask] __device__(size_type const idx) mutable { - if (!row_bitmask || cudf::bit_is_set(row_bitmask, idx)) { - return *set_ref.insert_and_find(idx).first; - } - return cudf::detail::CUDF_SIZE_TYPE_SENTINEL; - }); - return key_indices; -} - -/** - * @brief Compute aggregations and write the results directly to a dense output table. - * - * The target indices in the dense output table are computed by firstly inserting all keys into the - * hash set then extracting the unique keys and computing a transform map from the input key - * indices to the output key indices. This incurs some overhead, but it allows us to avoid - * allocating extra memory for the sparse table and gathering the results from the sparse table into - * the final dense table. - */ -template -std::pair, rmm::device_uvector> compute_aggs_dense_output( - bitmask_type const* row_bitmask, - table_view const& values, - SetType const& key_set, - host_span h_agg_kinds, - device_span d_agg_kinds, - std::span is_agg_intermediate, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - auto const num_rows = values.num_rows(); - auto [unique_keys, target_indices] = [&] { - auto matching_keys = - compute_matching_keys(row_bitmask, key_set.ref(cuco::op::insert_and_find), num_rows, stream); - auto unique_keys = extract_populated_keys(key_set, num_rows, stream, mr); - auto key_transform_map = compute_key_transform_map( - num_rows, unique_keys, stream, cudf::get_current_device_resource_ref()); - auto target_indices = compute_target_indices( - matching_keys, key_transform_map, stream, cudf::get_current_device_resource_ref()); - return std::pair{std::move(unique_keys), std::move(target_indices)}; - }(); - - auto const d_values = table_device_view::create(values, stream); - auto agg_results = create_results_table(static_cast(unique_keys.size()), - values, - h_agg_kinds, - is_agg_intermediate, - stream, - mr); - auto d_results_ptr = mutable_table_device_view::create(*agg_results, stream); - - thrust::for_each_n(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{0}, - num_rows * static_cast(h_agg_kinds.size()), - compute_single_pass_aggs_dense_output_fn{ - target_indices.begin(), d_agg_kinds.data(), *d_values, *d_results_ptr}); - - return {std::move(agg_results), std::move(unique_keys)}; -} - -/** - * @brief Compute aggregations and write the results to a sparse intermediate output table, then - * generate the final output table by gathering the relevant rows from it. - * - * During computing the aggregations, we write the results to a sparse intermediate output table - * using the target indices as indices of the input keys. Such input key indices are computed by - * inserting keys into the hash set on-the-fly. - */ -template -std::pair, rmm::device_uvector> compute_aggs_sparse_output_gather( - bitmask_type const* row_bitmask, - table_view const& values, - SetType const& key_set, - host_span h_agg_kinds, - device_span d_agg_kinds, - std::span is_agg_intermediate, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - auto const num_rows = values.num_rows(); - auto const d_values = table_device_view::create(values, stream); - auto agg_results = - create_results_table(num_rows, values, h_agg_kinds, is_agg_intermediate, stream, mr); - auto d_results_ptr = mutable_table_device_view::create(*agg_results, stream); - - thrust::for_each_n( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{0}, - num_rows, - compute_single_pass_aggs_sparse_output_fn{key_set.ref(cuco::op::insert_and_find), - row_bitmask, - d_agg_kinds.data(), - *d_values, - *d_results_ptr}); - - auto unique_keys = extract_populated_keys(key_set, num_rows, stream, mr); - auto dense_results = cudf::detail::gather(agg_results->view(), - unique_keys, - out_of_bounds_policy::DONT_CHECK, - cudf::negative_index_policy::NOT_ALLOWED, - stream, - mr); - return {std::move(dense_results), std::move(unique_keys)}; -} - -template -std::pair, rmm::device_uvector> compute_global_memory_aggs( - bitmask_type const* row_bitmask, - table_view const& values, - SetType const& key_set, - host_span h_agg_kinds, - device_span d_agg_kinds, - std::span is_agg_intermediate, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - return h_agg_kinds.size() > GROUPBY_DENSE_OUTPUT_THRESHOLD - ? compute_aggs_dense_output(row_bitmask, - values, - key_set, - h_agg_kinds, - d_agg_kinds, - is_agg_intermediate, - stream, - mr) - : compute_aggs_sparse_output_gather(row_bitmask, - values, - key_set, - h_agg_kinds, - d_agg_kinds, - is_agg_intermediate, - stream, - mr); -} - -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_global_memory_aggs.hpp b/cpp/src/groupby/hash/compute_global_memory_aggs.hpp deleted file mode 100644 index c3b3367ed79f..000000000000 --- a/cpp/src/groupby/hash/compute_global_memory_aggs.hpp +++ /dev/null @@ -1,32 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include -#include -#include -#include - -#include - -#include - -#include -#include -#include -#include - -namespace cudf::groupby::detail::hash { -template -std::pair, rmm::device_uvector> compute_global_memory_aggs( - bitmask_type const* row_bitmask, - table_view const& values, - SetType const& key_set, - host_span h_agg_kinds, - device_span d_agg_kinds, - std::span is_agg_intermediate, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_global_memory_aggs_null.cu b/cpp/src/groupby/hash/compute_global_memory_aggs_null.cu deleted file mode 100644 index a2359fb5581e..000000000000 --- a/cpp/src/groupby/hash/compute_global_memory_aggs_null.cu +++ /dev/null @@ -1,22 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "compute_global_memory_aggs.cuh" - -#include - -namespace cudf::groupby::detail::hash { - -template std::pair, rmm::device_uvector> -compute_global_memory_aggs(bitmask_type const* row_bitmask, - table_view const& values, - nullable_global_set_t const& key_set, - host_span h_agg_kinds, - device_span d_agg_kinds, - std::span is_agg_intermediate, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_groupby.cu b/cpp/src/groupby/hash/compute_groupby.cu index c06aaa6aa2f8..af4eef4f8640 100644 --- a/cpp/src/groupby/hash/compute_groupby.cu +++ b/cpp/src/groupby/hash/compute_groupby.cu @@ -5,44 +5,316 @@ #include "compute_groupby.hpp" #include "compute_single_pass_aggs.hpp" +#include "extract_single_pass_aggs.hpp" #include "groupby/common/utils.hpp" #include "hash_compound_agg_finalizer.hpp" +#include "hash_csr_kernels.cuh" #include "helpers.cuh" -#include "output_utils.hpp" -#include +#include #include +#include #include -#include +#include +#include #include +#include #include +#include +#include +#include +#include #include -#include -#include #include +#include #include #include -#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include namespace cudf::groupby::detail::hash { namespace { -// The number of columns in the keys table that will trigger caching of row hashes. -// This is a heuristic to reduce memory read when the keys table is hashed twice. -constexpr int HASH_CACHING_THRESHOLD = 4; +/// The keys grouped by the HashCSR build. +struct grouped_keys { + size_type num_groups; + size_type num_grouped_rows; + rmm::device_uvector key_rows; ///< One representative input row per group + rmm::device_uvector group_offsets; ///< `num_groups + 1` offsets into `grouped_rows` + rmm::device_uvector grouped_rows; ///< Input rows reordered so groups are contiguous +}; + +cuda::std::uint32_t hash_csr_capacity(size_type num_rows) +{ + auto const requested = + std::max(static_cast(num_rows) + 1, + std::ceil(static_cast(num_rows) / cudf::detail::CUCO_DESIRED_LOAD_FACTOR)); + CUDF_EXPECTS(requested <= std::numeric_limits::max(), + "HashCSR table capacity is not representable", + std::overflow_error); + return static_cast(requested); +} + +struct is_occupied_fn { + __device__ bool operator()(hash_table_entry_type entry) const + { + return entry != cudf::detail::CUDF_SIZE_TYPE_SENTINEL; + } +}; -int count_nested_columns(column_view const& input) +/** + * @brief Estimates the table capacity that fits the distinct keys of a large input. + * + * Every `stride`-th row is inserted into a small table, and the number of distinct keys among + * those rows is corrected for the keys the sample missed: `D` distinct keys show + * `D * (1 - exp(-s / D))` of themselves in a sample of `s` rows, which is solved for `D`. + * + * @return Four slots per estimated distinct key, or the maximum capacity when the estimate + * exceeds the representable capacity + */ +template +cuda::std::uint32_t estimate_capacity(size_type num_rows, + bitmask_type const* row_bitmask, + Equal const& d_row_equal, + Hash const& d_row_hash, + cuda::stream_ref stream) { - if (!is_nested(input.type())) { return 1; } + auto const temp_mr = cudf::get_current_device_resource_ref(); + rmm::device_uvector entries(hash_csr_sample_capacity, stream, temp_mr); + rmm::device_uvector counts(2, stream, temp_mr); + // Counts the valid rows among every `stride`-th row and the distinct keys among them. + auto const sample = [&](size_type stride) { + CUDF_CUDA_TRY(cudaMemsetAsync( + entries.data(), 0xff, entries.size() * sizeof(hash_table_entry_type), stream.get())); + CUDF_CUDA_TRY( + cudaMemsetAsync(counts.data(), 0, counts.size() * sizeof(size_type), stream.get())); + launch_hash_csr_sample_kernel( + num_rows, + stride, + row_bitmask, + hash_csr_table_ref{entries.data(), hash_csr_sample_capacity, hash_csr_sample_capacity}, + d_row_equal, + d_row_hash, + counts.data(), + stream); + auto const h_counts = + cudf::detail::make_pinned_vector(device_span{counts}, stream); + return std::pair{static_cast(h_counts[0]), static_cast(h_counts[1])}; + }; + // One row in 64 is sampled, fewer when that would fill more than half of the sample table. + auto const stride = std::max( + 64, cudf::util::div_rounding_up_safe(num_rows, hash_csr_sample_capacity / 2)); + auto const max_capacity = std::numeric_limits::max(); + auto [sampled, distinct] = sample(stride); + if (sampled == 0) { return hash_csr_min_estimated_capacity; } + + // The expected number of distinct keys seen grows with the population, so bisect on it. + auto const seen = [sampled](double population) { + return population * (1.0 - std::exp(-sampled / population)); + }; + auto low = distinct; + auto high = static_cast(num_rows); + for (int i = 0; i < 64; ++i) { + auto const mid = 0.5 * (low + high); + (seen(mid) < distinct ? low : high) = mid; + } + // Four slots per distinct key keep the probes short while the table stays small. + auto const estimate = 4.0 * high; + if (estimate >= static_cast(max_capacity)) { return max_capacity; } + return std::max(hash_csr_min_estimated_capacity, static_cast(estimate)); +} + +/** + * @brief Groups the input rows by key with a HashCSR build. + * + * Every valid row inserts its key into an open-addressed table and takes a rank within the slot + * it lands in. The occupied slots become the groups, a scan of their row counts gives the group + * offsets, and a scatter of the rows by slot offset plus rank yields the grouped row order. + */ +template +grouped_keys group_keys(size_type num_rows, + size_type key_bytes, + bitmask_type const* row_bitmask, + Equal const& d_row_equal, + Hash const& d_row_hash, + bool need_grouped_rows, + cuda::stream_ref stream) +{ + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto const policy = rmm::exec_policy_nosync(stream, temp_mr); + + // A table with a slot for every row would spread a few distinct keys over a table too large for + // the cache and make clearing and compacting its slots the dominant cost of low-cardinality + // inputs, so large inputs get a table sized from an estimate of their number of distinct keys. + // Should the estimate fall short, the build restarts with the table sized for every row. + auto const full_capacity = hash_csr_capacity(num_rows); + auto capacity = + num_rows < hash_csr_min_rows_to_estimate || key_bytes > hash_csr_max_estimated_key_bytes + ? full_capacity + : std::min(full_capacity, + estimate_capacity(num_rows, row_bitmask, d_row_equal, d_row_hash, stream)); + rmm::device_uvector entries(0, stream, temp_mr); + rmm::device_uvector slot_counts(0, stream, temp_mr); + rmm::device_uvector positions( + need_grouped_rows ? num_rows : 0, stream, temp_mr); + // Set by the build when the estimated table turns out to be too small. + std::optional> overflow; + if (capacity < full_capacity) { overflow.emplace(0, stream, temp_mr); } + // The occupied slots, in slot order, are the groups: without aggregations the slots hold the + // one row wanted for each group, otherwise the slot indices lead to the counts and rows. + rmm::device_uvector key_rows(0, stream, temp_mr); + rmm::device_uvector group_slots(0, stream, temp_mr); + bool count_by_representative{}; + while (true) { + auto const is_full_size = capacity == full_capacity; + count_by_representative = + need_grouped_rows && static_cast(num_rows) < capacity; + auto const count_capacity = + count_by_representative ? static_cast(num_rows) : capacity; + entries.resize(capacity, stream); + CUDF_CUDA_TRY(cudaMemsetAsync( + entries.data(), 0xff, entries.size() * sizeof(hash_table_entry_type), stream.get())); + if (need_grouped_rows) { + slot_counts.resize(count_capacity, stream); + if (count_capacity != 0) { + CUDF_CUDA_TRY(cudaMemsetAsync( + slot_counts.data(), 0, slot_counts.size() * sizeof(size_type), stream.get())); + } + } + auto const table = + hash_csr_table_ref{entries.data(), capacity, is_full_size ? capacity : hash_csr_max_probes}; + launch_hash_csr_build_kernel(num_rows, + row_bitmask, + need_grouped_rows ? positions.data() : nullptr, + need_grouped_rows ? slot_counts.data() : nullptr, + count_by_representative, + table, + d_row_equal, + d_row_hash, + is_full_size ? nullptr : overflow->data(), + stream); + if (count_by_representative) { + // Count indices identify representative rows, so selection no longer needs the table. + entries.resize(0, stream); + entries.shrink_to_fit(stream); + } + if (!need_grouped_rows) { + key_rows.resize(std::min(num_rows, capacity), stream); + auto const key_rows_end = + thrust::copy_if(policy, entries.begin(), entries.end(), key_rows.begin(), is_occupied_fn{}); + key_rows.resize(cuda::std::distance(key_rows.begin(), key_rows_end), stream); + } else { + group_slots.resize(std::min(num_rows, capacity), stream); + auto const group_slots_end = + thrust::copy_if(policy, + cuda::counting_iterator{0}, + cuda::counting_iterator{count_capacity}, + slot_counts.begin(), + group_slots.begin(), + [] __device__(size_type count) -> bool { return count > 0; }); + group_slots.resize(cuda::std::distance(group_slots.begin(), group_slots_end), stream); + } + // The compaction has just synchronized the stream, so reading the flag is cheap here. + if (is_full_size || overflow->value(stream) == 0) { break; } + // The retry overwrites the table, so release it instead of copying it while growing. + entries.resize(0, stream); + entries.shrink_to_fit(stream); + slot_counts.resize(0, stream); + slot_counts.shrink_to_fit(stream); + key_rows.resize(0, stream); + key_rows.shrink_to_fit(stream); + group_slots.resize(0, stream); + group_slots.shrink_to_fit(stream); + overflow.reset(); + capacity = full_capacity; + } + + if (!need_grouped_rows) { + auto const num_groups = static_cast(key_rows.size()); + return {num_groups, + 0, + std::move(key_rows), + rmm::device_uvector{0, stream, temp_mr}, + rmm::device_uvector{0, stream, temp_mr}}; + } + + auto const num_groups = static_cast(group_slots.size()); + // Every row is an included singleton group, so input order already forms a valid grouping. + if (num_groups == num_rows) { + entries.resize(0, stream); + entries.shrink_to_fit(stream); + slot_counts.resize(0, stream); + slot_counts.shrink_to_fit(stream); + positions.resize(0, stream); + positions.shrink_to_fit(stream); + group_slots.resize(0, stream); + group_slots.shrink_to_fit(stream); + overflow.reset(); - // Count the current column too. - return 1 + std::accumulate( - input.child_begin(), input.child_end(), 0, [](int count, column_view const& child) { - return count + count_nested_columns(child); - }); + key_rows.resize(num_rows, stream); + rmm::device_uvector group_offsets( + static_cast(num_rows) + 1, stream, temp_mr); + rmm::device_uvector grouped_rows(num_rows, stream, temp_mr); + thrust::sequence(policy, key_rows.begin(), key_rows.end(), size_type{0}); + thrust::sequence(policy, group_offsets.begin(), group_offsets.end(), size_type{0}); + thrust::sequence(policy, grouped_rows.begin(), grouped_rows.end(), size_type{0}); + return { + num_groups, num_rows, std::move(key_rows), std::move(group_offsets), std::move(grouped_rows)}; + } + + auto const entry_rows = entries.begin(); + key_rows.resize(num_groups, stream); + if (count_by_representative) { + thrust::copy(policy, group_slots.begin(), group_slots.end(), key_rows.begin()); + } else { + thrust::gather(policy, group_slots.begin(), group_slots.end(), entry_rows, key_rows.begin()); + } + entries.resize(0, stream); + entries.shrink_to_fit(stream); + + rmm::device_uvector group_offsets(group_slots.size() + 1, stream, temp_mr); + group_offsets.set_element_to_zero_async(0, stream); + auto const group_counts = + cuda::make_permutation_iterator(slot_counts.begin(), group_slots.begin()); + thrust::inclusive_scan( + policy, group_counts, group_counts + num_groups, group_offsets.begin() + 1); + auto const num_grouped_rows = + row_bitmask == nullptr ? num_rows : group_offsets.back_element(stream); + + // Reuse the slot counts to hold the start offset of the group of each occupied slot, then + // scatter every row to its group. + thrust::scatter(policy, + group_offsets.begin(), + group_offsets.begin() + num_groups, + group_slots.begin(), + slot_counts.begin()); + group_slots.resize(0, stream); + group_slots.shrink_to_fit(stream); + rmm::device_uvector grouped_rows(num_grouped_rows, stream, temp_mr); + launch_hash_csr_fill_kernel( + num_rows, positions.data(), slot_counts.data(), grouped_rows.data(), stream); + + return {num_groups, + num_grouped_rows, + std::move(key_rows), + std::move(group_offsets), + std::move(grouped_rows)}; } } // namespace @@ -57,7 +329,7 @@ std::unique_ptr compute_groupby(table_view const& keys, cuda::stream_ref stream, rmm::device_async_resource_ref mr) { - auto const num_keys = keys.num_rows(); + auto const num_rows = keys.num_rows(); [[maybe_unused]] auto [row_bitmask_data, row_bitmask] = skip_rows_with_nulls @@ -65,45 +337,17 @@ std::unique_ptr
compute_groupby(table_view const& keys, : std::pair{ rmm::device_buffer{0, stream, cudf::get_current_device_resource_ref()}, nullptr}; - auto const cached_hashes = [&]() -> rmm::device_uvector { - auto const num_columns = - std::accumulate(keys.begin(), keys.end(), 0, [](int count, column_view const& col) { - return count + count_nested_columns(col); - }); - - if (num_columns <= HASH_CACHING_THRESHOLD) { - return rmm::device_uvector{ - 0, stream, cudf::get_current_device_resource_ref()}; - } + // Bytes of one key row, with variable-width and nested columns counted as wide. + auto const key_bytes = std::accumulate( + keys.begin(), keys.end(), size_type{0}, [](size_type bytes, column_view const& col) { + return bytes + (cudf::is_fixed_width(col.type()) ? cudf::size_of(col.type()) : 64); + }); + auto const groups = group_keys( + num_rows, key_bytes, row_bitmask, d_row_equal, d_row_hash, !requests.empty(), stream); - rmm::device_uvector hashes( - num_keys, stream, cudf::get_current_device_resource_ref()); - thrust::tabulate(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - hashes.begin(), - hashes.end(), - [d_row_hash, row_bitmask] __device__(size_type const idx) { - if (!row_bitmask || cudf::bit_is_set(row_bitmask, idx)) { - return d_row_hash(idx); - } - return hash_value_type{0}; // dummy value, as it will be unused - }); - return hashes; - }(); - - auto set = - cuco::static_set{cuco::extent{static_cast(num_keys)}, - cudf::detail::CUCO_DESIRED_LOAD_FACTOR, // 50% load factor - cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL}, - d_row_equal, - probing_scheme_t{row_hasher_with_cache_t{d_row_hash, cached_hashes.data()}}, - cuco::thread_scope_device, - cuco::storage{}, - rmm::mr::polymorphic_allocator{}, - stream.get()}; - - auto const gather_keys = [&](auto const& gather_map) { + auto const gather_keys = [&] { return cudf::detail::gather(keys, - gather_map, + groups.key_rows, out_of_bounds_policy::DONT_CHECK, cudf::negative_index_policy::NOT_ALLOWED, stream, @@ -111,39 +355,26 @@ std::unique_ptr
compute_groupby(table_view const& keys, }; // In case of no requests, we still need to generate a set of unique keys. - if (requests.empty()) { - thrust::for_each_n( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{0}, - num_keys, - [set_ref = set.ref(cuco::op::insert), row_bitmask] __device__(size_type const idx) mutable { - if (!row_bitmask || cudf::bit_is_set(row_bitmask, idx)) { set_ref.insert(idx); } - }); - - rmm::device_uvector unique_key_indices( - num_keys, stream, cudf::get_current_device_resource_ref()); - auto const keys_end = set.retrieve_all(unique_key_indices.begin(), stream.get()); - auto const key_gather_map = device_span{ - unique_key_indices.data(), - static_cast(cuda::std::distance(unique_key_indices.begin(), keys_end))}; - return gather_keys(key_gather_map); - } + if (requests.empty()) { return gather_keys(); } // Compute all single pass aggs first. - auto const [key_gather_map, has_compound_aggs] = - compute_single_pass_aggs(set, row_bitmask, requests, cache, stream, mr); + auto const [values, agg_kinds, aggs, is_agg_intermediate, has_compound_aggs] = + extract_single_pass_aggs(requests, stream); + + auto const grouped = make_grouped_rows(groups.grouped_rows, groups.group_offsets, stream); + auto results = + compute_single_pass_aggs(values, agg_kinds, is_agg_intermediate, grouped, stream, mr); + for (std::size_t i = 0; i < results.size(); ++i) { + cache->add_result(values.column(i), *aggs[i], std::move(results[i])); + } if (has_compound_aggs) { for (auto const& request : requests) { auto const& agg_v = request.aggregations; auto const& col = request.values; - // The map to find the target output index for each input row is not always available due to - // minimizing overhead. As such, there is no way for the finalizers to perform additional - // aggregation operations. They can only compute their output using the previously computed - // single-pass aggregations with linear transformations such as addition/multiplication (e.g. - // for variance/stddev). In the future, if there are more compound aggregations that require - // additional aggregation steps, we can revisit this design. + // The finalizers only combine the single-pass results with linear transformations such as + // addition/multiplication (e.g. for variance/stddev); they do not aggregate further. auto const finalizer = hash_compound_agg_finalizer(col, cache, row_bitmask, stream, mr); for (auto&& agg : agg_v) { cudf::detail::aggregation_dispatcher(agg->kind, finalizer, *agg); @@ -151,7 +382,7 @@ std::unique_ptr
compute_groupby(table_view const& keys, } } - return gather_keys(key_gather_map); + return gather_keys(); } template std::unique_ptr
compute_groupby( @@ -173,4 +404,5 @@ template std::unique_ptr
compute_groupby>(); - -template void compute_mapping_indices>( - size_type grid_size, - size_type num_rows, - hash_set_ref_t global_set, - bitmask_type const* row_bitmask, - size_type* local_mapping_index, - size_type* global_mapping_index, - size_type* block_cardinality, - cuda::std::atomic_flag* needs_global_memory_fallback, - cuda::stream_ref stream); -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_mapping_indices.cuh b/cpp/src/groupby/hash/compute_mapping_indices.cuh deleted file mode 100644 index dff18f512834..000000000000 --- a/cpp/src/groupby/hash/compute_mapping_indices.cuh +++ /dev/null @@ -1,183 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include "compute_mapping_indices.hpp" -#include "helpers.cuh" - -#include -#include -#include -#include - -#include -#include -#include -#include - -#include - -namespace cudf::groupby::detail::hash { -template -__device__ void find_local_mapping(cooperative_groups::thread_block const& block, - size_type idx, - size_type num_input_rows, - SetType shared_set, - bitmask_type const* row_bitmask, - size_type* cardinality, - size_type* local_mapping_indices, - size_type* shared_set_indices) -{ - auto const is_valid_input = - idx < num_input_rows and (not row_bitmask or cudf::bit_is_set(row_bitmask, idx)); - auto const [result_idx, inserted] = [&]() { - if (is_valid_input) { - auto const result = shared_set.insert_and_find(idx); - auto const matched_idx = *result.first; - auto const inserted = result.second; - if (inserted) { // inserted a new element - auto const ref_cardinality = - cuda::atomic_ref{*cardinality}; - auto const shared_set_index = ref_cardinality.fetch_add(1, cuda::std::memory_order_relaxed); - - // The value of `shared_set_index` is before increment, thus if we have - // `shared_set_index == GROUPBY_CARDINALITY_THRESHOLD` the value of cardinality - // will be at least `GROUPBY_CARDINALITY_THRESHOLD + 1`. - // This will trigger fallback to global memory. - if (shared_set_index >= GROUPBY_CARDINALITY_THRESHOLD) { return cuda::std::pair{0, true}; } - - shared_set_indices[shared_set_index] = idx; - local_mapping_indices[idx] = shared_set_index; - } - return cuda::std::pair{matched_idx, inserted}; - } - return cuda::std::pair{0, false}; // dummy values - }(); - // Syncing the thread block is needed so that updates in `local_mapping_indices` are visible to - // all threads in the thread block. - block.sync(); - if (is_valid_input) { - // element was already in set - if (!inserted) { local_mapping_indices[idx] = local_mapping_indices[result_idx]; } - } -} - -template -__device__ void find_global_mapping(cooperative_groups::thread_block const& block, - size_type cardinality, - SetRef global_set, - size_type* shared_set_indices, - size_type* global_mapping_indices) -{ - // for all unique keys in shared memory hash set, stores their matches in - // global hash set to `global_mapping_indices` - for (auto idx = block.thread_rank(); idx < cardinality; idx += block.num_threads()) { - auto const input_idx = shared_set_indices[idx]; - auto const key_idx = *global_set.insert_and_find(input_idx).first; - - global_mapping_indices[block.group_index().x * GROUPBY_CARDINALITY_THRESHOLD + idx] = key_idx; - } -} - -/* - * @brief Inserts keys into the shared memory hash set, and stores the block-wise rank for a given - * row index in `local_mapping_indices`. If the number of unique keys found in a threadblock exceeds - * `GROUPBY_CARDINALITY_THRESHOLD`, the threads in that block will exit without updating - * `global_set` or setting `global_mapping_indices`. Else, we insert the unique keys found to the - * global hash set, and save the row index of the global sparse table in `global_mapping_indices`. - */ -template -CUDF_KERNEL void mapping_indices_kernel(size_type num_input_rows, - SetRef global_set, - bitmask_type const* row_bitmask, - size_type* local_mapping_indices, - size_type* global_mapping_indices, - size_type* block_cardinality, - cuda::std::atomic_flag* needs_global_memory_fallback) -{ - __shared__ size_type shared_set_indices[GROUPBY_CARDINALITY_THRESHOLD]; - - // Shared set initialization - __shared__ size_type slots[valid_extent.value()]; - - auto raw_set = cuco::static_set_ref{ - cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL}, - global_set.key_eq(), - probing_scheme_t{global_set.hash_function()}, - cuco::thread_scope_block, - cuco::bucket_storage_ref{valid_extent, - slots}}; - auto shared_set = raw_set.rebind_operators(cuco::insert_and_find); - - auto const block = cooperative_groups::this_thread_block(); - shared_set.initialize(block); - - __shared__ size_type cardinality; - if (block.thread_rank() == 0) { cardinality = 0; } - block.sync(); - - // All threads in the block will participate in the loop, and sync. - auto const stride = cudf::detail::grid_1d::grid_stride(); - for (auto idx = cudf::detail::grid_1d::global_thread_id(); - idx - block.thread_rank() < num_input_rows; - idx += stride) { - find_local_mapping(block, - idx, - num_input_rows, - shared_set, - row_bitmask, - &cardinality, - local_mapping_indices, - shared_set_indices); - - block.sync(); - if (cardinality > GROUPBY_CARDINALITY_THRESHOLD) { - if (block.thread_rank() == 0) { - needs_global_memory_fallback->test_and_set(cuda::std::memory_order_relaxed); - } - break; - } - } - - // Insert unique keys from shared to global hash set if block-cardinality - // doesn't exceed the threshold upper-limit - if (cardinality <= GROUPBY_CARDINALITY_THRESHOLD) { - find_global_mapping(block, cardinality, global_set, shared_set_indices, global_mapping_indices); - } - - if (block.thread_rank() == 0) { block_cardinality[block.group_index().x] = cardinality; } -} - -template -int32_t max_active_blocks_mapping_kernel() -{ - int32_t max_active_blocks{-1}; - CUDF_CUDA_TRY(cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &max_active_blocks, mapping_indices_kernel, GROUPBY_BLOCK_SIZE, 0)); - return max_active_blocks; -} - -template -void compute_mapping_indices(size_type grid_size, - size_type num_rows, - SetRef global_set, - bitmask_type const* row_bitmask, - size_type* local_mapping_indices, - size_type* global_mapping_indices, - size_type* block_cardinality, - cuda::std::atomic_flag* needs_global_memory_fallback, - cuda::stream_ref stream) -{ - mapping_indices_kernel<<>>( - num_rows, - global_set, - row_bitmask, - local_mapping_indices, - global_mapping_indices, - block_cardinality, - needs_global_memory_fallback); - CUDF_CUDA_TRY(cudaGetLastError()); -} -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_mapping_indices.hpp b/cpp/src/groupby/hash/compute_mapping_indices.hpp deleted file mode 100644 index be098275dcc7..000000000000 --- a/cpp/src/groupby/hash/compute_mapping_indices.hpp +++ /dev/null @@ -1,31 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include - -#include -#include - -namespace cudf::groupby::detail::hash { - -/* - * @brief Computes the maximum number of active blocks of the mapping indices kernel that can be - * executed on the underlying device. - */ -template -[[nodiscard]] int32_t max_active_blocks_mapping_kernel(); - -template -void compute_mapping_indices(size_type grid_size, - size_type num_rows, - SetRef global_set, - bitmask_type const* row_bitmask, - size_type* local_mapping_index, - size_type* global_mapping_index, - size_type* block_cardinality, - cuda::std::atomic_flag* needs_global_memory_fallback, - cuda::stream_ref stream); -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_mapping_indices_null.cu b/cpp/src/groupby/hash/compute_mapping_indices_null.cu deleted file mode 100644 index 38bd3647c420..000000000000 --- a/cpp/src/groupby/hash/compute_mapping_indices_null.cu +++ /dev/null @@ -1,23 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "compute_mapping_indices.cuh" -#include "compute_mapping_indices.hpp" - -namespace cudf::groupby::detail::hash { -template int32_t -max_active_blocks_mapping_kernel>(); - -template void compute_mapping_indices>( - size_type grid_size, - size_type num_rows, - nullable_hash_set_ref_t global_set, - bitmask_type const* row_bitmask, - size_type* local_mapping_index, - size_type* global_mapping_index, - size_type* block_cardinality, - cuda::std::atomic_flag* needs_global_memory_fallback, - cuda::stream_ref stream); -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_shared_memory_aggs.cu b/cpp/src/groupby/hash/compute_shared_memory_aggs.cu deleted file mode 100644 index b97de8e9b378..000000000000 --- a/cpp/src/groupby/hash/compute_shared_memory_aggs.cu +++ /dev/null @@ -1,410 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "compute_shared_memory_aggs.hpp" -#include "global_memory_aggregator.cuh" -#include "helpers.cuh" -#include "shared_memory_aggregator.cuh" -#include "single_pass_functors.cuh" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -#include -#include - -namespace cudf::groupby::detail::hash { -namespace { -/// Shared memory data alignment -CUDF_HOST_DEVICE cudf::size_type constexpr ALIGNMENT = 16; - -// Dictionary and nested value columns are rejected before this kernel is launched. -struct unsupported_shared_memory_type {}; - -template -struct dispatch_shared_memory_type { - using type = cuda::std::conditional_t>; -}; - -// Compound hash aggregations are decomposed into these simple aggregations before this kernel is -// launched. SUM_OVERFLOW is the only other simple hash aggregation and is explicitly rejected by -// is_shared_memory_compatible. -template -__device__ auto dispatch_shared_memory_aggregation(cudf::aggregation::Kind kind, - F&& f, - Ts&&... args) -{ - switch (kind) { - case cudf::aggregation::SUM: - return f.template operator()(cuda::std::forward(args)...); - case cudf::aggregation::PRODUCT: - return f.template operator()(cuda::std::forward(args)...); - case cudf::aggregation::MIN: - return f.template operator()(cuda::std::forward(args)...); - case cudf::aggregation::MAX: - return f.template operator()(cuda::std::forward(args)...); - case cudf::aggregation::COUNT_VALID: - return f.template operator()(cuda::std::forward(args)...); - case cudf::aggregation::COUNT_ALL: - return f.template operator()(cuda::std::forward(args)...); - case cudf::aggregation::SUM_OF_SQUARES: - return f.template operator()( - cuda::std::forward(args)...); - case cudf::aggregation::ARGMAX: - return f.template operator()(cuda::std::forward(args)...); - case cudf::aggregation::ARGMIN: - return f.template operator()(cuda::std::forward(args)...); - default: CUDF_UNREACHABLE("Unsupported shared memory aggregation."); - } -} - -template -struct dispatch_shared_memory_aggregation_fn { - template - __device__ auto operator()(F&& f, Ts&&... args) const - { - return f.template operator()(cuda::std::forward(args)...); - } -}; - -struct dispatch_shared_memory_source_fn { - template - __device__ auto operator()(cudf::aggregation::Kind kind, F&& f, Ts&&... args) const - { - if constexpr (cuda::std::is_same_v) { - CUDF_UNREACHABLE("Unsupported shared memory aggregation type."); - } else { - return dispatch_shared_memory_aggregation(kind, - dispatch_shared_memory_aggregation_fn{}, - cuda::std::forward(f), - cuda::std::forward(args)...); - } - } -}; - -template -__device__ auto dispatch_shared_memory_type_and_aggregation(cudf::data_type type, - cudf::aggregation::Kind kind, - F&& f, - Ts&&... args) -{ - return cudf::type_dispatcher(type, - dispatch_shared_memory_source_fn{}, - kind, - cuda::std::forward(f), - cuda::std::forward(args)...); -} - -// Allocates shared memory required for output columns. Exits if there is insufficient memory to -// perform shared memory aggregation for the current output column. -__device__ void calculate_columns_to_aggregate(cudf::size_type& col_start, - cudf::size_type& col_end, - cudf::mutable_table_device_view output_values, - cudf::size_type output_size, - cudf::size_type* shmem_agg_res_offsets, - cudf::size_type* shmem_agg_mask_offsets, - cudf::size_type cardinality, - cudf::size_type total_agg_size) -{ - col_start = col_end; - cudf::size_type bytes_allocated = 0; - - auto const valid_col_size = - cudf::util::round_up_safe(static_cast(sizeof(bool) * cardinality), ALIGNMENT); - - while (bytes_allocated < total_agg_size && col_end < output_size) { - auto const col_idx = col_end; - auto const next_col_size = - cudf::util::round_up_safe(cudf::type_dispatcher( - output_values.column(col_idx).type(), size_of_functor{}) * - cardinality, - ALIGNMENT); - auto const next_col_total_size = next_col_size + valid_col_size; - - if (bytes_allocated + next_col_total_size > total_agg_size) { break; } - - shmem_agg_res_offsets[col_end] = bytes_allocated; - shmem_agg_mask_offsets[col_end] = bytes_allocated + next_col_size; - - bytes_allocated += next_col_total_size; - ++col_end; - } -} - -// Each block initialize its own shared memory aggregation results -__device__ void initialize_shmem_aggregations(cooperative_groups::thread_block const& block, - cudf::size_type col_start, - cudf::size_type col_end, - cudf::mutable_table_device_view output_values, - cuda::std::byte* shmem_agg_storage, - cudf::size_type* shmem_agg_res_offsets, - cudf::size_type* shmem_agg_mask_offsets, - cudf::size_type cardinality, - cudf::aggregation::Kind const* d_agg_kinds) -{ - for (auto col_idx = col_start; col_idx < col_end; col_idx++) { - for (auto idx = block.thread_rank(); idx < cardinality; idx += block.num_threads()) { - auto target = - reinterpret_cast(shmem_agg_storage + shmem_agg_res_offsets[col_idx]); - auto target_mask = - reinterpret_cast(shmem_agg_storage + shmem_agg_mask_offsets[col_idx]); - dispatch_shared_memory_type_and_aggregation(output_values.column(col_idx).type(), - d_agg_kinds[col_idx], - initialize_shmem{}, - target, - target_mask, - idx); - } - } -} - -__device__ void compute_pre_aggregations(cudf::size_type col_start, - cudf::size_type col_end, - bitmask_type const* row_bitmask, - cudf::table_device_view source, - cudf::size_type num_input_rows, - cudf::size_type* local_mapping_index, - cuda::std::byte* shmem_agg_storage, - cudf::size_type* shmem_agg_res_offsets, - cudf::size_type* shmem_agg_mask_offsets, - cudf::aggregation::Kind const* d_agg_kinds, - cudf::size_type agg_location_offset) -{ - // Aggregates global memory sources to shared memory targets - for (auto source_idx = cudf::detail::grid_1d::global_thread_id(); source_idx < num_input_rows; - source_idx += cudf::detail::grid_1d::grid_stride()) { - if (not row_bitmask or cudf::bit_is_set(row_bitmask, source_idx)) { - auto const target_idx = local_mapping_index[source_idx] + agg_location_offset; - for (auto col_idx = col_start; col_idx < col_end; col_idx++) { - auto const source_col = source.column(col_idx); - - cuda::std::byte* target = - reinterpret_cast(shmem_agg_storage + shmem_agg_res_offsets[col_idx]); - bool* target_mask = - reinterpret_cast(shmem_agg_storage + shmem_agg_mask_offsets[col_idx]); - - dispatch_shared_memory_type_and_aggregation(source_col.type(), - d_agg_kinds[col_idx], - shmem_element_aggregator{}, - target, - target_mask, - target_idx, - source_col, - source_idx); - } - } - } -} - -__device__ void compute_final_aggregations(cooperative_groups::thread_block const& block, - cudf::size_type col_start, - cudf::size_type col_end, - cudf::table_device_view input_values, - cudf::mutable_table_device_view target, - cudf::size_type cardinality, - cudf::size_type num_agg_locations, - cudf::size_type* global_mapping_index, - cuda::std::byte* shmem_agg_storage, - cudf::size_type* agg_res_offsets, - cudf::size_type* agg_mask_offsets, - cudf::aggregation::Kind const* d_agg_kinds) -{ - // Aggregates shared memory sources to global memory targets - for (auto idx = block.thread_rank(); idx < num_agg_locations; idx += block.num_threads()) { - auto const target_idx = - global_mapping_index[(block.group_index().x * GROUPBY_CARDINALITY_THRESHOLD) + - (idx % cardinality)]; - for (auto col_idx = col_start; col_idx < col_end; col_idx++) { - auto target_col = target.column(col_idx); - - cuda::std::byte* source = - reinterpret_cast(shmem_agg_storage + agg_res_offsets[col_idx]); - bool* source_mask = reinterpret_cast(shmem_agg_storage + agg_mask_offsets[col_idx]); - - dispatch_shared_memory_type_and_aggregation(input_values.column(col_idx).type(), - d_agg_kinds[col_idx], - gmem_element_aggregator{}, - target_col, - target_idx, - input_values.column(col_idx), - source, - source_mask, - idx); - } - } -} - -/* Takes the local_mapping_index and global_mapping_index to compute - * pre (shared) and final (global) aggregates*/ -CUDF_KERNEL void single_pass_shmem_aggs_kernel(cudf::size_type num_rows, - bitmask_type const* row_bitmask, - cudf::size_type* local_mapping_index, - cudf::size_type* global_mapping_index, - cudf::size_type* block_cardinality, - cudf::table_device_view input_values, - cudf::mutable_table_device_view output_values, - cudf::aggregation::Kind const* d_agg_kinds, - cudf::size_type total_agg_size, - cudf::size_type offsets_size) -{ - auto const block = cooperative_groups::this_thread_block(); - auto const cardinality = block_cardinality[block.group_index().x]; - if (cardinality > GROUPBY_CARDINALITY_THRESHOLD or cardinality == 0) { return; } - - auto constexpr min_shmem_agg_locations = 32; - auto const multiplication_factor = min_shmem_agg_locations / cardinality; - auto const num_agg_locations = cuda::std::max(multiplication_factor, 1) * cardinality; - auto const agg_location_offset = - multiplication_factor > 1 ? (block.thread_rank() % multiplication_factor) * cardinality : 0; - - auto const num_cols = output_values.num_columns(); - - __shared__ cudf::size_type col_start; - __shared__ cudf::size_type col_end; - extern __shared__ cuda::std::byte shmem_agg_storage[]; - - cudf::size_type* shmem_agg_res_offsets = - reinterpret_cast(shmem_agg_storage + total_agg_size); - cudf::size_type* shmem_agg_mask_offsets = - reinterpret_cast(shmem_agg_storage + total_agg_size + offsets_size); - - if (block.thread_rank() == 0) { - col_start = 0; - col_end = 0; - } - // Workaround: use __syncthreads() instead of block.sync() throughout this - // kernel. cooperative_groups::thread_block::sync() does not properly fence - // shared memory on sm_120 with CUDA 13.2, causing init stores to be - // invisible to subsequent phases. - __syncthreads(); - - while (col_end < num_cols) { - __syncthreads(); - if (block.thread_rank() == 0) { - calculate_columns_to_aggregate(col_start, - col_end, - output_values, - num_cols, - shmem_agg_res_offsets, - shmem_agg_mask_offsets, - num_agg_locations, - total_agg_size); - } - __syncthreads(); - - initialize_shmem_aggregations(block, - col_start, - col_end, - output_values, - shmem_agg_storage, - shmem_agg_res_offsets, - shmem_agg_mask_offsets, - num_agg_locations, - d_agg_kinds); - __syncthreads(); - - compute_pre_aggregations(col_start, - col_end, - row_bitmask, - input_values, - num_rows, - local_mapping_index, - shmem_agg_storage, - shmem_agg_res_offsets, - shmem_agg_mask_offsets, - d_agg_kinds, - agg_location_offset); - __syncthreads(); - - compute_final_aggregations(block, - col_start, - col_end, - input_values, - output_values, - cardinality, - num_agg_locations, - global_mapping_index, - shmem_agg_storage, - shmem_agg_res_offsets, - shmem_agg_mask_offsets, - d_agg_kinds); - } -} -} // namespace - -size_type get_available_shared_memory_size(cudf::size_type grid_size) -{ - auto const active_blocks_per_sm = - cudf::util::div_rounding_up_safe(grid_size, cudf::detail::num_multiprocessors()); - - size_t dynamic_shmem_size = 0; - CUDF_CUDA_TRY(cudaOccupancyAvailableDynamicSMemPerBlock( - &dynamic_shmem_size, single_pass_shmem_aggs_kernel, active_blocks_per_sm, GROUPBY_BLOCK_SIZE)); - return cudf::util::round_down_safe(static_cast(0.5 * dynamic_shmem_size), - ALIGNMENT); -} - -int32_t max_active_blocks_shmem_aggs_kernel() -{ - int32_t max_active_blocks{-1}; - CUDF_CUDA_TRY(cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &max_active_blocks, single_pass_shmem_aggs_kernel, GROUPBY_BLOCK_SIZE, 0)); - return max_active_blocks; -} - -void compute_shared_memory_aggs(cudf::size_type grid_size, - size_type available_shmem_size, - cudf::size_type num_input_rows, - bitmask_type const* row_bitmask, - cudf::size_type* local_mapping_index, - cudf::size_type* global_mapping_index, - cudf::size_type* block_cardinality, - cudf::table_device_view input_values, - cudf::mutable_table_device_view output_values, - cudf::aggregation::Kind const* d_agg_kinds, - cuda::stream_ref stream) -{ - // For each aggregation, need one offset determining where the aggregation is - // performed, another indicating the validity of the aggregation - auto const offsets_size = compute_shmem_offsets_size(output_values.num_columns()); - // The rest of shmem is utilized for the actual arrays in shmem - CUDF_EXPECTS(available_shmem_size > offsets_size * 2, - "No enough space for shared memory aggregations"); - auto const shmem_agg_size = available_shmem_size - offsets_size * 2; - single_pass_shmem_aggs_kernel<<>>(num_input_rows, - row_bitmask, - local_mapping_index, - global_mapping_index, - block_cardinality, - input_values, - output_values, - d_agg_kinds, - shmem_agg_size, - offsets_size); - CUDF_CUDA_TRY(cudaGetLastError()); -} -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_shared_memory_aggs.hpp b/cpp/src/groupby/hash/compute_shared_memory_aggs.hpp deleted file mode 100644 index 9a9645008c33..000000000000 --- a/cpp/src/groupby/hash/compute_shared_memory_aggs.hpp +++ /dev/null @@ -1,41 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include -#include -#include - -#include - -#include - -namespace cudf::groupby::detail::hash { - -/* - * @brief Computes the maximum number of active blocks of the shared memory aggregation kernel that - * can be executed on the underlying device. - */ -int32_t max_active_blocks_shmem_aggs_kernel(); - -size_type get_available_shared_memory_size(size_type grid_size); - -size_type constexpr compute_shmem_offsets_size(size_type num_cols) -{ - return static_cast(sizeof(size_type) * num_cols); -} - -void compute_shared_memory_aggs(size_type grid_size, - size_type available_shmem_size, - size_type num_input_rows, - bitmask_type const* row_bitmask, - size_type* local_mapping_index, - size_type* global_mapping_index, - size_type* block_cardinality, - table_device_view input_values, - mutable_table_device_view output_values, - aggregation::Kind const* d_agg_kinds, - cuda::stream_ref stream); -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_single_pass_aggs.cu b/cpp/src/groupby/hash/compute_single_pass_aggs.cu index 33c745739bab..731dd27d7ca7 100644 --- a/cpp/src/groupby/hash/compute_single_pass_aggs.cu +++ b/cpp/src/groupby/hash/compute_single_pass_aggs.cu @@ -3,46 +3,270 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "compute_single_pass_aggs.cuh" -#include "compute_single_pass_aggs.hpp" -#include "single_pass_functors.cuh" +#include "single_pass_reductions.cuh" -#include +#include +#include +#include namespace cudf::groupby::detail::hash { +namespace single_pass { -std::pair is_shared_memory_compatible(host_span agg_kinds, - table_view const& values, - size_type grid_size) +/// A group is valid when any of its rows is valid. +std::pair reduce_group_validity(reduction_context const& ctx) { - // If any aggregation has values type is dictionary, or the aggregation is SUM_OVERFLOW, - // we should always use global memory code path. - for (std::size_t i = 0; i < agg_kinds.size(); ++i) { - if (is_dictionary(values.column(i).type()) || agg_kinds[i] == aggregation::SUM_OVERFLOW) { - return {false, 0}; - } + rmm::device_uvector group_valid( + ctx.num_groups, ctx.stream, cudf::get_current_device_resource_ref()); + reduce_groups(ctx.grouped, + ctx.grouped.packed_rows, + cuda::make_permutation_iterator(cudf::detail::make_validity_iterator(ctx.d_values), + ctx.grouped.rows.begin()), + group_valid.begin(), + cuda::std::logical_or{}, + false, + ctx.stream); + return cudf::detail::valid_if( + group_valid.begin(), group_valid.end(), cuda::std::identity{}, ctx.stream, ctx.mr); +} + +void set_group_null_mask(column& result, reduction_context const& ctx) +{ + if (!ctx.nullable || ctx.num_groups == 0) { return; } + auto [null_mask, null_count] = reduce_group_validity(ctx); + result.set_null_mask(std::move(null_mask), null_count); +} + +std::unique_ptr make_size_type_column(reduction_context const& ctx) +{ + return make_fixed_width_column(data_type{type_to_id()}, + ctx.num_groups, + mask_state::UNALLOCATED, + ctx.stream, + ctx.mr); +} + +std::unique_ptr count_groups(reduction_context const& ctx, bool valid_only) +{ + auto result = make_size_type_column(ctx); + if (ctx.num_groups == 0) { return result; } + + if (valid_only && ctx.values.has_nulls()) { + auto const valid_counts = cuda::transform_iterator{ + cuda::make_permutation_iterator(cudf::detail::make_validity_iterator(ctx.d_values), + ctx.grouped.rows.begin()), + [] __device__(bool valid) -> size_type { return static_cast(valid); }}; + reduce_groups(ctx.grouped, + ctx.grouped.packed_rows, + valid_counts, + result->mutable_view().begin(), + cuda::std::plus{}, + size_type{0}, + ctx.stream); + } else { + thrust::adjacent_difference( + rmm::exec_policy_nosync(ctx.stream, cudf::get_current_device_resource_ref()), + ctx.grouped.offsets.begin() + 1, + ctx.grouped.offsets.end(), + result->mutable_view().begin()); + } + return result; +} + +/// Calls `f.template operator()()` for the reduction kind `kind`. +template +auto dispatch_reduction_kind(aggregation::Kind kind, F&& f) +{ + switch (kind) { + case aggregation::SUM: return f.template operator()(); + case aggregation::PRODUCT: return f.template operator()(); + case aggregation::SUM_OF_SQUARES: return f.template operator()(); + case aggregation::MIN: return f.template operator()(); + case aggregation::MAX: return f.template operator()(); + case aggregation::ARGMIN: return f.template operator()(); + case aggregation::ARGMAX: return f.template operator()(); + case aggregation::SUM_OVERFLOW: return f.template operator()(); + default: CUDF_FAIL("Unsupported hash groupby aggregation"); + } +} + +struct compute_reduction_fn { + reduction_context const& ctx; + + template + std::unique_ptr operator()() const + { + return compute_reduction(ctx); + } +}; + +template +struct is_reduction_supported_fn { + template + bool operator()() const + { + return is_reduction_supported(); + } +}; + +struct is_reduction_kind_supported_fn { + data_type values_type; + + template + bool operator()() const + { + return type_dispatcher(values_type, is_reduction_supported_fn{}); } +}; - auto const available_shmem_size = get_available_shared_memory_size(grid_size); - auto const offsets_buffer_size = compute_shmem_offsets_size(values.num_columns()) * 2; - auto const data_buffer_size = available_shmem_size - offsets_buffer_size; +std::unique_ptr compute_aggregation(aggregation::Kind kind, reduction_context const& ctx) +{ + switch (kind) { + case aggregation::COUNT_VALID: return count_groups(ctx, true); + case aggregation::COUNT_ALL: return count_groups(ctx, false); + default: return dispatch_reduction_kind(kind, compute_reduction_fn{ctx}); + } +} + +} // namespace single_pass + +bool is_single_pass_agg_supported(data_type values_type, aggregation::Kind kind) +{ + // Values of STRUCT and LIST types are not aggregated by the hash groupby. + if (cudf::is_nested(values_type)) { return false; } + switch (kind) { + case aggregation::COUNT_VALID: + case aggregation::COUNT_ALL: return true; + case aggregation::SUM: + case aggregation::PRODUCT: + case aggregation::SUM_OF_SQUARES: + case aggregation::MIN: + case aggregation::MAX: + case aggregation::ARGMIN: + case aggregation::ARGMAX: + case aggregation::SUM_OVERFLOW: + return single_pass::dispatch_reduction_kind( + kind, single_pass::is_reduction_kind_supported_fn{values_type}); + default: return false; + } +} + +grouped_rows make_grouped_rows(device_span rows, + device_span offsets, + cuda::stream_ref stream) +{ + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto const num_rows = static_cast(rows.size()); + auto const num_groups = static_cast(offsets.size() - 1); + grouped_rows grouped{rows, + offsets, + rmm::device_uvector{0, stream, temp_mr}, + rmm::device_uvector{0, stream, temp_mr}}; + if (num_groups == 0) { return grouped; } + + // Small groups are packed several per block; every segment is then bounded by a shorter chunk + // so that no thread or sub-warp is left walking a long group alone. + auto const avg_rows = num_rows / num_groups; + auto const packed = avg_rows < single_pass::min_avg_rows_per_segment; + auto const chunk_rows = packed ? single_pass::packed_rows_per_chunk : single_pass::rows_per_chunk; + grouped.packed_rows = packed ? std::max(avg_rows, 1) : 0; + + auto const policy = rmm::exec_policy_nosync(stream, temp_mr); + auto const chunk_counts = cudf::detail::make_counting_transform_iterator( + 0, [offsets = offsets.begin(), chunk_rows] __device__(size_type group) -> size_type { + return cudf::util::div_rounding_up_safe(offsets[group + 1] - offsets[group], chunk_rows); + }); + grouped.group_chunks.resize(offsets.size(), stream); + grouped.group_chunks.set_element_to_zero_async(0, stream); + thrust::inclusive_scan( + policy, chunk_counts, chunk_counts + num_groups, grouped.group_chunks.begin() + 1); + auto const num_chunks = grouped.group_chunks.back_element(stream); + if (num_chunks == num_groups) { + // No group spans several chunks, so the groups are the segments. + grouped.group_chunks.resize(0, stream); + grouped.group_chunks.shrink_to_fit(stream); + return grouped; + } - auto const can_run_by_shared_mem_kernel = - std::all_of(values.begin(), values.end(), [&](auto const& col) { - // Ensure there is enough buffer space to store local aggregations up to the max - // cardinality for shared memory aggregations - auto const size = type_dispatcher(col.type(), size_of_functor{}); - return data_buffer_size >= size * GROUPBY_CARDINALITY_THRESHOLD; + grouped.chunk_offsets.resize(static_cast(num_chunks) + 1, stream); + thrust::tabulate( + policy, + grouped.chunk_offsets.begin(), + grouped.chunk_offsets.end(), + [offsets = offsets.begin(), + group_chunks = grouped.group_chunks.begin(), + group_chunks_end = grouped.group_chunks.end(), + num_chunks, + num_rows, + chunk_rows] __device__(size_type chunk) -> size_type { + if (chunk == num_chunks) { return num_rows; } + auto const group = static_cast( + cuda::std::upper_bound(group_chunks, group_chunks_end, chunk) - group_chunks - 1); + return offsets[group] + (chunk - group_chunks[group]) * chunk_rows; }); - return {can_run_by_shared_mem_kernel, available_shmem_size}; + return grouped; } -template std::pair, bool> compute_single_pass_aggs( - global_set_t& global_set, - bitmask_type const* row_bitmask, - std::span requests, - cudf::detail::result_cache* cache, +std::vector> compute_single_pass_aggs( + table_view const& values, + host_span agg_kinds, + std::span is_agg_intermediate, + grouped_rows const& grouped, cuda::stream_ref stream, - rmm::device_async_resource_ref mr); + rmm::device_async_resource_ref mr) +{ + CUDF_EXPECTS(values.num_columns() == static_cast(agg_kinds.size()), + "The number of values columns and aggregation kinds must be the same."); + CUDF_EXPECTS(values.num_columns() == static_cast(is_agg_intermediate.size()), + "The number of values columns and intermediate flags must be the same."); + + auto const num_groups = static_cast(grouped.offsets.size() - 1); + auto const num_aggs = agg_kinds.size(); + + // Returns one past the last of the consecutive additive aggregations on the column of `begin` + // that can be computed together with the aggregation at `begin`. + auto const fused_end = [&](std::size_t begin, data_type values_type) { + auto const& col = values.column(begin); + if (!single_pass::is_fusable_sum(agg_kinds[begin]) || + !is_single_pass_agg_supported(values_type, aggregation::SUM_OF_SQUARES)) { + return begin + 1; + } + auto end = begin + 1; + while (end < num_aggs && single_pass::is_fusable_sum(agg_kinds[end]) && + cudf::detail::is_shallow_equivalent(col, values.column(end)) && + std::find(agg_kinds.begin() + begin, agg_kinds.begin() + end, agg_kinds[end]) == + agg_kinds.begin() + end) { + ++end; + } + return end; + }; + + std::vector> results; + results.reserve(num_aggs); + for (std::size_t i = 0; i < num_aggs;) { + auto const& col = values.column(i); + auto const d_col = column_device_view::create(col, stream); + auto const values_type = + is_dictionary(col.type()) ? dictionary_column_view(col).keys().type() : col.type(); + auto const kind = agg_kinds[i]; + // Counts are never null, and intermediate results skip the null mask to avoid the extra work. + auto const nullable = !is_agg_intermediate[i] && kind != aggregation::COUNT_VALID && + kind != aggregation::COUNT_ALL && col.has_nulls(); + auto const ctx = single_pass::reduction_context{ + col, *d_col, values_type, grouped, num_groups, nullable, stream, mr}; + + auto const end = fused_end(i, values_type); + if (end > i + 1) { + auto fused = single_pass::compute_fused_sums( + ctx, + host_span{agg_kinds}.subspan(i, end - i), + is_agg_intermediate.subspan(i, end - i)); + std::move(fused.begin(), fused.end(), std::back_inserter(results)); + } else { + results.push_back(single_pass::compute_aggregation(kind, ctx)); + } + i = end; + } + return results; +} } // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_single_pass_aggs.cuh b/cpp/src/groupby/hash/compute_single_pass_aggs.cuh deleted file mode 100644 index d38acb4b6522..000000000000 --- a/cpp/src/groupby/hash/compute_single_pass_aggs.cuh +++ /dev/null @@ -1,165 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include "compute_global_memory_aggs.hpp" -#include "compute_mapping_indices.hpp" -#include "compute_shared_memory_aggs.hpp" -#include "compute_single_pass_aggs.hpp" -#include "extract_single_pass_aggs.hpp" -#include "helpers.cuh" -#include "output_utils.hpp" - -#include -#include -#include -#include - -#include -#include - -#include -#include -#include -#include - -namespace cudf::groupby::detail::hash { - -template -std::pair, bool> compute_single_pass_aggs( - SetType& global_set, - bitmask_type const* row_bitmask, - std::span requests, - cudf::detail::result_cache* cache, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - // Collect the single-pass aggregations that can be processed in this function. - // The compound aggregations that require multiple passes will be handled separately later on. - auto const [values, agg_kinds, aggs, is_agg_intermediate, has_compound_aggs] = - extract_single_pass_aggs(requests, stream); - auto const d_agg_kinds = cudf::detail::make_device_uvector_async( - agg_kinds, stream, cudf::get_current_device_resource_ref()); - auto const num_rows = values.num_rows(); - - // Performs naive global memory aggregations when the workload is not compatible with shared - // memory, such as when aggregating dictionary columns, when there is insufficient dynamic - // shared memory for shared memory aggregations, or when SUM_OVERFLOW aggregations are - // present. - auto const run_aggs_by_global_mem_kernel = [&] { - auto [agg_results, unique_key_indices] = compute_global_memory_aggs( - row_bitmask, values, global_set, agg_kinds, d_agg_kinds, is_agg_intermediate, stream, mr); - finalize_output(values, aggs, agg_results, cache, stream); - return std::pair{std::move(unique_key_indices), has_compound_aggs}; - }; - - // Grid size used for both index mapping and shared memory aggregation kernels. - auto const grid_size = [&] { - auto const max_blocks_mapping = - max_active_blocks_mapping_kernel>(); - auto const max_blocks_aggs = max_active_blocks_shmem_aggs_kernel(); - // We launch the same grid size for both kernels, thus we need to take the minimum of the two. - auto const max_blocks = std::min(max_blocks_mapping, max_blocks_aggs); - auto const max_grid_size = max_blocks * cudf::detail::num_multiprocessors(); - auto const num_blocks = cudf::util::div_rounding_up_safe(num_rows, GROUPBY_BLOCK_SIZE); - return std::min(max_grid_size, num_blocks); - }(); - - // grid_size is zero means the shared memory kernel cannot be launched, since input cannot be - // empty: empty input should already been handled before reaching here. - if (grid_size <= 0) { return run_aggs_by_global_mem_kernel(); } - - auto const [can_use_shared_mem_kernel, available_shmem_size] = - is_shared_memory_compatible(agg_kinds, values, grid_size); - - if (!can_use_shared_mem_kernel) { return run_aggs_by_global_mem_kernel(); } - - // Maps from the global row index of the input table to its block-wise rank. - rmm::device_uvector local_mapping_indices(num_rows, stream); - // Maps from the block-wise rank to the row index of result table. - rmm::device_uvector global_mapping_indices(grid_size * GROUPBY_CARDINALITY_THRESHOLD, - stream); - // Initialize it with a sentinel value, so later we can identify which ones are unused and which - // ones need to be updated. - thrust::uninitialized_fill( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - global_mapping_indices.begin(), - global_mapping_indices.end(), - cudf::detail::CUDF_SIZE_TYPE_SENTINEL); - // Compute the cardinality (the number of unique keys) encounter by each thread block. - rmm::device_uvector block_cardinality(grid_size, stream); - - // Flag indicating whether a global memory aggregation fallback is required or not. - rmm::device_uvector needs_global_memory_fallback(1, stream); - CUDF_CUDA_TRY(cudaMemsetAsync( - needs_global_memory_fallback.data(), 0, sizeof(cuda::std::atomic_flag), stream.get())); - - auto set_ref_insert = global_set.ref(cuco::op::insert_and_find); - compute_mapping_indices(grid_size, - num_rows, - set_ref_insert, - row_bitmask, - local_mapping_indices.data(), - global_mapping_indices.data(), - block_cardinality.data(), - needs_global_memory_fallback.data(), - stream); - - auto const needs_fallback = [&] { - cuda::std::atomic_flag h_needs_fallback; - // Cannot use a value-returning helper because atomic_flag is not copy-constructible; - // copy the raw bytes back to host instead. - CUDF_CUDA_TRY(cudf::detail::memcpy_async(&h_needs_fallback, - needs_global_memory_fallback.data(), - sizeof(cuda::std::atomic_flag), - stream)); - stream.sync(); - return h_needs_fallback.test(cuda::std::memory_order_relaxed); - }(); - if (needs_fallback) { return run_aggs_by_global_mem_kernel(); } - - auto unique_keys = extract_populated_keys(global_set, num_rows, stream, mr); - - // Now, update the target indices for computing aggregations using the shared memory kernel. - { - auto key_transform_map = compute_key_transform_map( - num_rows, unique_keys, stream, cudf::get_current_device_resource_ref()); - thrust::for_each_n( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{0}, - grid_size * GROUPBY_BLOCK_SIZE, - [key_transform_map = key_transform_map.begin(), - global_mapping_indices = global_mapping_indices.begin()] __device__(auto const idx) { - auto const block_id = idx / GROUPBY_BLOCK_SIZE; - auto const thread_rank = idx % GROUPBY_BLOCK_SIZE; - auto const mapping_idx = block_id * GROUPBY_CARDINALITY_THRESHOLD + thread_rank; - auto const old_idx = global_mapping_indices[mapping_idx]; - if (old_idx != cudf::detail::CUDF_SIZE_TYPE_SENTINEL) { - global_mapping_indices[mapping_idx] = key_transform_map[old_idx]; - } - }); - } - - auto const d_spass_values = table_device_view::create(values, stream); - auto agg_results = create_results_table( - static_cast(unique_keys.size()), values, agg_kinds, is_agg_intermediate, stream, mr); - auto d_results_ptr = mutable_table_device_view::create(*agg_results, stream); - compute_shared_memory_aggs(grid_size, - available_shmem_size, - num_rows, - row_bitmask, - local_mapping_indices.data(), - global_mapping_indices.data(), - block_cardinality.data(), - *d_spass_values, - *d_results_ptr, - d_agg_kinds.data(), - stream); - - finalize_output(values, aggs, agg_results, cache, stream); - return {std::move(unique_keys), has_compound_aggs}; -} -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/compute_single_pass_aggs.hpp b/cpp/src/groupby/hash/compute_single_pass_aggs.hpp index 122750ff5de0..aa2edc47d097 100644 --- a/cpp/src/groupby/hash/compute_single_pass_aggs.hpp +++ b/cpp/src/groupby/hash/compute_single_pass_aggs.hpp @@ -4,43 +4,80 @@ */ #pragma once -#include -#include +#include +#include +#include #include +#include #include #include #include +#include +#include +#include +#include + namespace cudf::groupby::detail::hash { /** - * @brief Determine if all of provided aggregations can be computed using shared memory kernels. + * @brief Whether the hash groupby can compute the single-pass aggregation `kind` on values of + * type `values_type` (the keys type for dictionary values). + */ +bool is_single_pass_agg_supported(data_type values_type, aggregation::Kind kind); + +/** + * @brief Input rows reordered so that the rows of every group are contiguous, together with the + * arrays of the reduction strategy chosen for the group size distribution. * - * @param agg_kinds The aggregation kinds to check - * @param values The input values table corresponding to the aggregation kinds - * @param grid_size The CUDA grid size to be used for launching the aggregation kernels - * @return A pair consisting of a boolean indicating if all aggregations can be computed using - * shared memory kernels, and the currently available shared memory size + * Every group is reduced as a segment: groups that are small on average are packed several per + * block, one thread or one sub-warp each, otherwise a block reduces each group. Groups spanning + * more than one chunk of rows are first reduced per chunk, so that a few long groups still occupy + * the whole device and no thread walks a long group alone. */ -std::pair is_shared_memory_compatible(host_span agg_kinds, - table_view const& values, - size_type grid_size); +struct grouped_rows { + device_span rows; ///< Input row index at each grouped position + device_span offsets; ///< `num_groups + 1` offsets delimiting the groups + rmm::device_uvector chunk_offsets; ///< `num_chunks + 1` chunk boundaries, only when + ///< some group spans several chunks + rmm::device_uvector group_chunks; ///< `num_groups + 1` offsets into the chunks, only + ///< when some group spans several chunks + size_type packed_rows = 0; ///< Average rows per group when the segments are packed several per + ///< block; 0 when every segment gets a block +}; /** - * @brief Computes all aggregations from `requests` that can run only a single pass over the data - * and stores the results in `cache`. + * @brief Chooses the reduction strategy for the grouped rows and builds its arrays. + * + * @param rows Input row index at each grouped position + * @param offsets `num_groups + 1` offsets delimiting the groups + * @param stream CUDA stream used for device memory operations and kernel launches + */ +grouped_rows make_grouped_rows(device_span rows, + device_span offsets, + cuda::stream_ref stream); + +/** + * @brief Computes one single-pass aggregation per values column as a reduction over the grouped + * rows. + * + * Results of aggregations that only feed a compound aggregation are created without a null mask. * - * @return A pair containing a gather map to collect the unique keys from the input keys table, and - * a boolean indicating if there are any compound aggregations to process further + * @param values One values column per aggregation + * @param agg_kinds The aggregation to compute on each values column + * @param is_agg_intermediate Whether each aggregation is only an intermediate result + * @param grouped The input rows grouped by key + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the result columns + * @return One result column per aggregation with one row per group */ -template -std::pair, bool> compute_single_pass_aggs( - SetType& global_set, - bitmask_type const* row_bitmask, - std::span requests, - cudf::detail::result_cache* cache, +std::vector> compute_single_pass_aggs( + table_view const& values, + host_span agg_kinds, + std::span is_agg_intermediate, + grouped_rows const& grouped, cuda::stream_ref stream, rmm::device_async_resource_ref mr); diff --git a/cpp/src/groupby/hash/compute_single_pass_aggs_null.cu b/cpp/src/groupby/hash/compute_single_pass_aggs_null.cu deleted file mode 100644 index 0181c2576678..000000000000 --- a/cpp/src/groupby/hash/compute_single_pass_aggs_null.cu +++ /dev/null @@ -1,17 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "compute_single_pass_aggs.cuh" -#include "compute_single_pass_aggs.hpp" - -namespace cudf::groupby::detail::hash { -template std::pair, bool> -compute_single_pass_aggs(nullable_global_set_t& global_set, - bitmask_type const* row_bitmask, - std::span requests, - cudf::detail::result_cache* cache, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/global_memory_aggregator.cuh b/cpp/src/groupby/hash/global_memory_aggregator.cuh deleted file mode 100644 index c1a360319e80..000000000000 --- a/cpp/src/groupby/hash/global_memory_aggregator.cuh +++ /dev/null @@ -1,239 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include -#include -#include -#include -#include -#include - -#include -#include - -namespace cudf::groupby::detail::hash { -template -struct update_target_element_gmem { - __device__ void operator()(cudf::mutable_column_device_view, - cudf::size_type, - cudf::column_device_view, - cuda::std::byte*, - cudf::size_type) const noexcept - { - CUDF_UNREACHABLE("Invalid source type and aggregation combination."); - } -}; - -template - requires(cudf::is_fixed_width() && - cudf::has_atomic_support>()) -struct update_target_element_gmem { - __device__ void operator()(cudf::mutable_column_device_view target, - cudf::size_type target_index, - cudf::column_device_view source_column, - cuda::std::byte* source, - cudf::size_type source_index) const noexcept - { - using DeviceType = - cudf::device_storage_type_t>; - DeviceType* source_casted = reinterpret_cast(source); - cudf::detail::atomic_min(&target.element(target_index), - static_cast(source_casted[source_index])); - } -}; - -template - requires(cudf::is_fixed_width() && - cudf::has_atomic_support>()) -struct update_target_element_gmem { - __device__ void operator()(cudf::mutable_column_device_view target, - cudf::size_type target_index, - cudf::column_device_view source_column, - cuda::std::byte* source, - cudf::size_type source_index) const noexcept - { - using DeviceType = - cudf::device_storage_type_t>; - DeviceType* source_casted = reinterpret_cast(source); - cudf::detail::atomic_max(&target.element(target_index), - static_cast(source_casted[source_index])); - } -}; - -template - requires(cudf::is_fixed_width() && !cudf::is_timestamp()) -struct update_target_element_gmem { - __device__ void operator()(cudf::mutable_column_device_view target, - cudf::size_type target_index, - cudf::column_device_view source_column, - cuda::std::byte* source, - cudf::size_type source_index) const noexcept - { - using DeviceType = - cudf::device_storage_type_t>; - DeviceType* source_casted = reinterpret_cast(source); - cudf::detail::atomic_add(&target.element(target_index), - static_cast(source_casted[source_index])); - } -}; - -// The shared memory will already have it squared -template - requires(cudf::detail::is_product_supported()) -struct update_target_element_gmem { - __device__ void operator()(cudf::mutable_column_device_view target, - cudf::size_type target_index, - cudf::column_device_view source_column, - cuda::std::byte* source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - - Target* source_casted = reinterpret_cast(source); - Target value = static_cast(source_casted[source_index]); - - cudf::detail::atomic_add(&target.element(target_index), value); - } -}; - -template - requires(cudf::detail::is_product_supported()) -struct update_target_element_gmem { - __device__ void operator()(cudf::mutable_column_device_view target, - cudf::size_type target_index, - cudf::column_device_view source_column, - cuda::std::byte* source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - - Target* source_casted = reinterpret_cast(source); - cudf::detail::atomic_mul(&target.element(target_index), - static_cast(source_casted[source_index])); - } -}; - -// Assuming that the target column of COUNT_VALID, COUNT_ALL would be using fixed_width column and -// non-fixed point column -template - requires(cudf::detail::is_valid_aggregation()) -struct update_target_element_gmem { - __device__ void operator()(cudf::mutable_column_device_view target, - cudf::size_type target_index, - cudf::column_device_view source_column, - cuda::std::byte* source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - - Target* source_casted = reinterpret_cast(source); - cudf::detail::atomic_add(&target.element(target_index), - static_cast(source_casted[source_index])); - } -}; - -template - requires(cudf::detail::is_valid_aggregation()) -struct update_target_element_gmem { - __device__ void operator()(cudf::mutable_column_device_view target, - cudf::size_type target_index, - cudf::column_device_view source_column, - cuda::std::byte* source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - - Target* source_casted = reinterpret_cast(source); - cudf::detail::atomic_add(&target.element(target_index), - static_cast(source_casted[source_index])); - } -}; - -template - requires(cudf::detail::is_valid_aggregation() && - cudf::is_relationally_comparable()) -struct update_target_element_gmem { - __device__ void operator()(cudf::mutable_column_device_view target, - cudf::size_type target_index, - cudf::column_device_view source_column, - cuda::std::byte* source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - Target* source_casted = reinterpret_cast(source); - auto source_argmax_index = source_casted[source_index]; - auto old = cudf::detail::atomic_cas( - &target.element(target_index), cudf::detail::ARGMAX_SENTINEL, source_argmax_index); - if (old != cudf::detail::ARGMAX_SENTINEL) { - while (source_column.element(source_argmax_index) > - source_column.element(old)) { - old = - cudf::detail::atomic_cas(&target.element(target_index), old, source_argmax_index); - } - } - } -}; - -template - requires(cudf::detail::is_valid_aggregation() && - cudf::is_relationally_comparable()) -struct update_target_element_gmem { - __device__ void operator()(cudf::mutable_column_device_view target, - cudf::size_type target_index, - cudf::column_device_view source_column, - cuda::std::byte* source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - Target* source_casted = reinterpret_cast(source); - auto source_argmin_index = source_casted[source_index]; - auto old = cudf::detail::atomic_cas( - &target.element(target_index), cudf::detail::ARGMIN_SENTINEL, source_argmin_index); - if (old != cudf::detail::ARGMIN_SENTINEL) { - while (source_column.element(source_argmin_index) < - source_column.element(old)) { - old = - cudf::detail::atomic_cas(&target.element(target_index), old, source_argmin_index); - } - } - } -}; - -/** - * @brief A functor that updates a single element in the target column stored in global memory by - * applying an aggregation operation to a corresponding element from a source column in shared - * memory. - * - * This functor can NOT be used for dictionary columns. - * - * This is a redundant copy replicating the behavior of `elementwise_aggregator` from - * `cudf/detail/aggregation/device_aggregators.cuh`. The key difference is that this functor accepts - * a pointer to raw bytes as the source, as `column_device_view` cannot yet be constructed from - * shared memory. - */ -struct gmem_element_aggregator { - template - __device__ void operator()(cudf::mutable_column_device_view target, - cudf::size_type target_index, - cudf::column_device_view source_column, - cuda::std::byte* source, - bool* source_mask, - cudf::size_type source_index) const noexcept - { - // Early exit for all aggregation kinds since shared memory aggregation of - // `COUNT_ALL` is always valid - if (!source_mask[source_index]) { return; } - - // The output for COUNT_VALID and COUNT_ALL is initialized to be all valid - if constexpr (!(k == cudf::aggregation::COUNT_VALID or k == cudf::aggregation::COUNT_ALL)) { - if (target.is_null(target_index)) { target.set_valid(target_index); } - } - - update_target_element_gmem{}( - target, target_index, source_column, source, source_index); - } -}; -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/groupby.cu b/cpp/src/groupby/hash/groupby.cu index 814b55300f64..bf20308e1dd0 100644 --- a/cpp/src/groupby/hash/groupby.cu +++ b/cpp/src/groupby/hash/groupby.cu @@ -4,6 +4,7 @@ */ #include "compute_groupby.hpp" +#include "compute_single_pass_aggs.hpp" #include "extract_single_pass_aggs.hpp" #include "groupby/common/utils.hpp" #include "helpers.cuh" @@ -13,12 +14,10 @@ #include #include #include -#include #include #include #include #include -#include #include #include @@ -59,48 +58,6 @@ std::unique_ptr
dispatch_groupby(table_view const& keys, } } -// check if the target_type of the aggregation/type pair supports atomic operations -struct can_use_hash_groupby_fn { - template - requires(cudf::is_nested()) - bool operator()() const - { - // Currently, input values (not keys) of STRUCT and LIST types are not supported in any of - // hash-based aggregations. For those situations, we fallback to sort-based aggregations. - return false; - } - - template - constexpr static bool uses_underlying_type() - { - return k == aggregation::MIN or k == aggregation::MAX or k == aggregation::SUM; - } - - template - requires(cudf::is_fixed_point()) - bool operator()() const - { - if constexpr (std::is_same_v && K == aggregation::SUM) { return true; } - - using TargetType = cudf::detail::target_type_t; - using DeviceTargetType = std:: - conditional_t(), cudf::device_storage_type_t, TargetType>; - if constexpr (not std::is_void_v) { - return cudf::has_atomic_support(); - } - return false; - } - - template - requires(not cudf::is_nested() and not cudf::is_fixed_point()) - bool operator()() const - { - using TargetType = cudf::detail::target_type_t; - if constexpr (not std::is_void_v) { return cudf::has_atomic_support(); } - return false; - } -}; - } // namespace /** @@ -124,7 +81,7 @@ bool can_use_hash_groupby(std::span requests) // compound aggregations are made up of simple aggregations auto const agg_kinds = get_simple_aggregations(*a, v_type); return std::all_of(agg_kinds.begin(), agg_kinds.end(), [v_type = v_type](auto k) { - return cudf::detail::dispatch_type_and_aggregation(v_type, k, can_use_hash_groupby_fn{}); + return is_single_pass_agg_supported(v_type, k); }); }); }); diff --git a/cpp/src/groupby/hash/hash_csr_kernels.cuh b/cpp/src/groupby/hash/hash_csr_kernels.cuh new file mode 100644 index 000000000000..e01be1134a86 --- /dev/null +++ b/cpp/src/groupby/hash/hash_csr_kernels.cuh @@ -0,0 +1,286 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cudf::groupby::detail::hash { + +/// One open-addressed slot: the first input row that claimed the slot. +using hash_table_entry_type = size_type; + +/// Each input row records its group count index and its rank among the rows of that group. +using build_position_type = cuco::pair; + +/// Slot recorded for rows that are excluded from the groupby because their keys contain nulls. +constexpr cuda::std::uint32_t hash_csr_no_slot = + cuda::std::numeric_limits::max(); + +constexpr thread_index_type hash_csr_block_size = 256; + +/// Inputs with at least this many rows size the table from an estimate of the number of distinct +/// keys instead of the row count; smaller inputs get a table for every row right away. +constexpr size_type hash_csr_min_rows_to_estimate = 1 << 21; +/// Keys wider than this (in bytes per row) always get a table for every row: every probe of an +/// occupied slot compares keys, which costs more than the smaller table saves for them. +constexpr size_type hash_csr_max_estimated_key_bytes = 32; +/// Slots of the table the distinct keys of a sample of the rows are counted in. +constexpr cuda::std::uint32_t hash_csr_sample_capacity = 1u << 20; +/// Fewest slots of a table sized from an estimate, so that a handful of hot slots still spread +/// over many cache lines. +constexpr cuda::std::uint32_t hash_csr_min_estimated_capacity = 1u << 18; +/// A probe this long is vanishingly unlikely at a load factor of one half, so it means the table +/// is (nearly) full and the build has to restart with a larger one. +constexpr cuda::std::uint32_t hash_csr_max_probes = 256; + +/// Device view of the linearly probed open-addressed table that maps each distinct key to a slot. +struct hash_csr_table_ref { + hash_table_entry_type* entries; + cuda::std::uint32_t capacity; + cuda::std::uint32_t max_probes; ///< Probes after which the table is declared full + + /** + * @brief Returns the slot owned by the key of `row`, claiming an empty slot when the key is new. + * + * Most groupby rows repeat a key that is already in the table, so each slot is read before any + * attempt to claim it and the compare-and-swap only runs on empty slots. + * + * @param representative Receives the row stored in the matching slot, or the empty sentinel + * @return The slot and whether `row` claimed it, or `capacity` when `max_probes` slots were + * probed without success + */ + template + __device__ cuda::std::pair insert_or_find( + size_type row, hash_value_type hash, Equal const& equal, size_type& representative) const + { + representative = cudf::detail::CUDF_SIZE_TYPE_SENTINEL; + auto slot = static_cast(hash % capacity); + for (cuda::std::uint32_t step = 0; step < max_probes; ++step) { + auto entry_ref = + cuda::atomic_ref{entries[slot]}; + auto current = entry_ref.load(cuda::memory_order_relaxed); + if (current == cudf::detail::CUDF_SIZE_TYPE_SENTINEL && + entry_ref.compare_exchange_strong(current, row, cuda::memory_order_relaxed)) { + representative = row; + return {slot, true}; + } + if (equal(row, current)) { + representative = current; + return {slot, false}; + } + slot = slot + 1 == capacity ? 0 : slot + 1; + } + return {capacity, false}; + } +}; + +/** + * @brief Inserts every valid row into the table and, when `positions` is given, records the slot + * of each row and its rank within that slot. + * + * Ranks are handed out by one atomic per distinct slot per warp: lanes that landed in the same + * slot combine their increments, which keeps low-cardinality inputs from serializing on a few + * counters. + * + * When `overflow` is given, a row whose probe runs out of `table.max_probes` slots sets it and + * the rows still to come are skipped: the caller then rebuilds with a larger table. + */ +template +CUDF_KERNEL void hash_csr_build_kernel(size_type num_rows, + bitmask_type const* valid_rows, + build_position_type* positions, + size_type* slot_counts, + bool count_by_representative, + hash_csr_table_ref table, + Equal equal, + Hasher hasher, + int* overflow) +{ + auto const lane = static_cast(threadIdx.x % cudf::detail::warp_size); + auto const stride = cudf::detail::grid_1d::grid_stride(); + // One thread checks the cross-block abort flag. A block that starts before an overflow + // finishes its bounded probes; the host still discards the partial build and retries. + if (overflow != nullptr) { + auto const stop = + threadIdx.x == 0 && cuda::atomic_ref{*overflow}.load( + cuda::memory_order_relaxed) != 0; + if (__syncthreads_or(stop)) { return; } + } + // Every lane of a warp runs the same number of iterations so the warp-wide match and shuffle + // below always see a converged warp; lanes past the last row simply do not participate. + for (auto first_row = cudf::detail::grid_1d::global_thread_id() - lane; first_row < num_rows; + first_row += stride) { + auto const row = first_row + lane; + auto const is_active = + row < num_rows && + (valid_rows == nullptr || cudf::bit_is_set(valid_rows, static_cast(row))); + auto slot = hash_csr_no_slot; + if (is_active) { + auto const index = static_cast(row); + size_type representative{}; + slot = table.insert_or_find(index, hasher(index), equal, representative).first; + if (slot == table.capacity) { + cuda::atomic_ref{*overflow}.store( + 1, cuda::memory_order_relaxed); + slot = hash_csr_no_slot; + } else if (count_by_representative) { + slot = static_cast(representative); + } + } + // Without aggregations only the distinct keys matter, and the table alone provides them. + if (positions == nullptr) { continue; } + + auto const has_slot = slot != hash_csr_no_slot; + auto const active_mask = __ballot_sync(0xffff'ffffu, has_slot); + if (has_slot) { + auto const peers = __match_any_sync(active_mask, slot); + auto const leader = cuda::std::countr_zero(peers); + size_type first_rank{}; + if (lane == static_cast(leader)) { + first_rank = + cuda::atomic_ref{slot_counts[slot]}.fetch_add( + static_cast(cuda::std::popcount(peers)), cuda::memory_order_relaxed); + } + first_rank = __shfl_sync(peers, first_rank, leader); + auto const rank_in_warp = + static_cast(cuda::std::popcount(peers & ((1u << lane) - 1u))); + positions[row] = {slot, first_rank + rank_in_warp}; + } else if (row < num_rows) { + positions[row] = {hash_csr_no_slot, cudf::detail::CUDF_SIZE_TYPE_SENTINEL}; + } + } +} + +/** + * @brief Inserts every `stride`-th valid row into the table and counts the sampled rows and the + * slots they claim. + * + * Rows are sampled one at a time rather than in runs, since neighboring rows often share a key. + * `counts[0]` receives the number of valid sampled rows and `counts[1]` the number of distinct + * keys among them. + */ +template +CUDF_KERNEL void hash_csr_sample_kernel(size_type num_rows, + size_type stride, + bitmask_type const* valid_rows, + hash_csr_table_ref table, + Equal equal, + Hasher hasher, + size_type* counts) +{ + auto const row = cudf::detail::grid_1d::global_thread_id() * stride; + auto const is_valid = + row < num_rows && + (valid_rows == nullptr || cudf::bit_is_set(valid_rows, static_cast(row))); + auto is_new = false; + if (is_valid) { + auto const index = static_cast(row); + size_type representative{}; + is_new = table.insert_or_find(index, hasher(index), equal, representative).second; + } + // Every thread of the block reaches both counts, which is what they require. + auto const num_valid = __syncthreads_count(is_valid); + auto const num_new = __syncthreads_count(is_new); + if (threadIdx.x == 0) { + cuda::atomic_ref{counts[0]}.fetch_add( + num_valid, cuda::memory_order_relaxed); + cuda::atomic_ref{counts[1]}.fetch_add( + num_new, cuda::memory_order_relaxed); + } +} + +/// Scatters each valid row to its group: the start offset of its slot plus its rank in the slot. +CUDF_KERNEL void hash_csr_fill_kernel(size_type num_rows, + build_position_type const* positions, + size_type const* slot_offsets, + size_type* grouped_rows) +{ + auto const stride = cudf::detail::grid_1d::grid_stride(); + for (auto row = cudf::detail::grid_1d::global_thread_id(); row < num_rows; row += stride) { + // The positions are read once, so they stream past the caches and leave them to the slot + // offsets, which every row looks up, and to the grouped rows, which the aggregations read next. + auto const position = cub::ThreadLoad(positions + row); + if (position.first == hash_csr_no_slot) { continue; } + grouped_rows[slot_offsets[position.first] + position.second] = static_cast(row); + } +} + +template +void launch_hash_csr_build_kernel(size_type num_rows, + bitmask_type const* valid_rows, + build_position_type* positions, + size_type* slot_counts, + bool count_by_representative, + hash_csr_table_ref table, + Equal equal, + Hasher hasher, + int* overflow, + cuda::stream_ref stream) +{ + if (num_rows == 0) { return; } + auto const config = cudf::detail::grid_1d{num_rows, hash_csr_block_size}; + hash_csr_build_kernel<<>>( + num_rows, + valid_rows, + positions, + slot_counts, + count_by_representative, + table, + equal, + hasher, + overflow); + CUDF_CUDA_TRY(cudaGetLastError()); +} + +template +void launch_hash_csr_sample_kernel(size_type num_rows, + size_type stride, + bitmask_type const* valid_rows, + hash_csr_table_ref table, + Equal equal, + Hasher hasher, + size_type* counts, + cuda::stream_ref stream) +{ + auto const num_samples = cudf::util::div_rounding_up_safe(num_rows, stride); + if (num_samples == 0) { return; } + auto const config = cudf::detail::grid_1d{num_samples, hash_csr_block_size}; + hash_csr_sample_kernel<<>>( + num_rows, stride, valid_rows, table, equal, hasher, counts); + CUDF_CUDA_TRY(cudaGetLastError()); +} + +inline void launch_hash_csr_fill_kernel(size_type num_rows, + build_position_type const* positions, + size_type const* slot_offsets, + size_type* grouped_rows, + cuda::stream_ref stream) +{ + if (num_rows == 0) { return; } + auto const config = cudf::detail::grid_1d{num_rows, hash_csr_block_size}; + hash_csr_fill_kernel<<>>( + num_rows, positions, slot_offsets, grouped_rows); + CUDF_CUDA_TRY(cudaGetLastError()); +} + +} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/helpers.cuh b/cpp/src/groupby/hash/helpers.cuh index fe536355ee83..b35a4c6b1ee2 100644 --- a/cpp/src/groupby/hash/helpers.cuh +++ b/cpp/src/groupby/hash/helpers.cuh @@ -1,19 +1,15 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include #include #include -#include #include #include -#include - namespace cudf::groupby::detail::hash { /// Number of threads to handle each input element CUDF_HOST_DEVICE auto constexpr GROUPBY_CG_SIZE = 1; @@ -21,59 +17,9 @@ CUDF_HOST_DEVICE auto constexpr GROUPBY_CG_SIZE = 1; /// Number of slots per thread CUDF_HOST_DEVICE auto constexpr GROUPBY_BUCKET_SIZE = 1; -/// Thread block size -CUDF_HOST_DEVICE auto constexpr GROUPBY_BLOCK_SIZE = 128; - -/// Threshold cardinality to switch between shared memory aggregations and global memory -/// aggregations -CUDF_HOST_DEVICE auto constexpr GROUPBY_CARDINALITY_THRESHOLD = 128; - -/// Threshold to switch between two strategies: one is to output the aggregation results directly to -/// the final dense output columns, the other is to output the results to sparse intermediate -/// buffers then gather to the final dense output columns. -auto constexpr GROUPBY_DENSE_OUTPUT_THRESHOLD = 2; - -// We add additional `block_size`, because after the number of elements in the local hash set -// exceeds the threshold, all threads in the thread block can still insert one more element. -/// The maximum number of elements handled per block -CUDF_HOST_DEVICE auto constexpr GROUPBY_SHM_MAX_ELEMENTS = - GROUPBY_CARDINALITY_THRESHOLD + GROUPBY_BLOCK_SIZE; - -// GROUPBY_SHM_MAX_ELEMENTS with 0.7 occupancy -/// Shared memory hash set extent type -using shmem_extent_t = - cuco::extent(static_cast(GROUPBY_SHM_MAX_ELEMENTS) * 1.43)>; - -/// Number of slots needed by each shared memory hash set -CUDF_HOST_DEVICE auto constexpr valid_extent = - cuco::make_valid_extent(shmem_extent_t{}); - using row_hash_t = cudf::detail::row::hash::device_row_hasher; -/// Adapter to cudf row hasher with caching support. -class row_hasher_with_cache_t { - row_hash_t hasher; - hash_value_type const* values; - - public: - row_hasher_with_cache_t(row_hash_t const& hasher, - hash_value_type const* values = nullptr) noexcept - : hasher(hasher), values(values) - { - } - - __device__ hash_value_type operator()(size_type const idx) const noexcept - { - if (values) { return values[idx]; } - return hasher(idx); - } -}; - -/// Probing scheme type used by groupby hash table -using probing_scheme_t = cuco::linear_probing; - using row_comparator_t = cudf::detail::row::equality::device_row_comparator< false, cudf::nullate::DYNAMIC, @@ -84,41 +30,4 @@ using nullable_row_comparator_t = cudf::detail::row::equality::device_row_compar cudf::nullate::DYNAMIC, cudf::detail::row::equality::nan_equal_physical_equality_comparator>; -using global_set_t = cuco::static_set, - cuda::thread_scope_device, - row_comparator_t, - probing_scheme_t, - rmm::mr::polymorphic_allocator, - cuco::storage>; - -using nullable_global_set_t = cuco::static_set, - cuda::thread_scope_device, - nullable_row_comparator_t, - probing_scheme_t, - rmm::mr::polymorphic_allocator, - cuco::storage>; - -template -using hash_set_ref_t = - cuco::static_set_ref>, - Op>; - -template -using nullable_hash_set_ref_t = - cuco::static_set_ref>, - Op>; } // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/output_utils.cu b/cpp/src/groupby/hash/output_utils.cu index 0f8e0d9265de..c8bc2f3d0be6 100644 --- a/cpp/src/groupby/hash/output_utils.cu +++ b/cpp/src/groupby/hash/output_utils.cu @@ -3,7 +3,6 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "helpers.cuh" #include "output_utils.hpp" #include @@ -17,14 +16,7 @@ #include #include -#include -#include - -#include -#include #include -#include -#include #include #include @@ -136,65 +128,6 @@ std::unique_ptr
create_results_table(size_type output_size, return result_table; } -template -rmm::device_uvector extract_populated_keys(SetType const& key_set, - size_type num_total_keys, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - rmm::device_uvector unique_key_indices(num_total_keys, stream, mr); - auto const keys_end = key_set.retrieve_all(unique_key_indices.begin(), stream.get()); - unique_key_indices.resize(std::distance(unique_key_indices.begin(), keys_end), stream); - return unique_key_indices; -} - -template rmm::device_uvector extract_populated_keys( - global_set_t const& key_set, - size_type num_total_keys, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - -template rmm::device_uvector extract_populated_keys( - nullable_global_set_t const& key_set, - size_type num_total_keys, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - -rmm::device_uvector compute_key_transform_map( - size_type num_total_keys, - device_span unique_key_indices, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - // Map from old key indices (index of the keys in the original input keys table) to new key - // indices (indices of the keys in the final output table, which contains only the extracted - // unique keys). Only these extracted unique keys are mapped. - rmm::device_uvector key_transform_map(num_total_keys, stream, mr); - thrust::scatter(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{0}, - cuda::counting_iterator{static_cast(unique_key_indices.size())}, - unique_key_indices.begin(), - key_transform_map.begin()); - - return key_transform_map; -} - -rmm::device_uvector compute_target_indices(device_span input, - device_span transform_map, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr) -{ - rmm::device_uvector target_indices(input.size(), stream, mr); - thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - input.begin(), - input.end(), - target_indices.begin(), - [new_indices = transform_map.begin()] __device__(size_type const idx) { - return idx == cudf::detail::CUDF_SIZE_TYPE_SENTINEL ? idx : new_indices[idx]; - }); - return target_indices; -} - void finalize_output(table_view const& values, std::vector> const& aggregations, std::unique_ptr
& agg_results, diff --git a/cpp/src/groupby/hash/output_utils.hpp b/cpp/src/groupby/hash/output_utils.hpp index 97f0dcf98e27..5d7ae132937d 100644 --- a/cpp/src/groupby/hash/output_utils.hpp +++ b/cpp/src/groupby/hash/output_utils.hpp @@ -12,8 +12,6 @@ #include #include -#include - #include #include @@ -46,60 +44,6 @@ std::unique_ptr
create_results_table(size_type output_size, cuda::stream_ref stream, rmm::device_async_resource_ref mr); -/** - * @brief Return an array containing indices of (unique) keys in `key_set`. - * - * @tparam SetType Type of the key hash set - * - * @param key_set Key hash set - * @param num_total_keys Number of total keys - * @param stream CUDA stream used for device memory operations and kernel launches - * @param mr Device memory resource used to allocate the returned array - * @return An array containing indices of unique keys retrieved from `key_set` - */ -template -rmm::device_uvector extract_populated_keys(SetType const& key_set, - size_type num_total_keys, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - -/** - * @brief Compute and return a mapping array that maps from the original input keys to their - * positions in the input array which contains indices of the unique keys. - * - * Note that the output mapping array only covers the keys with indices existing in the input array, - * leaving other keys with uninitialized mapping values. - * - * @param num_total_keys Number of total keys - * @param unique_key_indices Array containing indices of the unique keys - * @param stream CUDA stream used for device memory operations and kernel launches - * @param mr Device memory resource used to allocate the returned array - * @return An array mapping from the original input keys to their positions in the input array - */ -rmm::device_uvector compute_key_transform_map( - size_type num_total_keys, - device_span unique_key_indices, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - -/** - * @brief Transform from row indices of the keys in the input keys table into indices of these keys - * in the output unique keys table. - * - * Note that the positions (indices) of all output unique keys must be covered in the array - * `transform_map`. This is guaranteed as it was generated in `extract_populated_keys` function. - * - * @param input The indices of the keys to transform - * @param transform_map The mapping array from the input keys table to the output unique keys table - * @param stream CUDA stream used for device memory operations and kernel launches - * @param mr Device memory resource used to allocate the returned array - * @return A device vector mapping each input row to its output row index - */ -rmm::device_uvector compute_target_indices(device_span input, - device_span transform_map, - cuda::stream_ref stream, - rmm::device_async_resource_ref mr); - /** * @brief Perform some final computation for the aggregation results such as null count and move * the result columns into a `result_cache` object. diff --git a/cpp/src/groupby/hash/shared_memory_aggregator.cuh b/cpp/src/groupby/hash/shared_memory_aggregator.cuh deleted file mode 100644 index 3008c130338c..000000000000 --- a/cpp/src/groupby/hash/shared_memory_aggregator.cuh +++ /dev/null @@ -1,220 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include -#include -#include -#include -#include - -#include -#include - -namespace cudf::groupby::detail::hash { -template -struct update_target_element_shmem { - __device__ void operator()(cuda::std::byte*, - cudf::size_type, - cudf::column_device_view, - cudf::size_type) const - { - CUDF_UNREACHABLE("Invalid source type and aggregation combination."); - } -}; - -template - requires(cudf::is_fixed_width() && - cudf::has_atomic_support>()) -struct update_target_element_shmem { - __device__ void operator()(cuda::std::byte* target, - cudf::size_type target_index, - cudf::column_device_view source, - cudf::size_type source_index) const noexcept - { - using DeviceTarget = - cudf::device_storage_type_t>; - using DeviceSource = cudf::device_storage_type_t; - - DeviceTarget* target_casted = reinterpret_cast(target); - cudf::detail::atomic_min(&target_casted[target_index], - static_cast(source.element(source_index))); - } -}; - -template - requires(cudf::is_fixed_width() && - cudf::has_atomic_support>()) -struct update_target_element_shmem { - __device__ void operator()(cuda::std::byte* target, - cudf::size_type target_index, - cudf::column_device_view source, - cudf::size_type source_index) const noexcept - { - using DeviceTarget = - cudf::device_storage_type_t>; - using DeviceSource = cudf::device_storage_type_t; - - DeviceTarget* target_casted = reinterpret_cast(target); - cudf::detail::atomic_max(&target_casted[target_index], - static_cast(source.element(source_index))); - } -}; - -template - requires(cudf::is_fixed_width() && !cudf::is_timestamp()) -struct update_target_element_shmem { - __device__ void operator()(cuda::std::byte* target, - cudf::size_type target_index, - cudf::column_device_view source, - cudf::size_type source_index) const noexcept - { - using DeviceTarget = - cudf::device_storage_type_t>; - using DeviceSource = cudf::device_storage_type_t; - - DeviceTarget* target_casted = reinterpret_cast(target); - cudf::detail::atomic_add(&target_casted[target_index], - static_cast(source.element(source_index))); - } -}; - -template - requires(cudf::detail::is_product_supported()) -struct update_target_element_shmem { - __device__ void operator()(cuda::std::byte* target, - cudf::size_type target_index, - cudf::column_device_view source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - Target* target_casted = reinterpret_cast(target); - auto value = static_cast(source.element(source_index)); - cudf::detail::atomic_add(&target_casted[target_index], value * value); - } -}; - -template - requires(cudf::detail::is_product_supported()) -struct update_target_element_shmem { - __device__ void operator()(cuda::std::byte* target, - cudf::size_type target_index, - cudf::column_device_view source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - Target* target_casted = reinterpret_cast(target); - cudf::detail::atomic_mul(&target_casted[target_index], - static_cast(source.element(source_index))); - } -}; - -template - requires(cudf::detail::is_valid_aggregation()) -struct update_target_element_shmem { - __device__ void operator()(cuda::std::byte* target, - cudf::size_type target_index, - cudf::column_device_view source, - cudf::size_type source_index) const noexcept - { - // The nullability was checked prior to this call in the `shmem_element_aggregator` functor - using Target = cudf::detail::target_type_t; - Target* target_casted = reinterpret_cast(target); - cudf::detail::atomic_add(&target_casted[target_index], Target{1}); - } -}; - -template - requires(cudf::detail::is_valid_aggregation()) -struct update_target_element_shmem { - __device__ void operator()(cuda::std::byte* target, - cudf::size_type target_index, - cudf::column_device_view source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - Target* target_casted = reinterpret_cast(target); - cudf::detail::atomic_add(&target_casted[target_index], Target{1}); - } -}; - -template - requires(cudf::detail::is_valid_aggregation() && - cudf::is_relationally_comparable()) -struct update_target_element_shmem { - __device__ void operator()(cuda::std::byte* target, - cudf::size_type target_index, - cudf::column_device_view source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - Target* target_casted = reinterpret_cast(target); - auto old = cudf::detail::atomic_cas( - &target_casted[target_index], cudf::detail::ARGMAX_SENTINEL, source_index); - if (old != cudf::detail::ARGMAX_SENTINEL) { - while (source.element(source_index) > source.element(old)) { - old = cudf::detail::atomic_cas(&target_casted[target_index], old, source_index); - } - } - } -}; - -template - requires(cudf::detail::is_valid_aggregation() && - cudf::is_relationally_comparable()) -struct update_target_element_shmem { - __device__ void operator()(cuda::std::byte* target, - cudf::size_type target_index, - cudf::column_device_view source, - cudf::size_type source_index) const noexcept - { - using Target = cudf::detail::target_type_t; - Target* target_casted = reinterpret_cast(target); - auto old = cudf::detail::atomic_cas( - &target_casted[target_index], cudf::detail::ARGMIN_SENTINEL, source_index); - if (old != cudf::detail::ARGMIN_SENTINEL) { - while (source.element(source_index) < source.element(old)) { - old = cudf::detail::atomic_cas(&target_casted[target_index], old, source_index); - } - } - } -}; - -/** - * @brief A functor that updates a single element in the target column stored in shared memory by - * applying an aggregation operation to a corresponding element from a source column in global - * memory. - * - * This functor can NOT be used for dictionary columns. - * - * This is a redundant copy replicating the behavior of `elementwise_aggregator` from - * `cudf/detail/aggregation/device_aggregators.cuh`. The key difference is that this functor accepts - * a pointer to raw bytes as the source, as `column_device_view` cannot yet be constructed from - * shared memory. - */ -struct shmem_element_aggregator { - template - __device__ void operator()(cuda::std::byte* target, - bool* target_mask, - cudf::size_type target_index, - cudf::column_device_view source, - cudf::size_type source_index) const noexcept - { - // Check nullability for all aggregation kinds but `COUNT_ALL` - if constexpr (k != cudf::aggregation::COUNT_ALL) { - if (source.is_null(source_index)) { return; } - } - - // The output for COUNT_VALID and COUNT_ALL is initialized to be all valid - if constexpr (!(k == cudf::aggregation::COUNT_VALID or k == cudf::aggregation::COUNT_ALL)) { - if (!target_mask[target_index]) { - cudf::detail::atomic_max(target_mask + target_index, true); - } - } - - update_target_element_shmem{}(target, target_index, source, source_index); - } -}; -} // namespace cudf::groupby::detail::hash diff --git a/cpp/src/groupby/hash/single_pass_argminmax.cu b/cpp/src/groupby/hash/single_pass_argminmax.cu new file mode 100644 index 000000000000..413e813b9068 --- /dev/null +++ b/cpp/src/groupby/hash/single_pass_argminmax.cu @@ -0,0 +1,15 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "single_pass_reductions.cuh" + +namespace cudf::groupby::detail::hash::single_pass { + +template std::unique_ptr compute_reduction( + reduction_context const& ctx); +template std::unique_ptr compute_reduction( + reduction_context const& ctx); + +} // namespace cudf::groupby::detail::hash::single_pass diff --git a/cpp/src/groupby/hash/single_pass_functors.cuh b/cpp/src/groupby/hash/single_pass_functors.cuh index 165b7fa2bc41..0a1724a38cd3 100644 --- a/cpp/src/groupby/hash/single_pass_functors.cuh +++ b/cpp/src/groupby/hash/single_pass_functors.cuh @@ -1,65 +1,19 @@ /* - * 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 */ #pragma once -#include "helpers.cuh" - -#include +#include +#include #include -#include +#include +#include +#include -#include +#include namespace cudf::groupby::detail::hash { -/// Functor used by type dispatcher returning the size of the underlying C++ type -struct size_of_functor { - template - CUDF_HOST_DEVICE constexpr cudf::size_type operator()() - { - return sizeof(T); - } -}; - -template -struct initialize_target_element { - __device__ void operator()(cuda::std::byte* target, - bool* target_mask, - cudf::size_type idx) const noexcept - requires(not cudf::detail::is_identity_supported()) - { - CUDF_UNREACHABLE("Invalid source type and aggregation combination."); - } - - __device__ void operator()(cuda::std::byte* target, - bool* target_mask, - cudf::size_type idx) const noexcept - requires(cudf::detail::is_identity_supported()) - { - using DeviceType = cudf::device_storage_type_t; - DeviceType* target_casted = reinterpret_cast(target); - - target_casted[idx] = cudf::detail::get_identity(); - - target_mask[idx] = (k == cudf::aggregation::COUNT_ALL) or (k == cudf::aggregation::COUNT_VALID); - } -}; - -struct initialize_shmem { - template - __device__ void operator()(cuda::std::byte* target, - bool* target_mask, - cudf::size_type idx) const noexcept - { - initialize_target_element{}(target, target_mask, idx); - } -}; - -/** - * @brief Base struct to compute single-pass aggregations and store the results into an output - * table, executing for all input rows. - */ struct compute_single_pass_aggs_base_fn { aggregation::Kind const* aggs; table_device_view input_values; @@ -73,50 +27,6 @@ struct compute_single_pass_aggs_base_fn { } }; -/** - * @brief Functor to compute single-pass aggregations and store the results into an output table, - * executing for all input rows. - * - * This functor writes output to the sparse intermediate output table, using the target indices - * computed on-the-fly. In addition, aggregations are computed in serial order for each row. - * - * @tparam SetType Type of the key hash set - */ -template -struct compute_single_pass_aggs_sparse_output_fn : compute_single_pass_aggs_base_fn { - SetRef set_ref; - bitmask_type const* row_bitmask; - - compute_single_pass_aggs_sparse_output_fn(SetRef set_ref, - bitmask_type const* row_bitmask, - aggregation::Kind const* aggs, - table_device_view const& input_values, - mutable_table_device_view const& output_values) - : compute_single_pass_aggs_base_fn(aggs, input_values, output_values), - set_ref{set_ref}, - row_bitmask{row_bitmask} - { - } - - __device__ void operator()(size_type idx) - { - if (row_bitmask && !cudf::bit_is_set(row_bitmask, idx)) { return; } - auto const target_row_idx = *set_ref.insert_and_find(idx).first; - - for (size_type col_idx = 0; col_idx < input_values.num_columns(); ++col_idx) { - auto const& source_col = input_values.column(col_idx); - auto const& target_col = output_values.column(col_idx); - dispatch_type_and_aggregation(source_col.type(), - aggs[col_idx], - cudf::detail::element_aggregator{}, - target_col, - target_row_idx, - source_col, - idx); - } - } -}; - /** * @brief Functor to compute single-pass aggregations and store the results into an output table, * executing for all input rows. diff --git a/cpp/src/groupby/hash/single_pass_minmax.cu b/cpp/src/groupby/hash/single_pass_minmax.cu new file mode 100644 index 000000000000..befd295581d0 --- /dev/null +++ b/cpp/src/groupby/hash/single_pass_minmax.cu @@ -0,0 +1,13 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "single_pass_reductions.cuh" + +namespace cudf::groupby::detail::hash::single_pass { + +template std::unique_ptr compute_reduction(reduction_context const& ctx); +template std::unique_ptr compute_reduction(reduction_context const& ctx); + +} // namespace cudf::groupby::detail::hash::single_pass diff --git a/cpp/src/groupby/hash/single_pass_product.cu b/cpp/src/groupby/hash/single_pass_product.cu new file mode 100644 index 000000000000..d822613a71d6 --- /dev/null +++ b/cpp/src/groupby/hash/single_pass_product.cu @@ -0,0 +1,13 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "single_pass_reductions.cuh" + +namespace cudf::groupby::detail::hash::single_pass { + +template std::unique_ptr compute_reduction( + reduction_context const& ctx); + +} // namespace cudf::groupby::detail::hash::single_pass diff --git a/cpp/src/groupby/hash/single_pass_reductions.cuh b/cpp/src/groupby/hash/single_pass_reductions.cuh new file mode 100644 index 000000000000..37de7a19735c --- /dev/null +++ b/cpp/src/groupby/hash/single_pass_reductions.cuh @@ -0,0 +1,623 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "single_pass_reductions.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace cudf::groupby::detail::hash::single_pass { + +/// Reads a fixed-width element, going through the keys when the column is a dictionary. +template +struct value_accessor { + column_device_view col; + bool is_dictionary; + + __device__ T operator()(size_type row) const + { + if (is_dictionary) { + auto const keys = col.child(dictionary_column_view::keys_column_index); + return keys.element(static_cast(col.element(row))); + } + return col.element(row); + } +}; + +template +value_accessor reduction_context::accessor() const +{ + return {d_values, is_dictionary(values.type())}; +} + +/// Maps a grouped position to the value of the input row at that position, substituting +/// `null_value` for null rows and optionally squaring the value. +template +struct grouped_value_fn { + size_type const* grouped_rows; + value_accessor value; + Target null_value; + bool has_nulls; + + __device__ bool col_is_null(size_type row) const + { + return has_nulls && value.col.is_null_nocheck(row); + } + + __device__ Target compute(size_type row) const + { + auto const result = static_cast(value(row)); + if constexpr (Square) { return result * result; } + return result; + } + + __device__ Target operator()(size_type position) const + { + auto const row = grouped_rows[position]; + return col_is_null(row) ? null_value : compute(row); + } +}; + +/// A reduced value together with whether any of the reduced rows was valid, so that a nullable +/// aggregation and its null mask come out of one pass. +template +struct valid_value { + Result value; + bool valid; +}; + +template +struct valid_value_op { + __device__ valid_value operator()(valid_value const& lhs, + valid_value const& rhs) const + { + return {Op{}(lhs.value, rhs.value), lhs.valid || rhs.valid}; + } +}; + +/// Maps a grouped position to the value of the input row at that position and its validity; null +/// rows contribute the identity. +template +struct grouped_valid_value_fn { + grouped_value_fn value; + + __device__ valid_value operator()(size_type position) const + { + auto const row = value.grouped_rows[position]; + if (value.col_is_null(row)) { return {value.null_value, false}; } + return {value.compute(row), true}; + } +}; + +template +struct split_valid_value_fn { + __device__ cuda::std::tuple operator()(valid_value const& v) const + { + return {v.value, v.valid}; + } +}; + +/// Maps a grouped position to a SUM_OVERFLOW accumulator, treating nulls as a zero contribution. +template +struct grouped_sum_overflow_fn { + size_type const* grouped_rows; + value_accessor value; + bool has_nulls; + + __device__ cudf::reduction::detail::sum_overflow_result operator()( + size_type position) const + { + auto const row = grouped_rows[position]; + if (has_nulls && value.col.is_null_nocheck(row)) { return {DeviceType{0}, 0}; } + return {value(row), 0}; + } +}; + +/// Splits a reduced accumulator into the sum and overflow-flag children of the output struct. +template +struct split_sum_overflow_fn { + __device__ cuda::std::tuple operator()( + cudf::reduction::detail::sum_overflow_result const& accumulator) const + { + return {accumulator.sum, accumulator.wraps != 0}; + } +}; + +/// Sums accumulated together when several additive aggregations are requested on one column. +template +struct fused_sums { + Result sum; + Result sum_of_squares; + size_type count; +}; + +template +struct fused_sums_plus { + __device__ fused_sums operator()(fused_sums const& lhs, + fused_sums const& rhs) const + { + return {lhs.sum + rhs.sum, lhs.sum_of_squares + rhs.sum_of_squares, lhs.count + rhs.count}; + } +}; + +/// Maps a grouped position to the sums contributed by the input row at that position. +template +struct grouped_fused_sums_fn { + size_type const* grouped_rows; + value_accessor value; + bool has_nulls; + + __device__ fused_sums operator()(size_type position) const + { + auto const row = grouped_rows[position]; + if (has_nulls && value.col.is_null_nocheck(row)) { return {Result{0}, Result{0}, 0}; } + auto const result = static_cast(value(row)); + return {result, result * result, 1}; + } +}; + +/// Splits the reduced sums into the SUM, SUM_OF_SQUARES and COUNT_VALID outputs. +template +struct split_fused_sums_fn { + __device__ cuda::std::tuple operator()( + fused_sums const& sums) const + { + return {sums.sum, sums.sum_of_squares, sums.count}; + } +}; + +constexpr bool is_fusable_sum(aggregation::Kind kind) +{ + return kind == aggregation::SUM || kind == aggregation::SUM_OF_SQUARES || + kind == aggregation::COUNT_VALID; +} + +/// Groups this small on average are reduced as segments packed several per block: one block per +/// segment would leave most of the device idle. +constexpr size_type min_avg_rows_per_segment = 128; + +/// Groups longer than this are reduced chunk by chunk so that every block has a bounded range. +constexpr size_type rows_per_chunk = 1 << 14; + +/// Chunk length when the segments are packed several per block, so that a thread or a sub-warp +/// never walks a long group alone. +constexpr size_type packed_rows_per_chunk = 1 << 10; + +/// Threads that share one segment on the packed path. +constexpr int packed_threads_per_group = 8; + +// The packed path calls `cub::detail::segmented_reduce::dispatch` directly, an internal entry point +// whose namespace, signature and hint semantics were verified in CCCL 3.5 only. +static_assert(CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION == 5, + "re-verify cub::detail::segmented_reduce::dispatch and the max_segment_size hint " + "(dispatch_segmented_reduce.cuh, kernels/kernel_segmented_reduce.cuh) for this CCCL"); + +/** + * @brief Policy selector for cub::DeviceSegmentedReduce whose medium path hands each segment to + * `packed_threads_per_group` threads instead of a full warp. + * + * The large (one block per segment) and small (one thread per segment) policies are those of + * CUB's own selector, so the large path reduces in exactly the order it does today. Only the + * medium tile changes: `packed_threads_per_group * items_per_thread` rows instead of + * `32 * items_per_thread`, and `threads_per_block / packed_threads_per_group` segments per block. + * + * Verified against cub/device/dispatch/tuning/tuning_segmented_reduce.cuh:23-63 (the policy + * structs), :101-140 (the default selector) and cub/util_device.cuh:865 (the selector concept: + * stateless, `operator()(cuda::compute_capability) -> cub::SegmentedReducePolicy`). The kernel + * evaluates it as a constant expression for `__launch_bounds__` (kernel_segmented_reduce.cuh:113), + * hence `constexpr` and host/device. + */ +template +struct subwarp_segmented_reduce_policy_selector { + [[nodiscard]] CUDF_HOST_DEVICE constexpr cub::SegmentedReducePolicy operator()( + cuda::compute_capability cc) const + { + auto const base = + cub::detail::segmented_reduce::policy_selector_from_types{}(cc); + auto const& large = base.large_reduce; + return cub::SegmentedReducePolicy{large, + cub::SegmentedReduceWarpReducePolicy{large.threads_per_block, + packed_threads_per_group, + large.items_per_thread, + large.vec_size, + large.load_modifier}, + base.small_reduce}; + } +}; + +/// The stream and temporary-storage resource handed to the CUB algorithms. Spelled out without +/// class template argument deduction, which the host compiler rejects for these types outside of +/// a template. +using cub_stream_prop_t = cuda::std::execution::prop; +using cub_mr_prop_t = + cuda::std::execution::prop; +using cub_env_t = cuda::std::execution::env; + +inline cub_env_t make_cub_env(cuda::stream_ref stream) +{ + return cub_env_t{ + cub_stream_prop_t{cuda::get_stream_t{}, stream}, + cub_mr_prop_t{cuda::mr::get_memory_resource_t{}, cudf::get_current_device_resource_ref()}}; +} + +/** + * @brief Reduces segments with CUB, packing several per block when `avg_rows` is positive. + * + * `cub::DeviceSegmentedReduce::Reduce` always launches one block per segment. Its dispatch also + * accepts a segment size hint, which is what makes it hand every segment to one thread (hint + * within the small tile) or to one sub-warp (hint within the medium tile) instead; the kernel's + * agents loop over as many tiles as a segment has, so a segment longer than the hint is still + * reduced correctly, only by that thread or sub-warp alone. Segments averaging at most the small + * tile of the accumulator get a thread each, longer ones a sub-warp each. + * + * @param avg_rows Average number of rows per segment; zero selects one block per segment + */ +template +void reduce_segments(device_span offsets, + size_type avg_rows, + ValueIterator values, + OutputIterator output, + Op op, + T init, + cuda::stream_ref stream) +{ + using offset_type = cub::detail::common_iterator_value_t; + using accum_type = cuda::std::__accumulator_t, T>; + using selector_type = subwarp_segmented_reduce_policy_selector; + + std::size_t hint = 0; + if (avg_rows > 0) { + cuda::compute_capability cc{}; + CUDF_CUDA_TRY(cub::detail::ptx_compute_cap(cc)); + auto const policy = selector_type{}(cc); + auto const small_tile = policy.small_reduce.items_per_tile(); + auto const medium_tile = policy.medium_reduce.items_per_tile(); + CUDF_EXPECTS(small_tile + 1 <= medium_tile, "Unexpected segmented reduce tuning"); + hint = avg_rows <= small_tile ? std::size_t{1} : static_cast(small_tile) + 1; + } + CUDF_CUDA_TRY(cub::detail::dispatch_with_env( + make_cub_env(stream), + [&](auto, void* d_temp_storage, std::size_t& temp_storage_bytes, cudaStream_t launch_stream) { + return cub::detail::segmented_reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + values, + output, + static_cast(offsets.size() - 1), + offsets.begin(), + offsets.begin() + 1, + op, + init, + hint, + launch_stream, + selector_type{}); + })); +} + +/// Reduces the grouped values of every group into one output element per group. `packed_rows` +/// estimates the rows requiring value loads per segment; zero selects one block per segment. +template +void reduce_groups(grouped_rows const& grouped, + size_type packed_rows, + ValueIterator values, + OutputIterator output, + Op op, + T init, + cuda::stream_ref stream) +{ + if (grouped.group_chunks.is_empty()) { + reduce_segments(grouped.offsets, packed_rows, values, output, op, init, stream); + return; + } + rmm::device_uvector partials( + grouped.chunk_offsets.size() - 1, stream, cudf::get_current_device_resource_ref()); + reduce_segments(grouped.chunk_offsets, packed_rows, values, partials.begin(), op, init, stream); + reduce_segments(grouped.group_chunks, + packed_rows > 0 ? packed_rows_per_chunk : 0, + partials.begin(), + output, + op, + init, + stream); +} + +/// The device representation of a column element. Chrono and fixed-point columns reduce as their +/// integer reps, so the reduction kernels are only instantiated once per representation. +template +struct rep_type { + using type = device_storage_type_t; +}; + +template + requires(cudf::is_chrono()) +struct rep_type { + using type = typename T::rep; +}; + +template +using rep_type_t = typename rep_type::type; + +template +constexpr bool is_reduction_supported() +{ + switch (K) { + case aggregation::SUM: + return cudf::is_numeric() || cudf::is_duration() || cudf::is_fixed_point(); + case aggregation::PRODUCT: + case aggregation::SUM_OF_SQUARES: return cudf::detail::is_product_supported(); + case aggregation::MIN: + case aggregation::MAX: return cudf::is_fixed_width() && is_relationally_comparable(); + case aggregation::ARGMIN: + case aggregation::ARGMAX: return is_relationally_comparable(); + case aggregation::SUM_OVERFLOW: return cudf::detail::sum_overflow_supported; + default: return false; + } +} + +template +struct reduce_fn { + template + requires(is_reduction_supported() && + (K == aggregation::SUM || K == aggregation::PRODUCT || + K == aggregation::SUM_OF_SQUARES || K == aggregation::MIN || K == aggregation::MAX)) + std::unique_ptr operator()(reduction_context const& ctx) const + { + using Source = rep_type_t; + using Result = rep_type_t>; + using Op = cudf::detail::corresponding_operator_t; + + auto result = make_fixed_width_column(cudf::detail::target_type(ctx.values_type, K), + ctx.num_groups, + mask_state::UNALLOCATED, + ctx.stream, + ctx.mr); + if (ctx.num_groups == 0) { return result; } + + using value_fn = grouped_value_fn; + auto const identity = Op::template identity(); + auto const value = + value_fn{ctx.grouped.rows.data(), ctx.accessor(), identity, ctx.values.has_nulls()}; + auto const output = result->mutable_view().begin(); + if (!ctx.nullable) { + reduce_groups(ctx.grouped, + ctx.grouped.packed_rows, + cudf::detail::make_counting_transform_iterator(0, value), + output, + Op{}, + identity, + ctx.stream); + return result; + } + + // The validity of a group (any valid row) rides along with its value in one pass. + rmm::device_uvector group_valid( + ctx.num_groups, ctx.stream, cudf::get_current_device_resource_ref()); + auto const values = cudf::detail::make_counting_transform_iterator( + 0, grouped_valid_value_fn < Source, Result, K == aggregation::SUM_OF_SQUARES > {value}); + auto const outputs = cuda::transform_output_iterator{ + cuda::make_zip_iterator(output, group_valid.begin()), split_valid_value_fn{}}; + auto packed_rows = ctx.grouped.packed_rows; + if (packed_rows > 0 && ctx.grouped.group_chunks.is_empty()) { + // Null rows still read their validity, but skip the value load. Use the expected number + // of valid rows per group when choosing how many threads share a segment. + auto const valid_rows = + static_cast(ctx.values.size() - ctx.values.null_count()); + auto const avg_valid_rows = cudf::util::div_rounding_up_safe( + packed_rows * valid_rows, static_cast(ctx.values.size())); + packed_rows = static_cast(std::max(avg_valid_rows, 1)); + } + reduce_groups(ctx.grouped, + packed_rows, + values, + outputs, + valid_value_op{}, + valid_value{identity, false}, + ctx.stream); + auto [null_mask, null_count] = cudf::detail::valid_if( + group_valid.begin(), group_valid.end(), cuda::std::identity{}, ctx.stream, ctx.mr); + result->set_null_mask(std::move(null_mask), null_count); + return result; + } + + template + requires(is_reduction_supported() && + (K == aggregation::ARGMIN || K == aggregation::ARGMAX)) + std::unique_ptr operator()(reduction_context const& ctx) const + { + auto result = make_size_type_column(ctx); + if (ctx.num_groups == 0) { return result; } + + // The grouped rows are the input row indices themselves, so reducing them with the + // element comparator yields the input index of each group's extremum. The sentinel identity + // loses against every valid row and is left in place for all-null groups. + constexpr auto is_argmin = K == aggregation::ARGMIN; + reduce_groups(ctx.grouped, + ctx.grouped.packed_rows, + ctx.grouped.rows.begin(), + result->mutable_view().begin(), + cudf::detail::element_argminmax_fn>{ + ctx.d_values, ctx.values.has_nulls(), is_argmin}, + is_argmin ? cudf::detail::ARGMIN_SENTINEL : cudf::detail::ARGMAX_SENTINEL, + ctx.stream); + set_group_null_mask(*result, ctx); + return result; + } + + template + requires(is_reduction_supported() && K == aggregation::SUM_OVERFLOW) + std::unique_ptr operator()(reduction_context const& ctx) const + { + using Source = rep_type_t; + using accumulator = cudf::reduction::detail::sum_overflow_result; + + auto sum_child = make_fixed_width_column( + ctx.values_type, ctx.num_groups, mask_state::UNALLOCATED, ctx.stream, ctx.mr); + auto overflow_child = make_fixed_width_column( + data_type{type_id::BOOL8}, ctx.num_groups, mask_state::UNALLOCATED, ctx.stream, ctx.mr); + if (ctx.num_groups > 0) { + auto const values = cudf::detail::make_counting_transform_iterator( + 0, + grouped_sum_overflow_fn{ + ctx.grouped.rows.data(), ctx.accessor(), ctx.values.has_nulls()}); + auto const children = cuda::transform_output_iterator{ + cuda::make_zip_iterator(sum_child->mutable_view().begin(), + overflow_child->mutable_view().begin()), + split_sum_overflow_fn{}}; + reduce_groups(ctx.grouped, + ctx.grouped.packed_rows, + values, + children, + cudf::reduction::detail::overflow_sum_op{}, + accumulator{}, + ctx.stream); + } + + auto [null_mask, null_count] = ctx.nullable && ctx.num_groups > 0 + ? reduce_group_validity(ctx) + : std::pair{rmm::device_buffer{}, size_type{0}}; + std::vector> children; + children.push_back(std::move(sum_child)); + children.push_back(std::move(overflow_child)); + return create_structs_hierarchy( + ctx.num_groups, std::move(children), null_count, std::move(null_mask), ctx.stream, ctx.mr); + } + + template + requires(!is_reduction_supported()) + std::unique_ptr operator()(reduction_context const&) const + { + CUDF_FAIL("Unsupported type for hash groupby aggregation"); + } +}; + +/// Computes the SUM, SUM_OF_SQUARES and COUNT_VALID aggregations requested on one column, as +/// extracted for MEAN, M2, VARIANCE and STD, with a single segmented reduction. +struct fused_sums_fn { + template + requires(cudf::detail::is_product_supported()) + std::vector> operator()(reduction_context const& ctx, + host_span kinds, + std::span is_intermediate) const + { + using Source = rep_type_t; + using Result = rep_type_t>; + static_assert( + cuda::std:: + is_same_v>>); + + // Every sum is reduced even when it is not requested; those land in temporary columns. + auto const make_output = [&](aggregation::Kind kind) { + auto const requested = std::find(kinds.begin(), kinds.end(), kind) != kinds.end(); + return make_fixed_width_column(cudf::detail::target_type(ctx.values_type, kind), + ctx.num_groups, + mask_state::UNALLOCATED, + ctx.stream, + requested ? ctx.mr : cudf::get_current_device_resource_ref()); + }; + auto sum = make_output(aggregation::SUM); + auto sum_of_squares = make_output(aggregation::SUM_OF_SQUARES); + auto count = make_output(aggregation::COUNT_VALID); + auto const counts = count->view().template begin(); + if (ctx.num_groups > 0) { + auto const values = cudf::detail::make_counting_transform_iterator( + 0, + grouped_fused_sums_fn{ + ctx.grouped.rows.data(), ctx.accessor(), ctx.values.has_nulls()}); + auto const outputs = cuda::transform_output_iterator{ + cuda::make_zip_iterator(sum->mutable_view().template begin(), + sum_of_squares->mutable_view().template begin(), + count->mutable_view().template begin()), + split_fused_sums_fn{}}; + reduce_groups(ctx.grouped, + ctx.grouped.packed_rows, + values, + outputs, + fused_sums_plus{}, + fused_sums{Result{0}, Result{0}, 0}, + ctx.stream); + } + + std::vector> results; + for (std::size_t i = 0; i < kinds.size(); ++i) { + auto result = kinds[i] == aggregation::SUM ? std::move(sum) + : kinds[i] == aggregation::SUM_OF_SQUARES ? std::move(sum_of_squares) + : std::move(count); + // A sum is null when its group has no valid row, which the valid count already tells. + auto const nullable = + !is_intermediate[i] && kinds[i] != aggregation::COUNT_VALID && ctx.values.has_nulls(); + if (nullable && ctx.num_groups > 0) { + auto [null_mask, null_count] = cudf::detail::valid_if( + counts, + counts + ctx.num_groups, + [] __device__(size_type count) { return count > 0; }, + ctx.stream, + ctx.mr); + result->set_null_mask(std::move(null_mask), null_count); + } + results.push_back(std::move(result)); + } + return results; + } + + template + requires(!cudf::detail::is_product_supported()) + std::vector> operator()(reduction_context const&, + host_span, + std::span) const + { + CUDF_FAIL("Unsupported type for fused hash groupby sums"); + } +}; + +template +std::unique_ptr compute_reduction(reduction_context const& ctx) +{ + return type_dispatcher(ctx.values_type, reduce_fn{}, ctx); +} + +} // namespace cudf::groupby::detail::hash::single_pass diff --git a/cpp/src/groupby/hash/single_pass_reductions.hpp b/cpp/src/groupby/hash/single_pass_reductions.hpp new file mode 100644 index 000000000000..4555cc7c1dda --- /dev/null +++ b/cpp/src/groupby/hash/single_pass_reductions.hpp @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "compute_single_pass_aggs.hpp" + +#include + +#include + +#include + +namespace cudf::groupby::detail::hash::single_pass { + +template +struct value_accessor; + +struct reduction_context { + column_view const& values; + column_device_view const& d_values; + data_type values_type; ///< Type of the values, or of the keys for dictionary values + grouped_rows const& grouped; + size_type num_groups; + bool nullable; ///< Whether the result carries a null mask + cuda::stream_ref stream; + rmm::device_async_resource_ref mr; + + template + value_accessor accessor() const; +}; + +// Shared host helpers are defined only in the frontend, keeping their reduction kernels unique. +std::pair reduce_group_validity(reduction_context const& ctx); +void set_group_null_mask(column& result, reduction_context const& ctx); +std::unique_ptr make_size_type_column(reduction_context const& ctx); + +// Kind-specific TUs explicitly instantiate this bridge; the frontend needs no reducer definition. +template +std::unique_ptr compute_reduction(reduction_context const& ctx); + +// Suppress implicit instantiation in the frontend and other reducer translation units. +extern template std::unique_ptr compute_reduction( + reduction_context const& ctx); +extern template std::unique_ptr compute_reduction( + reduction_context const& ctx); +extern template std::unique_ptr compute_reduction( + reduction_context const& ctx); +extern template std::unique_ptr compute_reduction( + reduction_context const& ctx); +extern template std::unique_ptr compute_reduction( + reduction_context const& ctx); +extern template std::unique_ptr compute_reduction( + reduction_context const& ctx); +extern template std::unique_ptr compute_reduction( + reduction_context const& ctx); +extern template std::unique_ptr compute_reduction( + reduction_context const& ctx); + +std::vector> compute_fused_sums(reduction_context const& ctx, + host_span kinds, + std::span is_intermediate); + +} // namespace cudf::groupby::detail::hash::single_pass diff --git a/cpp/src/groupby/hash/single_pass_sum_overflow.cu b/cpp/src/groupby/hash/single_pass_sum_overflow.cu new file mode 100644 index 000000000000..e47a953910ce --- /dev/null +++ b/cpp/src/groupby/hash/single_pass_sum_overflow.cu @@ -0,0 +1,13 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "single_pass_reductions.cuh" + +namespace cudf::groupby::detail::hash::single_pass { + +template std::unique_ptr compute_reduction( + reduction_context const& ctx); + +} // namespace cudf::groupby::detail::hash::single_pass diff --git a/cpp/src/groupby/hash/single_pass_sums.cu b/cpp/src/groupby/hash/single_pass_sums.cu new file mode 100644 index 000000000000..898356d4c2ec --- /dev/null +++ b/cpp/src/groupby/hash/single_pass_sums.cu @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "single_pass_reductions.cuh" + +namespace cudf::groupby::detail::hash::single_pass { + +template std::unique_ptr compute_reduction(reduction_context const& ctx); +template std::unique_ptr compute_reduction( + reduction_context const& ctx); + +std::vector> compute_fused_sums(reduction_context const& ctx, + host_span kinds, + std::span is_intermediate) +{ + return type_dispatcher(ctx.values_type, fused_sums_fn{}, ctx, kinds, is_intermediate); +} + +} // namespace cudf::groupby::detail::hash::single_pass diff --git a/cpp/src/groupby/streaming_groupby/common.cuh b/cpp/src/groupby/streaming_groupby/common.cuh index 5f68ceb35866..cf0f47269f11 100644 --- a/cpp/src/groupby/streaming_groupby/common.cuh +++ b/cpp/src/groupby/streaming_groupby/common.cuh @@ -27,6 +27,7 @@ #include #include +#include #include #include #include @@ -333,7 +334,8 @@ struct streaming_groupby::impl { * built once and reused on every aggregate() / merge() call rather than rebuilt * (which requires a host-to-device copy of the column metadata). */ - std::unique_ptr _d_agg_results; + std::unique_ptr> + _d_agg_results; std::vector _value_col_indices; std::unique_ptr> _d_agg_kinds; diff --git a/cpp/src/groupby/streaming_groupby/impl.cu b/cpp/src/groupby/streaming_groupby/impl.cu index 19b94b7ce132..68c93567c26a 100644 --- a/cpp/src/groupby/streaming_groupby/impl.cu +++ b/cpp/src/groupby/streaming_groupby/impl.cu @@ -15,21 +15,25 @@ #include #include #include -#include #include +#include +#include #include #include #include #include #include +#include #include #include #include +#include #include #include +#include #include #include #include @@ -38,6 +42,32 @@ namespace cudf::groupby { namespace { +// Streaming still uses element_aggregator, whose atomic requirements are independent of the +// reductions used by ordinary hash groupby. +struct is_atomic_aggregation_supported { + template + bool operator()() const + { + if constexpr (cudf::is_nested()) { + return false; + } else if constexpr (std::is_same_v && K == aggregation::SUM) { + // The existing decimal128 SUM implementation provides its own atomic addition. + return true; + } else { + using Target = cudf::detail::target_type_t; + constexpr auto uses_storage = + cudf::is_fixed_point() && + (K == aggregation::MIN || K == aggregation::MAX || K == aggregation::SUM); + using DeviceTarget = + std::conditional_t, Target>; + if constexpr (!std::is_void_v) { + return cudf::has_atomic_support(); + } + return false; + } + } +}; + void validate_requests(host_span requests) { for (auto const& req : requests) { @@ -86,10 +116,7 @@ streaming_groupby::impl::impl(host_span key_indices, size_type max_distinct_keys, null_policy null_handling, cuda::mr::any_resource mr) - : _max_distinct_keys{max_distinct_keys}, - _null_handling{null_handling}, - _mr{std::move(mr)}, - _d_agg_results{nullptr, +[](mutable_table_device_view*) {}} + : _max_distinct_keys{max_distinct_keys}, _null_handling{null_handling}, _mr{std::move(mr)} { CUDF_EXPECTS(max_distinct_keys > 0, "max_distinct_keys must be positive.", std::invalid_argument); if (!key_indices.empty()) { _key_indices.assign(key_indices.begin(), key_indices.end()); } @@ -125,18 +152,6 @@ void streaming_groupby::impl::initialize(table_view const& data, cuda::stream_re auto agg_requests = build_aggregation_requests(_requests_clone, data); - // TODO: streaming aggregation reuses the cudf hash-groupby element_aggregator, - // so it inherits the same atomic-support requirement. In particular, decimal128 - // MIN/MAX/SUM falls through to CUDF_UNREACHABLE because __int128 is not - // lock-free atomic. Stateless cudf::groupby falls back to sort-based groupby - // in that case; streaming has no such fallback. Until streaming has a - // non-atomic aggregator path (or 128-bit atomics gain hardware support), gate - // by the same predicate to fail loudly instead of silently producing garbage. - CUDF_EXPECTS(detail::hash::can_use_hash_groupby(agg_requests), - "streaming_groupby does not support this combination of value type and " - "aggregation kind (e.g. decimal128 MIN/MAX/SUM require 128-bit atomics).", - std::invalid_argument); - auto [values_view, agg_kinds_hv, agg_objects, is_intermediate, has_compound] = detail::hash::extract_single_pass_aggs(agg_requests, stream); @@ -146,7 +161,8 @@ void streaming_groupby::impl::initialize(table_view const& data, cuda::stream_re _has_compound_aggs = has_compound; // Reject aggregation kinds that are unsupported in streaming after decomposition. - for (auto k : _agg_kinds) { + for (std::size_t i = 0; i < _agg_kinds.size(); ++i) { + auto const k = _agg_kinds[i]; CUDF_EXPECTS(k != aggregation::ARGMIN && k != aggregation::ARGMAX, "Streaming groupby does not support MIN/MAX on variable-width types " "(internally decomposed to ARGMIN/ARGMAX).", @@ -155,6 +171,15 @@ void streaming_groupby::impl::initialize(table_view const& data, cuda::stream_re "Streaming groupby does not support SUM_OVERFLOW " "(struct intermediate cannot be merged across batches).", std::invalid_argument); + auto const& values = values_view.column(i); + auto const values_type = cudf::is_dictionary(values.type()) + ? cudf::dictionary_column_view(values).keys().type() + : values.type(); + CUDF_EXPECTS(cudf::detail::dispatch_type_and_aggregation( + values_type, k, is_atomic_aggregation_supported{}), + "streaming_groupby does not support this combination of value type and " + "aggregation kind.", + std::invalid_argument); } _agg_results = detail::hash::create_results_table( @@ -163,11 +188,7 @@ void streaming_groupby::impl::initialize(table_view const& data, cuda::stream_re // Cache the mutable_table_device_view once; the underlying table is fixed-size and // never reallocated, so the device-side descriptor stays valid for the whole // lifetime of this impl. - { - auto raii = mutable_table_device_view::create(*_agg_results, stream); - _d_agg_results = - decltype(_d_agg_results){raii.release(), +[](mutable_table_device_view* t) { t->destroy(); }}; - } + _d_agg_results = mutable_table_device_view::create(*_agg_results, stream); _d_agg_kinds = std::make_unique>( cudf::detail::make_device_uvector_async(_agg_kinds, stream, mr)); @@ -404,8 +425,8 @@ bool is_streaming_groupby_supported(data_type values_type, aggregation::Kind kin break; default: return false; } - // decimal128 SUM/MIN/MAX needs 128-bit atomics, which aren't supported. - if ((kind == aggregation::SUM || kind == aggregation::MIN || kind == aggregation::MAX) && + // decimal128 MIN/MAX needs unsupported 128-bit atomics; SUM has its own atomic addition. + if ((kind == aggregation::MIN || kind == aggregation::MAX) && values_type.id() == type_id::DECIMAL128) { return false; } diff --git a/cpp/tests/groupby/keys_tests.cpp b/cpp/tests/groupby/keys_tests.cpp index f97ee317dbb1..b5df0b3c2b10 100644 --- a/cpp/tests/groupby/keys_tests.cpp +++ b/cpp/tests/groupby/keys_tests.cpp @@ -1,16 +1,23 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include #include +#include #include +#include #include #include #include +#include + +#include + +#include using namespace cudf::test::iterators; @@ -107,6 +114,68 @@ TYPED_TEST(groupby_keys_test, include_null_keys) cudf::null_policy::INCLUDE); } +TYPED_TEST(groupby_keys_test, distinct_keys_with_nullable_values) +{ + using K = TypeParam; + cudf::test::fixed_width_column_wrapper keys({3, 1, 0, 2}, {1, 1, 0, 1}); + cudf::test::fixed_width_column_wrapper vals({30, 99, 7, 20}, {1, 0, 1, 1}); + cudf::test::fixed_width_column_wrapper included_counts{1, 0, 1, 1}; + + cudf::test::fixed_width_column_wrapper excluded_keys({1, 2, 3}, no_nulls()); + cudf::test::fixed_width_column_wrapper excluded_vals({0, 20, 30}, {0, 1, 1}); + cudf::test::fixed_width_column_wrapper excluded_counts{0, 1, 1}; + + // Including the null key yields one group per input row; excluding it must omit that row. + for (auto null_handling : {cudf::null_policy::INCLUDE, cudf::null_policy::EXCLUDE}) { + auto const include_null = null_handling == cudf::null_policy::INCLUDE; + auto const& expect_keys = include_null ? keys : excluded_keys; + auto const& expect_vals = include_null ? vals : excluded_vals; + auto const& expect_counts = include_null ? included_counts : excluded_counts; + test_single_agg(keys, + vals, + expect_keys, + expect_vals, + cudf::make_max_aggregation(), + force_use_sort_impl::NO, + null_handling); + test_single_agg( + keys, + vals, + expect_keys, + expect_counts, + cudf::make_count_aggregation(cudf::null_policy::EXCLUDE), + force_use_sort_impl::NO, + null_handling); + } +} + +TYPED_TEST(groupby_keys_test, one_duplicate_key_with_nullable_values) +{ + using K = TypeParam; + // Exactly one repeated key gives one fewer group than input rows, including the null key. + cudf::test::fixed_width_column_wrapper keys({3, 1, 0, 3}, {1, 1, 0, 1}); + cudf::test::fixed_width_column_wrapper vals({30, 99, 7, 20}, {1, 0, 1, 1}); + cudf::test::fixed_width_column_wrapper expect_keys({1, 3, 0}, {1, 1, 0}); + cudf::test::fixed_width_column_wrapper expect_vals({0, 30, 7}, {0, 1, 1}); + cudf::test::fixed_width_column_wrapper expect_counts{0, 2, 1}; + + test_single_agg(keys, + vals, + expect_keys, + expect_vals, + cudf::make_max_aggregation(), + force_use_sort_impl::NO, + cudf::null_policy::INCLUDE); + test_single_agg( + keys, + vals, + expect_keys, + expect_counts, + cudf::make_count_aggregation(cudf::null_policy::EXCLUDE), + force_use_sort_impl::NO, + cudf::null_policy::INCLUDE); +} + TYPED_TEST(groupby_keys_test, pre_sorted_keys) { using K = TypeParam; @@ -407,3 +476,121 @@ TEST_F(groupby_cache_test, duplicate_columns) cudf::make_nth_element_aggregation(0)); EXPECT_NO_THROW(gb_obj.aggregate(requests)); } + +using groupby_sampling_test = groupby_keys_test; + +TEST_F(groupby_sampling_test, NearlyDistinctSampleUnderestimatesPopulation) +{ + constexpr cudf::size_type num_rows = 1 << 21; + constexpr cudf::size_type stride = 64; + constexpr cudf::size_type sample_keys = 31'000; + constexpr cudf::size_type num_samples = num_rows / stride; + + // The periodic sample is almost entirely distinct, but still has far fewer keys than the + // complete input. Every row outside the sample has a unique key. An undersized table must + // restart the build without dropping rows or duplicating groups. + std::vector keys_data(num_rows); + std::vector expected_keys; + std::vector expected_counts; + std::vector expected_maxima; + expected_keys.reserve(num_rows - num_samples + sample_keys); + expected_counts.reserve(num_rows - num_samples + sample_keys); + expected_maxima.reserve(num_rows - num_samples + sample_keys); + for (cudf::size_type key = 0; key < sample_keys; ++key) { + expected_keys.push_back(key); + expected_counts.push_back(num_samples / sample_keys + (key < num_samples % sample_keys)); + auto const last_sample = key + ((num_samples - 1 - key) / sample_keys) * sample_keys; + expected_maxima.push_back(last_sample * stride); + } + for (cudf::size_type row = 0; row < num_rows; ++row) { + if (row % stride == 0) { + keys_data[row] = (row / stride) % sample_keys; + } else { + keys_data[row] = sample_keys + row; + expected_keys.push_back(keys_data[row]); + expected_counts.push_back(1); + expected_maxima.push_back(row); + } + } + + auto const keys = + cudf::test::fixed_width_column_wrapper(keys_data.begin(), keys_data.end()); + auto const expect_keys = + cudf::test::fixed_width_column_wrapper(expected_keys.begin(), expected_keys.end()); + auto const expect_counts = cudf::test::fixed_width_column_wrapper( + expected_counts.begin(), expected_counts.end()); + test_single_agg(keys, + keys, + expect_keys, + expect_counts, + cudf::make_count_aggregation()); + + // COUNT only needs the group offsets. MAX of the row indices also verifies that the retry + // rebuilt the row positions and filled the grouped row order correctly. + auto const values = cudf::test::fixed_width_column_wrapper( + cuda::counting_iterator{0}, cuda::counting_iterator{num_rows}); + auto const expect_maxima = + cudf::test::fixed_width_column_wrapper(expected_maxima.begin(), expected_maxima.end()); + test_single_agg(keys, + values, + expect_keys, + expect_maxima, + cudf::make_max_aggregation()); + + // Without requests the retry rebuilds the representative key rows instead of group slots. + cudf::groupby::groupby gb_obj(cudf::table_view({keys})); + auto const result = gb_obj.aggregate({}, cudf::test::get_default_stream()); + auto const sorted_keys = + cudf::sort(result.first->view(), {}, {}, cudf::test::get_default_stream()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expect_keys, sorted_keys->view().column(0)); + EXPECT_TRUE(result.second.empty()); +} + +TEST_F(groupby_sampling_test, NullableMaxConsumesFullPackedSegment) +{ + constexpr cudf::size_type num_groups = 1'000; + constexpr cudf::size_type long_group_rows = 1'024; + constexpr cudf::size_type short_group_rows = 5; + constexpr cudf::size_type num_rows = long_group_rows + (num_groups - 1) * short_group_rows; + + // Groups average six rows, but fewer than one in four values is valid. The density hint + // chooses one thread per segment; it must still consume all 1,024 rows of the longest group. + std::vector keys_data; + std::vector values_data; + std::vector validity; + std::vector expected_maxima; + std::vector expected_validity; + keys_data.reserve(num_rows); + values_data.reserve(num_rows); + validity.reserve(num_rows); + expected_maxima.reserve(num_groups); + expected_validity.reserve(num_groups); + + for (cudf::size_type group = 0; group < num_groups; ++group) { + auto const group_rows = group == 0 ? long_group_rows : short_group_rows; + auto const maximum = static_cast(group * long_group_rows + group_rows - 1); + expected_maxima.push_back(maximum); + expected_validity.push_back(group != 1); + for (cudf::size_type row = 0; row < group_rows; ++row) { + auto const valid = group != 1 && row == group_rows - 1; + keys_data.push_back(group); + // Null payloads exceed every valid maximum, so loading one as valid also fails the test. + values_data.push_back(valid ? maximum : 1.0e9); + validity.push_back(valid); + } + } + + auto const keys = + cudf::test::fixed_width_column_wrapper(keys_data.begin(), keys_data.end()); + auto const values = cudf::test::fixed_width_column_wrapper( + values_data.begin(), values_data.end(), validity.begin()); + auto const expect_keys = cudf::test::fixed_width_column_wrapper( + cuda::counting_iterator{0}, cuda::counting_iterator{num_groups}); + auto const expect_maxima = cudf::test::fixed_width_column_wrapper( + expected_maxima.begin(), expected_maxima.end(), expected_validity.begin()); + test_single_agg(keys, + values, + expect_keys, + expect_maxima, + cudf::make_max_aggregation()); +} diff --git a/cpp/tests/groupby/streaming_groupby_test.cpp b/cpp/tests/groupby/streaming_groupby_test.cpp index 343ca6663ef2..c60b6d6c56ad 100644 --- a/cpp/tests/groupby/streaming_groupby_test.cpp +++ b/cpp/tests/groupby/streaming_groupby_test.cpp @@ -6,12 +6,14 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include #include @@ -24,6 +26,7 @@ #include #include +#include #include #include @@ -153,6 +156,33 @@ TEST_F(StreamingGroupbyTest, MemoryResource) EXPECT_GT(mr.get_bytes_counter().peak, 0); } +TEST_F(StreamingGroupbyTest, ReleasesDeviceMemory) +{ + cudf::test::fixed_width_column_wrapper keys{1, 2, 3, 1}; + cudf::test::fixed_width_column_wrapper values{10, 20, 30, 40}; + cudf::table_view batch{{keys, values}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + reqs.push_back(make_req(1, cudf::make_max_aggregation())); + auto statistics_mr = rmm::mr::statistics_resource_adaptor(mr()); + cudf::test::scoped_current_device_resource resource_scope{statistics_mr}; + + for (int iteration = 0; iteration < 3; ++iteration) { + SCOPED_TRACE(iteration); + { + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch); + streaming_agg.aggregate(batch); + auto [result_keys, result_values] = streaming_agg.finalize(); + EXPECT_EQ(result_keys->num_rows(), 3); + EXPECT_EQ(result_values.size(), reqs.size()); + } + cudf::test::get_default_stream().sync(); + EXPECT_EQ(statistics_mr.get_bytes_counter().value, 0); + } + EXPECT_GT(statistics_mr.get_bytes_counter().peak, 0); +} + TEST_F(StreamingGroupbyTest, SumTwoBatches) { using K = int32_t; @@ -728,6 +758,49 @@ TEST_F(StreamingGroupbyTest, UnsupportedAggThrows) std::invalid_argument); } +TEST_F(StreamingGroupbyTest, Decimal128MinMaxRejected) +{ + cudf::test::fixed_width_column_wrapper keys{1, 1, 2}; + cudf::test::fixed_point_column_wrapper<__int128_t> values{{1, 2, 3}, numeric::scale_type{-2}}; + auto const batch = cudf::table_view{{keys, values}}; + + for (auto const kind : {cudf::aggregation::MIN, cudf::aggregation::MAX}) { + SCOPED_TRACE(static_cast(kind)); + EXPECT_FALSE(cudf::groupby::is_streaming_groupby_supported( + cudf::data_type{cudf::type_id::DECIMAL128}, kind)); + auto reqs = single_agg_req(1, + kind == cudf::aggregation::MIN + ? cudf::make_min_aggregation() + : cudf::make_max_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + EXPECT_THROW(streaming_agg.aggregate(batch), std::invalid_argument); + } +} + +TEST_F(StreamingGroupbyTest, Decimal128SumTwoBatches) +{ + EXPECT_TRUE(cudf::groupby::is_streaming_groupby_supported( + cudf::data_type{cudf::type_id::DECIMAL128}, cudf::aggregation::SUM)); + auto constexpr scale = numeric::scale_type{-2}; + auto constexpr large = __int128_t{1} << 80; + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1}; + cudf::test::fixed_width_column_wrapper keys2{2, 1, 3}; + cudf::test::fixed_point_column_wrapper<__int128_t> values1{{large, -large, 7}, scale}; + cudf::test::fixed_point_column_wrapper<__int128_t> values2{{3, 5, large}, scale}; + cudf::test::fixed_width_column_wrapper expected_keys{1, 2, 3}; + cudf::test::fixed_point_column_wrapper<__int128_t> expected_values{ + {large + 12, -large + 3, large}, scale}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(cudf::table_view{{keys1, values1}}); + streaming_agg.aggregate(cudf::table_view{{keys2, values2}}); + auto [keys, results] = streaming_agg.finalize(); + ASSERT_EQ(results.size(), 1); + ASSERT_EQ(results[0].results.size(), 1); + check(keys, results, cudf::table_view{{expected_keys}}, {expected_values}); +} + TEST_F(StreamingGroupbyTest, BatchExceedsMaxDistinctKeysThrows) { using K = int32_t;