Merge pull request #5983 from masterleinad/fix_unordered_map_m_size

UnorderedMap: Fix mutable m_size corruption bug
sandialabs · Mar 15, 2023 · 65bf47c · 65bf47c
2 parents 0eeb3a4 + 3cb200c
commit 65bf47c
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 6 deletions.
diff --git a/containers/src/Kokkos_UnorderedMap.hpp b/containers/src/Kokkos_UnorderedMap.hpp
@@ -275,7 +275,7 @@ class UnorderedMap {
       : m_bounded_insert(true),
         m_hasher(hasher),
         m_equal_to(equal_to),
-        m_size(),
+        m_size(std::make_shared<size_type>()),
         m_available_indexes(calculate_capacity(capacity_hint)),
         m_hash_lists(view_alloc(WithoutInitializing, "UnorderedMap hash list"),
                      Impl::find_hash_size(capacity())),
@@ -315,7 +315,7 @@ class UnorderedMap {
       Kokkos::deep_copy(m_keys, tmp);
     }
     Kokkos::deep_copy(m_scalars, 0);
-    m_size = 0;
+    *m_size = 0;
   }
 
   KOKKOS_INLINE_FUNCTION constexpr bool is_allocated() const {
@@ -369,10 +369,10 @@ class UnorderedMap {
   size_type size() const {
     if (capacity() == 0u) return 0u;
     if (modified()) {
-      m_size = m_available_indexes.count();
+      *m_size = m_available_indexes.count();
       reset_flag(modified_idx);
     }
-    return m_size;
+    return *m_size;
   }
 
   /// \brief The current number of failed insert() calls.
@@ -725,7 +725,7 @@ class UnorderedMap {
       tmp.m_bounded_insert    = src.m_bounded_insert;
       tmp.m_hasher            = src.m_hasher;
       tmp.m_equal_to          = src.m_equal_to;
-      tmp.m_size              = src.size();
+      tmp.m_size              = src.m_size;
       tmp.m_available_indexes = bitset_type(src.capacity());
       tmp.m_hash_lists        = size_type_view(
           view_alloc(WithoutInitializing, "UnorderedMap hash list"),
@@ -818,7 +818,7 @@ class UnorderedMap {
   bool m_bounded_insert;
   hasher_type m_hasher;
   equal_to_type m_equal_to;
-  mutable size_type m_size;
+  std::shared_ptr<size_type> m_size;
   bitset_type m_available_indexes;
   size_type_view m_hash_lists;
   size_type_view m_next_index;

diff --git a/containers/unit_tests/TestUnorderedMap.hpp b/containers/unit_tests/TestUnorderedMap.hpp
@@ -54,7 +54,12 @@ struct TestInsert {
       }
     } while (rehash_on_fail && failed_count > 0u);
 
+    // Trigger the m_size mutable bug.
+    typename map_type::HostMirror map_h;
     execution_space().fence();
+    Kokkos::deep_copy(map_h, map);
+    execution_space().fence();
+    ASSERT_EQ(map_h.size(), map.size());
   }
 
   KOKKOS_INLINE_FUNCTION
@@ -328,6 +333,25 @@ TEST(TEST_CATEGORY, UnorderedMap_clear_zero_size) {
   ASSERT_EQ(0u, m.size());
 }
 
+TEST(TEST_CATEGORY, UnorderedMap_consistent_size) {
+  using Map =
+      Kokkos::UnorderedMap<int, void, Kokkos::DefaultHostExecutionSpace>;
+
+  Map m(11);
+  m.insert(7);
+  ;
+  ASSERT_EQ(1u, m.size());
+
+  {
+    auto m2 = m;
+    m2.insert(2);
+    // This line triggers modified flags to be cleared in both m and m2
+    [[maybe_unused]] auto sz = m2.size();
+  }
+
+  ASSERT_EQ(2u, m.size());
+}
+
 }  // namespace Test
 
 #endif  // KOKKOS_TEST_UNORDERED_MAP_HPP