/
ArborX_DetailsDendrogram.hpp
121 lines (94 loc) · 4.05 KB
/
ArborX_DetailsDendrogram.hpp
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
115
116
117
118
119
120
121
/****************************************************************************
* Copyright (c) 2017-2023 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 *
****************************************************************************/
#ifndef ARBORX_DETAILS_DENDROGRAM_HPP
#define ARBORX_DETAILS_DENDROGRAM_HPP
#include <ArborX_DetailsKokkosExtStdAlgorithms.hpp> // iota
#include <ArborX_DetailsUnionFind.hpp>
#include <Kokkos_Core.hpp>
namespace ArborX::Details
{
struct UnweightedEdge
{
int source;
int target;
};
template <typename Edges, typename Parents>
void dendrogramUnionFindHost(Edges sorted_edges_host, Parents &parents_host)
{
Kokkos::Profiling::pushRegion(
"ArborX::Dendrogram::dendrogram_union_find::union_find_host");
using ExecutionSpace = Kokkos::DefaultHostExecutionSpace;
ExecutionSpace host_space;
int const num_edges = sorted_edges_host.size();
int const num_vertices = num_edges + 1;
auto const vertices_offset = num_edges;
Kokkos::View<int *, Kokkos::HostSpace> labels(
Kokkos::view_alloc(Kokkos::WithoutInitializing,
"ArborX::Dendrogram::dendrogram_union_find::labels"),
num_vertices);
KokkosExt::iota(host_space, labels);
UnionFind<Kokkos::HostSpace, true> union_find(labels);
Kokkos::View<int *, Kokkos::HostSpace> set_edges_host(
Kokkos::view_alloc(host_space, Kokkos::WithoutInitializing,
"ArborX::Dendrogram::labels"),
num_vertices);
constexpr int UNDEFINED = -1;
Kokkos::deep_copy(host_space, set_edges_host, UNDEFINED);
// Fence all execution spaces to make sure all the data is ready
Kokkos::fence("ArborX::Dendrogram::dendrogramUnionFindHost"
" (global fence before performing union-find on host)");
Kokkos::Profiling::pushRegion(
"ArborX::Dendrogram::dendrogram_union_find::union_find");
for (int e = 0; e < num_edges; ++e)
{
int i = union_find.representative(sorted_edges_host(e).source);
int j = union_find.representative(sorted_edges_host(e).target);
for (int k : {i, j})
{
auto edge_child = set_edges_host(k);
if (edge_child != UNDEFINED)
parents_host(edge_child) = e;
else
parents_host(vertices_offset + k) = e;
}
union_find.merge(i, j);
set_edges_host(union_find.representative(i)) = e;
}
parents_host(num_edges - 1) = -1; // root
Kokkos::Profiling::popRegion();
Kokkos::Profiling::popRegion();
}
template <typename ExecutionSpace, typename MemorySpace>
void dendrogramUnionFind(
ExecutionSpace const &exec_space,
Kokkos::View<UnweightedEdge const *, MemorySpace> sorted_edges,
Kokkos::View<int *, MemorySpace> &parents)
{
Kokkos::Profiling::pushRegion("ArborX::Dendrogram::dendrogram_union_find");
Kokkos::Profiling::pushRegion(
"ArborX::Dendrogram::dendrogram_union_find::copy_to_host");
auto sorted_edges_host =
Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, sorted_edges);
auto parents_host = Kokkos::create_mirror_view(
Kokkos::view_alloc(Kokkos::WithoutInitializing), parents);
Kokkos::Profiling::popRegion();
Kokkos::Profiling::pushRegion(
"ArborX::Dendrogram::dendrogram_union_find::union_find");
dendrogramUnionFindHost(sorted_edges_host, parents_host);
Kokkos::Profiling::popRegion();
Kokkos::Profiling::pushRegion(
"ArborX::Dendrogram::dendrogram_union_find::copy_to_device");
Kokkos::deep_copy(exec_space, parents, parents_host);
Kokkos::Profiling::popRegion();
Kokkos::Profiling::popRegion();
}
} // namespace ArborX::Details
#endif