Skip to content

Commit

Permalink
Merge pull request #906 from dalg24/tool_exec_instance
Browse files Browse the repository at this point in the history
Check that work is enqueued on the user-provided execution space instance
  • Loading branch information
aprokop authored Sep 6, 2024
2 parents c16c321 + 7c28f20 commit 50a3858
Show file tree
Hide file tree
Showing 3 changed files with 199 additions and 30 deletions.
1 change: 1 addition & 0 deletions test/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,7 @@ list(APPEND ARBORX_TEST_QUERY_TREE_SOURCES
tstQueryTreeTraversalPolicy.cpp
tstQueryTreeIntersectsKDOP.cpp
tstKokkosToolsAnnotations.cpp
tstKokkosToolsExecutionSpaceInstances.cpp
utf_main.cpp
)
add_executable(ArborX_Test_QueryTree.exe ${ARBORX_TEST_QUERY_TREE_SOURCES})
Expand Down
90 changes: 60 additions & 30 deletions test/Search_UnitTestHelpers.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -146,122 +146,152 @@ makeDistributedTree(MPI_Comm comm, std::vector<ArborX::Box> const &b)
#endif

template <typename DeviceType>
auto makeIntersectsBoxQueries(std::vector<ArborX::Box> const &boxes)
auto makeIntersectsBoxQueries(
std::vector<ArborX::Box> const &boxes,
typename DeviceType::execution_space const &exec_space = {})
{
int const n = boxes.size();
Kokkos::View<decltype(ArborX::intersects(ArborX::Box{})) *, DeviceType>
queries("Testing::intersecting_with_box_predicates", n);
auto queries_host = Kokkos::create_mirror_view(queries);
queries(Kokkos::view_alloc(Kokkos::WithoutInitializing,
"Testing::intersecting_with_box_predicates"),
n);
auto queries_host =
Kokkos::create_mirror_view(Kokkos::WithoutInitializing, queries);
for (int i = 0; i < n; ++i)
queries_host(i) = ArborX::intersects(boxes[i]);
Kokkos::deep_copy(queries, queries_host);
Kokkos::deep_copy(exec_space, queries, queries_host);
return queries;
}

template <typename DeviceType, typename Data>
auto makeIntersectsBoxWithAttachmentQueries(
std::vector<ArborX::Box> const &boxes, std::vector<Data> const &data)
std::vector<ArborX::Box> const &boxes, std::vector<Data> const &data,
typename DeviceType::execution_space const &exec_space = {})
{
int const n = boxes.size();
Kokkos::View<decltype(ArborX::attach(ArborX::intersects(ArborX::Box{}),
Data{})) *,
DeviceType>
queries("Testing::intersecting_with_box_with_attachment_predicates", n);
auto queries_host = Kokkos::create_mirror_view(queries);
queries(Kokkos::view_alloc(
Kokkos::WithoutInitializing,
"Testing::intersecting_with_box_with_attachment_predicates"),
n);
auto queries_host =
Kokkos::create_mirror_view(Kokkos::WithoutInitializing, queries);
for (int i = 0; i < n; ++i)
queries_host(i) = ArborX::attach(ArborX::intersects(boxes[i]), data[i]);
Kokkos::deep_copy(queries, queries_host);
Kokkos::deep_copy(exec_space, queries, queries_host);
return queries;
}

template <typename DeviceType>
auto makeNearestQueries(
std::vector<std::pair<ArborX::Point, int>> const &points)
std::vector<std::pair<ArborX::Point, int>> const &points,
typename DeviceType::execution_space const &exec_space = {})
{
// NOTE: `points` is not a very descriptive name here. It stores both the
// actual point and the number k of neighbors to query for.
int const n = points.size();
Kokkos::View<ArborX::Nearest<ArborX::Point> *, DeviceType> queries(
"Testing::nearest_queries", n);
auto queries_host = Kokkos::create_mirror_view(queries);
Kokkos::view_alloc(Kokkos::WithoutInitializing,
"Testing::nearest_queries"),
n);
auto queries_host =
Kokkos::create_mirror_view(Kokkos::WithoutInitializing, queries);
for (int i = 0; i < n; ++i)
queries_host(i) = ArborX::nearest(points[i].first, points[i].second);
Kokkos::deep_copy(queries, queries_host);
Kokkos::deep_copy(exec_space, queries, queries_host);
return queries;
}

template <typename DeviceType>
auto makeBoxNearestQueries(
std::vector<std::tuple<ArborX::Point, ArborX::Point, int>> const &boxes)
std::vector<std::tuple<ArborX::Point, ArborX::Point, int>> const &boxes,
typename DeviceType::execution_space const &exec_space = {})
{
// NOTE: `boxes` is not a very descriptive name here. It stores both the
// corners of the boxe and the number k of neighbors to query for.
int const n = boxes.size();
Kokkos::View<ArborX::Nearest<ArborX::Box> *, DeviceType> queries(
"Testing::nearest_queries", n);
auto queries_host = Kokkos::create_mirror_view(queries);
Kokkos::view_alloc(Kokkos::WithoutInitializing,
"Testing::nearest_queries"),
n);
auto queries_host =
Kokkos::create_mirror_view(Kokkos::WithoutInitializing, queries);
for (int i = 0; i < n; ++i)
queries_host(i) = ArborX::nearest(
ArborX::Box{std::get<0>(boxes[i]), std::get<1>(boxes[i])},
std::get<2>(boxes[i]));
Kokkos::deep_copy(queries, queries_host);
Kokkos::deep_copy(exec_space, queries, queries_host);
return queries;
}

template <typename DeviceType>
auto makeSphereNearestQueries(
std::vector<std::tuple<ArborX::Point, float, int>> const &spheres)
std::vector<std::tuple<ArborX::Point, float, int>> const &spheres,
typename DeviceType::execution_space const &exec_space = {})
{
// NOTE: `sphere` is not a very descriptive name here. It stores both the
// center and the radius of the sphere and the number k of neighbors to query
// for.
int const n = spheres.size();
Kokkos::View<ArborX::Nearest<ArborX::Sphere> *, DeviceType> queries(
"Testing::nearest_queries", n);
auto queries_host = Kokkos::create_mirror_view(queries);
Kokkos::view_alloc(Kokkos::WithoutInitializing,
"Testing::nearest_queries"),
n);
auto queries_host =
Kokkos::create_mirror_view(Kokkos::WithoutInitializing, queries);
for (int i = 0; i < n; ++i)
queries_host(i) = ArborX::nearest(
ArborX::Sphere{std::get<0>(spheres[i]), std::get<1>(spheres[i])},
std::get<2>(spheres[i]));
Kokkos::deep_copy(queries, queries_host);
Kokkos::deep_copy(exec_space, queries, queries_host);
return queries;
}

template <typename DeviceType, typename Data>
auto makeNearestWithAttachmentQueries(
std::vector<std::pair<ArborX::Point, int>> const &points,
std::vector<Data> const &data)
std::vector<Data> const &data,
typename DeviceType::execution_space const &exec_space = {})
{
// NOTE: `points` is not a very descriptive name here. It stores both the
// actual point and the number k of neighbors to query for.
int const n = points.size();
Kokkos::View<decltype(ArborX::attach(ArborX::Nearest<ArborX::Point>{},
Data{})) *,
DeviceType>
queries("Testing::nearest_queries", n);
auto queries_host = Kokkos::create_mirror_view(queries);
queries(Kokkos::view_alloc(Kokkos::WithoutInitializing,
"Testing::nearest_queries"),
n);
auto queries_host =
Kokkos::create_mirror_view(Kokkos::WithoutInitializing, queries);
for (int i = 0; i < n; ++i)
queries_host(i) = ArborX::attach(
ArborX::nearest(points[i].first, points[i].second), data[i]);
Kokkos::deep_copy(queries, queries_host);
Kokkos::deep_copy(exec_space, queries, queries_host);
return queries;
}

template <typename DeviceType>
Kokkos::View<decltype(ArborX::intersects(ArborX::Sphere{})) *, DeviceType>
makeIntersectsSphereQueries(
std::vector<std::pair<ArborX::Point, float>> const &points)
auto makeIntersectsSphereQueries(
std::vector<std::pair<ArborX::Point, float>> const &points,
typename DeviceType::execution_space const &exec_space = {})
{
// NOTE: `points` is not a very descriptive name here. It stores both the
// actual point and the radius for the search around that point.
int const n = points.size();
Kokkos::View<decltype(ArborX::intersects(ArborX::Sphere{})) *, DeviceType>
queries("Testing::intersecting_with_sphere_predicates", n);
auto queries_host = Kokkos::create_mirror_view(queries);
queries(
Kokkos::view_alloc(Kokkos::WithoutInitializing,
"Testing::intersecting_with_sphere_predicates"),
n);
auto queries_host =
Kokkos::create_mirror_view(Kokkos::WithoutInitializing, queries);
for (int i = 0; i < n; ++i)
queries_host(i) =
ArborX::intersects(ArborX::Sphere{points[i].first, points[i].second});
Kokkos::deep_copy(queries, queries_host);
Kokkos::deep_copy(exec_space, queries, queries_host);
return queries;
}

Expand Down
138 changes: 138 additions & 0 deletions test/tstKokkosToolsExecutionSpaceInstances.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,138 @@
/****************************************************************************
* 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 *
****************************************************************************/

#include "ArborX_EnableDeviceTypes.hpp" // ARBORX_DEVICE_TYPES
#include <ArborX_LinearBVH.hpp>

#include <boost/test/unit_test.hpp>

#include <string>

#include "Search_UnitTestHelpers.hpp"

BOOST_AUTO_TEST_SUITE(KokkosToolsExecutionSpaceInstances)

namespace tt = boost::test_tools;

namespace
{
// Lambdas can only be converted to function pointers if they do not capture.
// Using a global non-static variable in an unnamed namespace to "capture" the
// device id.
uint32_t arborx_test_device_id = -1;
uint32_t arborx_test_root_device_id = -1;

void arborx_test_parallel_x_callback(char const *label, uint32_t device_id,
uint64_t * /*kernel_id*/)
{
std::string label_str(label);

for (std::string s : {"Kokkos::View::destruction []"})
if (label_str.find(s) != std::string::npos)
return;

BOOST_TEST(device_id == arborx_test_device_id,
"\"" << label
<< "\" kernel not on the right execution space instance: "
<< device_id << " != " << arborx_test_device_id);
}

template <class ExecutionSpace>
void arborx_test_set_tools_callbacks(ExecutionSpace exec)
{
arborx_test_device_id = Kokkos::Tools::Experimental::device_id(exec);
arborx_test_root_device_id =
Kokkos::Tools::Experimental::device_id_root<ExecutionSpace>();

Kokkos::Tools::Experimental::set_begin_parallel_for_callback(
arborx_test_parallel_x_callback);
Kokkos::Tools::Experimental::set_begin_parallel_reduce_callback(
arborx_test_parallel_x_callback);
Kokkos::Tools::Experimental::set_begin_parallel_scan_callback(
arborx_test_parallel_x_callback);
}

void arborx_test_unset_tools_callbacks()
{
Kokkos::Tools::Experimental::set_begin_parallel_for_callback(nullptr);
Kokkos::Tools::Experimental::set_begin_parallel_reduce_callback(nullptr);
Kokkos::Tools::Experimental::set_begin_parallel_scan_callback(nullptr);
arborx_test_device_id = -1;
arborx_test_root_device_id = -1;
}

} // namespace

BOOST_AUTO_TEST_CASE_TEMPLATE(bvh_bvh_execution_space_instance, DeviceType,
ARBORX_DEVICE_TYPES)
{
using Tree = ArborX::BVH<typename DeviceType::memory_space>;
using ExecutionSpace = typename DeviceType::execution_space;

auto exec = Kokkos::Experimental::partition_space(ExecutionSpace{}, 1)[0];
arborx_test_set_tools_callbacks(exec);

{ // default constructed
Tree tree;
}

{ // empty
auto tree = make<Tree>(exec, {});
}

{ // one leaf
auto tree = make<Tree>(exec, {
{{{0, 0, 0}}, {{1, 1, 1}}},
});
}

{ // two leaves
auto tree = make<Tree>(exec, {
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
});
}

arborx_test_unset_tools_callbacks();
}

BOOST_AUTO_TEST_CASE_TEMPLATE(bvh_query_execution_space_instance, DeviceType,
ARBORX_DEVICE_TYPES)
{
using ExecutionSpace = typename DeviceType::execution_space;

auto tree = make<ArborX::BVH<typename DeviceType::memory_space>>(
ExecutionSpace{}, {
{{{0, 0, 0}}, {{1, 1, 1}}},
{{{0, 0, 0}}, {{1, 1, 1}}},
});

auto exec = Kokkos::Experimental::partition_space(ExecutionSpace{}, 1)[0];
arborx_test_set_tools_callbacks(exec);

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

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

arborx_test_unset_tools_callbacks();
}

BOOST_AUTO_TEST_SUITE_END()

0 comments on commit 50a3858

Please sign in to comment.