Expose window_extent to device ref (#400)

This is the first step to migrate shared memory code from the legacy map to the new map. It exposes the `window_extent` ref API thus it can be used to construct container device ref on shared memory.
NVIDIA · Dec 2, 2023 · 7fae640 · 7fae640
1 parent 7404bd2
commit 7fae640
Show file tree

Hide file tree

Showing 7 changed files with 96 additions and 0 deletions.
diff --git a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh
@@ -241,6 +241,16 @@ class open_addressing_ref_impl {
     return storage_ref_.capacity();
   }
 
+  /**
+   * @brief Gets the window extent of the current storage.
+   *
+   * @return The window extent.
+   */
+  [[nodiscard]] __host__ __device__ constexpr extent_type window_extent() const noexcept
+  {
+    return storage_ref_.window_extent();
+  }
+
   /**
    * @brief Returns a const_iterator to one past the last slot.
    *

diff --git a/include/cuco/detail/static_map/static_map_ref.inl b/include/cuco/detail/static_map/static_map_ref.inl
@@ -115,6 +115,26 @@ static_map_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>
   return impl_.capacity();
 }
 
+template <typename Key,
+          typename T,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr static_map_ref<Key,
+                                             T,
+                                             Scope,
+                                             KeyEqual,
+                                             ProbingScheme,
+                                             StorageRef,
+                                             Operators...>::extent_type
+static_map_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::window_extent()
+  const noexcept
+{
+  return impl_.window_extent();
+}
+
 template <typename Key,
           typename T,
           cuda::thread_scope Scope,

diff --git a/include/cuco/detail/static_set/static_set_ref.inl b/include/cuco/detail/static_set/static_set_ref.inl
@@ -99,6 +99,24 @@ static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::c
   return impl_.capacity();
 }
 
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr static_set_ref<Key,
+                                             Scope,
+                                             KeyEqual,
+                                             ProbingScheme,
+                                             StorageRef,
+                                             Operators...>::extent_type
+static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::window_extent()
+  const noexcept
+{
+  return impl_.window_extent();
+}
+
 template <typename Key,
           cuda::thread_scope Scope,
           typename KeyEqual,

diff --git a/include/cuco/static_map_ref.cuh b/include/cuco/static_map_ref.cuh
@@ -145,6 +145,13 @@ class static_map_ref
    */
   [[nodiscard]] __host__ __device__ constexpr auto capacity() const noexcept;
 
+  /**
+   * @brief Gets the window extent of the current storage.
+   *
+   * @return The window extent.
+   */
+  [[nodiscard]] __host__ __device__ constexpr extent_type window_extent() const noexcept;
+
   /**
    * @brief Gets the sentinel value used to represent an empty key slot.
    *

diff --git a/include/cuco/static_set_ref.cuh b/include/cuco/static_set_ref.cuh
@@ -130,6 +130,13 @@ class static_set_ref
    */
   [[nodiscard]] __host__ __device__ constexpr auto capacity() const noexcept;
 
+  /**
+   * @brief Gets the window extent of the current storage.
+   *
+   * @return The window extent.
+   */
+  [[nodiscard]] __host__ __device__ constexpr extent_type window_extent() const noexcept;
+
   /**
    * @brief Gets the sentinel value used to represent an empty key slot.
    *

diff --git a/tests/static_map/capacity_test.cu b/tests/static_map/capacity_test.cu
@@ -73,6 +73,26 @@ TEST_CASE("Static map capacity", "")
 
   constexpr std::size_t num_keys{400};
 
+  SECTION("Static window extent can be evaluated at build time.")
+  {
+    std::size_t constexpr gold_extent = 211;
+
+    using extent_type = cuco::experimental::extent<std::size_t, num_keys>;
+    cuco::experimental::static_map<Key,
+                                   T,
+                                   extent_type,
+                                   cuda::thread_scope_device,
+                                   Equal,
+                                   ProbeT,
+                                   AllocatorT,
+                                   StorageT>
+      map{extent_type{}, cuco::empty_key<Key>{-1}, cuco::empty_value<T>{-1}};
+
+    auto ref               = map.ref(cuco::experimental::insert);
+    auto const num_windows = ref.window_extent();
+    STATIC_REQUIRE(static_cast<std::size_t>(num_windows) == gold_extent);
+  }
+
   SECTION("Dynamic extent is evaluated at run time.")
   {
     auto constexpr gold_capacity = 422;  // 211 x 2

diff --git a/tests/static_set/capacity_test.cu b/tests/static_set/capacity_test.cu
@@ -60,6 +60,20 @@ TEST_CASE("Static set capacity", "")
 
   constexpr std::size_t num_keys{400};
 
+  SECTION("Static window extent can be evaluated at build time.")
+  {
+    std::size_t constexpr gold_extent = 211;
+
+    using extent_type = cuco::experimental::extent<std::size_t, num_keys>;
+    cuco::experimental::
+      static_set<Key, extent_type, cuda::thread_scope_device, Equal, ProbeT, AllocatorT, StorageT>
+        set{extent_type{}, cuco::empty_key<Key>{-1}};
+
+    auto ref               = set.ref(cuco::experimental::insert);
+    auto const num_windows = ref.window_extent();
+    STATIC_REQUIRE(static_cast<std::size_t>(num_windows) == gold_extent);
+  }
+
   SECTION("Dynamic extent is evaluated at run time.")
   {
     auto constexpr gold_capacity = 422;  // 211 x 2