diff --git a/containers/src/Kokkos_UnorderedMap.hpp b/containers/src/Kokkos_UnorderedMap.hpp index e45ff61da19..fd8e94c5c08 100644 --- a/containers/src/Kokkos_UnorderedMap.hpp +++ b/containers/src/Kokkos_UnorderedMap.hpp @@ -414,23 +414,33 @@ class UnorderedMap { /// This is not a device function; it may not 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 f(tmp, *this); - f.apply(); + Impl::UnorderedMapRehash 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 f(*this); - f.apply(); - execution_space().fence( - "Kokkos::UnorderedMap::end_erase: fence after erasing"); + Impl::UnorderedMapErase 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 index a979ee40d8c..e9f8d7bdfa9 100644 --- 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>; 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>; 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,8 +222,6 @@ 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); @@ -232,7 +229,7 @@ struct UnorderedMapPrint { 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 index f63f1c6afe3..c80c1661bf6 100644 --- 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 f(map); + Kokkos::Impl::UnorderedMapPrint f{map}; f.apply(); } diff --git a/core/src/Kokkos_Macros.hpp b/core/src/Kokkos_Macros.hpp index a77e50b65b3..69dc9df4959 100644 --- 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