Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Store dendrogram in chain form when computing MST #1187

Draft
wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

Store dendrogram in chain form when computing MST #1187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 21 additions & 1 deletion src/cluster/ArborX_MinimumSpanningTree.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@ struct MinimumSpanningTree
Kokkos::View<WeightedEdge *, MemorySpace> edges;
Kokkos::View<int *, MemorySpace> dendrogram_parents;
Kokkos::View<float *, MemorySpace> dendrogram_parent_heights;
Kokkos::View<int *, MemorySpace> _chain_offsets;
Kokkos::View<int *, MemorySpace> _chain_levels;

template <class ExecutionSpace, class Primitives>
MinimumSpanningTree(ExecutionSpace const &space, Primitives const &primitives,
Expand All @@ -46,6 +48,8 @@ struct MinimumSpanningTree
AccessTraits<Primitives, PrimitivesTag>::size(primitives) - 1)
, dendrogram_parents("ArborX::MST::dendrogram_parents", 0)
, dendrogram_parent_heights("ArborX::MST::dendrogram_parent_heights", 0)
, _chain_offsets("ArborX::MST::chain_offsets", 0)
, _chain_levels("ArborX::MST::chain_levels", 0)
{
Kokkos::Profiling::pushRegion("ArborX::MST::MST");

Expand Down Expand Up @@ -182,6 +186,8 @@ struct MinimumSpanningTree
int num_components = n;
[[maybe_unused]] int edges_start = 0;
[[maybe_unused]] int edges_end = 0;
std::vector<int> edge_offsets;
edge_offsets.push_back(0);
do
{
Kokkos::Profiling::pushRegion("ArborX::Boruvka_" +
Expand Down Expand Up @@ -224,6 +230,8 @@ struct MinimumSpanningTree
Kokkos::deep_copy(space, num_edges_host, num_edges);
space.fence();

edge_offsets.push_back(num_edges_host);

if constexpr (Mode == BoruvkaMode::HDBSCAN)
{
Kokkos::parallel_for(
Expand Down Expand Up @@ -277,7 +285,19 @@ struct MinimumSpanningTree
std::make_pair(edges_start, edges_end)),
ROOT_CHAIN_VALUE);

computeParents(space, edges, sided_parents, dendrogram_parents);
Kokkos::View<int *, MemorySpace> edge_hierarchy_offsets(
Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
"ArborX::MST::edge_hierarchy_offsets"),
edge_offsets.size());
Kokkos::deep_copy(
space, edge_hierarchy_offsets,
Kokkos::View<int *, Kokkos::HostSpace, Kokkos::MemoryUnmanaged>{
edge_offsets.data(), edge_offsets.size()});

computeParentsAndReorderEdges(space, edges, edge_hierarchy_offsets,
sided_parents, dendrogram_parents,
_chain_offsets, _chain_levels);
Kokkos::resize(sided_parents, 0);

KokkosExt::reallocWithoutInitializing(space, dendrogram_parent_heights,
n - 1);
Expand Down
96 changes: 75 additions & 21 deletions src/cluster/detail/ArborX_BoruvkaHelpers.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@
#include <detail/ArborX_HappyTreeFriends.hpp>
#include <detail/ArborX_WeightedEdge.hpp>
#include <kokkos_ext/ArborX_KokkosExtArithmeticTraits.hpp>
#include <misc/ArborX_Utils.hpp>
#include <kokkos_ext/ArborX_KokkosExtKernelStdAlgorithms.hpp>
#include <misc/ArborX_SortUtils.hpp>

#include <Kokkos_Core.hpp>
#include <Kokkos_Profiling_ScopedRegion.hpp>
Expand Down Expand Up @@ -547,10 +548,14 @@ void assignVertexParents(ExecutionSpace const &space, Labels const &labels,
});
}

template <typename ExecutionSpace, typename Edges, typename SidedParents,
typename Parents>
void computeParents(ExecutionSpace const &space, Edges const &edges,
SidedParents const &sided_parents, Parents &parents)
template <typename ExecutionSpace, typename Edges,
typename EdgeHierarchyOffsets, typename SidedParents,
typename Parents, typename ChainOffsets, typename ChainLevels>
void computeParentsAndReorderEdges(
ExecutionSpace const &space, Edges &edges,
EdgeHierarchyOffsets const &edge_hierarchy_offsets,
SidedParents const &sided_parents, Parents &parents,
ChainOffsets &chain_offsets, ChainLevels &chain_levels)
{
Kokkos::Profiling::ScopedRegion guard("ArborX::MST::compute_edge_parents");

Expand Down Expand Up @@ -654,28 +659,77 @@ void computeParents(ExecutionSpace const &space, Edges const &edges,
}
});

auto rev_permute =
KokkosExt::cloneWithoutInitializingNorCopying(space, permute);
Kokkos::parallel_for(
"ArborX::MST::compute_rev_permute",
Kokkos::RangePolicy(space, 0, num_edges),
KOKKOS_LAMBDA(int const i) { rev_permute(permute(i)) = i; });

Kokkos::parallel_for(
"ArborX::MST::update_vertex_parents",
Kokkos::RangePolicy(space, num_edges, parents.size()),
KOKKOS_LAMBDA(int const i) { parents(i) = rev_permute(parents(i)); });

KokkosExt::reallocWithoutInitializing(space, chain_offsets, num_edges + 1);
int num_chains;
Kokkos::parallel_scan(
"ArborX::MST::compute_parents", Kokkos::RangePolicy(space, 0, num_edges),
KOKKOS_LAMBDA(int const i) {
int e = permute(i);
if (i == num_edges - 1)
{
// The parent of the root node is set to -1
parents(e) = -1;
}
else if ((keys(i) >> shift) == (keys(i + 1) >> shift))
KOKKOS_LAMBDA(int const i, int &update, bool final_pass) {
if (i == num_edges - 1 || (keys(i) >> shift) != (keys(i + 1) >> shift))
++update;

if (final_pass)
{
// For the edges belonging to the same chain, assign the parent of an
// edge to the edge with the next larger value
parents(e) = permute(i + 1);
if (i == num_edges - 1)
{
// The parent of the root node is set to -1
parents(i) = -1;
chain_offsets(update) = num_edges;
chain_offsets(0) = 0;
}
else if ((keys(i) >> shift) == (keys(i + 1) >> shift))
{
// For the edges belonging to the same chain, assign the parent of
// an edge to the edge with the next larger value
parents(i) = i + 1;
}
else
{
// For an edge which points to the root of a chain, assign edge's
// parent to be that root
parents(i) = rev_permute((keys(i) >> shift) / 2);
chain_offsets(update) = i + 1;
}
}
else
},
num_chains);
Kokkos::resize(Kokkos::WithoutInitializing, chain_offsets, num_chains + 1);

Details::applyPermutation(space, permute, edges);

Kokkos::resize(Kokkos::WithoutInitializing, chain_levels, num_chains + 1);
int num_levels = 0;
Kokkos::parallel_scan(
"ArborX::MST::compute_chain_levels",
Kokkos::RangePolicy<ExecutionSpace>(space, 0, num_chains),
KOKKOS_LAMBDA(int i, int &level, bool final_pass) {
auto upper_bound = [&v = edge_hierarchy_offsets](int x) {
return KokkosExt::upper_bound(v.data(), v.data() + v.size(), x);
};

if (i == num_chains - 1 ||
upper_bound(permute(chain_offsets(i))) !=
upper_bound(permute(chain_offsets(i + 1))))
{
// For an edge which points to the root of a chain, assign edge's
// parent to be that root
parents(e) = (keys(i) >> shift) / 2;
++level;
if (final_pass)
chain_levels(level) = i + 1;
}
});
},
num_levels);
++num_levels;
Kokkos::resize(Kokkos::WithoutInitializing, chain_levels, num_levels);
}

// Compute upper bound on the shortest edge of each component.
Expand Down
Loading