Skip to content

Commit

Permalink
Allow strongly typed index-rank pairs in parallel tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@masterleinad
masterleinad committed Jan 11, 2023
1 parent 44ff8aa commit 78fd0b6
Show file tree
Hide file tree
Showing 6 changed files with 131 additions and 107 deletions.
14 changes: 6 additions & 8 deletions test/ArborX_BoostRTreeHelpers.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -180,8 +180,7 @@ static auto translate(ArborX::Nearest<Geometry> const &query)
return boost::geometry::index::nearest(geometry, k);
}

template <typename Indexable, typename InputView,
typename OutputView = Kokkos::View<int *, Kokkos::HostSpace>>
template <typename OutputView, typename Indexable, typename InputView>
static std::tuple<OutputView, OutputView>
performQueries(RTree<Indexable> const &rtree, InputView const &queries)
{
Expand All @@ -206,10 +205,8 @@ performQueries(RTree<Indexable> const &rtree, InputView const &queries)
}

#ifdef ARBORX_ENABLE_MPI
template <typename Indexable, typename InputView,
typename OutputView1 =
Kokkos::View<Kokkos::pair<int, int> *, Kokkos::HostSpace>,
typename OutputView2 = Kokkos::View<int *, Kokkos::HostSpace>>
template <typename OutputView1, typename OutputView2, typename Indexable,
typename InputView>
static std::tuple<OutputView2, OutputView1>
performQueries(ParallelRTree<Indexable> const &rtree, InputView const &queries)
{
Expand Down Expand Up @@ -267,7 +264,7 @@ class RTree
static_assert(Kokkos::is_execution_space<ExecutionSpace>::value);

std::tie(offset, indices) =
BoostRTreeHelpers::performQueries(_tree, predicates);
BoostRTreeHelpers::performQueries<InputView>(_tree, predicates);
}

template <typename ExecutionSpace, typename Predicates, typename Callback,
Expand Down Expand Up @@ -311,7 +308,8 @@ class ParallelRTree
static_assert(Kokkos::is_execution_space<ExecutionSpace>::value);

std::tie(offset, indices) =
BoostRTreeHelpers::performQueries(_tree, predicates);
BoostRTreeHelpers::performQueries<InputView1, InputView2, Indexable,
Predicates>(_tree, predicates);
}

private:
Expand Down
47 changes: 28 additions & 19 deletions test/Search_UnitTestHelpers.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,14 +67,13 @@ void query(ArborX::DistributedTree<MemorySpace> const &tree,
}
#endif

template <typename ExecutionSpace, typename Tree, typename Queries>
template <typename ValueType, typename ExecutionSpace, typename Tree,
typename Queries>
auto query(ExecutionSpace const &exec_space, Tree const &tree,
Queries const &queries)
{
using memory_space = typename Tree::memory_space;
using value_type =
std::conditional_t<is_distributed<Tree>{}, Kokkos::pair<int, int>, int>;
Kokkos::View<value_type *, memory_space> values("Testing::values", 0);
Kokkos::View<ValueType *, memory_space> values("Testing::values", 0);
Kokkos::View<int *, memory_space> offsets("Testing::offsets", 0);
tree.query(exec_space, queries, values, offsets);
return make_compressed_storage(
Expand All @@ -83,8 +82,11 @@ auto query(ExecutionSpace const &exec_space, Tree const &tree,
}

#define ARBORX_TEST_QUERY_TREE(exec_space, tree, queries, reference) \
BOOST_TEST(query(exec_space, tree, queries) == (reference), \
boost::test_tools::per_element());
{ \
using value_type = typename decltype(reference)::value_type; \
BOOST_TEST(query<value_type>(exec_space, tree, queries) == (reference), \
boost::test_tools::per_element()); \
}

template <typename ValueType, typename ExecutionSpace, typename Tree,
typename Queries, typename Callback>
Expand All @@ -102,24 +104,26 @@ auto query(ExecutionSpace const &exec_space, Tree const &tree,

#define ARBORX_TEST_QUERY_TREE_CALLBACK(exec_space, tree, queries, callback, \
reference) \
using value_type = typename decltype(reference)::value_type; \
BOOST_TEST(query<value_type>(exec_space, tree, queries, callback) == \
(reference), \
boost::test_tools::per_element());
{ \
using value_type = typename decltype(reference)::value_type; \
BOOST_TEST(query<value_type>(exec_space, tree, queries, callback) == \
(reference), \
boost::test_tools::per_element()); \
}

// Workaround for NVCC that complains that the enclosing parent function
// (query_with_distance) for an extended __host__ __device__ lambda must not
// have deduced return type
template <typename DeviceType, typename ExecutionSpace>
Kokkos::View<Kokkos::pair<Kokkos::pair<int, int>, float> *, DeviceType>
template <typename TupleIndexRankDistanceType, typename DeviceType,
typename ExecutionSpace>
Kokkos::View<TupleIndexRankDistanceType *, DeviceType>
zip(ExecutionSpace const &space, Kokkos::View<int *, DeviceType> indices,
Kokkos::View<int *, DeviceType> ranks,
Kokkos::View<float *, DeviceType> distances)
{
auto const n = indices.extent(0);
Kokkos::View<Kokkos::pair<Kokkos::pair<int, int>, float> *, DeviceType>
values(Kokkos::view_alloc(Kokkos::WithoutInitializing, "Testing::values"),
n);
Kokkos::View<TupleIndexRankDistanceType *, DeviceType> values(
Kokkos::view_alloc(Kokkos::WithoutInitializing, "Testing::values"), n);
Kokkos::parallel_for(
"ArborX:UnitTestSupport:zip",
Kokkos::RangePolicy<ExecutionSpace>(space, 0, n), KOKKOS_LAMBDA(int i) {
Expand All @@ -129,7 +133,8 @@ zip(ExecutionSpace const &space, Kokkos::View<int *, DeviceType> indices,
}

#ifdef ARBORX_ENABLE_MPI
template <typename ExecutionSpace, typename Tree, typename Queries>
template <typename ValueType, typename ExecutionSpace, typename Tree,
typename Queries>
auto query_with_distance(ExecutionSpace const &exec_space, Tree const &tree,
Queries const &queries,
std::enable_if_t<is_distributed<Tree>{}> * = nullptr)
Expand All @@ -145,7 +150,7 @@ auto query_with_distance(ExecutionSpace const &exec_space, Tree const &tree,
ArborX::Details::NearestPredicateTag{}, tree, exec_space, queries,
indices, offsets, ranks, &distances);

auto values = zip(exec_space, indices, ranks, distances);
auto values = zip<ValueType>(exec_space, indices, ranks, distances);

return make_compressed_storage(
Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, offsets),
Expand All @@ -155,8 +160,12 @@ auto query_with_distance(ExecutionSpace const &exec_space, Tree const &tree,

#define ARBORX_TEST_QUERY_TREE_WITH_DISTANCE(exec_space, tree, queries, \
reference) \
BOOST_TEST(query_with_distance(exec_space, tree, queries) == (reference), \
boost::test_tools::per_element());
{ \
using value_type = typename decltype(reference)::value_type; \
BOOST_TEST(query_with_distance<value_type>(exec_space, tree, queries) == \
(reference), \
boost::test_tools::per_element()); \
}

template <typename Tree, typename ExecutionSpace>
auto make(ExecutionSpace const &exec_space, std::vector<ArborX::Box> const &b)
Expand Down
45 changes: 30 additions & 15 deletions test/tstDistributedTree.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,26 @@

namespace tt = boost::test_tools;

using PairIndexRank = Kokkos::pair<int, int>;
using TupleIndexRankDistance = Kokkos::pair<Kokkos::pair<int, int>, float>;
struct PairIndexRank
{
int index;
int rank;
friend bool operator==(PairIndexRank lhs, PairIndexRank rhs)
{
return lhs.index == rhs.index && lhs.rank == rhs.rank;
}
friend bool operator<(PairIndexRank lhs, PairIndexRank rhs)
{
return lhs.index < rhs.index ||
(lhs.index == rhs.index && lhs.rank < rhs.rank);
}
friend std::ostream &operator<<(std::ostream &stream,
PairIndexRank const &pair)
{
return stream << '[' << pair.index << ',' << pair.rank << ']';
}
};
using TupleIndexRankDistance = Kokkos::pair<PairIndexRank, float>;

BOOST_AUTO_TEST_CASE_TEMPLATE(hello_world, DeviceType, ARBORX_DEVICE_TYPES)
{
Expand Down Expand Up @@ -96,11 +114,11 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(hello_world, DeviceType, ARBORX_DEVICE_TYPES)
values.reserve(n + 1);
for (int i = 0; i < n; ++i)
{
values.emplace_back(n - 1 - i, comm_size - 1 - comm_rank);
values.push_back({n - 1 - i, comm_size - 1 - comm_rank});
}
if (comm_rank > 0)
{
values.emplace_back(0, comm_size - comm_rank);
values.push_back({0, comm_size - comm_rank});
ARBORX_TEST_QUERY_TREE(ExecutionSpace{}, tree, queries,
make_reference_solution(values, {0, n + 1}));
}
Expand Down Expand Up @@ -317,19 +335,15 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(one_leaf_per_rank, DeviceType,

if (comm_rank > 0)
{
std::vector<PairIndexRank> values;
values.reserve(comm_size);
for (int i = 0; i < comm_size; ++i)
values.push_back({0, i});
ARBORX_TEST_QUERY_TREE(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0., 0., 0.}}, comm_rank * comm_size},
}),
make_reference_solution(
[comm_size]() {
std::vector<PairIndexRank> values;
values.reserve(comm_size);
for (int i = 0; i < comm_size; ++i)
values.emplace_back(0, i);
return values;
}(),
{0, comm_size}));
make_reference_solution(values, {0, comm_size}));
}
else
{
Expand Down Expand Up @@ -812,8 +826,9 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(boost_comparison, DeviceType, ARBORX_DEVICE_TYPES)
BoostExt::ParallelRTree<ArborX::Box> rtree(comm, ExecutionSpace{},
bounding_boxes_host);

ARBORX_TEST_QUERY_TREE(ExecutionSpace{}, distributed_tree, within_queries,
query(ExecutionSpace{}, rtree, within_queries_host));
ARBORX_TEST_QUERY_TREE(
ExecutionSpace{}, distributed_tree, within_queries,
query<PairIndexRank>(ExecutionSpace{}, rtree, within_queries_host));
}

template <typename MemorySpace>
Expand Down
60 changes: 30 additions & 30 deletions test/tstKokkosToolsAnnotations.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,18 +100,18 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(bvh_query_allocations_prefixed, DeviceType,
});

// spatial predicates
query(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));
query<int>(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));

// nearest predicates
query(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));
query<int>(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));

Kokkos::Tools::Experimental::set_allocate_data_callback(nullptr);
}
Expand Down Expand Up @@ -161,18 +161,18 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(kernels_prefixed, DeviceType, ARBORX_DEVICE_TYPES)
});

// spatial predicates
query(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));
query<int>(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));

// nearest predicates
query(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));
query<int>(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));

Kokkos::Tools::Experimental::set_begin_parallel_for_callback(nullptr);
Kokkos::Tools::Experimental::set_begin_parallel_scan_callback(nullptr);
Expand Down Expand Up @@ -221,18 +221,18 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(regions_prefixed, DeviceType, ARBORX_DEVICE_TYPES)
});

// spatial predicates
query(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));
query<int>(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));

// nearest predicates
query(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));
query<int>(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));

Kokkos::Tools::Experimental::set_push_region_callback(nullptr);
}
Expand Down
60 changes: 30 additions & 30 deletions test/tstKokkosToolsDistributedAnnotations.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,18 +84,18 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(
});

// spatial predicates
query(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));
query<Kokkos::pair<int, int>>(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));

// nearest predicates
query(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));
query<Kokkos::pair<int, int>>(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));

Kokkos::Tools::Experimental::set_allocate_data_callback(nullptr);
}
Expand All @@ -122,18 +122,18 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(kernels_prefixed, DeviceType, ARBORX_DEVICE_TYPES)
});

// spatial predicates
query(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));
query<Kokkos::pair<int, int>>(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));

// nearest predicates
query(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));
query<Kokkos::pair<int, int>>(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));

Kokkos::Tools::Experimental::set_begin_parallel_for_callback(nullptr);
Kokkos::Tools::Experimental::set_begin_parallel_scan_callback(nullptr);
Expand Down Expand Up @@ -165,18 +165,18 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(regions_prefixed, DeviceType, ARBORX_DEVICE_TYPES)
});

// spatial predicates
query(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));
query<Kokkos::pair<int, int>>(ExecutionSpace{}, tree,
makeIntersectsBoxQueries<DeviceType>({
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
}));

// nearest predicates
query(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));
query<Kokkos::pair<int, int>>(ExecutionSpace{}, tree,
makeNearestQueries<DeviceType>({
{{{0, 0, 0}}, 1},
{{{0, 0, 0}}, 2},
}));

Kokkos::Tools::Experimental::set_push_region_callback(nullptr);
}
Expand Down
Loading

0 comments on commit 78fd0b6

Please sign in to comment.