Skip to content

Commit

Permalink
Add UnorderedMapInsertOps for coo2crs (#5877)
Browse files Browse the repository at this point in the history 
* Add UnorderedMapInsertOps for coo2crs

* Apply suggestions from code review

Co-authored-by: Daniel Arndt <arndtd@ornl.gov>

* Implement PR feedback
Fix OpenMP test failures
Fix CI build error
Increase size of expected_values

* Create host mirrors for checking expected values
Ensure m_insert_op gets copied too

* Address PR feedback. Depends on #5899.

* Fixups from rebase and clang 15

* cleanup

* Ensure if evaluates at runtime with constexpr

* Do not allow insert ops on sets

* Silence intel error 869

* Reduce unit test runtime

* Update containers/src/Kokkos_UnorderedMap.hpp

Co-authored-by: Daniel Arndt <arndtd@ornl.gov>

* Fix typo

---------

Co-authored-by: Daniel Arndt <arndtd@ornl.gov>
  • Loading branch information
@e10harvey @masterleinad
e10harvey and masterleinad committed Mar 30, 2023
1 parent 0476985 commit fdb089b
Show file tree
Hide file tree
Showing 2 changed files with 192 additions and 25 deletions.
60 changes: 52 additions & 8 deletions containers/src/Kokkos_UnorderedMap.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,36 @@ class UnorderedMapInsertResult {
uint32_t m_status;
};

/// \class UnorderedMapInsertOpTypes
///
/// \brief Operations applied to the values array upon subsequent insertions.
///
/// The default behavior when a k,v pair already exists in the UnorderedMap is
/// to perform no operation. Alternatively, the caller may select to
/// instantiate the UnorderedMap with the AtomicAdd insert operator such that
/// duplicate keys accumulate values into the given values array entry.
/// \tparam ValueTypeView The UnorderedMap value array type.
/// \tparam ValuesIdxType The index type for lookups in the value array.
///
/// Supported operations:
/// NoOp: the first key inserted stores the associated value.
/// AtomicAdd: duplicate key insertions sum values together.
template <class ValueTypeView, class ValuesIdxType>
struct UnorderedMapInsertOpTypes {
using value_type = typename ValueTypeView::non_const_value_type;
struct NoOp {
KOKKOS_FUNCTION
void op(ValueTypeView, ValuesIdxType, const value_type) const {}
};
struct AtomicAdd {
KOKKOS_FUNCTION
void op(ValueTypeView values, ValuesIdxType values_idx,
const value_type v) const {
Kokkos::atomic_add(values.data() + values_idx, v);
}
};
};

/// \class UnorderedMap
/// \brief Thread-safe, performance-portable lookup table.
///
Expand Down Expand Up @@ -186,7 +216,6 @@ class UnorderedMap {
public:
//! \name Public types and constants
//@{

// key_types
using declared_key_type = Key;
using key_type = std::remove_const_t<declared_key_type>;
Expand Down Expand Up @@ -232,7 +261,6 @@ class UnorderedMap {
UnorderedMap<Key, Value, host_mirror_space, Hasher, EqualTo>;

using histogram_type = Impl::UnorderedMapHistogram<const_map_type>;

//@}

private:
Expand Down Expand Up @@ -263,13 +291,17 @@ class UnorderedMap {
public:
//! \name Public member functions
//@{
using default_op_type =
typename UnorderedMapInsertOpTypes<value_type_view, uint32_t>::NoOp;

/// \brief Constructor
///
/// \param capacity_hint [in] Initial guess of how many unique keys will be
/// inserted into the map \param hash [in] Hasher function for \c Key
/// instances. The
/// default value usually suffices.
/// inserted into the map.
/// \param hash [in] Hasher function for \c Key instances. The
/// default value usually suffices.
/// \param equal_to [in] The operator used for determining if two
/// 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),
Expand Down Expand Up @@ -442,9 +474,18 @@ class UnorderedMap {
/// \param v [in] The corresponding value to attempt to insert. If
/// using this class as a set (with Value = void), then you need not
/// provide this value.
KOKKOS_INLINE_FUNCTION
insert_result insert(key_type const &k,
impl_value_type const &v = impl_value_type()) const {
/// \param insert_op [in] The operator used for combining values if a
/// key already exists. See
/// Kokkos::UnorderedMapInsertOpTypes for more ops.
template <typename InsertOpType = default_op_type>
KOKKOS_INLINE_FUNCTION insert_result
insert(key_type const &k, impl_value_type const &v = impl_value_type(),
[[maybe_unused]] InsertOpType arg_insert_op = InsertOpType()) const {
if constexpr (is_set) {
static_assert(std::is_same_v<InsertOpType, default_op_type>,
"Insert Operations are not supported on sets.");
}

insert_result result;

if (!is_insertable_map || capacity() == 0u ||
Expand Down Expand Up @@ -532,6 +573,9 @@ class UnorderedMap {
}

result.set_existing(curr, free_existing);
if constexpr (!is_set) {
arg_insert_op.op(m_values, curr, v);
}
not_done = false;
}
//------------------------------------------------------------
Expand Down
157 changes: 140 additions & 17 deletions containers/unit_tests/TestUnorderedMap.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,18 +25,31 @@ namespace Test {

namespace Impl {

template <typename MapType, bool Near = false>
template <typename MapType,
typename InsertOp = typename MapType::default_op_type,
bool Near = false, bool CheckValues = false>
struct TestInsert {
using map_type = MapType;
using execution_space = typename map_type::execution_space;
using value_type = uint32_t;

struct ExpectedValues {
unsigned map_idx;
typename map_type::value_type v;
};
using expected_values_type = Kokkos::View<ExpectedValues *, execution_space>;
expected_values_type expected_values;

map_type map;
uint32_t inserts;
uint32_t collisions;
InsertOp insert_op;

TestInsert(map_type arg_map, uint32_t arg_inserts, uint32_t arg_collisions)
: map(arg_map), inserts(arg_inserts), collisions(arg_collisions) {}
: map(arg_map), inserts(arg_inserts), collisions(arg_collisions) {
auto len = map.capacity() > arg_inserts ? map.capacity() : arg_inserts;
expected_values = expected_values_type("ExpectedValues", len);
}

void testit(bool rehash_on_fail = true) {
execution_space().fence();
Expand All @@ -60,6 +73,18 @@ struct TestInsert {
Kokkos::deep_copy(map_h, map);
execution_space().fence();
ASSERT_EQ(map_h.size(), map.size());

if (!rehash_on_fail && CheckValues) {
typename expected_values_type::HostMirror expected_values_h =
create_mirror_view(expected_values);
Kokkos::deep_copy(expected_values_h, expected_values);
for (unsigned i = 0; i < map_h.size(); i++) {
auto map_idx = expected_values_h(i).map_idx;
if (map_idx != static_cast<unsigned>(~0)) {
ASSERT_EQ(expected_values_h(map_idx).v, map_h.value_at(map_idx));
}
}
}
}

KOKKOS_INLINE_FUNCTION
Expand All @@ -70,10 +95,47 @@ struct TestInsert {
failed_count += count;
}

template <typename UmapOpType = InsertOp>
KOKKOS_FORCEINLINE_FUNCTION bool is_op_noop() const {
using vt = typename map_type::value_type;
using Device = typename map_type::device_type;
using UmapOpTypeArg1 = Kokkos::View<
std::remove_const_t<std::conditional_t<std::is_void_v<vt>, int, vt>> *,
Device>;
return std::is_base_of_v<
InsertOp, typename Kokkos::UnorderedMapInsertOpTypes<UmapOpTypeArg1,
uint32_t>::NoOp>;
}

template <typename UmapOpType = InsertOp>
KOKKOS_FORCEINLINE_FUNCTION bool is_op_atomic_add() const {
using vt = typename map_type::value_type;
using Device = typename map_type::device_type;
using UmapOpTypeArg1 = Kokkos::View<
std::remove_const_t<std::conditional_t<std::is_void_v<vt>, int, vt>> *,
Device>;
return std::is_base_of_v<UmapOpType,
typename Kokkos::UnorderedMapInsertOpTypes<
UmapOpTypeArg1, uint32_t>::AtomicAdd>;
}

KOKKOS_INLINE_FUNCTION
void operator()(uint32_t i, value_type &failed_count) const {
const uint32_t key = Near ? i / collisions : i % (inserts / collisions);
if (map.insert(key, i).failed()) ++failed_count;
auto ret = map.insert(key, i, insert_op);
if (ret.failed()) {
++failed_count;
expected_values(i).map_idx = static_cast<unsigned>(~0);
} else if (CheckValues) {
auto map_idx = map.find(key);
expected_values(map_idx).map_idx = map_idx;
auto ptr = expected_values.data();
if (is_op_atomic_add()) {
Kokkos::atomic_add(&((ptr + map_idx)[0].v), i);
} else if (ret.success() && is_op_noop()) {
Kokkos::atomic_store(&((ptr + map_idx)[0].v), i);
}
}
}
};

Expand Down Expand Up @@ -154,26 +216,26 @@ struct TestFind {
// MSVC reports a syntax error for this test.
// WORKAROUND MSVC
#ifndef _WIN32
template <typename Device>
template <typename Device, class map_type, class const_map_type,
class insert_op_type, bool check_values = false>
void test_insert(uint32_t num_nodes, uint32_t num_inserts,
uint32_t num_duplicates, bool near) {
using map_type = Kokkos::UnorderedMap<uint32_t, uint32_t, Device>;
using const_map_type =
Kokkos::UnorderedMap<const uint32_t, const uint32_t, Device>;

const uint32_t expected_inserts =
(num_inserts + num_duplicates - 1u) / num_duplicates;
typename map_type::size_type arg_capacity_hint = 0;
typename map_type::hasher_type arg_hasher;
typename map_type::equal_to_type arg_equal_to;

map_type map;
map_type map(arg_capacity_hint, arg_hasher, arg_equal_to);
map.rehash(num_nodes, false);

if (near) {
Impl::TestInsert<map_type, true> test_insert(map, num_inserts,
num_duplicates);
Impl::TestInsert<map_type, insert_op_type, true, check_values> test_insert(
map, num_inserts, num_duplicates);
test_insert.testit();
} else {
Impl::TestInsert<map_type, false> test_insert(map, num_inserts,
num_duplicates);
Impl::TestInsert<map_type, insert_op_type, false, check_values> test_insert(
map, num_inserts, num_duplicates);
test_insert.testit();
}

Expand All @@ -191,8 +253,7 @@ void test_insert(uint32_t num_nodes, uint32_t num_inserts,

{
uint32_t find_errors = 0;
Impl::TestFind<const_map_type> test_find(map, num_inserts,
num_duplicates);
Impl::TestFind<map_type> test_find(map, num_inserts, num_duplicates);
test_find.testit(find_errors);
EXPECT_EQ(0u, find_errors);
}
Expand All @@ -204,6 +265,64 @@ void test_insert(uint32_t num_nodes, uint32_t num_inserts,
map.end_erase();
EXPECT_EQ(0u, map.size());
}

// Check the values from the insert operation
{
Impl::TestInsert<map_type, insert_op_type, true> test_insert(
map, num_inserts, num_duplicates);
test_insert.testit(false);
}
}

template <typename Device>
void test_inserts(uint32_t num_nodes, uint32_t num_inserts,
uint32_t num_duplicates, bool near) {
using key_type = uint32_t;
using value_type = uint32_t;
using value_view_type = Kokkos::View<value_type *, Device>;
using size_type = uint32_t;
using hasher_type = typename Kokkos::pod_hash<key_type>;
using equal_to_type = typename Kokkos::pod_equal_to<key_type>;

using map_op_type =
Kokkos::UnorderedMapInsertOpTypes<value_view_type, size_type>;
using noop_type = typename map_op_type::NoOp;

using map_type = Kokkos::UnorderedMap<key_type, value_type, Device,
hasher_type, equal_to_type>;
using const_map_type =
Kokkos::UnorderedMap<const key_type, const value_type, Device,
hasher_type, equal_to_type>;

test_insert<Device, map_type, const_map_type, noop_type>(
num_nodes, num_inserts, num_duplicates, near);
}

template <typename Device>
void test_all_insert_ops(uint32_t num_nodes, uint32_t num_inserts,
uint32_t num_duplicates, bool near) {
using key_type = uint32_t;
using value_type = uint32_t;
using value_view_type = Kokkos::View<value_type *, Device>;
using size_type = uint32_t;
using hasher_type = typename Kokkos::pod_hash<key_type>;
using equal_to_type = typename Kokkos::pod_equal_to<key_type>;

using map_op_type =
Kokkos::UnorderedMapInsertOpTypes<value_view_type, size_type>;
using noop_type = typename map_op_type::NoOp;
using atomic_add_type = typename map_op_type::AtomicAdd;

using map_type = Kokkos::UnorderedMap<key_type, value_type, Device,
hasher_type, equal_to_type>;
using const_map_type =
Kokkos::UnorderedMap<const key_type, const value_type, Device,
hasher_type, equal_to_type>;

test_insert<Device, map_type, const_map_type, noop_type, true>(
num_nodes, num_inserts, num_duplicates, near);
test_insert<Device, map_type, const_map_type, atomic_add_type, true>(
num_nodes, num_inserts, num_duplicates, near);
}
#endif

Expand Down Expand Up @@ -279,8 +398,12 @@ TEST(TEST_CATEGORY, UnorderedMap_insert) {
}
#endif
for (int i = 0; i < 500; ++i) {
test_insert<TEST_EXECSPACE>(100000, 90000, 100, true);
test_insert<TEST_EXECSPACE>(100000, 90000, 100, false);
test_inserts<TEST_EXECSPACE>(100000, 90000, 100, true);
test_inserts<TEST_EXECSPACE>(100000, 90000, 100, false);
}
for (int i = 0; i < 5; ++i) {
test_all_insert_ops<TEST_EXECSPACE>(1000, 900, 10, true);
test_all_insert_ops<TEST_EXECSPACE>(1000, 900, 10, false);
}
}
#endif
Expand Down

0 comments on commit fdb089b

Please sign in to comment.