UnorderedMap(space instance): proposal for kokkos#6067

containers/src/Kokkos_Bitset.hpp

@@ -28,6 +28,23 @@
 
 namespace Kokkos {
 
+namespace Impl {
+//! Either append to the label if the property already exists, or set it.
+template <typename... P>
+auto with_updated_label(const ViewCtorProp<P...>& view_ctor_prop,
+                        const std::string& label) {
+  //! If the label property is already set, append. Otherwise, set label.
+  if constexpr (ViewCtorProp<P...>::has_label) {
+    auto new_ctor_props(view_ctor_prop);
+    static_cast<ViewCtorProp<void, std::string>&>(new_ctor_props)
+        .value.append(label);
+    return new_ctor_props;
+  } else {
+    return Impl::with_properties_if_unset(view_ctor_prop, label);
+  }
+}
+}  // namespace Impl
+
 template <typename Device = Kokkos::DefaultExecutionSpace>
 class Bitset;
 
@@ -70,13 +87,32 @@ class Bitset {
     block_shift = Kokkos::Impl::integral_power_of_two(block_size)
   };
 
+  //! Type of @ref m_blocks.
+  using block_view_type = View<unsigned*, Device, MemoryTraits<RandomAccess>>;
+
  public:
   /// constructor
   /// arg_size := number of bit in set
-  Bitset(unsigned arg_size = 0u)
-      : m_size(arg_size),
-        m_last_block_mask(0u),
-        m_blocks("Bitset", ((m_size + block_mask) >> block_shift)) {
+  Bitset(unsigned arg_size = 0u) : Bitset(Kokkos::view_alloc(), arg_size) {}
+
+  template <class... P>
+  Bitset(const Impl::ViewCtorProp<P...>& arg_prop, unsigned arg_size)
+      : m_size(arg_size), m_last_block_mask(0u) {
+    //! Ensure that allocation properties are consistent.
+    using alloc_prop_t = std::decay_t<decltype(arg_prop)>;
+    static_assert(alloc_prop_t::initialize,
+                  "Allocation property 'initialize' should be true.");
+    static_assert(
+        !alloc_prop_t::has_pointer,
+        "Allocation properties should not contain the 'pointer' property.");
+
+    //! Update 'label' property and allocate.
+    const auto prop_copy = Kokkos::Impl::with_updated_label(
+        Impl::with_properties_if_unset(arg_prop, std::string("Bitset")),
+        " - blocks");
+    m_blocks =
+        block_view_type(prop_copy, ((m_size + block_mask) >> block_shift));
+
     for (int i = 0, end = static_cast<int>(m_size & block_mask); i < end; ++i) {
       m_last_block_mask |= 1u << i;
     }
@@ -105,7 +141,7 @@ class Bitset {
   /// number of bits which are set to 1
   /// can only be called from the host
   unsigned count() const {
-    Impl::BitsetCount<Bitset<Device> > f(*this);
+    Impl::BitsetCount<Bitset<Device>> f(*this);
     return f.apply();
   }
 
@@ -275,7 +311,7 @@ class Bitset {
  private:
   unsigned m_size;
   unsigned m_last_block_mask;
-  View<unsigned*, Device, MemoryTraits<RandomAccess> > m_blocks;
+  block_view_type m_blocks;
 
  private:
   template <typename DDevice>
@@ -302,6 +338,7 @@ class ConstBitset {
  public:
   using execution_space = typename Device::execution_space;
   using size_type       = unsigned int;
+  using block_view_type = typename Bitset<Device>::block_view_type::const_type;
 
  private:
   enum { block_size = static_cast<unsigned>(sizeof(unsigned) * CHAR_BIT) };
@@ -340,7 +377,7 @@ class ConstBitset {
   unsigned size() const { return m_size; }
 
   unsigned count() const {
-    Impl::BitsetCount<ConstBitset<Device> > f(*this);
+    Impl::BitsetCount<ConstBitset<Device>> f(*this);
     return f.apply();
   }
 
@@ -356,7 +393,7 @@ class ConstBitset {
 
  private:
   unsigned m_size;
-  View<const unsigned*, Device, MemoryTraits<RandomAccess> > m_blocks;
+  block_view_type m_blocks;
 
  private:
   template <typename DDevice>

containers/src/Kokkos_UnorderedMap.hpp

@@ -34,6 +34,7 @@
 
 #include <impl/Kokkos_Traits.hpp>
 #include <impl/Kokkos_UnorderedMap_impl.hpp>
+#include <impl/Kokkos_ViewCtor.hpp>
 
 #include <cstdint>
 
@@ -302,28 +303,76 @@ class UnorderedMap {
   ///                           keys are equal.
   UnorderedMap(size_type capacity_hint = 0, hasher_type hasher = hasher_type(),
                equal_to_type equal_to = equal_to_type())
-      : m_bounded_insert(true),
-        m_hasher(hasher),
-        m_equal_to(equal_to),
-        m_size("UnorderedMap size"),
-        m_available_indexes(calculate_capacity(capacity_hint)),
-        m_hash_lists(view_alloc(WithoutInitializing, "UnorderedMap hash list"),
-                     Impl::find_hash_size(capacity())),
-        m_next_index(view_alloc(WithoutInitializing, "UnorderedMap next index"),
-                     capacity() + 1)  // +1 so that the *_at functions can
-                                      // always return a valid reference
-        ,
-        m_keys("UnorderedMap keys", capacity()),
-        m_values("UnorderedMap values", (is_set ? 0 : capacity())),
-        m_scalars("UnorderedMap scalars") {
+      : UnorderedMap(Kokkos::view_alloc(), capacity_hint, hasher, equal_to) {}
+
+  template <class... P>
+  UnorderedMap(const Impl::ViewCtorProp<P...> &arg_prop,
+               size_type capacity_hint = 0, hasher_type hasher = hasher_type(),
+               equal_to_type equal_to = equal_to_type())
+      : m_bounded_insert(true), m_hasher(hasher), m_equal_to(equal_to) {
     if (!is_insertable_map) {
       Kokkos::Impl::throw_runtime_exception(
           "Cannot construct a non-insertable (i.e. const key_type) "
           "unordered_map");
     }
 
-    Kokkos::deep_copy(m_hash_lists, invalid_index);
-    Kokkos::deep_copy(m_next_index, invalid_index);
+    //! Ensure that allocation properties are consistent.
+    using alloc_prop_t = std::decay_t<decltype(arg_prop)>;
+    static_assert(alloc_prop_t::initialize,
+                  "Allocation property 'initialize' should be true.");
+    static_assert(
+        !alloc_prop_t::has_pointer,
+        "Allocation properties should not contain the 'pointer' property.");
+
+    /// Update allocation properties with 'label' and 'without initializing'
+    /// properties.
+    const auto prop_copy =
+        Impl::with_properties_if_unset(arg_prop, std::string("UnorderedMap"));
+    const auto prop_copy_noinit =
+        Impl::with_properties_if_unset(prop_copy, Kokkos::WithoutInitializing);
+
+    //! Initialize member views.
+    m_size = shared_size_t(Kokkos::view_alloc(
+        Kokkos::DefaultHostExecutionSpace{},
+        Impl::get_property<Impl::LabelTag>(prop_copy) + " - size"));
+
+    m_available_indexes =
+        bitset_type(Kokkos::Impl::with_updated_label(prop_copy, " - bitset"),
+                    calculate_capacity(capacity_hint));
+
+    m_hash_lists = size_type_view(
+        Kokkos::Impl::with_updated_label(prop_copy_noinit, " - hash list"),
+        Impl::find_hash_size(capacity()));
+
+    m_next_index = size_type_view(
+        Kokkos::Impl::with_updated_label(prop_copy_noinit, " - next index"),
+        capacity() + 1);  // +1 so that the *_at functions can always return a
+                          // valid reference
+
+    m_keys = key_type_view(
+        Kokkos::Impl::with_updated_label(prop_copy, " - keys"), capacity());
+
+    m_values = value_type_view(
+        Kokkos::Impl::with_updated_label(prop_copy, " - values"),
+        is_set ? 0 : capacity());
+
+    m_scalars =
+        scalars_view(Kokkos::Impl::with_updated_label(prop_copy, " - scalars"));
+
+    /**
+     * Deep copies should also be done using the space instance if given.
+     * Instead of the if/else we could use the
+     * @c get_property_or_default, but giving even the default execution space
+     * instance will change the behavior of @c deep_copy.
+     */
+    if constexpr (alloc_prop_t::has_execution_space) {
+      const auto &space = Impl::get_property<Impl::ExecutionSpaceTag>(arg_prop);
+      Kokkos::deep_copy(space, m_hash_lists, invalid_index);
+      Kokkos::deep_copy(space, m_next_index, invalid_index);
+    } else {
+      Kokkos::deep_copy(m_hash_lists, invalid_index);
+      Kokkos::deep_copy(m_next_index, invalid_index);
+    }
   }
 
   void reset_failed_insert_flag() { reset_flag(failed_insert_idx); }
@@ -860,7 +909,8 @@ class UnorderedMap {
   bool m_bounded_insert;
   hasher_type m_hasher;
   equal_to_type m_equal_to;
-  View<size_type, HostSpace> m_size;
+  using shared_size_t = View<size_type, Kokkos::DefaultHostExecutionSpace>;
+  shared_size_t m_size;
   bitset_type m_available_indexes;
   size_type_view m_hash_lists;
   size_type_view m_next_index;

containers/unit_tests/TestUnorderedMap.hpp

@@ -429,17 +429,57 @@ TEST(TEST_CATEGORY, UnorderedMap_valid_empty) {
   ASSERT_TRUE(n.is_allocated());
 }
 
+/**
+ * This helper is needed because NVCC does not like extended lambdas
+ * in private member functions.
+ * Google Test bodies are private member functions. So it is incompatible.
+ * See also https://github.com/google/googletest/issues/4104.
+ */
+template <typename map_type>
+struct UnorderedMapInsert {
+  //! Type of range-for policy and its index type.
+  using range_policy_t =
+      Kokkos::RangePolicy<typename map_type::execution_space,
+                          Kokkos::IndexType<unsigned short int>>;
+  using index_t = typename range_policy_t::index_type;
+
+  const map_type m_map;
+
+  //! Ensure shared ownership of @ref m_map.
+  UnorderedMapInsert(map_type map) : m_map(std::move(map)) {}
+
+  //! Insert a single value.
+  template <typename T>
+  void insert_single(const T &arg) const {
+    Kokkos::parallel_for(
+        Kokkos::RangePolicy<typename map_type::execution_space>(0, 1),
+        // NOLINTNEXTLINE(kokkos-implicit-this-capture)
+        KOKKOS_CLASS_LAMBDA(const index_t) { m_map.insert(arg); });
+  }
+
+  //! Insert multiple values.
+  template <typename... Args>
+  void insert(Args &&... args) const {
+    static_assert(sizeof...(Args) > 1, "Prefer the single value version");
+    constexpr size_t size = sizeof...(Args);
+    Kokkos::Array<typename map_type::key_type, size> values{
+        std::forward<Args>(args)...};
+    Kokkos::parallel_for(
+        Kokkos::RangePolicy<typename map_type::execution_space>(0, size),
+        // NOLINTNEXTLINE(kokkos-implicit-this-capture)
+        KOKKOS_CLASS_LAMBDA(const index_t i) { m_map.insert(values[i]); });
+  }
+};
+
 TEST(TEST_CATEGORY, UnorderedMap_clear_zero_size) {
-  using Map =
-      Kokkos::UnorderedMap<int, void, Kokkos::DefaultHostExecutionSpace>;
+  using map_type = Kokkos::UnorderedMap<int, void, TEST_EXECSPACE>;
+
+  map_type m(11);
 
-  Map m(11);
   ASSERT_EQ(0u, m.size());
 
-  m.insert(2);
-  m.insert(3);
-  m.insert(5);
-  m.insert(7);
+  UnorderedMapInsert<map_type>(m).insert(2, 3, 5, 7);
+
   ASSERT_EQ(4u, m.size());
   m.rehash(0);
   ASSERT_EQ(128u, m.capacity());
@@ -450,19 +490,22 @@ TEST(TEST_CATEGORY, UnorderedMap_clear_zero_size) {
 }
 
 TEST(TEST_CATEGORY, UnorderedMap_consistent_size) {
-  using Map =
-      Kokkos::UnorderedMap<int, void, Kokkos::DefaultHostExecutionSpace>;
+  using map_type = Kokkos::UnorderedMap<int, void, TEST_EXECSPACE>;
+
+  map_type m(11);
+  UnorderedMapInsert<map_type> inserter(m);
+
+  inserter.insert_single(7);
 
-  Map m(11);
-  m.insert(7);
-  ;
   ASSERT_EQ(1u, m.size());
 
   {
-    auto m2 = m;
-    m2.insert(2);
+    auto m_copy = m;
+    UnorderedMapInsert<decltype(m_copy)> inserter_copy(m_copy);
+    inserter_copy.insert_single(2);
     // This line triggers modified flags to be cleared in both m and m2
-    [[maybe_unused]] auto sz = m2.size();
+    const auto sz = m_copy.size();
+    ASSERT_EQ(2u, sz);
   }
 
   ASSERT_EQ(2u, m.size());
@@ -507,6 +550,18 @@ TEST(TEST_CATEGORY, UnorderedMap_lambda_capturable) {
 }
 #endif
 
+/**
+ * @test This test ensures that an @ref UnorderedMap can be built
+ *       with an execution space instance (using @ref view_alloc).
+ */
+TEST(TEST_CATEGORY, UnorderedMap_constructor_view_alloc) {
+  using map_type = Kokkos::UnorderedMap<size_t, void, TEST_EXECSPACE>;
+  map_type map(Kokkos::view_alloc(TEST_EXECSPACE{}, "test umap"), 150);
+  ASSERT_EQ(map.size(), 0u);
+  ASSERT_GE(map.capacity(), 150u);
+  ASSERT_TRUE(map.is_allocated());
+}
+
 }  // namespace Test
 
 #endif  // KOKKOS_TEST_UNORDERED_MAP_HPP