Fix a circular inclusion

Author: PointKernel

cmake/header_testing.cmake

@@ -41,15 +41,13 @@ function(cuco_add_header_test label definitions)
 
   # List of headers that have known issues or are not meant to be included directly
   set(excluded_headers
-    # Headers with circular dependencies that are not meant to be included directly
-    cuco/detail/static_map/helpers.cuh
+    # Add any headers that should be excluded from testing here
+    # Example: cuco/internal_header.cuh
   )
 
   # Remove excluded headers
   if(excluded_headers)
     list(REMOVE_ITEM headers ${excluded_headers})
-    list(LENGTH headers headers_count_after_exclusion)
-    message(STATUS "After exclusion: ${headers_count_after_exclusion} headers remaining")
   endif()
 
   # Only test with CUDA compiler since cuco is device-only

include/cuco/detail/static_map/helpers.cuh

@@ -15,12 +15,13 @@
  */
 
 #include <cuco/detail/bitwise_compare.cuh>
-#include <cuco/detail/static_map/helpers.cuh>
 #include <cuco/detail/static_map/kernels.cuh>
 #include <cuco/detail/utility/cuda.hpp>
 #include <cuco/detail/utils.hpp>
+#include <cuco/extent.cuh>
 #include <cuco/operator.hpp>
 #include <cuco/static_map_ref.cuh>
+#include <cuco/storage.cuh>
 
 #include <cuda/std/tuple>
 #include <cuda/stream_ref>
@@ -352,7 +353,7 @@ void static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Stora
 {
   auto constexpr has_init = false;
   auto const init = this->empty_value_sentinel();  // use empty_sentinel as unused init value
-  detail::static_map_ns::dispatch_insert_or_apply<has_init, cg_size, Allocator>(
+  this->dispatch_insert_or_apply<has_init, cg_size, Allocator>(
     first, last, init, op, ref(op::insert_or_apply), stream);
 }
 
@@ -370,7 +371,7 @@ void static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Stora
     InputIt first, InputIt last, Init init, Op op, cuda::stream_ref stream) noexcept
 {
   auto constexpr has_init = true;
-  detail::static_map_ns::dispatch_insert_or_apply<has_init, cg_size, Allocator>(
+  this->dispatch_insert_or_apply<has_init, cg_size, Allocator>(
     first, last, init, op, ref(op::insert_or_apply), stream);
 }
 
@@ -960,4 +961,92 @@ auto static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Stora
                                     cuda_thread_scope<Scope>{},
                                     impl_->storage_ref()};
 }
+
+template <class Key,
+          class T,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+template <bool HasInit,
+          int32_t CGSize,
+          typename AllocatorType,
+          typename InputIt,
+          typename InitType,
+          typename OpType,
+          typename RefType>
+void static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::
+  dispatch_insert_or_apply(
+    InputIt first, InputIt last, InitType init, OpType op, RefType ref, cuda::stream_ref stream)
+{
+  auto const num = cuco::detail::distance(first, last);
+  if (num == 0) { return; }
+
+  int32_t const default_grid_size = cuco::detail::grid_size(num, CGSize);
+
+  if constexpr (CGSize == 1) {
+    using shmem_size_type = int32_t;
+
+    int32_t constexpr shmem_block_size                = 1024;
+    shmem_size_type constexpr cardinality_threshold   = shmem_block_size;
+    shmem_size_type constexpr shared_map_num_elements = cardinality_threshold + shmem_block_size;
+    float constexpr load_factor                       = 0.7;
+    shmem_size_type constexpr shared_map_size =
+      static_cast<shmem_size_type>((1.0 / load_factor) * shared_map_num_elements);
+
+    using extent_type     = cuco::extent<shmem_size_type, shared_map_size>;
+    using shared_map_type = cuco::static_map<typename RefType::key_type,
+                                             typename RefType::mapped_type,
+                                             extent_type,
+                                             cuda::thread_scope_block,
+                                             typename RefType::key_equal,
+                                             typename RefType::probing_scheme_type,
+                                             AllocatorType,
+                                             cuco::storage<1>>;
+
+    using shared_map_ref_type = typename shared_map_type::template ref_type<>;
+    auto constexpr bucket_extent =
+      cuco::make_valid_extent<typename shared_map_ref_type::probing_scheme_type,
+                              typename shared_map_ref_type::storage_ref_type>(extent_type{});
+
+    auto insert_or_apply_shmem_fn_ptr =
+      cuco::detail::static_map_ns::insert_or_apply_shmem<HasInit,
+                                                         CGSize,
+                                                         shmem_block_size,
+                                                         shared_map_ref_type,
+                                                         InputIt,
+                                                         InitType,
+                                                         OpType,
+                                                         RefType>;
+
+    int32_t const max_op_grid_size =
+      cuco::detail::max_occupancy_grid_size(shmem_block_size, insert_or_apply_shmem_fn_ptr);
+
+    int32_t const shmem_default_grid_size =
+      cuco::detail::grid_size(num, CGSize, cuco::detail::default_stride(), shmem_block_size);
+
+    auto const shmem_grid_size         = std::min(shmem_default_grid_size, max_op_grid_size);
+    auto const num_elements_per_thread = num / (shmem_grid_size * shmem_block_size);
+
+    // use shared_memory only if each thread has atleast 3 elements to process
+    if (num_elements_per_thread > 2) {
+      cuco::detail::static_map_ns::
+        insert_or_apply_shmem<HasInit, CGSize, shmem_block_size, shared_map_ref_type>
+        <<<shmem_grid_size, shmem_block_size, 0, stream.get()>>>(
+          first, num, init, op, ref, bucket_extent);
+    } else {
+      cuco::detail::static_map_ns::
+        insert_or_apply<HasInit, CGSize, cuco::detail::default_block_size()>
+        <<<default_grid_size, cuco::detail::default_block_size(), 0, stream.get()>>>(
+          first, num, init, op, ref);
+    }
+  } else {
+    cuco::detail::static_map_ns::
+      insert_or_apply<HasInit, CGSize, cuco::detail::default_block_size()>
+      <<<default_grid_size, cuco::detail::default_block_size(), 0, stream.get()>>>(
+        first, num, init, op, ref);
+  }
+}
 }  // namespace cuco

include/cuco/detail/static_map/static_map.inl

@@ -1215,6 +1215,36 @@ class static_map {
   [[nodiscard]] auto ref(Operators... ops) const noexcept;
 
  private:
+  /**
+   * @brief Dispatches to shared memory map kernel if `num_elements_per_thread > 2`, else
+   * fallbacks to global memory map kernel.
+   *
+   * @tparam HasInit Boolean to dispatch based on init parameter
+   * @tparam CGSize Number of threads in each CG
+   * @tparam Allocator Allocator type used to created shared_memory map
+   * @tparam InputIt Device accessible input iterator whose `value_type` is
+   * convertible to the `value_type` of the data structure
+   * @tparam Init Type of init value convertible to payload type
+   * @tparam Op Callable type used to peform `apply` operation.
+   * @tparam Ref Type of non-owning device ref allowing access to storage
+   *
+   * @param first Beginning of the sequence of input elements
+   * @param last End of the sequence of input elements
+   * @param init The init value of the `op`
+   * @param op Callable object to perform apply operation.
+   * @param ref Non-owning container device ref used to access the slot storage
+   * @param stream CUDA stream used for insert_or_apply operation
+   */
+  template <bool HasInit,
+            int32_t CGSize,
+            typename AllocatorType,
+            typename InputIt,
+            typename InitType,
+            typename OpType,
+            typename RefType>
+  void dispatch_insert_or_apply(
+    InputIt first, InputIt last, InitType init, OpType op, RefType ref, cuda::stream_ref stream);
+
   std::unique_ptr<impl_type> impl_;   ///< Static map implementation
   mapped_type empty_value_sentinel_;  ///< Sentinel value that indicates an empty payload
 };