Skip to content

Commit

Permalink
Merge pull request #1180 from aprokop/remove_batched_queries
Browse files Browse the repository at this point in the history 
Refactor space filling curves surrounding code
  • Loading branch information
@aprokop
aprokop authored Jan 27, 2025
2 parents 7dc00bf + 8a7d92b commit b007f37
Show file tree
Hide file tree
Showing 10 changed files with 111 additions and 144 deletions.
8 changes: 4 additions & 4 deletions examples/viz/tree_visualization.cpp
View file
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
27 changes: 15 additions & 12 deletions src/spatial/ArborX_LinearBVH.hpp
View file
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 @@ -227,21 +226,24 @@ 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>;
using Details::returnCentroid;
using LinearOrderingValueType = std::invoke_result_t<
SpaceFillingCurve, decltype(scene_bounding_box),
std::decay_t<decltype(returnCentroid(indexables(0)))>>;
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);
Details::projectOntoSpaceFillingCurve(
space, indexables, curve, scene_bounding_box, linear_ordering_indices);

Kokkos::Profiling::popRegion();
Kokkos::Profiling::pushRegion("ArborX::BVH::BVH::sort_linearized_order");
Expand Down Expand Up @@ -308,15 +310,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

This file was deleted.

11 changes: 4 additions & 7 deletions src/spatial/detail/ArborX_CrsGraphWrapperImpl.hpp
View file
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 @@ -342,9 +342,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 @@ -387,9 +384,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
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 @@ -207,6 +223,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
72 changes: 48 additions & 24 deletions src/spatial/detail/ArborX_SpaceFillingCurves.hpp
View file
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, typename 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 @@ -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, typename 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,48 @@ struct Morton64
namespace Details
{

template <typename ExecutionSpace, typename Values, typename SpaceFillingCurve,
typename Box, typename LinearOrdering>
inline void
projectOntoSpaceFillingCurve(ExecutionSpace const &space, Values const &values,
SpaceFillingCurve const &curve,
Box const &scene_bounding_box,
LinearOrdering &linear_ordering_indices)
{
using Point = std::decay_t<decltype(returnCentroid(values(0)))>;
static_assert(GeometryTraits::is_point_v<Point>);
static_assert(GeometryTraits::is_box_v<Box>);
ARBORX_ASSERT(linear_ordering_indices.size() == values.size());
static_assert(std::is_same_v<typename LinearOrdering::value_type,
decltype(curve(scene_bounding_box, values(0)))>);

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

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)
{
using Point = std::decay_t<decltype(returnCentroid(values(0)))>;
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"),
values.size());
projectOntoSpaceFillingCurve(space, values, curve, scene_bounding_box,
linear_ordering_indices);
return sortObjects(space, linear_ordering_indices);
}

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
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

0 comments on commit b007f37

Please sign in to comment.