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

Refactor space filling curves surrounding code #1180

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Refactor space filling curves surrounding code #1180

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
43a55fd
Fix getGeometry for prvalues
aprokop Oct 11, 2024
aaff2d0
Get rid of BatchedQueries header and unify space filling curve projec…
aprokop Oct 9, 2024
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
8 changes: 4 additions & 4 deletions examples/viz/tree_visualization.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,10 +161,10 @@ void viz(std::string const &prefix, std::string const &infile, int n_neighbors)
performQueries(prefix + "shuffled_", suffix);

// Sort them
auto permute = ArborX::Details::BatchedQueries<DeviceType>::
sortPredicatesAlongSpaceFillingCurve(ExecutionSpace{},
ArborX::Experimental::Morton32(),
bvh.bounds(), queries);
auto permute = ArborX::Details::computeSpaceFillingCurvePermutation(
ExecutionSpace{},
ArborX::Details::PredicateIndexables<decltype(queries)>{queries},
ArborX::Experimental::Morton32{}, bvh.bounds());
ArborX::Details::applyPermutation(ExecutionSpace{}, permute, queries);
performQueries(prefix + "sorted_", suffix);
}
Expand Down
29 changes: 11 additions & 18 deletions src/spatial/ArborX_LinearBVH.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
#include <ArborX_CrsGraphWrapper.hpp>
#include <detail/ArborX_AccessTraits.hpp>
#include <detail/ArborX_AttachIndices.hpp>
#include <detail/ArborX_BatchedQueries.hpp>
#include <detail/ArborX_Callbacks.hpp>
#include <detail/ArborX_CrsGraphWrapperImpl.hpp>
#include <detail/ArborX_IndexableGetter.hpp>
Expand Down Expand Up @@ -245,26 +244,19 @@ BoundingVolumeHierarchy<MemorySpace, Value, IndexableGetter, BoundingVolume>::
"ArborX::BVH::BVH::calculate_scene_bounding_box");

// determine the bounding box of the scene
Box<DIM, typename GeometryTraits::coordinate_type_t<BoundingVolume>> bbox{};
Box<DIM, typename GeometryTraits::coordinate_type_t<BoundingVolume>>
scene_bounding_box{};
Details::TreeConstruction::calculateBoundingBoxOfTheScene(space, indexables,
bbox);
scene_bounding_box);
Kokkos::Profiling::popRegion();
Kokkos::Profiling::pushRegion("ArborX::BVH::BVH::compute_linear_ordering");

// Map indexables from multidimensional domain to one-dimensional interval
using LinearOrderingValueType =
std::invoke_result_t<SpaceFillingCurve, decltype(bbox), indexable_type>;
Kokkos::View<LinearOrderingValueType *, MemorySpace> linear_ordering_indices(
Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
"ArborX::BVH::BVH::linear_ordering"),
size());
Details::TreeConstruction::projectOntoSpaceFillingCurve(
space, indexables, curve, bbox, linear_ordering_indices);
auto linear_ordering_indices = Details::projectOntoSpaceFillingCurve(
space, indexables, curve, scene_bounding_box);

Kokkos::Profiling::popRegion();
Kokkos::Profiling::pushRegion("ArborX::BVH::BVH::sort_linearized_order");

// Compute the ordering of the indexables along the space-filling curve
auto permutation_indices =
Details::sortObjects(space, linear_ordering_indices);

Expand Down Expand Up @@ -325,15 +317,16 @@ void BoundingVolumeHierarchy<
if (policy._sort_predicates)
{
Kokkos::Profiling::pushRegion(profiling_prefix + "::compute_permutation");
using DeviceType = Kokkos::Device<ExecutionSpace, MemorySpace>;
using Details::expand;
Box<GeometryTraits::dimension_v<bounding_volume_type>,
typename GeometryTraits::coordinate_type_t<bounding_volume_type>>
scene_bounding_box{};
using namespace Details;
expand(scene_bounding_box, bounds());
auto permute = Details::BatchedQueries<DeviceType>::
sortPredicatesAlongSpaceFillingCurve(space, Experimental::Morton32(),
scene_bounding_box, predicates);

auto permute = Details::computeSpaceFillingCurvePermutation(
space, Details::PredicateIndexables<Predicates>{predicates},
Experimental::Morton32{}, scene_bounding_box);

Kokkos::Profiling::popRegion();

using PermutedPredicates =
Expand Down
74 changes: 0 additions & 74 deletions src/spatial/detail/ArborX_BatchedQueries.hpp
View file
Open in desktop

This file was deleted.

11 changes: 4 additions & 7 deletions src/spatial/detail/ArborX_CrsGraphWrapperImpl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,10 @@
#define ARBORX_DETAIL_CRS_GRAPH_WRAPPER_IMPL_HPP

#include <ArborX_Box.hpp>
#include <detail/ArborX_BatchedQueries.hpp>
#include <detail/ArborX_Callbacks.hpp>
#include <detail/ArborX_PermutedData.hpp>
#include <detail/ArborX_Predicates.hpp>
#include <detail/ArborX_SpaceFillingCurves.hpp>
#include <detail/ArborX_TraversalPolicy.hpp>
#include <kokkos_ext/ArborX_KokkosExtStdAlgorithms.hpp>
#include <kokkos_ext/ArborX_KokkosExtViewHelpers.hpp>
Expand Down Expand Up @@ -343,9 +343,6 @@ queryDispatch(Tag, Tree const &tree, ExecutionSpace const &space,
Experimental::TraversalPolicy const &policy =
Experimental::TraversalPolicy())
{
using MemorySpace = typename Tree::memory_space;
using DeviceType = Kokkos::Device<ExecutionSpace, MemorySpace>;

check_valid_callback<typename Tree::value_type>(callback, predicates, out);

std::string profiling_prefix = "ArborX::CrsGraphWrapper::query::";
Expand Down Expand Up @@ -388,9 +385,9 @@ queryDispatch(Tag, Tree const &tree, ExecutionSpace const &space,
scene_bounding_box{};
using namespace Details;
expand(scene_bounding_box, tree.bounds());
auto permute = Details::BatchedQueries<DeviceType>::
sortPredicatesAlongSpaceFillingCurve(space, Experimental::Morton32(),
scene_bounding_box, predicates);
auto permute = computeSpaceFillingCurvePermutation(
space, PredicateIndexables<Predicates>{predicates},
Experimental::Morton32{}, scene_bounding_box);
Kokkos::Profiling::popRegion();

queryImpl(space, tree, predicates, callback, out, offset, permute,
Expand Down
34 changes: 34 additions & 0 deletions src/spatial/detail/ArborX_Predicates.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,20 +147,36 @@ getGeometry(Nearest<Geometry> const &pred)
{
return pred._geometry;
}
template <typename Geometry>
KOKKOS_INLINE_FUNCTION Geometry getGeometry(Nearest<Geometry> &&pred)
{
return pred._geometry;
}

template <typename Geometry>
KOKKOS_INLINE_FUNCTION Geometry const &
getGeometry(Intersects<Geometry> const &pred)
{
return pred._geometry;
}
template <typename Geometry>
KOKKOS_INLINE_FUNCTION Geometry getGeometry(Intersects<Geometry> &&pred)
{
return pred._geometry;
}

template <typename Geometry>
KOKKOS_INLINE_FUNCTION Geometry const &
getGeometry(Experimental::OrderedSpatial<Geometry> const &pred)
{
return pred._geometry;
}
template <typename Geometry>
KOKKOS_INLINE_FUNCTION Geometry
getGeometry(Experimental::OrderedSpatial<Geometry> &&pred)
{
return pred._geometry;
}

template <typename Predicate, typename Data>
struct PredicateWithAttachment : Predicate
Expand Down Expand Up @@ -199,6 +215,24 @@ KOKKOS_INLINE_FUNCTION constexpr auto attach(Predicate &&pred, Data &&data)
std::forward<Predicate>(pred), std::forward<Data>(data)};
}

namespace Details
{
template <typename Predicates>
struct PredicateIndexables
{
Predicates _predicates;

using memory_space = typename Predicates::memory_space;

KOKKOS_FUNCTION decltype(auto) operator()(int i) const
{
return getGeometry(_predicates(i));
}

KOKKOS_FUNCTION auto size() const { return _predicates.size(); }
};
} // namespace Details

} // namespace ArborX

#endif
74 changes: 50 additions & 24 deletions src/spatial/detail/ArborX_SpaceFillingCurves.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,11 @@
#ifndef ARBORX_SPACE_FILLING_CURVES_HPP
#define ARBORX_SPACE_FILLING_CURVES_HPP

#include <ArborX_Box.hpp>
#include <ArborX_Point.hpp>
#include <ArborX_GeometryTraits.hpp>
#include <algorithms/ArborX_Centroid.hpp>
#include <algorithms/ArborX_TranslateAndScale.hpp>
#include <detail/ArborX_MortonCode.hpp>
#include <misc/ArborX_SortUtils.hpp>

#include <Kokkos_DetectionIdiom.hpp>
#include <Kokkos_Macros.hpp>
Expand All @@ -30,20 +30,11 @@ namespace Experimental

struct Morton32
{
template <typename Box, typename Point,
std::enable_if_t<GeometryTraits::is_box_v<Box> &&
GeometryTraits::is_point_v<Point>> * = nullptr>
KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box, Point p) const
{
Details::translateAndScale(p, p, scene_bounding_box);
return Details::morton32(p);
}
template <typename Box, typename Geometry,
std::enable_if_t<GeometryTraits::is_box_v<Box> &&
!GeometryTraits::is_point_v<Geometry>> * = nullptr>
template <typename Box, class Geometry>
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do not mix typename and class

KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box,
Geometry const &geometry) const
{
static_assert(GeometryTraits::is_box_v<Box>);
using Details::returnCentroid;
auto p = returnCentroid(geometry);
Details::translateAndScale(p, p, scene_bounding_box);
Expand All @@ -53,20 +44,11 @@ struct Morton32

struct Morton64
{
template <typename Box, typename Point,
std::enable_if_t<GeometryTraits::is_box_v<Box> &&
GeometryTraits::is_point_v<Point>> * = nullptr>
KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box, Point p) const
{
Details::translateAndScale(p, p, scene_bounding_box);
return Details::morton64(p);
}
template <typename Box, class Geometry,
std::enable_if_t<GeometryTraits::is_box_v<Box> &&
!GeometryTraits::is_point_v<Geometry>> * = nullptr>
template <typename Box, class Geometry>
KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box,
Geometry const &geometry) const
{
static_assert(GeometryTraits::is_box_v<Box>);
using Details::returnCentroid;
auto p = returnCentroid(geometry);
Details::translateAndScale(p, p, scene_bounding_box);
Expand All @@ -79,6 +61,50 @@ struct Morton64
namespace Details
{

template <typename ExecutionSpace, typename Values, typename SpaceFillingCurve,
typename Box>
inline Kokkos::View<std::invoke_result_t<SpaceFillingCurve, Box,
std::decay_t<decltype(returnCentroid(
std::declval<Values>()(0)))>> *,
typename Values::memory_space>
projectOntoSpaceFillingCurve(ExecutionSpace const &space, Values const &values,
SpaceFillingCurve const &curve,
Box const &scene_bounding_box)
{
using Point = std::decay_t<decltype(returnCentroid(values(0)))>;
static_assert(GeometryTraits::is_point_v<Point>);
static_assert(GeometryTraits::is_box_v<Box>);

auto const n = values.size();

using LinearOrderingValueType =
std::invoke_result_t<SpaceFillingCurve, Box, Point>;
Kokkos::View<LinearOrderingValueType *, typename Values::memory_space>
linear_ordering_indices(
Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
"ArborX::SpaceFillingCurve::linear_ordering"),
n);
Comment on lines +82 to +86
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think that moving the allocation into that function and returning the Morton codes is a good change

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't like it much either. However, our boilerplate code for determining the type of indices is really several lines at every call site.

Alternatively, we could introduce a requirement for a space filling curve to provide index_type that it will return. It would simplify things a lot.

Kokkos::parallel_for(
"ArborX::SpaceFillingCurve::project_onto_space_filling_curve",
Kokkos::RangePolicy(space, 0, n), KOKKOS_LAMBDA(int i) {
linear_ordering_indices(i) = curve(scene_bounding_box, values(i));
});

return linear_ordering_indices;
}

template <typename ExecutionSpace, typename Values, typename SpaceFillingCurve,
typename Box>
inline auto computeSpaceFillingCurvePermutation(ExecutionSpace const &space,
Values const &values,
SpaceFillingCurve const &curve,
Box const &scene_bounding_box)
{
auto linear_ordering_indices =
projectOntoSpaceFillingCurve(space, values, curve, scene_bounding_box);
return sortObjects(space, linear_ordering_indices);
}
Comment on lines +96 to +106
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Saves little code to bundle two distinct operations.
I am not convinced it is a good idea to introduce it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's fair. It does a bit more than that in that it also reduces the amount of used memory as linear_ordering_indices are deallocated automatically once this function returns. This behavior is similar to the original sortPredicatesAlongSpaceFillingCurve.


template <int DIM, class SpaceFillingCurve>
void check_valid_space_filling_curve(SpaceFillingCurve const &)
{
Expand Down
22 changes: 0 additions & 22 deletions src/spatial/detail/ArborX_TreeConstruction.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
#include <algorithms/ArborX_Expand.hpp>
#include <algorithms/ArborX_Reducer.hpp>
#include <detail/ArborX_Node.hpp> // makeLeafNode
#include <detail/ArborX_SpaceFillingCurves.hpp>
#include <kokkos_ext/ArborX_KokkosExtArithmeticTraits.hpp>
#include <misc/ArborX_Exception.hpp>

Expand All @@ -39,27 +38,6 @@ inline void calculateBoundingBoxOfTheScene(ExecutionSpace const &space,
GeometryReducer<Box>(scene_bounding_box));
}

template <typename ExecutionSpace, typename Indexables,
typename SpaceFillingCurve, typename Box, typename LinearOrdering>
inline void projectOntoSpaceFillingCurve(ExecutionSpace const &space,
Indexables const &indexables,
SpaceFillingCurve const &curve,
Box const &scene_bounding_box,
LinearOrdering linear_ordering_indices)
{
size_t const n = indexables.size();
ARBORX_ASSERT(linear_ordering_indices.extent(0) == n);
static_assert(
std::is_same_v<typename LinearOrdering::value_type,
decltype(curve(scene_bounding_box, indexables(0)))>);

Kokkos::parallel_for(
"ArborX::TreeConstruction::project_primitives_onto_space_filling_curve",
Kokkos::RangePolicy(space, 0, n), KOKKOS_LAMBDA(int i) {
linear_ordering_indices(i) = curve(scene_bounding_box, indexables(i));
});
}

template <typename ExecutionSpace, typename Values, typename IndexableGetter,
typename Nodes, typename BoundingVolume>
inline void
Expand Down
Loading
Loading