-
Notifications
You must be signed in to change notification settings - Fork 40
ArborX::AccessTraits
Defined in header <ArborX_AccessTraits.hpp>
template <typename T, typename Tag, typename Enable=void>
struct AccessTraits;
The struct template ArborX::AccessTraits
tells the ArborX::BVH
constructor how much data to index and how to access that data,
or ArborX::BVH::query()
what predicates to query for.
ArborX provides the following partial specialization for Kokkos views
template <typename View, typename Tag>
AccessTraits<View, Tag, std::enable_if_t<Kokkos::is_view<View>{}>>;
struct PrimitivesTag { };
struct PredicatesTag { };
Empty class types used to indicate the data category.
T
: User data type.
Tag
: Tag to distinguish between primitives and predicates.
Member type | Description |
---|---|
memory_space |
A valid Kokkos memory space |
size |
returns the number of elements |
get |
access the specified element |
#include <ArborX.hpp>
#include <Kokkos_Core.hpp>
#include <array>
#include <iostream>
#include <numeric>
struct PointCloud
{
float *d_x;
float *d_y;
float *d_z;
int N;
};
struct NearestToOrigin
{
int k;
};
template <>
struct ArborX::AccessTraits<PointCloud, ArborX::PrimitivesTag>
{
static KOKKOS_FUNCTION std::size_t size(PointCloud const &cloud)
{
return cloud.N;
}
static KOKKOS_FUNCTION static ArborX::Point get(PointCloud const &cloud,
std::size_t i)
{
return {{cloud.d_x[i], cloud.d_y[i], cloud.d_z[i]}};
}
using memory_space = Kokkos::CudaSpace;
};
template <>
struct ArborX::AccessTraits<NearestToOrigin, ArborX::PredicatesTag>
{
static KOKKOS_FUNCTION std::size_t size(NearestToOrigin const &)
{
return 1;
}
static KOKKOS_FUNCTION auto get(NearestToOrigin const &d, std::size_t)
{
return ArborX::nearest(ArborX::Point{{0, 0, 0}}, d.k);
}
using memory_space = Kokkos::CudaSpace;
};
int main(int argc, char *argv[])
{
Kokkos::ScopeGuard guard(argc, argv);
constexpr std::size_t N = 1000;
std::array<float, N> a;
float *d_a;
cudaMalloc(&d_a, sizeof(float) * N);
std::iota(std::begin(a), std::end(a), 1.0);
cudaMemcpy(d_a, a.data(), sizeof(float) * N, cudaMemcpyHostToDevice);
using device_type = Kokkos::Cuda::device_type;
ArborX::BVH<Kokkos::CudaSpace> bvh{Kokkos::Cuda{}, PointCloud{d_a, d_a, d_a, N}};
Kokkos::View<int *, device_type> indices("indices", 0);
Kokkos::View<int *, device_type> offset("offset", 0);
bvh.query(Kokkos::Cuda{}, NearestToOrigin{5}, indices, offset);
Kokkos::parallel_for(1, KOKKOS_LAMBDA(int i) {
for (int j = offset(i); j < offset(i + 1); ++j)
{
printf("%i %i\n", i, indices(j));
}
});
return 0;
}
Output
0 0
0 1
0 2
0 3
0 4
static KOKKOS_FUNCTION size_type size(T const& data);
data
: data argument passed to the ArborX::BVH<MemorySpace>::BVH
constructor or ArborX::BVH<MemorySpace>::query()
The number of primitives if the second template argument is ArborX::PrimitivesTag
or the number of queries to perform if it is ArborX::PredicatesTag
.
static KOKKOS_FUNCTION context_dependent get(T const& data, size_type pos);
If std::is_same<Tag,PrimitimesTag>::value
is true
, the return value type must either decay to ArborX::Point
or to ArborX::Box
.
If std::is_same<Tag,PredicatesTag>::value
is true
, the return value type must decay to a valid ArborX predicate, e.g. a predicate generated by ArborX::nearest()
or ArborX::intersects()
.
get()
needs to be marked with the KOKKOS_FUNCTION
or KOKKOS_INLINE_FUNCTION
macro because it is called from the device.
data
: data argument passed to the ArborX::BVH<MemorySpace>::BVH
constructor or ArborX::BVH<MemorySpace>::query()
pos
: position of the primitive from which to return the centroid or the smallest bounding volume