Merge remote-tracking branch 'upstream/dev' into fix-large-retrieve-all

NVIDIA · Aug 16, 2024 · dc8235a · dc8235a
2 parents e0ed502 + 6eaed1b
commit dc8235a
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 12 deletions.
diff --git a/include/cuco/detail/static_map/kernels.cuh b/include/cuco/detail/static_map/kernels.cuh
@@ -202,8 +202,8 @@ CUCO_KERNEL __launch_bounds__(BlockSize) void insert_or_apply_shmem(
 
   auto shared_map     = SharedMapRefType{cuco::empty_key<Key>{ref.empty_key_sentinel()},
                                      cuco::empty_value<Value>{ref.empty_value_sentinel()},
-                                         {},
-                                         {},
+                                     ref.key_eq(),
+                                     ref.probing_scheme(),
                                          {},
                                      storage};
   auto shared_map_ref = std::move(shared_map).with(cuco::op::insert_or_apply);

diff --git a/include/cuco/detail/static_map/static_map.inl b/include/cuco/detail/static_map/static_map.inl
@@ -347,19 +347,11 @@ template <class Key,
           class ProbingScheme,
           class Allocator,
           class Storage>
-template <typename InputIt, typename Init, typename Op>
+template <typename InputIt, typename Init, typename Op, typename>
 void static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::
   insert_or_apply_async(
     InputIt first, InputIt last, Init init, Op op, cuda::stream_ref stream) noexcept
 {
-  using shared_map_type   = cuco::static_map<Key,
-                                             T,
-                                             int32_t,
-                                             cuda::thread_scope_block,
-                                             KeyEqual,
-                                             ProbingScheme,
-                                             Allocator,
-                                             cuco::storage<1>>;
   auto constexpr has_init = true;
   static_map_ns::detail::dispatch_insert_or_apply<has_init, cg_size, Allocator>(
     first, last, init, op, ref(op::insert_or_apply), stream);

diff --git a/include/cuco/detail/storage/aow_storage.inl b/include/cuco/detail/storage/aow_storage.inl
@@ -74,6 +74,8 @@ template <typename T, int32_t WindowSize, typename Extent, typename Allocator>
 void aow_storage<T, WindowSize, Extent, Allocator>::initialize_async(
   value_type key, cuda::stream_ref stream) noexcept
 {
+  if (this->num_windows() == 0) { return; }
+
   auto constexpr cg_size = 1;
   auto constexpr stride  = 4;
   auto const grid_size   = cuco::detail::grid_size(this->num_windows(), cg_size, stride);

diff --git a/include/cuco/static_map.cuh b/include/cuco/static_map.cuh
@@ -37,6 +37,7 @@
 
 #include <cstddef>
 #include <memory>
+#include <type_traits>
 #include <utility>
 
 namespace cuco {
@@ -564,7 +565,10 @@ class static_map {
    * @param op Callable object to perform apply operation.
    * @param stream CUDA stream used for insert
    */
-  template <typename InputIt, typename Init, typename Op>
+  template <typename InputIt,
+            typename Init,
+            typename Op,
+            typename = std::enable_if_t<std::is_convertible_v<Init, T>>>
   void insert_or_apply_async(
     InputIt first, InputIt last, Init init, Op op, cuda::stream_ref stream = {}) noexcept;
 

diff --git a/tests/utility/storage_test.cu b/tests/utility/storage_test.cu
@@ -37,6 +37,15 @@ TEMPLATE_TEST_CASE_SIG("Storage tests",
   using allocator_type = cuco::cuda_allocator<char>;
   auto allocator       = allocator_type{};
 
+  SECTION("Initialize empty storage is allowed.")
+  {
+    auto s = cuco::
+      aow_storage<cuco::pair<Key, Value>, window_size, cuco::extent<std::size_t>, allocator_type>{
+        cuco::extent<std::size_t>{0}, allocator};
+
+    s.initialize(cuco::pair<Key, Value>{1, 1});
+  }
+
   SECTION("Allocate array of pairs with AoS storage.")
   {
     auto s = cuco::