API v2

[aprokop]

API v2 design idea

TODO

Remaining items that probably need to be addressed:

• Internal node callbacks
• Storing data associated with internal nodes
• Customizing traversal routines

Range adaptor

// There is no need to differentiate between Values and Predicates, RangeAdaptor should work for both
template<typename Container>
struct RangeAdaptorArchetypeExpression {
  // memory_space is needed to understand where the return type of `get()` resides
  using memory_space = typename Container::memory_space;
  using size_type = typename Container::size_type;
  using value_type = typename Container::value_type;

  KOKKOS_FUNCTION static size_type size(Container const& container);
  KOKKOS_FUNCTION static value_type get(Container const& container, size_type index);
};

template<typename Container, typename Enable = void>
struct RangeAdaptor
{
  using not_specialized = void;
};

// Specialization for 1D Kokkos View
template <typename View>
struct RangeAdaptor<View, std::enable_if_t<Kokkos::is_view<View>{} && View::rank == 1>>
{
  using memory_space = typename View::memory_space;
  using size_type = typename View::size_type;
  using value_type = typename View::const_value_type;

  KOKKOS_FUNCTION static size_type size(View const &v) { return v.extent(0); }
  KOKKOS_FUNCTION static value_type &get(View const &v, size_type index) { return v(i); }
};

Indexable getter

template <typename ValueType, typename Indexable>
struct IndexableGetterArchetypeExpression
{
  using value_type = ValueType;
  using indexable_type = Indexable;

  // May be called multiple times during execution
  KOKKOS_FUNCTION Indexable operator()(ValueType const &value) const;
};

// TODO: maybe it should be not specialized, with Point and Box having specializations instead
template <typename ValueType, typename Indexable = ValueType>
struct DefaultIndexableGetter
{
  using value_type = ValueType;
  using indexable_type = Indexable;

  KOKKOS_FUNCTION static Indexable const& operator()(ValueType const &value) { return value; }
}
template <>
struct DefaultIndexableGetter<Sphere, Box>
{
  using value_type = Sphere;
  using indexable_type = Box;

  KOKKOS_FUNCTION static Box const& operator()(Sphere const &value)
  {
    auto const& c = sphere.centroid();
    auto const r = sphere.radius();
    return Box{{c[0] - r, c[1] - r, c[2] - r}, {c[0] + r, c[1] + r, c[2] + r}};
  }
}

Callbacks

// return_type can be
// - void or CallbackTreeTraversalControl (for spatial)
// - void (for knn)
struct CallbackArchetypeExpression
{
  template <typename Predicate, typename ValueType>
  KOKKOS_FUNCTION return_type operator()(Predicate const &predicate, ValueType const& &value) const;
};

struct CallbackWithOutputArchetypeExpression
{
  template <typename Predicate, typename ValueType, typename OutputFunctor>
  KOKKOS_FUNCTION return_type operator()(Predicate const &predicate, ValueType const& value, OutputFunctor const& output) const;
};

LBVH

template <typename MemorySpace, typename ValueType, typename IndexableGetter, typename Parameters...>
class LBVH
{
public:
  using memory_space = MemorySpace;
  using size_type = typename MemorySpace::size_type;
  using value_type = ValueType;
  using bounding_volume_type = ; // get from the Parameters

private:
  // An example demonstrating a straightforward implementation
  struct LeafNode
  {
    int rope;
    value_type value;
  };
  struct InternalNode
  {
    int left_child;
    int rope;
    bounding_volume_type bounding_volume;
  };

  size_type _size;
  bounding_volume_type _bounds;
  Kokkos::View<LeafNode *, MemorySpace> _leaf_nodes;
  Kokkos::View<InternalNode *, MemorySpace> _internal_nodes;
  IndexableGetter _indexable_getter;

public:
  LBVH() = default;

  template <typename ExecutionSpace, template Values>
  LBVH(ExecutionSpace const& exec_space, Values const& values, IndexableGetter const& indexable_getter);

  template <typename ExecutionSpace, typename Predicates, typename Callback>
  void query(ExecutionSpace const &space, Predicates const &predicates, Callback const &callback) const;

  KOKKOS_FUNCTION size_type size() const noexcept;
  KOKKOS_FUNCTION bool empty() const noexcept;
  KOKKOS_FUNCTION bounding_volume_type bounds() const noexcept;
};

template <typename MemorySpace, typename IndexableGetter, typename Parameters...>
template <typename ExecutionSpace, typename Values>
LBVH<MemorySpace, IndexableGetter, Parameters>::
    LBVH(ExecutionSpace const &exec_space, Values const& values, IndexableGetter const& indexable_getter)
    : _size(RangeAdaptor<Values>::size(values))
    , _leaf_nodes("leaf_nodes", _size > 0 ? _size - 1 : 0)
    , _internal_nodes("internal_nodes", _size)
    , _indexable_getter(indexable_getter)
{
  using Range = RangeAdaptor<Values>;
  static_assert(std::is_same<typename Range::memory_space, typename IndexableGetter::memory_space>{}, "");
  static_assert(KokkosExt::is_accessible_from<typename Range::memory_space,
                                              ExecutionSpace>::value, "");

  auto const n = Range::size(values);

  // Example: calculating bounding volume of the scene
  bounding_volume_type scene_bounding_volume;
  Kokkos::parallel_reduce(Kokkos::RangePolicy<ExecutionSpace>(exec_space, 0, n),
    KOKKOS_LAMBDA(int i, bounding_volume_type &update) {
      update += indexable_getter(Range::get(values, i));
  }, scene_bounding_box);

  // Example: constructing leaf nodes
  Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(exec_space, 0, n),
    KOKKOS_LAMBDA(int i) {
      _leaf_nodes(i) = {ROPE_SENTINEL, Range::get(values, i)};
  });

  // Constructing hierarchy
  Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(exec_space, 0, n),
    KOKKOS_LAMBDA(int i) {
      bounding_volume_type bounding_folume{};
      bounding_volume = _indexable_getter(_leaf_nodes(i).value);
      ...
  });
}

template <typename MemorySpace, typeanme IndexableGetter, typename Parameters...>
template <typename ExecutionSpace, typename Predicates, typename Callback>
void LBVH<MemorySpace, IndexableGetter, Parameters>::
query(ExecutionSpace const &exec_space, Predicates const &predicates, Callback const &callback) const
{
  using Range = RangeAdaptor<Predicates>;
  auto const n = Range::size(predicates);

  Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(exec_space, 0, n),
    KOKKOS_LAMBDA(int const queryIndex) {
      auto const &predicate = Range::get(predicates, queryIndex);

      int node_index = _bvh.getRootIndex();
      do
      {
        auto const& node = getNode(node_index);
        if (!isLeaf(node_index))
        {
          if (predicate(node.bounding_volume))
            node_index = node.left_child;
          else
            node_index = node.rope;
        }
        else
        {
          if (predicate(_indexable_getter(node.value)))
            callback(predicate, node.value);
          node_index = node.rope;
        }

      } while (node_index != ROPE_SENTINEL);
  });
}

API v1 to API v2 adaptation

v1 is equivalent to v2 with hardcoded ValueType = int and Indexable = Box.

// The purpose of the LegacyIndexableGetter is to convert Primitives to Boxes that we then hold on.
template<typename MemorySpace>
struct LegacyIndexableGetter
{
  using value_type = int;
  using indexable_type = Box;

  Kokkos::View<Box*, MemorySpace> _boxes;

  template <typename ExecutionSpace, typename Primitives>
  LegacyIndexableGetter(ExecutionSpace const& exec_space, Primitives const& primitives)
    : _boxes(AccessTraits<Primitives, PrimitivesTag>::size(primitives))
  {
    using Access = AccessTraits<Primitives, PrimitivesTag>;
    using Type = std::decay_t<declval(Access::get(std::declval<Primitives>(), std::declval<size_t>()))>;

    if (Kokkos::is_view<Primitives>::value && std::is_same<Type, Box>{})
    {
      // Shorcut if Primitives is a View of boxes
      _boxes = primitives;
    }
    else
    {
      auto const n = Access::size(primitives);

      Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(0, n),
        KOKKOS_LAMBDA(int const i) {
          _boxes(i) += Access::get(primitives, i);
      });
    }
  }

  KOKKOS_FUNCTION indexable_type const& operator()(int i) { return _boxes(i); }
};

// The purpose of LegacyIota is to be able to specialize RangeAdaptor
class LegacyIota
{
private:
  using value_type = int;
  using size_type = size_t;

  size_type _size;

public:
  LegacyIota(size_type size) : _size(size) { }

  KOKKOS_FUNCTION size_type size() const { return _size; }
  KOKKOS_FUNCTION value_type operator()(size_type index) const { return index; }
};

template <>
RangeAdaptor<LegacyIota>
{
  using Container = LegacyIota;

  using memory_space = typename Container::memory_space;
  using value_type = typename Container::value_type;
  using size_type = typename Container::size_type;

  KOKKOS_FUNCTION static size_type size(Container const& container) { return container.size(); }
  KOKKOS_FUNCTION static value_type get(Container const& container, size_type index) { return container(index); }
};

template <typename MemorySpace>
class BoundingVolumeHierarchy : public LBVH<MemorySpace, int, LegacyIndexableGetter<MemorySpace>> {
private:
  using base_type = LBVH<MemorySpace, int, LegacyIndexableGetter<MemorySpace>>;
public:
  template <typename ExecutionSpace, typename Primitives>
  BoundingVolumeHierarchy(ExecutionSpace const& exec_space, Primitives const &primitives)
    : base_type(exec_space,
                LegacyIota(AccessTraits<Primitives, PrimitivesTag>::size(primitives)),
                LegacyIndexableGetter(primitives))
  { }

  template <typename ExecutionSpace, typename Predicates, typename Callback>
  void query(ExecutionSpace const &exec_space, Predicates const &predicates, Callback const &callback)
  {
    base_type::query(exec_space, predicates, callback);
  }
};

Examples

Examples demonstrating how one would use the proposed interface in certain situations

Find all points within certain distance

Description: Given a cloud of points and a radius r, for every point find all points within distance r from it.

Applications: HACC, LAMPPS

Find a cell that point belongs to

Description: Given a mesh and a cloud of points, for every point find a mesh cell it belongs to.

Applications: DTK, ExaWind

Find k closest neighbors and distances to them

Description: Given two clouds of points, for ever point find k nearest
neighbors and compute a quantity based on values of those k points weighted by
distances.

Applications: DTK, machine learning

Force calculation

Description: Given a cloud of points, compute a force for every point. Computation of force to distant nodes is replaced by computing a force to internal nodes (that contains the center and mass of its descendants).

Applications: HACC