From b169b9bace81d3a1e0428474e4f9444626156298 Mon Sep 17 00:00:00 2001 From: Bruno Turcksin Date: Tue, 15 Feb 2022 19:04:09 +0000 Subject: [PATCH] Add support for ArborX::DistributedTree --- include/deal.II/arborx/distributed_tree.h | 134 ++++++++++++++++++++++ 1 file changed, 134 insertions(+) create mode 100644 include/deal.II/arborx/distributed_tree.h diff --git a/include/deal.II/arborx/distributed_tree.h b/include/deal.II/arborx/distributed_tree.h new file mode 100644 index 0000000000..2dba72b548 --- /dev/null +++ b/include/deal.II/arborx/distributed_tree.h @@ -0,0 +1,134 @@ +// --------------------------------------------------------------------- +// +// Copyright (C) 2022 by the deal.II authors +// +// This file is part of the deal.II library. +// +// The deal.II library is free software; you can use it, redistribute +// it, and/or modify it under the terms of the GNU Lesser General +// Public License as published by the Free Software Foundation; either +// version 2.1 of the License, or (at your option) any later version. +// The full text of the license can be found in the file LICENSE.md at +// the top level directory of deal.II. +// +// --------------------------------------------------------------------- + +#ifndef dealii_arborx_distributed_tree_h +#define dealii_arborx_distributed_tree_h + +#include + +#if defined(DEAL_II_ARBORX_WITH_MPI) && defined(DEAL_II_WITH_MPI) +# include + +# include +# include + +DEAL_II_NAMESPACE_OPEN + +namespace ArborXWrappers +{ + /** + * This class implements a wrapper around ArborX::DistributedTree, the + * distributed version of ArborX:BVH. + * + * Because ArborX uses Kokkos, Kokkos needs to be initialized and finalized + * before using this class. + */ + class DistributedTree + { + public: + /** + * Constructor. Use a vector of BoundingBox @p bounding_boxes as primitives. + * The BoundingBox objects in @p bounding_boxes are local to the MPI process. + */ + template + DistributedTree( + const MPI_Comm & comm, + const std::vector> &bounding_boxes); + + /** + * Constructor. Use a vector of @p points as primitives. The Point objects + * in @p points are local to the MPI process. + */ + template + DistributedTree(const MPI_Comm & comm, + const std::vector> &points); + + /** + * Return the indices and the MPI ranks of those BoundingBox objects that + * satisfy the @p queries. + * Because @p queries can contain multiple queries, the function returns a pair + * of indices and ranks and the associated offsets. + * + * Valid `QueryType` classes are + * ArborXWrappers::BoundingBoxIntersectPredicate, + * ArborXWrappers::BoundingBoxNearestPredicate, + * ArborXWrappers::PointIntersectPredicate, and + * ArborXWrappers::PointNearestPredicate, + */ + template + std::pair>, std::vector> + query(const QueryType &queries); + + private: + /** + * Underlying ArborX object. + */ + ArborX::DistributedTree distributed_tree; + }; + + + + template + DistributedTree::DistributedTree( + const MPI_Comm & comm, + const std::vector> &bounding_boxes) + : distributed_tree(comm, + Kokkos::DefaultHostExecutionSpace{}, + bounding_boxes) + {} + + + + template + DistributedTree::DistributedTree( + const MPI_Comm & comm, + const std::vector> &points) + : distributed_tree(comm, Kokkos::DefaultHostExecutionSpace{}, points) + {} + + + + template + std::pair>, std::vector> + DistributedTree::query(const QueryType &queries) + { + Kokkos::View offsets("offsets", 0); + Kokkos::View *, Kokkos::HostSpace> indices_ranks( + "indices_ranks", 0); + distributed_tree.query(Kokkos::DefaultHostExecutionSpace{}, + queries, + indices_ranks, + offsets); + + std::vector> indices_ranks_vector; + for (unsigned int i = 0; i < indices_ranks.extent(0); ++i) + { + indices_ranks_vector.emplace_back(indices_ranks(i).first, + indices_ranks(i).second); + } + + std::vector offsets_vector; + offsets_vector.insert(offsets_vector.begin(), + offsets.data(), + offsets.data() + offsets.extent(0)); + + return {indices_ranks_vector, offsets_vector}; + } +} // namespace ArborXWrappers + +DEAL_II_NAMESPACE_CLOSE + +#endif +#endif -- 2.39.5