Merge pull request #251 from DARMA-tasking/230-kokkos-unordered-map-s…

…erializer

230 implement Kokkos::UnorderedMap serializer

src/checkpoint/checkpoint.h

@@ -65,6 +65,8 @@
 #include "checkpoint/container/unique_ptr_serialize.h"
 #include "checkpoint/container/view_serialize.h"
 
+#include "checkpoint/container/kokkos_unordered_map_serialize.h"
+
 #include "checkpoint/checkpoint_api.h"
 #include "checkpoint/checkpoint_api.impl.h"
 

src/checkpoint/container/container_serialize.h

@@ -61,6 +61,16 @@ serializeContainerSize(Serializer& s, ContainerT& cont) {
   return cont_size;
 }
 
+template <typename Serializer, typename ContainerT>
+inline typename ContainerT::size_type
+serializeContainerCapacity(Serializer& s, ContainerT& cont) {
+  typename ContainerT::size_type cont_capacity = cont.capacity();
+  if (!s.isFootprinting()) {
+    s | cont_capacity;
+  }
+  return cont_capacity;
+}
+
 template <typename Serializer, typename ContainerT>
 inline void serializeContainerElems(Serializer& s, ContainerT& cont) {
   for (auto&& elm : cont) {

src/checkpoint/container/kokkos_unordered_map_serialize.h

@@ -0,0 +1,161 @@
+/*
+//@HEADER
+// *****************************************************************************
+//
+//                      kokkos_unordered_map_serialize.h
+//                 DARMA/checkpoint => Serialization Library
+//
+// Copyright 2019 National Technology & Engineering Solutions of Sandia, LLC
+// (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S.
+// Government retains certain rights in this software.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// * Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+//
+// * Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// * Neither the name of the copyright holder nor the names of its
+//   contributors may be used to endorse or promote products derived from this
+//   software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+// POSSIBILITY OF SUCH DAMAGE.
+//
+// Questions? Contact [email protected]
+//
+// *****************************************************************************
+//@HEADER
+*/
+
+#if !defined INCLUDED_CHECKPOINT_CONTAINER_KOKKOS_UNORDERED_MAP_SERIALIZE_H
+#define INCLUDED_CHECKPOINT_CONTAINER_KOKKOS_UNORDERED_MAP_SERIALIZE_H
+
+#include "checkpoint/common.h"
+#include "checkpoint/serializers/serializers_headers.h"
+#include "checkpoint/dispatch/allocator.h"
+#include "checkpoint/dispatch/dispatch.h"
+#include "checkpoint/dispatch/reconstructor.h"
+
+#if KOKKOS_ENABLED_CHECKPOINT
+
+#include <Kokkos_UnorderedMap.hpp>
+
+namespace checkpoint {
+
+template <
+  typename Serializer, typename Key, typename Value, typename Device,
+  typename Hasher, typename EqualTo
+>
+void serializeKokkosUnorderedMapElems(
+  Serializer& s,
+  Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>& map
+) {
+  using mapSizeType =
+    typename Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>::size_type;
+
+  for (mapSizeType i = 0; i < map.capacity(); i++) {
+    Key keyAtI = map.key_at(i);
+    if (map.exists(keyAtI)) {
+      Value val = map.value_at(map.find(keyAtI));
+
+      s | keyAtI;
+      s | val;
+    }
+  }
+}
+
+template <
+  typename Serializer, typename Key, typename Value, typename Device,
+  typename Hasher, typename EqualTo
+>
+void deserializeInsertElems(
+  Serializer& s, Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>& map,
+  typename Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>::size_type
+    map_size,
+  typename Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>::size_type
+    map_capacity
+) {
+  using KeyAlloc = dispatch::Allocator<Key>;
+  using ValueAlloc = dispatch::Allocator<Value>;
+  using KeyReconstructor =
+    dispatch::Reconstructor<typename dispatch::CleanType<Key>::CleanT>;
+  using ValReconstructor =
+    dispatch::Reconstructor<typename dispatch::CleanType<Value>::CleanT>;
+
+  // resize unordered map
+  map.rehash(map_capacity);
+
+  KeyAlloc keyAllocated;
+  ValueAlloc valAllocated;
+  for (SerialSizeType i = 0; i < map_size; i++) {
+    auto* key = KeyReconstructor::construct(keyAllocated.buf);
+    auto* val = ValReconstructor::construct(valAllocated.buf);
+
+    s | *key;
+    s | *val;
+
+    map.insert(std::move(*key), std::move(*val));
+  }
+}
+
+template <
+  typename SerializerT, typename Key, typename Value, typename Device,
+  typename Hasher, typename EqualTo
+>
+typename std::enable_if_t<
+  not std::is_same<SerializerT, checkpoint::Footprinter>::value, void
+> serialize(
+  SerializerT& s,
+  Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>& map
+) {
+  using UnorderedMapType =
+    Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>;
+
+  typename UnorderedMapType::size_type size = serializeContainerSize(s, map);
+  typename UnorderedMapType::size_type capacity = serializeContainerCapacity(s, map);
+
+  if (s.isUnpacking()) {
+    deserializeInsertElems(s, map, size, capacity);
+  } else {
+    serializeKokkosUnorderedMapElems(s, map);
+  }
+}
+
+template <
+  typename SerializerT, typename Key, typename Value, typename Device,
+  typename Hasher, typename EqualTo
+>
+typename std::enable_if_t<
+  std::is_same<SerializerT, checkpoint::Footprinter>::value, void
+> serialize(
+  SerializerT& s,
+  Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>& map
+) {
+  using UnorderedMapType =
+    Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>;
+
+  typename UnorderedMapType::size_type size = serializeContainerSize(s, map);
+  typename UnorderedMapType::size_type capacity = serializeContainerCapacity(s, map);
+
+  s.addBytes(capacity * sizeof(Key) + size * sizeof(Value));
+}
+
+} // namespace checkpoint
+
+#endif /*KOKKOS_ENABLED_CHECKPOINT*/
+
+#endif /*INCLUDED_CHECKPOINT_CONTAINER_KOKKOS_UNORDERED_MAP_SERIALIZE_H*/

tests/unit/test_footprinter.cc

@@ -612,4 +612,32 @@ TEST_F(TestFootprinter, test_no_serialize) {
     sizeof(m) + m.size() * (sizeof(p) + sizeof(p.first) + sizeof(p.second))
   );
 }
+
+#if KOKKOS_ENABLED_CHECKPOINT
+TEST_F(TestFootprinter, test_kokkos_unordered_map) {
+  // empty map
+  {
+    auto mapEmpty = Kokkos::UnorderedMap<int, int>();
+
+    auto expected_size =
+      sizeof(mapEmpty) + mapEmpty.capacity() * sizeof(int);
+
+    EXPECT_EQ(checkpoint::getMemoryFootprint(mapEmpty), expected_size);
+  }
+
+  // map with some elements
+  {
+    auto mapIntLong = Kokkos::UnorderedMap<int, long>(3);
+    mapIntLong.insert(4, 50);
+    mapIntLong.insert(5, 60);
+    mapIntLong.insert(6, 70);
+
+    auto expected_size =
+      sizeof(mapIntLong) + mapIntLong.capacity() * sizeof(int) + 3 * sizeof(long);
+
+    EXPECT_EQ(checkpoint::getMemoryFootprint(mapIntLong), expected_size);
+  }
+}
+#endif /*KOKKOS_ENABLED_CHECKPOINT*/
+
 }}} // end namespace checkpoint::tests::unit

tests/unit/test_kokkos_serialize_unordered_map.cc

@@ -0,0 +1,112 @@
+/*
+//@HEADER
+// *****************************************************************************
+//
+//                  test_kokkos_serialize_unordered_map.cc
+//                 DARMA/checkpoint => Serialization Library
+//
+// Copyright 2019 National Technology & Engineering Solutions of Sandia, LLC
+// (NTESS). Under the terms of Contract DE-NA0003525 with NTESS, the U.S.
+// Government retains certain rights in this software.
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// * Redistributions of source code must retain the above copyright notice,
+//   this list of conditions and the following disclaimer.
+//
+// * Redistributions in binary form must reproduce the above copyright notice,
+//   this list of conditions and the following disclaimer in the documentation
+//   and/or other materials provided with the distribution.
+//
+// * Neither the name of the copyright holder nor the names of its
+//   contributors may be used to endorse or promote products derived from this
+//   software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+// POSSIBILITY OF SUCH DAMAGE.
+//
+// Questions? Contact [email protected]
+//
+// *****************************************************************************
+//@HEADER
+*/
+
+#if KOKKOS_ENABLED_CHECKPOINT
+
+#include "test_commons.h"
+
+namespace checkpoint { namespace tests { namespace unit {
+
+struct KokkosUnorderedMapTest : virtual testing::Test { };
+
+template <
+  typename Key, typename Value, typename Device, typename Hasher,
+  typename EqualTo
+>
+static void test_kokkos_unordered_map(
+  Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>& refMap
+) {
+  using UnorderedMapType =
+    Kokkos::UnorderedMap<Key, Value, Device, Hasher, EqualTo>;
+
+  ASSERT_FALSE(refMap.failed_insert());
+
+  auto serialized = checkpoint::serialize<UnorderedMapType>(refMap);
+  auto deserialized =
+    checkpoint::deserialize<UnorderedMapType>(serialized->getBuffer());
+  auto& outMap = *deserialized;
+
+  ASSERT_FALSE(outMap.failed_insert());
+  ASSERT_EQ(refMap.size(), outMap.size());
+  ASSERT_LE(refMap.capacity(), outMap.capacity());
+
+  // check all keys and values
+  for (typename UnorderedMapType::size_type i = 0; i < refMap.capacity(); i++) {
+    Key refKey = refMap.key_at(i);
+    if (refMap.exists(refKey)) {
+      ASSERT_TRUE(outMap.exists(refKey));
+
+      Value refVal = refMap.value_at(refMap.find(refKey));
+      Value val = outMap.value_at(outMap.find(refKey));
+      ASSERT_EQ(refVal, val);
+    }
+  }
+}
+
+TEST_F(KokkosUnorderedMapTest, test_kokkos_unordered_map) {
+  int bigSize = 1000;
+  auto mapIntIntBig = Kokkos::UnorderedMap<int, int>(bigSize);
+  for(int i = 0; i < bigSize; i++) {
+    mapIntIntBig.insert(i + bigSize, i * bigSize);
+  }
+  test_kokkos_unordered_map(mapIntIntBig);
+
+  auto mapIntDouble = Kokkos::UnorderedMap<int, double>(2);
+  mapIntDouble.insert(3, 123.34);
+  mapIntDouble.insert(4, 10.1112);
+  test_kokkos_unordered_map(mapIntDouble);
+
+  auto mapShortString = Kokkos::UnorderedMap<short, std::string>(2);
+  mapShortString.insert(5, "123");
+  mapShortString.insert(27, "101112");
+  test_kokkos_unordered_map(mapShortString);
+
+  auto mapUnsignedVector = Kokkos::UnorderedMap<unsigned, std::vector<int>>(2);
+  mapUnsignedVector.insert(11, std::vector<int>{1, 2, 3});
+  mapUnsignedVector.insert(22, std::vector<int>{9, 8, 7});
+  test_kokkos_unordered_map(mapUnsignedVector);
+}
+
+}}} // namespace checkpoint::tests::unit
+
+#endif /*KOKKOS_ENABLED_CHECKPOINT*/