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