unorderedmap: adding missing 'space' argument for map erasure

uliegecsm · Dec 1, 2023 · 4d9517d · 4d9517d
1 parent ce7b6a6
commit 4d9517d
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 31 deletions.
diff --git a/containers/src/Kokkos_UnorderedMap.hpp b/containers/src/Kokkos_UnorderedMap.hpp
@@ -414,23 +414,33 @@ class UnorderedMap {
   /// This is <i>not</i> a device function; it may <i>not</i> be
   /// called in a parallel kernel.
   bool rehash(size_type requested_capacity = 0) {
+    return rehash(execution_space{}, requested_capacity);
+  }
+
+  bool rehash(const execution_space &space, size_type requested_capacity = 0) {
     const bool bounded_insert = (capacity() == 0) || (size() == 0u);
-    return rehash(requested_capacity, bounded_insert);
+    return rehash(space, requested_capacity, bounded_insert);
   }
 
   bool rehash(size_type requested_capacity, bool bounded_insert) {
-    if (!is_insertable_map) return false;
+    return rehash(execution_space{}, requested_capacity, bounded_insert);
+  }
+
+  bool rehash(const execution_space &space, size_type requested_capacity,
+              bool bounded_insert) {
+    if constexpr (!is_insertable_map) return false;
 
     const size_type curr_size = size();
     requested_capacity =
         (requested_capacity < curr_size) ? curr_size : requested_capacity;
 
-    insertable_map_type tmp(requested_capacity, m_hasher, m_equal_to);
+    insertable_map_type tmp(Kokkos::view_alloc(space, "rehash"), requested_capacity,
+                            m_hasher, m_equal_to);
 
     if (curr_size) {
       tmp.m_bounded_insert = false;
-      Impl::UnorderedMapRehash<insertable_map_type> f(tmp, *this);
-      f.apply();
+      Impl::UnorderedMapRehash<insertable_map_type> f{tmp, *this};
+      f.apply(space);
     }
     tmp.m_bounded_insert = bounded_insert;
 
@@ -467,25 +477,33 @@ class UnorderedMap {
   }
 
   bool begin_erase() {
+    execution_space().fence(
+        "Kokkos::UnorderedMap::begin_erase: fence before setting erasable "
+        "flag");
+    begin_erase(execution_space{});
+  }
+
+  bool begin_erase(const execution_space &space) {
     bool result = !erasable();
     if (is_insertable_map && result) {
-      execution_space().fence(
-          "Kokkos::UnorderedMap::begin_erase: fence before setting erasable "
-          "flag");
       set_flag(erasable_idx);
     }
     return result;
   }
 
   bool end_erase() {
+    execution_space().fence(
+        "Kokkos::UnorderedMap::end_erase: fence before erasing");
+    end_erase(execution_space{});
+    execution_space().fence(
+        "Kokkos::UnorderedMap::end_erase: fence after erasing");
+  }
+
+  bool end_erase(const execution_space &space) {
     bool result = erasable();
     if (is_insertable_map && result) {
-      execution_space().fence(
-          "Kokkos::UnorderedMap::end_erase: fence before erasing");
-      Impl::UnorderedMapErase<declared_map_type> f(*this);
-      f.apply();
-      execution_space().fence(
-          "Kokkos::UnorderedMap::end_erase: fence after erasing");
+      Impl::UnorderedMapErase<declared_map_type> f{*this};
+      f.apply(space);
       reset_flag(erasable_idx);
     }
     return result;

diff --git a/containers/src/impl/Kokkos_UnorderedMap_impl.hpp b/containers/src/impl/Kokkos_UnorderedMap_impl.hpp
@@ -47,16 +47,15 @@ struct UnorderedMapRehash {
   using const_map_type  = typename map_type::const_map_type;
   using execution_space = typename map_type::execution_space;
   using size_type       = typename map_type::size_type;
+  using policy_type =
+      Kokkos::RangePolicy<execution_space, IndexType<size_type>>;
 
   map_type m_dst;
   const_map_type m_src;
 
-  UnorderedMapRehash(map_type const& dst, const_map_type const& src)
-      : m_dst(dst), m_src(src) {}
-
-  void apply() const {
-    parallel_for("Kokkos::Impl::UnorderedMapRehash::apply", m_src.capacity(),
-                 *this);
+  void apply(const execution_space& space) const {
+    parallel_for("Kokkos::Impl::UnorderedMapRehash::apply",
+                 policy_type(space, 0, m_src.capacity()), *this);
   }
 
   KOKKOS_INLINE_FUNCTION
@@ -76,19 +75,19 @@ struct UnorderedMapErase {
   using size_type       = typename map_type::size_type;
   using key_type        = typename map_type::key_type;
   using value_type      = typename map_type::impl_value_type;
+  using policy_type =
+      Kokkos::RangePolicy<execution_space, IndexType<size_type>>;
 
   map_type m_map;
 
-  UnorderedMapErase(map_type const& map) : m_map(map) {}
-
-  void apply() const {
+  void apply(const execution_space& space) const {
     parallel_for("Kokkos::Impl::UnorderedMapErase::apply",
-                 m_map.m_hash_lists.extent(0), *this);
+                 policy_type(space, 0, m_map.m_hash_lists.extent(0)), *this);
   }
 
   KOKKOS_INLINE_FUNCTION
   void operator()(size_type i) const {
-    const size_type invalid_index = map_type::invalid_index;
+    constexpr size_type invalid_index = map_type::invalid_index;
 
     size_type curr = m_map.m_hash_lists(i);
     size_type next = invalid_index;
@@ -188,7 +187,7 @@ struct UnorderedMapHistogram {
 
   KOKKOS_INLINE_FUNCTION
   void operator()(size_type i) const {
-    const size_type invalid_index = map_type::invalid_index;
+    constexpr size_type invalid_index = map_type::invalid_index;
 
     uint32_t length     = 0;
     size_type min_index = ~0u, max_index = 0;
@@ -223,16 +222,14 @@ struct UnorderedMapPrint {
 
   map_type m_map;
 
-  UnorderedMapPrint(map_type const& map) : m_map(map) {}
-
   void apply() {
     parallel_for("Kokkos::Impl::UnorderedMapPrint::apply",
                  m_map.m_hash_lists.extent(0), *this);
   }
 
   KOKKOS_INLINE_FUNCTION
   void operator()(size_type i) const {
-    const size_type invalid_index = map_type::invalid_index;
+    constexpr size_type invalid_index = map_type::invalid_index;
 
     uint32_t list = m_map.m_hash_lists(i);
     for (size_type curr = list, ii = 0; curr != invalid_index;

diff --git a/containers/unit_tests/TestUnorderedMap.hpp b/containers/unit_tests/TestUnorderedMap.hpp
@@ -246,7 +246,7 @@ void test_insert(uint32_t num_nodes, uint32_t num_inserts,
 
   const bool print_list = false;
   if (print_list) {
-    Kokkos::Impl::UnorderedMapPrint<map_type> f(map);
+    Kokkos::Impl::UnorderedMapPrint<map_type> f{map};
     f.apply();
   }
 

diff --git a/core/src/Kokkos_Macros.hpp b/core/src/Kokkos_Macros.hpp
@@ -532,7 +532,11 @@ static constexpr bool kokkos_omp_on_host() { return false; }
 #define KOKKOS_ENABLE_CUDA_LDG_INTRINSIC
 #endif
 
-#define KOKKOS_INVALID_INDEX (~std::size_t(0))
+//! Invalid index of a given @p __type__.
+#define KOKKOS_INVALID_INDEX_TYPE(__type__) ~static_cast<__type__>(0)
+
+//! Invalid index.
+#define KOKKOS_INVALID_INDEX KOKKOS_INVALID_INDEX_TYPE(std::size_t)
 
 #define KOKKOS_IMPL_CTOR_DEFAULT_ARG KOKKOS_INVALID_INDEX