broadphase_arborx.hpp

Go to the documentation of this file.

 
#ifndef DEM_BROADPHASE_ARBORX_HPP
#define DEM_BROADPHASE_ARBORX_HPP
 
#include <ArborX.hpp>
#include <Kokkos_Core.hpp>
#include <utility>
 
namespace peclet::dem {
 
using BpExec = Kokkos::DefaultExecutionSpace;
using BpMem = BpExec::memory_space;
 
// View params are templated so callers can pass managed Views (tests) OR unmanaged Views wrapping
// raw device buffers (the CUDA dem module's float4*/int2* arrays). Element types/layout must be
// the usual (positions [N][3], radii [N], pairs [N][2], count scalar).
template <class PosV, class RadV, class PairsV, class CountV>
inline int findCollisionsArborX(PosV pos, RadV rad, int numParticles, int numReal, float margin,
                                PairsV outPairs, CountV outCount, float boxCap = 0.0f) {
  BpExec space;
  using Box = ArborX::Box<3>;
 
  // Sanitize the AABB half-width + centre: a non-finite (NaN/inf) or absurdly large box — a
  // particle blown up by a bad coupling/DEM transient — would otherwise overlap ~every other box
  // and make the ArborX query allocate billions of intersections (observed: a 36 GiB "bp::values"
  // OOM). SANE_BOX parks such a particle in a degenerate far box that matches nothing, so one bad
  // particle can't crash the run. boxCap<=0 disables the size clamp (NaN/inf are always caught).
#define PECLET_DEM_SANE_BOX(i, px, py, pz, b)                                                    \
  float px = pos(i, 0), py = pos(i, 1), pz = pos(i, 2), b = rad(i) + margin;                     \
  if (!(b == b) || !(px == px) || !(py == py) || !(pz == pz) || (boxCap > 0.0f && b > boxCap)) { \
    px = py = pz = -1.0e30f;                                                                     \
    b = 0.0f;                                                                                    \
  }
 
  // AABBs over all particles (these are the BVH primitives).
  Kokkos::View<Box*, BpMem> boxes(
      Kokkos::view_alloc(space, "peclet::dem::bp::boxes", Kokkos::WithoutInitializing),
      numParticles);
  Kokkos::parallel_for(
      "peclet::dem::bp::aabb", Kokkos::RangePolicy<BpExec>(space, 0, numParticles),
      KOKKOS_LAMBDA(int i) {
        PECLET_DEM_SANE_BOX(i, px, py, pz, b);
        boxes(i) = Box{{px - b, py - b, pz - b}, {px + b, py + b, pz + b}};
      });
 
  ArborX::BoundingVolumeHierarchy const tree(space, ArborX::Experimental::attach_indices(boxes));
 
  // One intersection query per real particle (box of the same half-width).
  using Predicate = decltype(ArborX::intersects(std::declval<Box>()));
  Kokkos::View<Predicate*, BpMem> preds(
      Kokkos::view_alloc(space, "peclet::dem::bp::preds", Kokkos::WithoutInitializing), numReal);
  Kokkos::parallel_for(
      "peclet::dem::bp::preds", Kokkos::RangePolicy<BpExec>(space, 0, numReal),
      KOKKOS_LAMBDA(int i) {
        PECLET_DEM_SANE_BOX(i, px, py, pz, b);
        preds(i) = ArborX::intersects(Box{{px - b, py - b, pz - b}, {px + b, py + b, pz + b}});
      });
#undef PECLET_DEM_SANE_BOX
 
  Kokkos::View<typename decltype(tree)::value_type*, BpMem> values("peclet::dem::bp::values", 0);
  Kokkos::View<int*, BpMem> offsets("peclet::dem::bp::offsets", 0);
  tree.query(space, preds, values, offsets);
 
  // Compact the (i, j) results with i<j into the output buffer (matches cuBQL's i<j filter).
  Kokkos::deep_copy(space, outCount, 0);
  const int maxPairs = static_cast<int>(outPairs.extent(0));
  Kokkos::View<int* [2], BpMem> pairs = outPairs;
  Kokkos::View<int, BpMem> cnt = outCount;
  Kokkos::parallel_for(
      "peclet::dem::bp::emit", Kokkos::RangePolicy<BpExec>(space, 0, numReal),
      KOKKOS_LAMBDA(int i) {
        for (int k = offsets(i); k < offsets(i + 1); ++k) {
          const int j = values(k).index;
          if (i < j) {
            const int slot = Kokkos::atomic_fetch_add(&cnt(), 1);
            if (slot < maxPairs) {
              pairs(slot, 0) = i;
              pairs(slot, 1) = j;
            }
          }
        }
      });
  space.fence();
 
  int h_count = 0;
  Kokkos::deep_copy(h_count, outCount);
  return h_count;
}
 
}  // namespace peclet::dem
 
#endif  // DEM_BROADPHASE_ARBORX_HPP