From 04b2e02ddbedfa0840f75cc5c25b942be16648b3 Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Tue, 5 Nov 2024 10:29:44 -0800
Subject: [PATCH] Add multimap size

---
 .../detail/static_multimap/static_multimap.inl    | 15 +++++++++++++++
 include/cuco/static_multimap.cuh                  | 10 ++++++++++
 2 files changed, 25 insertions(+)

diff --git a/include/cuco/detail/static_multimap/static_multimap.inl b/include/cuco/detail/static_multimap/static_multimap.inl
index 2df33b50b..9649df3e5 100644
--- a/include/cuco/detail/static_multimap/static_multimap.inl
+++ b/include/cuco/detail/static_multimap/static_multimap.inl
@@ -456,6 +456,21 @@ void static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator,
   this->impl_->rehash_async(extent, *this, stream);
 }
 
+template <class Key,
+          class T,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::size_type
+static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::size(
+  cuda::stream_ref stream) const
+{
+  return impl_->size(stream);
+}
+
 template <class Key,
           class T,
           class Extent,
diff --git a/include/cuco/static_multimap.cuh b/include/cuco/static_multimap.cuh
index 04fb03211..bf3e5998f 100644
--- a/include/cuco/static_multimap.cuh
+++ b/include/cuco/static_multimap.cuh
@@ -657,6 +657,16 @@ class static_multimap {
    */
   void rehash_async(size_type capacity, cuda::stream_ref stream = {});
 
+  /**
+   * @brief Gets the number of elements in the container.
+   *
+   * @note This function synchronizes the given stream.
+   *
+   * @param stream CUDA stream used to get the number of inserted elements
+   * @return The number of elements in the container
+   */
+  [[nodiscard]] size_type size(cuda::stream_ref stream = {}) const;
+
   /**
    * @brief Gets the maximum number of elements the hash map can hold.
    *