Clean up hash-based groupby aggregation, reducing overhead and memory usage (#20658)

This does some minor cleanup to the hash-based groupby aggregation code, removing some code path that seems never (or very rarely) executed and reducing a little of memory usage.

Authors:
  - Nghia Truong (https://github.com/ttnghia)

Approvers:
  - Yunsong Wang (https://github.com/PointKernel)
  - Bradley Dice (https://github.com/bdice)

URL: #20658

Author: ttnghia

cpp/src/groupby/hash/compute_mapping_indices.cu

@@ -17,5 +17,6 @@ template void compute_mapping_indices<hash_set_ref_t<cuco::insert_and_find_tag>>
   size_type* local_mapping_index,
   size_type* global_mapping_index,
   size_type* block_cardinality,
+  cuda::std::atomic_flag* needs_global_memory_fallback,
   rmm::cuda_stream_view stream);
 }  // namespace cudf::groupby::detail::hash

cpp/src/groupby/hash/compute_mapping_indices.cuh

@@ -42,6 +42,13 @@ __device__ void find_local_mapping(cooperative_groups::thread_block const& block
         auto const ref_cardinality =
           cuda::atomic_ref<size_type, cuda::thread_scope_block>{*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;
       }
@@ -71,7 +78,7 @@ __device__ void find_global_mapping(cooperative_groups::thread_block const& bloc
     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_SHM_MAX_ELEMENTS + idx] = key_idx;
+    global_mapping_indices[block.group_index().x * GROUPBY_CARDINALITY_THRESHOLD + idx] = key_idx;
   }
 }
 
@@ -88,9 +95,10 @@ CUDF_KERNEL void mapping_indices_kernel(size_type num_input_rows,
                                         bitmask_type const* row_bitmask,
                                         size_type* local_mapping_indices,
                                         size_type* global_mapping_indices,
-                                        size_type* block_cardinality)
+                                        size_type* block_cardinality,
+                                        cuda::std::atomic_flag* needs_global_memory_fallback)
 {
-  __shared__ size_type shared_set_indices[GROUPBY_SHM_MAX_ELEMENTS];
+  __shared__ size_type shared_set_indices[GROUPBY_CARDINALITY_THRESHOLD];
 
   // Shared set initialization
   __shared__ size_type slots[valid_extent.value()];
@@ -126,12 +134,17 @@ CUDF_KERNEL void mapping_indices_kernel(size_type num_input_rows,
                        shared_set_indices);
 
     block.sync();
-    if (cardinality >= GROUPBY_CARDINALITY_THRESHOLD) { break; }
+    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) {
+  if (cardinality <= GROUPBY_CARDINALITY_THRESHOLD) {
     find_global_mapping(block, cardinality, global_set, shared_set_indices, global_mapping_indices);
   }
 
@@ -155,13 +168,16 @@ void compute_mapping_indices(size_type grid_size,
                              size_type* local_mapping_indices,
                              size_type* global_mapping_indices,
                              size_type* block_cardinality,
+                             cuda::std::atomic_flag* needs_global_memory_fallback,
                              rmm::cuda_stream_view stream)
 {
-  mapping_indices_kernel<<<grid_size, GROUPBY_BLOCK_SIZE, 0, stream>>>(num_rows,
-                                                                       global_set,
-                                                                       row_bitmask,
-                                                                       local_mapping_indices,
-                                                                       global_mapping_indices,
-                                                                       block_cardinality);
+  mapping_indices_kernel<<<grid_size, GROUPBY_BLOCK_SIZE, 0, stream>>>(
+    num_rows,
+    global_set,
+    row_bitmask,
+    local_mapping_indices,
+    global_mapping_indices,
+    block_cardinality,
+    needs_global_memory_fallback);
 }
 }  // namespace cudf::groupby::detail::hash

cpp/src/groupby/hash/compute_mapping_indices.hpp

@@ -27,5 +27,6 @@ void compute_mapping_indices(size_type grid_size,
                              size_type* local_mapping_index,
                              size_type* global_mapping_index,
                              size_type* block_cardinality,
+                             cuda::std::atomic_flag* needs_global_memory_fallback,
                              rmm::cuda_stream_view stream);
 }  // namespace cudf::groupby::detail::hash

cpp/src/groupby/hash/compute_mapping_indices_null.cu

@@ -18,5 +18,6 @@ template void compute_mapping_indices<nullable_hash_set_ref_t<cuco::insert_and_f
   size_type* local_mapping_index,
   size_type* global_mapping_index,
   size_type* block_cardinality,
+  cuda::std::atomic_flag* needs_global_memory_fallback,
   rmm::cuda_stream_view stream);
 }  // namespace cudf::groupby::detail::hash

cpp/src/groupby/hash/compute_shared_memory_aggs.cu

@@ -145,7 +145,7 @@ __device__ void compute_final_aggregations(cooperative_groups::thread_block cons
   // 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_SHM_MAX_ELEMENTS) +
+      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);
@@ -182,7 +182,7 @@ CUDF_KERNEL void single_pass_shmem_aggs_kernel(cudf::size_type num_rows,
 {
   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; }
+  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;

cpp/src/groupby/hash/compute_single_pass_aggs.cu

@@ -5,12 +5,10 @@
 
 #include "compute_single_pass_aggs.cuh"
 #include "compute_single_pass_aggs.hpp"
+#include "single_pass_functors.cuh"
 
-#include <rmm/device_scalar.hpp>
 #include <rmm/device_uvector.hpp>
 
-#include <cub/device/device_select.cuh>
-
 namespace cudf::groupby::detail::hash {
 
 std::pair<bool, size_type> is_shared_memory_compatible(host_span<aggregation::Kind const> agg_kinds,
@@ -39,39 +37,6 @@ std::pair<bool, size_type> is_shared_memory_compatible(host_span<aggregation::Ki
   return {can_run_by_shared_mem_kernel, available_shmem_size};
 }
 
-std::pair<size_type, rmm::device_uvector<size_type>> find_fallback_blocks(
-  size_type grid_size, size_type const* block_cardinality, rmm::cuda_stream_view stream)
-{
-  rmm::device_uvector<size_type> fallback_block_ids(grid_size, stream);
-  rmm::device_scalar<size_type> d_num_fallback_blocks(stream);
-
-  std::size_t storage_bytes = 0;
-  auto const select_pred    = [block_cardinality] __device__(auto const idx) {
-    return block_cardinality[idx] >= GROUPBY_CARDINALITY_THRESHOLD;
-  };
-  auto const exec_copy_if = [&](auto const storage_ptr) {
-    cub::DeviceSelect::If(storage_ptr,
-                          storage_bytes,
-                          thrust::make_counting_iterator(0),
-                          fallback_block_ids.begin(),
-                          d_num_fallback_blocks.data(),
-                          grid_size,
-                          select_pred,
-                          stream.value());
-  };
-
-  exec_copy_if(nullptr);
-  rmm::device_buffer tmp_storage(storage_bytes, stream);
-  exec_copy_if(tmp_storage.data());
-
-  auto const num_fallback_blocks = d_num_fallback_blocks.value(stream);
-  if (num_fallback_blocks > 0) { fallback_block_ids.resize(num_fallback_blocks, stream); }
-
-  return {num_fallback_blocks,
-          num_fallback_blocks > 0 ? std::move(fallback_block_ids)
-                                  : rmm::device_uvector<size_type>{0, stream}};
-}
-
 template std::pair<rmm::device_uvector<size_type>, bool> compute_single_pass_aggs<global_set_t>(
   global_set_t& global_set,
   bitmask_type const* row_bitmask,