/
tstUnionFind.cpp
114 lines (95 loc) · 4.6 KB
/
tstUnionFind.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
/****************************************************************************
* Copyright (c) 2017-2022 by the ArborX authors *
* All rights reserved. *
* *
* This file is part of the ArborX library. ArborX is *
* distributed under a BSD 3-clause license. For the licensing terms see *
* the LICENSE file in the top-level directory. *
* *
* SPDX-License-Identifier: BSD-3-Clause *
****************************************************************************/
#include "ArborX_EnableDeviceTypes.hpp" // ARBORX_DEVICE_TYPES
#include "ArborX_EnableViewComparison.hpp"
#include <ArborX_DetailsKokkosExtStdAlgorithms.hpp>
#include <ArborX_DetailsUnionFind.hpp>
#include <ArborX_DetailsUtils.hpp> // iota
#include "BoostTest_CUDA_clang_workarounds.hpp"
#include <boost/test/unit_test.hpp>
BOOST_AUTO_TEST_SUITE(UnionFind)
template <typename ExecutionSpace, typename UnionFind>
Kokkos::View<int *, Kokkos::HostSpace>
build_representatives(ExecutionSpace const &space, UnionFind union_find)
{
using MemorySpace = typename UnionFind::memory_space;
auto const n = union_find.size();
Kokkos::View<int *, MemorySpace> representatives(
Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
"Test::representatives"),
n);
Kokkos::View<int *, MemorySpace> map2smallest(
Kokkos::view_alloc(space, Kokkos::WithoutInitializing, "Test::map"), n);
Kokkos::deep_copy(space, map2smallest, INT_MAX);
Kokkos::parallel_for(
"Test::find_representatives",
Kokkos::RangePolicy<ExecutionSpace>(space, 0, n), KOKKOS_LAMBDA(int i) {
auto r = union_find.representative(i);
Kokkos::atomic_min(&map2smallest(r), i);
representatives(i) = r;
});
// We want the representative values to not depend on a specific
// implementation of the union-find (e.g., not relying them being equal to
// the smallest index in the set), so we explicitly remap them.
Kokkos::parallel_for(
"Test::remap_representatives",
Kokkos::RangePolicy<ExecutionSpace>(space, 0, n), KOKKOS_LAMBDA(int i) {
representatives(i) = map2smallest(representatives(i));
});
return Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{},
representatives);
}
template <typename ExecutionSpace, typename UnionFind>
void merge(ExecutionSpace const &space, UnionFind &union_find, int i, int j)
{
Kokkos::parallel_for(
"Test::merge", Kokkos::RangePolicy<ExecutionSpace>(space, 0, 1),
KOKKOS_LAMBDA(int) { union_find.merge(i, j); });
}
#define ARBORX_TEST_UNION_FIND_REPRESENTATIVES(space, union_find, ref) \
BOOST_TEST(build_representatives(space, union_find) == ref, \
boost::test_tools::per_element());
BOOST_AUTO_TEST_CASE_TEMPLATE(union_find, DeviceType, ARBORX_DEVICE_TYPES)
{
using ExecutionSpace = typename DeviceType::execution_space;
using MemorySpace = typename DeviceType::memory_space;
#ifdef KOKKOS_ENABLE_SERIAL
using UnionFind = ArborX::Details::UnionFind<
MemorySpace,
/*DoSerial=*/std::is_same_v<ExecutionSpace, Kokkos::Serial>>;
#else
using UnionFind = ArborX::Details::UnionFind<MemorySpace>;
#endif
ExecutionSpace space;
constexpr int n = 5;
Kokkos::View<int *, MemorySpace> labels(
Kokkos::view_alloc(space, Kokkos::WithoutInitializing, "Test::labels"),
n);
ArborX::Details::KokkosExt::iota(space, labels);
UnionFind union_find(labels);
ARBORX_TEST_UNION_FIND_REPRESENTATIVES(space, union_find,
(std::vector<int>{0, 1, 2, 3, 4}));
merge(space, union_find, 1, 1);
ARBORX_TEST_UNION_FIND_REPRESENTATIVES(space, union_find,
(std::vector<int>{0, 1, 2, 3, 4}));
merge(space, union_find, 3, 0);
ARBORX_TEST_UNION_FIND_REPRESENTATIVES(space, union_find,
(std::vector<int>{0, 1, 2, 0, 4}));
merge(space, union_find, 1, 2);
merge(space, union_find, 4, 1);
merge(space, union_find, 1, 1);
ARBORX_TEST_UNION_FIND_REPRESENTATIVES(space, union_find,
(std::vector<int>{0, 1, 1, 0, 1}));
merge(space, union_find, 0, 1);
ARBORX_TEST_UNION_FIND_REPRESENTATIVES(space, union_find,
(std::vector<int>{0, 0, 0, 0, 0}));
}
BOOST_AUTO_TEST_SUITE_END()