DistributedClosestPoint GPU-related optimizations

src/axom/core/execution/internal/cuda_exec.hpp

@@ -55,6 +55,8 @@ struct execution_space<CUDA_EXEC<BLOCK_SIZE, SYNCHRONOUS>>
   using atomic_policy = RAJA::cuda_atomic;
   using sync_policy = RAJA::cuda_synchronize;
 
+  static constexpr int BlockSize = BLOCK_SIZE;
+
   static constexpr MemorySpace memory_space = MemorySpace::Device;
 
   AXOM_HOST_DEVICE static constexpr bool async() noexcept { return false; }
@@ -97,6 +99,8 @@ struct execution_space<CUDA_EXEC<BLOCK_SIZE, ASYNC>>
   using atomic_policy = RAJA::cuda_atomic;
   using sync_policy = RAJA::cuda_synchronize;
 
+  static constexpr int BlockSize = BLOCK_SIZE;
+
   static constexpr MemorySpace memory_space = MemorySpace::Device;
 
   AXOM_HOST_DEVICE static constexpr bool async() noexcept { return true; }

src/axom/core/execution/internal/hip_exec.hpp

@@ -53,6 +53,8 @@ struct execution_space<HIP_EXEC<BLOCK_SIZE, SYNCHRONOUS>>
   using atomic_policy = RAJA::hip_atomic;
   using sync_policy = RAJA::hip_synchronize;
 
+  static constexpr int BlockSize = BLOCK_SIZE;
+
   static constexpr MemorySpace memory_space = MemorySpace::Device;
 
   AXOM_HOST_DEVICE static constexpr bool async() noexcept { return false; }
@@ -95,6 +97,8 @@ struct execution_space<HIP_EXEC<BLOCK_SIZE, ASYNC>>
   using atomic_policy = RAJA::hip_atomic;
   using sync_policy = RAJA::hip_synchronize;
 
+  static constexpr int BlockSize = BLOCK_SIZE;
+
   static constexpr MemorySpace memory_space = MemorySpace::Device;
 
   AXOM_HOST_DEVICE static constexpr bool async() noexcept { return true; }

src/axom/core/execution/internal/omp_exec.hpp

@@ -41,6 +41,8 @@ struct execution_space<OMP_EXEC>
   using atomic_policy = RAJA::omp_atomic;
   using sync_policy = RAJA::omp_synchronize;
 
+  static constexpr int BlockSize = 1;
+
 #ifdef AXOM_USE_UMPIRE
   static constexpr MemorySpace memory_space = MemorySpace::Host;
 #else

src/axom/core/execution/internal/seq_exec.hpp

@@ -51,6 +51,8 @@ struct execution_space<SEQ_EXEC>
 
   using sync_policy = void;
 
+  static constexpr int BlockSize = 1;
+
 #ifdef AXOM_USE_UMPIRE
   static constexpr MemorySpace memory_space = MemorySpace::Host;
 #else

src/axom/primal/operators/squared_distance.hpp

@@ -85,16 +85,28 @@ AXOM_HOST_DEVICE inline double squared_distance(const Point<T, NDIMS>& P,
     return axom::numerics::floating_point_limits<T>::max();
   }
 
-  if(B.contains(P))
+  // compute closest point to the box
+  Point<T, NDIMS> cp;
+  if constexpr(std::is_floating_point_v<T>)
   {
-    return 0;
+    for(int i = 0; i < NDIMS; ++i)
+    {
+      cp[i] = fmax(B.getMin()[i], fmin(P[i], B.getMax()[i]));
+    }
+  }
+  else
+  {
+    for(int i = 0; i < NDIMS; ++i)
+    {
+      cp[i] = clampVal(P[i], B.getMin()[i], B.getMax()[i]);
+    }
   }
 
-  // compute closest point to the box
-  Point<T, NDIMS> cp;
-  for(int i = 0; i < NDIMS; ++i)
+  // if clamped point is the same as the original point, our point
+  // was already in the bounding box
+  if(cp == P)
   {
-    cp[i] = clampVal(P[i], B.getMin()[i], B.getMax()[i]);
+    return 0;
   }
 
   // return squared distance to the closest point

src/axom/quest/detail/DistributedClosestPointImpl.hpp

@@ -8,8 +8,6 @@
 
 #include "axom/config.hpp"
 #include "axom/core.hpp"
-#include "axom/core/NumericLimits.hpp"
-#include "axom/core/execution/runtime_policy.hpp"
 #include "axom/slic.hpp"
 #include "axom/primal.hpp"
 #include "axom/spin.hpp"
@@ -1058,6 +1056,8 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl
         /// Create an ArrayView in ExecSpace that is compatible with queryPts
         PointArray execPoints(queryPts, m_allocatorID);
         auto query_pts = execPoints.view();
+        auto query_order = mortonSortQueryPoints(query_pts, qPtCount);
+        auto query_order_view = query_order.view();
         const double sqDistThreshold = m_sqDistanceThreshold;
         auto it = m_bvh->getTraverser();
         const int rank = m_rank;
@@ -1074,7 +1074,8 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl
           axom::ReduceMax<ExecSpace, double> maxSqDistance(currentMaxSqDistance);
           axom::for_all<ExecSpace>(
             qPtCount,
-            AXOM_LAMBDA(std::int32_t idx) mutable {
+            AXOM_LAMBDA(std::int32_t sorted_idx) {
+              const auto idx = query_order_view[sorted_idx];
               PointType qpt = query_pts[idx];
 
               MinCandidate curr_min {};
@@ -1106,7 +1107,7 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl
                 return sqDist <= curr_min.sqDist && sqDist <= sqDistThreshold;
               };
 
-              it.traverse_tree(qpt, checkMinDist, traversePredicate);
+              it.template traverseTreeShared<ExecSpace>(qpt, checkMinDist, traversePredicate);
 
               if(curr_min.rank == rank)
               {
@@ -1131,7 +1132,8 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl
           AXOM_ANNOTATE_SCOPE("ComputeClosestPoints");
           axom::for_all<ExecSpace>(
             qPtCount,
-            AXOM_LAMBDA(std::int32_t idx) mutable {
+            AXOM_LAMBDA(std::int32_t sorted_idx) {
+              const auto idx = query_order_view[sorted_idx];
               PointType qpt = query_pts[idx];
 
               MinCandidate curr_min {};
@@ -1165,7 +1167,7 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl
               };
 
               // Traverse the tree, searching for the point with minimum distance.
-              it.traverse_tree(qpt, checkMinDist, traversePredicate);
+              it.template traverseTreeShared<ExecSpace>(qpt, checkMinDist, traversePredicate);
 
               // If modified, update the fields that changed
               if(curr_min.rank == rank)
@@ -1210,6 +1212,60 @@ class DistributedClosestPointExec : public DistributedClosestPointImpl
   }
 
 private:
+  /*! \brief Returns query point indices ordered by their Morton codes. */
+  axom::Array<axom::IndexType> mortonSortQueryPoints(const axom::ArrayView<PointType>& queryPoints,
+                                                     axom::IndexType queryPointCount) const
+  {
+    axom::Array<axom::IndexType> queryOrder(queryPointCount, queryPointCount, m_allocatorID);
+    if(queryPointCount == 0)
+    {
+      return queryOrder;
+    }
+
+    PointType minPoint;
+    PointType inverseExtent;
+    for(int dim = 0; dim < DIM; ++dim)
+    {
+      axom::ReduceMin<ExecSpace, double> minCoord(axom::numeric_limits<double>::max());
+      axom::ReduceMax<ExecSpace, double> maxCoord(axom::numeric_limits<double>::lowest());
+      axom::for_all<ExecSpace>(
+        queryPointCount,
+        AXOM_LAMBDA(axom::IndexType idx) {
+          minCoord.min(queryPoints[idx][dim]);
+          maxCoord.max(queryPoints[idx][dim]);
+        });
+
+      minPoint[dim] = minCoord.get();
+      const double extent = maxCoord.get() - minPoint[dim];
+      inverseExtent[dim] = extent > 0.0 ? 1.0 / extent : 0.0;
+    }
+
+    axom::Array<std::uint32_t> mortonCodes(queryPointCount, queryPointCount, m_allocatorID);
+    auto morton_codes = mortonCodes.view();
+    auto query_order = queryOrder.view();
+    axom::for_all<ExecSpace>(
+      queryPointCount,
+      AXOM_LAMBDA(axom::IndexType idx) {
+        constexpr int bits_per_dimension = 32 / DIM;
+        constexpr double coordinate_scale = 1 << bits_per_dimension;
+        constexpr double coordinate_max = coordinate_scale - 1.0;
+
+        primal::Point<std::int32_t, DIM> gridPoint;
+        for(int dim = 0; dim < DIM; ++dim)
+        {
+          const double coordinate = (queryPoints[idx][dim] - minPoint[dim]) * inverseExtent[dim];
+          gridPoint[dim] = static_cast<std::int32_t>(
+            axom::utilities::clampVal(coordinate * coordinate_scale, 0.0, coordinate_max));
+        }
+
+        morton_codes[idx] = spin::convertPointToMorton<std::uint32_t>(gridPoint);
+        query_order[idx] = idx;
+      });
+
+    axom::stable_sort_pairs<ExecSpace>(morton_codes, query_order);
+    return queryOrder;
+  }
+
   /*!
     @brief Object point coordindates array.
 

src/axom/spin/CMakeLists.txt

@@ -31,6 +31,7 @@ set( spin_headers
      UniformGrid.hpp
 
      ## internal
+     internal/linear_bvh/BVHNode.hpp
      internal/linear_bvh/RadixTree.hpp
      internal/linear_bvh/build_radix_tree.hpp
      internal/linear_bvh/bvh_traverse.hpp

src/axom/spin/internal/linear_bvh/BVHNode.hpp

@@ -0,0 +1,39 @@
+// Copyright (c) 2017-2025, Lawrence Livermore National Security, LLC and
+// other Axom Project Developers. See the top-level LICENSE file for details.
+//
+// SPDX-License-Identifier: (BSD-3-Clause)
+
+#ifndef Axom_Spin_BVHNode_HH
+#define Axom_Spin_BVHNode_HH
+
+#include "axom/primal/geometry/BoundingBox.hpp"
+
+namespace axom
+{
+namespace spin
+{
+namespace internal
+{
+namespace linear_bvh
+{
+
+/*!
+ * \brief Node structure for a 2-wide BVH tree.
+ */
+template <typename FloatType, int NDIMS>
+struct BVH2Node
+{
+  using BoxType = primal::BoundingBox<FloatType, NDIMS>;
+
+  BoxType left;
+  BoxType right;
+  std::int32_t left_child;
+  std::int32_t right_child;
+};
+
+}  // namespace linear_bvh
+}  // namespace internal
+}  // namespace spin
+}  // namespace axom
+
+#endif