Skip to content

Commit

Permalink
Merge pull request #1157 from aprokop/arborx-2.0-no-apiv1-distributed
Browse files Browse the repository at this point in the history 
Get rid of APIv1 in DistributedTree
  • Loading branch information
@aprokop
aprokop authored Sep 25, 2024
2 parents 0ca5049 + 67c945c commit 745f3ea
Show file tree
Hide file tree
Showing 7 changed files with 147 additions and 289 deletions.
80 changes: 7 additions & 73 deletions src/ArborX_DistributedTree.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,14 +118,13 @@ class DistributedTreeBase

// NOTE: query() must be called as collective over all processes in the
// communicator passed to the constructor
template <
typename MemorySpace, typename Value = Details::LegacyDefaultTemplateValue,
typename IndexableGetter = Details::DefaultIndexableGetter,
typename BoundingVolume =
Box<GeometryTraits::dimension_v<
std::decay_t<std::invoke_result_t<IndexableGetter, Value>>>,
typename GeometryTraits::coordinate_type_t<
std::decay_t<std::invoke_result_t<IndexableGetter, Value>>>>>
template <typename MemorySpace, typename Value,
typename IndexableGetter = Details::DefaultIndexableGetter,
typename BoundingVolume = Box<
GeometryTraits::dimension_v<
std::decay_t<std::invoke_result_t<IndexableGetter, Value>>>,
typename GeometryTraits::coordinate_type_t<
std::decay_t<std::invoke_result_t<IndexableGetter, Value>>>>>
class DistributedTree
: public DistributedTreeBase<BoundingVolumeHierarchy<
MemorySpace, Value, IndexableGetter, BoundingVolume>>
Expand All @@ -150,71 +149,6 @@ class DistributedTree
{}
};

template <typename MemorySpace>
class DistributedTree<MemorySpace, Details::LegacyDefaultTemplateValue,
Details::DefaultIndexableGetter, Box<3, float>>
: public DistributedTreeBase<BoundingVolumeHierarchy<MemorySpace>>
{
using base_type = DistributedTreeBase<BoundingVolumeHierarchy<MemorySpace>>;

public:
using memory_space = MemorySpace;
using value_type = int;
using bounding_volume_type = typename base_type::bounding_volume_type;

DistributedTree() = default; // build an empty tree

template <typename ExecutionSpace, typename Primitives>
DistributedTree(MPI_Comm comm, ExecutionSpace const &space,
Primitives const &primitives)
: base_type(comm, space, primitives)
{}

template <typename ExecutionSpace, typename UserPredicates,
typename IndicesAndRanks, typename OffsetView>
void query(ExecutionSpace const &space, UserPredicates const &user_predicates,
IndicesAndRanks &&indices_and_ranks, OffsetView &&offset) const
{
namespace KokkosExt = Details::KokkosExt;

static_assert(
KokkosExt::is_accessible_from<MemorySpace, ExecutionSpace>::value);

using Predicates = Details::AccessValues<UserPredicates, PredicatesTag>;
static_assert(
KokkosExt::is_accessible_from<typename Predicates::memory_space,
ExecutionSpace>::value,
"Predicates must be accessible from the execution space");

Predicates predicates{user_predicates}; // NOLINT

int comm_rank = -1;
if (base_type::getComm() != MPI_COMM_NULL)
MPI_Comm_rank(base_type::getComm(), &comm_rank);

base_type::query(space, predicates,
Details::DefaultCallbackWithRank{comm_rank},
std::forward<IndicesAndRanks>(indices_and_ranks),
std::forward<OffsetView>(offset));
}

template <typename ExecutionSpace, typename UserPredicates, typename Callback>
void query(ExecutionSpace const &space, UserPredicates const &user_predicates,
Callback &&callback) const
{
base_type::query(space, user_predicates, std::forward<Callback>(callback));
}

template <typename ExecutionSpace, typename UserPredicates, typename Callback,
typename Indices, typename Offset>
void query(ExecutionSpace const &space, UserPredicates const &user_predicates,
Callback &&callback, Indices &&out, Offset &&offset) const
{
base_type::query(space, user_predicates, std::forward<Callback>(callback),
std::forward<Indices>(out), std::forward<Offset>(offset));
}
};

template <typename BottomTree>
template <typename ExecutionSpace, typename... Args>
DistributedTreeBase<BottomTree>::DistributedTreeBase(
Expand Down
101 changes: 1 addition & 100 deletions src/details/ArborX_DetailsDistributedTreeNearestHelpers.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,8 @@
#define ARBORX_DETAILS_DISTRIBUTED_TREE_NEAREST_HELPERS_HPP

#include <ArborX_AccessTraits.hpp>
#include <ArborX_Box.hpp>
#include <ArborX_DetailsHappyTreeFriends.hpp>
#include <ArborX_LinearBVH.hpp>
#include <ArborX_Callbacks.hpp>
#include <ArborX_Point.hpp>
#include <ArborX_Ray.hpp>
#include <ArborX_Sphere.hpp>

#include <Kokkos_BitManipulation.hpp>
Expand Down Expand Up @@ -54,18 +51,6 @@ auto declare_callback_constrained(Callback const &callback)
namespace Details
{

struct DefaultCallbackWithRank
{
int _rank;

template <typename Predicate, typename Value, typename OutputFunctor>
KOKKOS_FUNCTION void operator()(Predicate const &, Value const &value,
OutputFunctor const &out) const
{
out({value, _rank});
}
};

template <class Callback>
struct is_constrained_callback : std::false_type
{};
Expand All @@ -79,11 +64,6 @@ struct is_constrained_callback<
template <>
struct is_constrained_callback<DefaultCallback> : std::true_type
{};
// We need DefaultCallbackWithRank specialization for the case without a
// user-provided callback and using APIv1
template <>
struct is_constrained_callback<DefaultCallbackWithRank> : std::true_type
{};

template <class Callback>
inline constexpr bool is_constrained_callback_v =
Expand Down Expand Up @@ -236,85 +216,6 @@ struct CallbackWithDistance
}
};

template <typename MemorySpace, typename Callback, typename OutValue,
bool UseValues>
struct CallbackWithDistance<
BoundingVolumeHierarchy<MemorySpace, Details::LegacyDefaultTemplateValue,
Details::DefaultIndexableGetter, Box<3, float>>,
Callback, OutValue, UseValues>
{
using Tree =
BoundingVolumeHierarchy<MemorySpace, Details::LegacyDefaultTemplateValue,
Details::DefaultIndexableGetter, Box<3, float>>;

Tree _tree;
Callback _callback;
Kokkos::View<unsigned int *, typename Tree::memory_space> _rev_permute;

template <typename ExecutionSpace>
CallbackWithDistance(ExecutionSpace const &exec_space, Tree const &tree,
Callback const &callback)
: _tree(tree)
, _callback(callback)
{
// NOTE cannot have extended __host__ __device__ lambda in constructor with
// NVCC
computeReversePermutation(exec_space);
}

template <typename ExecutionSpace>
void computeReversePermutation(ExecutionSpace const &exec_space)
{
if (_tree.empty())
return;

auto const n = _tree.size();

_rev_permute = Kokkos::View<unsigned int *, typename Tree::memory_space>(
Kokkos::view_alloc(
Kokkos::WithoutInitializing,
"ArborX::DistributedTree::query::nearest::reverse_permutation"),
n);
Kokkos::parallel_for(
"ArborX::DistributedTree::query::nearest::"
"compute_reverse_permutation",
Kokkos::RangePolicy<ExecutionSpace>(exec_space, 0, n),
KOKKOS_CLASS_LAMBDA(int const i) {
_rev_permute(HappyTreeFriends::getValue(_tree, i).index) = i;
});
}

template <typename Query, typename OutputFunctor>
KOKKOS_FUNCTION void operator()(Query const &query, int index,
OutputFunctor const &out) const
{
// TODO: This breaks the abstraction of the distributed Tree not knowing
// the details of the local tree. Right now, this is the only way. Will
// need to be fixed with a proper callback abstraction.
int const leaf_node_index = _rev_permute(index);
auto const &leaf_node_bounding_volume =
HappyTreeFriends::getIndexable(_tree, leaf_node_index);
if constexpr (UseValues)
{
OutValue out_value;
[[maybe_unused]] int count = 0;
_callback(query, index, [&](OutValue const &ov) {
out_value = ov;
++count;
});
// If the user callback produces no output, we have nothing to attach the
// distance to, which is problematic as we would not be able to do the
// final filtering. If there are multiple outputs, it will currently
// break our communication routines and filtering. We rely on 3-phase
// nearest implementation for these cases.
KOKKOS_ASSERT(count == 1);
out({out_value, distance(getGeometry(query), leaf_node_bounding_volume)});
}
else
out(distance(getGeometry(query), leaf_node_bounding_volume));
}
};

} // namespace Details
} // namespace ArborX

Expand Down
1 change: 0 additions & 1 deletion src/details/ArborX_DetailsDistributedTreeSpatial.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
#include <ArborX_Callbacks.hpp>
#include <ArborX_DetailsDistributedTreeImpl.hpp>
#include <ArborX_DetailsDistributedTreeUtils.hpp>
#include <ArborX_DetailsLegacy.hpp>
#include <ArborX_Predicates.hpp>

#include <Kokkos_Core.hpp>
Expand Down
64 changes: 38 additions & 26 deletions test/Search_UnitTestHelpers.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,9 @@ struct is_distributed : std::false_type
{};

#ifdef ARBORX_ENABLE_MPI
template <typename D>
struct is_distributed<ArborX::DistributedTree<D>> : std::true_type
template <typename MemorySpace, typename Value, typename... Args>
struct is_distributed<ArborX::DistributedTree<MemorySpace, Value, Args...>>
: std::true_type
{};

template <typename I>
Expand All @@ -55,19 +56,6 @@ auto make_reference_solution(std::vector<T> const &values,
return make_compressed_storage(offsets, values);
}

#ifdef ARBORX_ENABLE_MPI
// FIXME This is a temporary workaround until we reconcile interfaces of
// DistributedTree and BVH
template <typename ExecutionSpace, typename MemorySpace, typename Queries,
typename Values, typename Offsets>
void query(ArborX::DistributedTree<MemorySpace> const &tree,
ExecutionSpace const &space, Queries const &queries,
Values const &values, Offsets const &offsets)
{
tree.query(space, queries, values, offsets);
}
#endif

template <typename ExecutionSpace, typename Tree, typename Queries>
auto query(ExecutionSpace const &exec_space, Tree const &tree,
Queries const &queries)
Expand Down Expand Up @@ -130,22 +118,46 @@ auto make(ExecutionSpace const &exec_space,
}

#ifdef ARBORX_ENABLE_MPI
template <typename DeviceType, int DIM = 3, typename Coordinate = float>
auto makeDistributedTree(MPI_Comm comm,
std::vector<ArborX::Box<DIM, Coordinate>> const &b)
template <typename MemorySpace, typename Geometry>
struct PairIndexRankIndexableGetter
{
Kokkos::View<Geometry *, MemorySpace> _geometries;

KOKKOS_FUNCTION auto const &operator()(ArborXTest::PairIndexRank p) const
{
return _geometries(p.index);
}
};

template <typename DeviceType, typename Geometry>
auto makeDistributedTree(MPI_Comm comm, typename DeviceType::execution_space,
std::vector<Geometry> const &g)
{
using ExecutionSpace = typename DeviceType::execution_space;
using MemorySpace = typename DeviceType::memory_space;

int const n = b.size();
Kokkos::View<ArborX::Box<DIM, Coordinate> *, DeviceType> boxes(
"Testing::boxes", n);
auto boxes_host = Kokkos::create_mirror_view(boxes);
using PairIndexRank = ArborXTest::PairIndexRank;

int comm_rank;
MPI_Comm_rank(comm, &comm_rank);

int const n = g.size();
Kokkos::View<Geometry *, DeviceType> geometries("Testing::geometries", n);
Kokkos::View<PairIndexRank *, DeviceType> pairs("Testing::geometry_pairs", n);
auto geometries_host = Kokkos::create_mirror_view(geometries);
auto pairs_host = Kokkos::create_mirror_view(pairs);
for (int i = 0; i < n; ++i)
boxes_host(i) = b[i];
Kokkos::deep_copy(boxes, boxes_host);
{
geometries_host(i) = g[i];
pairs_host(i) = {i, comm_rank};
}
Kokkos::deep_copy(geometries, geometries_host);
Kokkos::deep_copy(pairs, pairs_host);

using IndexableGetter = PairIndexRankIndexableGetter<MemorySpace, Geometry>;

return ArborX::DistributedTree<typename DeviceType::memory_space>(
comm, ExecutionSpace{}, boxes);
return ArborX::DistributedTree<MemorySpace, PairIndexRank, IndexableGetter>(
comm, ExecutionSpace{}, pairs, IndexableGetter{geometries});
}
#endif

Expand Down
Loading

0 comments on commit 745f3ea

Please sign in to comment.