feat: Add clear function

NVIDIA · Jul 29, 2023 · 966da5c · 966da5c
1 parent f158748
commit 966da5c
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/include/cuco/detail/static_map.inl b/include/cuco/detail/static_map.inl
@@ -80,6 +80,16 @@ static_map<Key, Value, Scope, Allocator>::static_map(
       slots_, empty_key_sentinel_, empty_value_sentinel_, capacity_);
 }
 
+template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
+void static_map<Key, Value, Scope, Allocator>::clear(cudaStream_t stream) {
+  auto constexpr block_size = 256;
+  auto constexpr stride     = 4;
+  auto const grid_size      = (capacity_ + stride * block_size - 1) / (stride * block_size);
+  detail::initialize<block_size, atomic_key_type, atomic_mapped_type>
+    <<<grid_size, block_size, 0, stream>>>(
+      slots_, empty_key_sentinel_, empty_value_sentinel_, capacity_);
+}
+
 template <typename Key, typename Value, cuda::thread_scope Scope, typename Allocator>
 static_map<Key, Value, Scope, Allocator>::~static_map()
 {

diff --git a/include/cuco/static_map.cuh b/include/cuco/static_map.cuh
@@ -226,6 +226,8 @@ class static_map {
              Allocator const& alloc = Allocator{},
              cudaStream_t stream    = 0);
 
+  void clear(cudaStream_t stream    = 0);
+
   /**
    * @brief Destroys the map and frees its contents.
    *