diff --git a/examples/viz/tree_visualization.cpp b/examples/viz/tree_visualization.cpp index c66a401c4..e085502f5 100644 --- a/examples/viz/tree_visualization.cpp +++ b/examples/viz/tree_visualization.cpp @@ -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:: - sortPredicatesAlongSpaceFillingCurve(ExecutionSpace{}, - ArborX::Experimental::Morton32(), - bvh.bounds(), queries); + auto permute = ArborX::Details::computeSpaceFillingCurvePermutation( + ExecutionSpace{}, + ArborX::Details::PredicateIndexables{queries}, + ArborX::Experimental::Morton32{}, bvh.bounds()); ArborX::Details::applyPermutation(ExecutionSpace{}, permute, queries); performQueries(prefix + "sorted_", suffix); } diff --git a/src/spatial/ArborX_LinearBVH.hpp b/src/spatial/ArborX_LinearBVH.hpp index e6f2cafff..d7055bfcc 100644 --- a/src/spatial/ArborX_LinearBVH.hpp +++ b/src/spatial/ArborX_LinearBVH.hpp @@ -16,7 +16,6 @@ #include #include #include -#include #include #include #include @@ -245,26 +244,19 @@ BoundingVolumeHierarchy:: "ArborX::BVH::BVH::calculate_scene_bounding_box"); // determine the bounding box of the scene - Box> bbox{}; + Box> + 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; - Kokkos::View 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); @@ -325,15 +317,16 @@ void BoundingVolumeHierarchy< if (policy._sort_predicates) { Kokkos::Profiling::pushRegion(profiling_prefix + "::compute_permutation"); - using DeviceType = Kokkos::Device; + using Details::expand; Box, typename GeometryTraits::coordinate_type_t> scene_bounding_box{}; - using namespace Details; expand(scene_bounding_box, bounds()); - auto permute = Details::BatchedQueries:: - sortPredicatesAlongSpaceFillingCurve(space, Experimental::Morton32(), - scene_bounding_box, predicates); + + auto permute = Details::computeSpaceFillingCurvePermutation( + space, Details::PredicateIndexables{predicates}, + Experimental::Morton32{}, scene_bounding_box); + Kokkos::Profiling::popRegion(); using PermutedPredicates = diff --git a/src/spatial/detail/ArborX_BatchedQueries.hpp b/src/spatial/detail/ArborX_BatchedQueries.hpp deleted file mode 100644 index cf5650149..000000000 --- a/src/spatial/detail/ArborX_BatchedQueries.hpp +++ /dev/null @@ -1,74 +0,0 @@ -/**************************************************************************** - * 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 * - ****************************************************************************/ - -#ifndef ARBORX_BATCHED_QUERIES_HPP -#define ARBORX_BATCHED_QUERIES_HPP - -#include -#include -#include // sortObjects - -#include - -#include - -namespace ArborX::Details -{ -template -struct BatchedQueries -{ -public: - // BatchedQueries defines functions for sorting queries along the Z-order - // space-filling curve in order to minimize data divergence. The goal is - // to increase correlation between traversal decisions made by nearby - // threads and thereby increase performance. - // - // NOTE: sortQueriesAlongZOrderCurve() does not actually apply the sorting - // order, it returns the permutation indices. applyPermutation() was added - // in that purpose. reversePermutation() is able to restore the initial - // order on the results that are in "compressed row storage" format. You - // may notice it is not used anymore in the code that performs the batched - // queries. We found that it was slightly more performant to add a level of - // indirection when recording results rather than using that function at - // the end. We decided to keep reversePermutation around for now. - - template - static Kokkos::View - sortPredicatesAlongSpaceFillingCurve(ExecutionSpace const &space, - SpaceFillingCurve const &curve, - Box const &scene_bounding_box, - Predicates const &predicates) - { - auto const n_queries = predicates.size(); - - using Point = - std::decay_t; - using LinearOrderingValueType = - std::invoke_result_t; - Kokkos::View linear_ordering_indices( - Kokkos::view_alloc(space, Kokkos::WithoutInitializing, - "ArborX::BVH::query::linear_ordering"), - n_queries); - Kokkos::parallel_for( - "ArborX::BatchedQueries::project_predicates_onto_space_filling_curve", - Kokkos::RangePolicy(space, 0, n_queries), KOKKOS_LAMBDA(int i) { - linear_ordering_indices(i) = curve( - scene_bounding_box, returnCentroid(getGeometry(predicates(i)))); - }); - - return sortObjects(space, linear_ordering_indices); - } -}; - -} // namespace ArborX::Details - -#endif diff --git a/src/spatial/detail/ArborX_CrsGraphWrapperImpl.hpp b/src/spatial/detail/ArborX_CrsGraphWrapperImpl.hpp index 7f17f99b7..0b3ee3c90 100644 --- a/src/spatial/detail/ArborX_CrsGraphWrapperImpl.hpp +++ b/src/spatial/detail/ArborX_CrsGraphWrapperImpl.hpp @@ -13,10 +13,10 @@ #define ARBORX_DETAIL_CRS_GRAPH_WRAPPER_IMPL_HPP #include -#include #include #include #include +#include #include #include #include @@ -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; - check_valid_callback(callback, predicates, out); std::string profiling_prefix = "ArborX::CrsGraphWrapper::query::"; @@ -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:: - sortPredicatesAlongSpaceFillingCurve(space, Experimental::Morton32(), - scene_bounding_box, predicates); + auto permute = computeSpaceFillingCurvePermutation( + space, PredicateIndexables{predicates}, + Experimental::Morton32{}, scene_bounding_box); Kokkos::Profiling::popRegion(); queryImpl(space, tree, predicates, callback, out, offset, permute, diff --git a/src/spatial/detail/ArborX_Predicates.hpp b/src/spatial/detail/ArborX_Predicates.hpp index df3180bd6..7746c3930 100644 --- a/src/spatial/detail/ArborX_Predicates.hpp +++ b/src/spatial/detail/ArborX_Predicates.hpp @@ -215,6 +215,24 @@ KOKKOS_INLINE_FUNCTION constexpr auto attach(Predicate &&pred, Data &&data) std::forward(pred), std::forward(data)}; } +namespace Details +{ +template +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 diff --git a/src/spatial/detail/ArborX_SpaceFillingCurves.hpp b/src/spatial/detail/ArborX_SpaceFillingCurves.hpp index 80d8e0814..db35e0454 100644 --- a/src/spatial/detail/ArborX_SpaceFillingCurves.hpp +++ b/src/spatial/detail/ArborX_SpaceFillingCurves.hpp @@ -12,11 +12,11 @@ #ifndef ARBORX_SPACE_FILLING_CURVES_HPP #define ARBORX_SPACE_FILLING_CURVES_HPP -#include -#include +#include #include #include #include +#include #include #include @@ -30,20 +30,11 @@ namespace Experimental struct Morton32 { - template && - GeometryTraits::is_point_v> * = 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 && - !GeometryTraits::is_point_v> * = nullptr> + template KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box, Geometry const &geometry) const { + static_assert(GeometryTraits::is_box_v); using Details::returnCentroid; auto p = returnCentroid(geometry); Details::translateAndScale(p, p, scene_bounding_box); @@ -53,20 +44,11 @@ struct Morton32 struct Morton64 { - template && - GeometryTraits::is_point_v> * = 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 && - !GeometryTraits::is_point_v> * = nullptr> + template KOKKOS_FUNCTION auto operator()(Box const &scene_bounding_box, Geometry const &geometry) const { + static_assert(GeometryTraits::is_box_v); using Details::returnCentroid; auto p = returnCentroid(geometry); Details::translateAndScale(p, p, scene_bounding_box); @@ -79,6 +61,50 @@ struct Morton64 namespace Details { +template +inline Kokkos::View()(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; + static_assert(GeometryTraits::is_point_v); + static_assert(GeometryTraits::is_box_v); + + auto const n = values.size(); + + using LinearOrderingValueType = + std::invoke_result_t; + Kokkos::View + linear_ordering_indices( + Kokkos::view_alloc(space, Kokkos::WithoutInitializing, + "ArborX::SpaceFillingCurve::linear_ordering"), + n); + 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 +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); +} + template void check_valid_space_filling_curve(SpaceFillingCurve const &) { diff --git a/src/spatial/detail/ArborX_TreeConstruction.hpp b/src/spatial/detail/ArborX_TreeConstruction.hpp index 84c1f6d25..76e19af5d 100644 --- a/src/spatial/detail/ArborX_TreeConstruction.hpp +++ b/src/spatial/detail/ArborX_TreeConstruction.hpp @@ -15,7 +15,6 @@ #include #include #include // makeLeafNode -#include #include #include @@ -39,27 +38,6 @@ inline void calculateBoundingBoxOfTheScene(ExecutionSpace const &space, GeometryReducer(scene_bounding_box)); } -template -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); - - 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 inline void diff --git a/test/tstDetailsTreeConstruction.cpp b/test/tstDetailsTreeConstruction.cpp index f093d4136..da70ea091 100644 --- a/test/tstDetailsTreeConstruction.cpp +++ b/test/tstDetailsTreeConstruction.cpp @@ -14,6 +14,7 @@ #include #include // expandBits, morton32 #include // ROPE SENTINEL +#include #include #include #include // sortObjects @@ -92,10 +93,8 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(assign_morton_codes, DeviceType, BOOST_TEST(ArborX::Details::equals( scene_host, {{{0.0, 0.0, 0.0}}, {{(float)N, (float)N, (float)N}}})); - Kokkos::View morton_codes("morton_codes", - n); - ArborX::Details::TreeConstruction::projectOntoSpaceFillingCurve( - space, boxes, ArborX::Experimental::Morton64(), scene_host, morton_codes); + auto morton_codes = ArborX::Details::projectOntoSpaceFillingCurve( + space, boxes, ArborX::Experimental::Morton64(), scene_host); auto morton_codes_host = Kokkos::create_mirror_view(morton_codes); Kokkos::deep_copy(morton_codes_host, morton_codes); BOOST_TEST(morton_codes_host == ref, tt::per_element()); diff --git a/test/tstKokkosToolsAnnotations.cpp b/test/tstKokkosToolsAnnotations.cpp index 0b62e45ed..3b12d109b 100644 --- a/test/tstKokkosToolsAnnotations.cpp +++ b/test/tstKokkosToolsAnnotations.cpp @@ -45,6 +45,7 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(bvh_bvh_allocations_prefixed, DeviceType, std::regex re("^(Testing::" "|ArborX::BVH::" "|ArborX::Sorting::" + "|ArborX::SpaceFillingCurve::" "|Kokkos::SortImpl::BinSortFunctor::" "|Kokkos::Serial::" // unsure what's going on ").*"); @@ -94,6 +95,7 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(bvh_query_allocations_prefixed, DeviceType, void const * /*ptr*/, uint64_t /*size*/) { std::regex re("^(Testing::" "|ArborX::BVH::query::" + "|ArborX::SpaceFillingCurve::" "|ArborX::NearestBufferProvider::" "|ArborX::TreeTraversal::spatial::" "|ArborX::TreeTraversal::nearest::" diff --git a/test/tstKokkosToolsDistributedAnnotations.cpp b/test/tstKokkosToolsDistributedAnnotations.cpp index 9be59b236..5a7d0edbb 100644 --- a/test/tstKokkosToolsDistributedAnnotations.cpp +++ b/test/tstKokkosToolsDistributedAnnotations.cpp @@ -42,6 +42,7 @@ BOOST_AUTO_TEST_CASE_TEMPLATE( std::regex re("^(Testing::" "|ArborX::DistributedTree::" "|ArborX::BVH::" + "|ArborX::SpaceFillingCurve::" "|ArborX::Sorting::" ").*"); BOOST_TEST(std::regex_match(label, re), @@ -77,6 +78,7 @@ BOOST_AUTO_TEST_CASE_TEMPLATE( void const * /*ptr*/, uint64_t /*size*/) { std::regex re("^(Testing::" "|ArborX::DistributedTree::query::" + "|ArborX::SpaceFillingCurve::" "|ArborX::Distributor::" "|ArborX::BVH::query::" "|ArborX::NearestBufferProvider::"