[WIP] Implement DiscretizedBox for reducing memory size of nodes

examples/simple_intersection/example_intersection.cpp

@@ -40,9 +40,9 @@ class Boxes
     float Lx = 100.0;
     float Ly = 100.0;
     float Lz = 100.0;
-    int nx = 11;
-    int ny = 11;
-    int nz = 11;
+    int nx = 3;
+    int ny = 3;
+    int nz = 3;
     int n = nx * ny * nz;
     float hx = Lx / (nx - 1);
     float hy = Ly / (ny - 1);
@@ -56,6 +56,8 @@ class Boxes
         Kokkos::view_alloc(Kokkos::WithoutInitializing, "boxes"), n);
     auto boxes_host = Kokkos::create_mirror_view(_boxes);
 
+    ArborX::Box global_box{{0, 0, 0}, {Lx, Ly, Lz}};
+
     for (int i = 0; i < nx; ++i)
       for (int j = 0; j < ny; ++j)
         for (int k = 0; k < nz; ++k)
@@ -65,6 +67,16 @@ class Boxes
           ArborX::Point p_upper{
               {(i + .25) * hx, (j + .25) * hy, (k + .25) * hz}};
           boxes_host[index(i, j, k)] = {p_lower, p_upper};
+          ArborX::DiscretizedBox discrete_box(boxes_host[index(i, j, k)],
+                                              global_box);
+          auto recovered_box = discrete_box.to_box(global_box);
+          std::cout << i << ' ' << j << ' ' << ' ' << k << ": "
+                    << recovered_box.minCorner()[0] << ' '
+                    << recovered_box.minCorner()[1] << ' '
+                    << recovered_box.minCorner()[2] << ' '
+                    << recovered_box.maxCorner()[0] << ' '
+                    << recovered_box.maxCorner()[1] << ' '
+                    << recovered_box.maxCorner()[2] << std::endl;
         }
     Kokkos::deep_copy(execution_space, _boxes, boxes_host);
   }
@@ -154,7 +166,11 @@ int main()
       Kokkos::abort("Wrong dimensions for the offsets View!\n");
     for (int i = 0; i < static_cast<int>(n + 1); ++i)
       if (offsets_host(i) != i)
+      {
+        std::cout << "offsets_host(" << i << ") =" << offsets_host(i)
+                  << std::endl;
         Kokkos::abort("Wrong entry in the offsets View!\n");
+      }
 
     if (indices_host.size() != n)
       Kokkos::abort("Wrong dimensions for the indices View!\n");

src/ArborX_LinearBVH.hpp

@@ -46,6 +46,7 @@ class BasicBoundingVolumeHierarchy
   static_assert(Kokkos::is_memory_space<MemorySpace>::value, "");
   using size_type = typename MemorySpace::size_type;
   using bounding_volume_type = BoundingVolume;
+  using discretized_bounding_volume_type = ArborX::DiscretizedBox;
 
   BasicBoundingVolumeHierarchy() = default; // build an empty tree
 
@@ -81,7 +82,7 @@ class BasicBoundingVolumeHierarchy
                   std::forward<View>(view), std::forward<Args>(args)...);
   }
 
-private:
+  // private:
   friend struct Details::HappyTreeFriends;
 
 #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP)
@@ -93,48 +94,54 @@ class BasicBoundingVolumeHierarchy
                    Kokkos::CudaSpace
 #else
                    Kokkos::Experimental::HIPSpace
+#endif
+                   >{},
+      Details::NodeWithLeftChildAndRope<discretized_bounding_volume_type>,
+      Details::NodeWithTwoChildren<discretized_bounding_volume_type>>;
+#else
+  using node_type =
+      Details::NodeWithTwoChildren<discretized_bounding_volume_type>;
+#endif
+
+#if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP)
+  // Ropes based traversal is only used for CUDA, as it was found to be slower
+  // than regular one for Power9 on Summit.  It is also used with HIP.
+  using leaf_node_type = std::conditional_t<
+      std::is_same<MemorySpace,
+#if defined(KOKKOS_ENABLE_CUDA)
+                   Kokkos::CudaSpace
+#else
+                   Kokkos::Experimental::HIPSpace
 #endif
                    >{},
       Details::NodeWithLeftChildAndRope<bounding_volume_type>,
       Details::NodeWithTwoChildren<bounding_volume_type>>;
 #else
-  using node_type = Details::NodeWithTwoChildren<bounding_volume_type>;
+  using leaf_node_type = Details::NodeWithTwoChildren<bounding_volume_type>;
 #endif
 
   Kokkos::View<node_type *, MemorySpace> getInternalNodes()
   {
     assert(!empty());
-    return Kokkos::subview(_internal_and_leaf_nodes,
-                           std::make_pair(size_type{0}, size() - 1));
+    return _internal_nodes;
   }
 
-  Kokkos::View<node_type *, MemorySpace> getLeafNodes()
+  Kokkos::View<leaf_node_type *, MemorySpace> getLeafNodes()
   {
     assert(!empty());
-    return Kokkos::subview(_internal_and_leaf_nodes,
-                           std::make_pair(size() - 1, 2 * size() - 1));
+    return _leaf_nodes;
   }
-  Kokkos::View<node_type const *, MemorySpace> getLeafNodes() const
+  Kokkos::View<leaf_node_type const *, MemorySpace> getLeafNodes() const
   {
     assert(!empty());
-    return Kokkos::subview(_internal_and_leaf_nodes,
-                           std::make_pair(size() - 1, 2 * size() - 1));
-  }
-
-  KOKKOS_FUNCTION
-  bounding_volume_type const *getRootBoundingVolumePtr() const
-  {
-    // Need address of the root node's bounding box to copy it back on the host,
-    // but can't access _internal_and_leaf_nodes elements from the constructor
-    // since the data is on the device.
-    assert(Details::HappyTreeFriends::getRoot(*this) == 0 &&
-           "workaround below assumes root is stored as first element");
-    return &_internal_and_leaf_nodes.data()->bounding_volume;
+    return _leaf_nodes;
   }
 
   size_t _size;
   bounding_volume_type _bounds;
-  Kokkos::View<node_type *, MemorySpace> _internal_and_leaf_nodes;
+  Kokkos::View<node_type *, MemorySpace> _internal_nodes;
+  Kokkos::View<leaf_node_type *, MemorySpace> _leaf_nodes;
+  Kokkos::View<Box, MemorySpace> _scene_bounding_box;
 };
 
 template <typename DeviceType>
@@ -200,10 +207,13 @@ BasicBoundingVolumeHierarchy<MemorySpace, BoundingVolume, Enable>::
                                  Primitives const &primitives,
                                  SpaceFillingCurve const &curve)
     : _size(AccessTraits<Primitives, PrimitivesTag>::size(primitives))
-    , _internal_and_leaf_nodes(
-          Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
-                             "ArborX::BVH::internal_and_leaf_nodes"),
-          _size > 0 ? 2 * _size - 1 : 0)
+    , _internal_nodes(Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
+                                         "ArborX::BVH::internal_nodes"),
+                      _size > 0 ? _size - 1 : 0)
+    , _leaf_nodes(Kokkos::view_alloc(space, Kokkos::WithoutInitializing,
+                                     "ArborX::BVH::leaf_nodes"),
+                  _size)
+    , _scene_bounding_box("ArborX::BVH::scene_bounding_box")
 {
   static_assert(
       KokkosExt::is_accessible_from<MemorySpace, ExecutionSpace>::value, "");
@@ -228,19 +238,16 @@ BasicBoundingVolumeHierarchy<MemorySpace, BoundingVolume, Enable>::
   Box bbox{};
   Details::TreeConstruction::calculateBoundingBoxOfTheScene(space, primitives,
                                                             bbox);
+  _bounds = bounding_volume_type{};
+  _bounds += bbox;
 
+  Kokkos::deep_copy(space, _scene_bounding_box, bbox);
   Kokkos::Profiling::popRegion();
 
   if (size() == 1)
   {
-    Details::TreeConstruction::initializeSingleLeafNode(
-        space, primitives, _internal_and_leaf_nodes);
-    Kokkos::deep_copy(
-        space,
-        Kokkos::View<BoundingVolume, Kokkos::HostSpace,
-                     Kokkos::MemoryUnmanaged>(&_bounds),
-        Kokkos::View<BoundingVolume const, MemorySpace,
-                     Kokkos::MemoryUnmanaged>(getRootBoundingVolumePtr()));
+    Details::TreeConstruction::initializeSingleLeafNode(space, primitives,
+                                                        _leaf_nodes, bbox);
     return;
   }
 
@@ -270,14 +277,7 @@ BasicBoundingVolumeHierarchy<MemorySpace, BoundingVolume, Enable>::
   // generate bounding volume hierarchy
   Details::TreeConstruction::generateHierarchy(
       space, primitives, permutation_indices, linear_ordering_indices,
-      getLeafNodes(), getInternalNodes());
-
-  Kokkos::deep_copy(
-      space,
-      Kokkos::View<BoundingVolume, Kokkos::HostSpace, Kokkos::MemoryUnmanaged>(
-          &_bounds),
-      Kokkos::View<BoundingVolume const, MemorySpace, Kokkos::MemoryUnmanaged>(
-          getRootBoundingVolumePtr()));
+      getLeafNodes(), getInternalNodes(), _scene_bounding_box);
 
   Kokkos::Profiling::popRegion();
 }

src/details/ArborX_Box.hpp

@@ -18,6 +18,8 @@
 
 #include <Kokkos_Macros.hpp>
 
+#include <iostream>
+
 namespace ArborX
 {
 /**
@@ -115,6 +117,128 @@ struct Box
   }
 #endif
 };
+
+struct DiscretizedBox
+{
+  static constexpr int N = (1 << 10);
+
+  // Initialize all max coordinates with 0 and all min coordinates with max.
+  KOKKOS_FUNCTION constexpr DiscretizedBox()
+      : _data(single_max + (0lu << 10) + (single_max << 20) + (0lu << 30) +
+              (single_max << 40) + (0lu << 50))
+  {
+  }
+
+  KOKKOS_FUNCTION constexpr double step(float min, float max) const
+  {
+    return ((double)max - min) / (N - 1);
+  }
+
+  KOKKOS_FUNCTION constexpr float restore(float min, int i, double h) const
+  {
+    return min + i * h;
+  }
+
+  KOKKOS_FUNCTION DiscretizedBox(Box const &local_box, Box const &scene_box)
+  {
+    const Point &scene_min_corner = scene_box.minCorner();
+    const Point &scene_max_corner = scene_box.maxCorner();
+    const Point &local_min_corner = local_box.minCorner();
+    const Point &local_max_corner = local_box.maxCorner();
+
+    _data = 0;
+
+    int shift = 0;
+    for (int d = 0; d < 3; ++d)
+    {
+      const auto h = step(scene_min_corner[d], scene_max_corner[d]);
+      int min = std::floor((local_min_corner[d] - scene_min_corner[d]) / h);
+      if (restore(scene_min_corner[d], min, h) > local_min_corner[d])
+      {
+        printf("[%d] adjusting min down\n", d);
+        --min;
+      }
+      _data |= (((std::uint64_t)min) << shift);
+      shift += 10;
+
+      int max = std::ceil((local_max_corner[d] - scene_min_corner[d]) / h);
+      if (restore(scene_min_corner[d], max, h) < local_max_corner[d])
+      {
+        printf("[%d] adjusting max up\n", d);
+        ++max;
+      }
+
+      _data |= (((std::uint64_t)max) << shift);
+      shift += 10;
+    }
+  }
+
+  KOKKOS_FUNCTION constexpr Box to_box(Box const &scene_box) const
+  {
+    const Point &scene_min_corner = scene_box.minCorner();
+    const Point &scene_max_corner = scene_box.maxCorner();
+
+    Box box;
+    int shift = 0;
+    for (int d = 0; d < 3; ++d)
+    {
+      const auto h = step(scene_min_corner[d], scene_max_corner[d]);
+
+      box.minCorner()[d] =
+          restore(scene_min_corner[d], (_data >> shift) & single_max, h);
+      shift += 10;
+      box.maxCorner()[d] =
+          restore(scene_min_corner[d], (_data >> shift) & single_max, h);
+      shift += 10;
+    }
+
+    return box;
+  }
+
+  KOKKOS_FUNCTION constexpr DiscretizedBox &
+  operator+=(DiscretizedBox const &other)
+  {
+    using KokkosExt::max;
+    using KokkosExt::min;
+
+    _data = min(_data & min_x_mask, other._data & min_x_mask) |
+            max(_data & max_x_mask, other._data & max_x_mask) |
+            min(_data & min_y_mask, other._data & min_y_mask) |
+            max(_data & max_y_mask, other._data & max_y_mask) |
+            min(_data & min_z_mask, other._data & min_z_mask) |
+            max(_data & max_z_mask, other._data & max_z_mask);
+
+    return *this;
+  }
+
+private:
+  std::uint64_t _data;
+  static constexpr std::uint64_t single_max = (1lu << 10) - 1;
+  static constexpr std::uint64_t min_x_mask = single_max;
+  static constexpr std::uint64_t max_x_mask = single_max << 10;
+  static constexpr std::uint64_t min_y_mask = single_max << 20;
+  static constexpr std::uint64_t max_y_mask = single_max << 30;
+  static constexpr std::uint64_t min_z_mask = single_max << 40;
+  static constexpr std::uint64_t max_z_mask = single_max << 50;
+};
+
+template <typename T>
+KOKKOS_INLINE_FUNCTION Box convert_to_box(const T &t, const Box &global_box);
+
+template <>
+KOKKOS_INLINE_FUNCTION Box convert_to_box<Box>(const Box &local_box,
+                                               const Box &)
+{
+  return local_box;
+}
+
+template <>
+KOKKOS_INLINE_FUNCTION Box convert_to_box<DiscretizedBox>(
+    const DiscretizedBox &local_box, const Box &global_box)
+{
+  return local_box.to_box(global_box);
+}
+
 } // namespace ArborX
 
 #endif

src/details/ArborX_DetailsAlgorithms.hpp

@@ -169,6 +169,12 @@ KOKKOS_INLINE_FUNCTION void expand(BOX &box, BOX const &other)
   box += other;
 }
 
+KOKKOS_INLINE_FUNCTION void expand(DiscretizedBox &box,
+                                   DiscretizedBox const &other)
+{
+  box += other;
+}
+
 // expand an axis-aligned bounding box to include a sphere
 KOKKOS_INLINE_FUNCTION
 void expand(Box &box, Sphere const &sphere)
@@ -184,6 +190,33 @@ void expand(Box &box, Sphere const &sphere)
   }
 }
 
+template <typename T1, typename T2>
+KOKKOS_INLINE_FUNCTION void expand_helper(T1 &t1, T2 const &t2,
+                                          Box const & /*global_box*/)
+{
+  expand(t1, t2);
+}
+
+KOKKOS_INLINE_FUNCTION void
+expand_helper(DiscretizedBox &box, Point const &other, Box const &global_box)
+{
+  box += DiscretizedBox(Box(other, other), global_box);
+}
+
+KOKKOS_INLINE_FUNCTION void expand_helper(DiscretizedBox &box, Box const &other,
+                                          Box const &global_box)
+{
+  box += DiscretizedBox(other, global_box);
+}
+
+KOKKOS_INLINE_FUNCTION
+constexpr Box
+convert_from_discretized_box(const DiscretizedBox &discretized_box,
+                             const Box &global_box)
+{
+  return discretized_box.to_box(global_box);
+}
+
 // check if two axis-aligned bounding boxes intersect
 KOKKOS_INLINE_FUNCTION
 constexpr bool intersects(Box const &box, Box const &other)

src/details/ArborX_DetailsDistributedTreeImpl.hpp

@@ -498,8 +498,9 @@ struct CallbackWithDistance
     // the details of the local tree. Right now, this is the only way. Will
     // need to be fixed with a proper callback abstraction.
     int const leaf_node_index = _rev_permute(index);
-    auto const &leaf_node_bounding_volume =
-        HappyTreeFriends::getBoundingVolume(_tree, leaf_node_index);
+    auto const &leaf_node_bounding_volume = convert_from_discretized_box(
+        HappyTreeFriends::getBoundingVolume(_tree, leaf_node_index),
+        _tree._scene_bounding_box());
     out({index, distance(getGeometry(query), leaf_node_bounding_volume)});
   }
 };