Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enable generation of a 2D neighbor list using ArborX #248

Merged
merged 10 commits into from
Jun 5, 2020
Merged

Enable generation of a 2D neighbor list using ArborX #248

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
82db633
Fixup let size() member function in access traits specialization be a…
dalg24 Jun 4, 2020
d126eb0
Fixup use MemorySpace type alias in CrsGraph specialization of Neighb…
dalg24 Jun 4, 2020
5d528c8
Remove redundant bound check in NeighborList<CrsGraph>::getNeighbor()>
dalg24 Jun 4, 2020
41f2902
Abstract away the logic to filter collisions to get full/half neighbo…
dalg24 Jun 4, 2020
a8cb345
Add couple new custom callbacks to populate a 2D neighbor list
dalg24 Jun 4, 2020
1f261a6
Add Experimental::make2DNeighborList() that uses the bounding volume …
dalg24 Jun 4, 2020
6974ab4
Add new specialization of NeighborList for the 2D data structure
dalg24 Jun 4, 2020
150552d
Add unit tests for experimental 2D neighbor list with ArborX
dalg24 Jun 4, 2020
2053903
Fixup typo in assert count -> counts
dalg24 Jun 5, 2020
c79f7ac
Comment that counter increment should be atomic
dalg24 Jun 5, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
154 changes: 138 additions & 16 deletions core/src/Cabana_Experimental_NeighborList.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ struct Access<Slice, PrimitivesTag,
{
using memory_space = typename Slice::memory_space;
using size_type = typename Slice::size_type;
static size_type size( Slice const &x ) { return x.size(); }
static KOKKOS_FUNCTION size_type size( Slice const &x ) { return x.size(); }
static KOKKOS_FUNCTION Point get( Slice const &x, size_type i )
{
return {static_cast<float>( x( i, 0 ) ),
Expand Down Expand Up @@ -114,40 +114,86 @@ namespace Experimental
namespace Impl
{

// Custom callback for ArborX::BVH::query()
template <typename Tag>
struct NeighborDiscriminatorCallback;
struct CollisionFilter;

template <>
struct CollisionFilter<FullNeighborTag>
{
KOKKOS_FUNCTION bool static keep( int i, int j ) noexcept
{
return i != j; // discard self-collision
}
};

template <>
struct NeighborDiscriminatorCallback<Cabana::FullNeighborTag>
struct CollisionFilter<HalfNeighborTag>
{
KOKKOS_FUNCTION static bool keep( int i, int j ) noexcept { return i > j; }
};

// Custom callback for ArborX::BVH::query()
template <typename Tag>
struct NeighborDiscriminatorCallback
{
using tag = ArborX::Details::InlineCallbackTag;
template <typename Predicate, typename OutputFunctor>
KOKKOS_FUNCTION void operator()( Predicate const &predicate, int i,
KOKKOS_FUNCTION void operator()( Predicate const &predicate,
int primitive_index,
OutputFunctor const &out ) const
{
if ( getData( predicate ) != i ) // discard self-collision
int const predicate_index = getData( predicate );
if ( CollisionFilter<Tag>::keep( predicate_index, primitive_index ) )
{
out( i );
out( primitive_index );
}
}
};

template <>
struct NeighborDiscriminatorCallback<Cabana::HalfNeighborTag>
// Count in the first pass
template <typename Counts, typename Tag>
struct NeighborDiscriminatorCallback2D_FirstPass
{
Counts counts;
using tag = ArborX::Details::InlineCallbackTag;
template <typename Predicate, typename OutputFunctor>
KOKKOS_FUNCTION void operator()( Predicate const &predicate, int i,
OutputFunctor const &out ) const
KOKKOS_FUNCTION void operator()( Predicate const &predicate,
int primitive_index,
OutputFunctor const & ) const
{
int const predicate_index = getData( predicate );
if ( CollisionFilter<Tag>::keep( predicate_index, primitive_index ) )
{
++counts( predicate_index ); // WARNING see below**
}
}
};

// Fill in the second pass
template <typename Counts, typename Neighbors, typename Tag>
struct NeighborDiscriminatorCallback2D_SecondPass
{
Counts counts;
Neighbors neighbors;
using tag = ArborX::Details::InlineCallbackTag;
template <typename Predicate, typename OutputFunctor>
KOKKOS_FUNCTION void operator()( Predicate const &predicate,
int primitive_index,
OutputFunctor const & ) const
{
if ( getData( predicate ) > i )
int const predicate_index = getData( predicate );
if ( CollisionFilter<Tag>::keep( predicate_index, primitive_index ) )
{
out( i );
assert( counts( predicate_index ) < (int)neighbors.extent( 1 ) );
neighbors( predicate_index, counts( predicate_index )++ ) =
primitive_index; // WARNING see below**
}
}
};

// NOTE** Taking advantage of the knowledge that one predicate is processed by a
// single thread. Count increment should be atomic otherwise.

} // namespace Impl

template <typename MemorySpace, typename Tag>
Expand Down Expand Up @@ -179,8 +225,59 @@ auto makeNeighborList( Tag, Slice const &coordinate_slice,
Impl::makePredicates( coordinate_slice, first, last, radius ),
Impl::NeighborDiscriminatorCallback<Tag>{}, indices, offset );

return CrsGraph<typename DeviceType::memory_space, Tag>{
std::move( indices ), std::move( offset ), first, bvh.size()};
return CrsGraph<MemorySpace, Tag>{std::move( indices ), std::move( offset ),
first, bvh.size()};
}

template <typename MemorySpace, typename Tag>
struct Dense
{
Kokkos::View<int *, MemorySpace> cnt;
Kokkos::View<int **, MemorySpace> val;
typename MemorySpace::size_type shift;
typename MemorySpace::size_type total;
};

template <typename DeviceType, typename Slice, typename Tag>
auto make2DNeighborList( Tag, Slice const &coordinate_slice,
typename Slice::size_type first,
typename Slice::size_type last,
typename Slice::value_type radius )
{
using MemorySpace = typename DeviceType::memory_space;
using ExecutionSpace = typename DeviceType::execution_space;
ExecutionSpace space{};

ArborX::BVH<MemorySpace> bvh( space, coordinate_slice );

Kokkos::View<int *, DeviceType> indices( "indices", 0 );
Kokkos::View<int *, DeviceType> offset( "offset", 0 );

auto const predicates =
Impl::makePredicates( coordinate_slice, first, last, radius );

auto const n_queries = ArborX::Traits::
Access<decltype( predicates ), ArborX::Traits::PredicatesTag>::size(
predicates );

Kokkos::View<int *, DeviceType> counts( "counts", n_queries );
bvh.query(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK so you do your own 2-pass here - do you force ArborX to do the queries in 1-pass mode so this is effectively still just a 2-pass process as you are filling your own data structure?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess you must check if there is enough space allocated by the user but you never fill the native data structure so ArborX will always do 1-pass because the internal count never goes past 0?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct. We detect that you never called the output functor and do a single pass. Note that we are working on adding a new BVH::query() overload that does only takes the callback (no offsets and values).

space, predicates,
Impl::NeighborDiscriminatorCallback2D_FirstPass<decltype( counts ),
Tag>{counts},
indices, offset );

Kokkos::View<int **, DeviceType> neighbors(
Kokkos::view_alloc( "neighbors", Kokkos::WithoutInitializing ),
n_queries, ArborX::max( space, counts ) );
Kokkos::deep_copy( counts, 0 ); // reset counts to zero
bvh.query(
space, predicates,
Impl::NeighborDiscriminatorCallback2D_SecondPass<
decltype( counts ), decltype( neighbors ), Tag>{counts, neighbors},
indices, offset );

return Dense<MemorySpace, Tag>{counts, neighbors, first, bvh.size()};
}

} // namespace Experimental
Expand All @@ -205,13 +302,38 @@ class NeighborList<Experimental::CrsGraph<MemorySpace, Tag>>
static KOKKOS_FUNCTION size_type
getNeighbor( crs_graph_type const &crs_graph, size_type p, size_type n )
{
assert( (int)p >= 0 && p < crs_graph.total );
assert( n < numNeighbor( crs_graph, p ) );
p -= crs_graph.shift;
return crs_graph.col_ind( crs_graph.row_ptr( p ) + n );
}
};

template <typename MemorySpace, typename Tag>
class NeighborList<Experimental::Dense<MemorySpace, Tag>>
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like this Dense type name

{
using size_type = std::size_t;
using specialization_type = Experimental::Dense<MemorySpace, Tag>;

public:
using memory_space = MemorySpace;
static KOKKOS_FUNCTION size_type numNeighbor( specialization_type const &d,
size_type p )
{
assert( (int)p >= 0 && p < d.total );
p -= d.shift;
if ( (int)p < 0 || p >= d.cnt.size() )
return 0;
return d.cnt( p );
}
static KOKKOS_FUNCTION size_type getNeighbor( specialization_type const &d,
size_type p, size_type n )
{
assert( n < numNeighbor( d, p ) );
p -= d.shift;
return d.val( p, n );
}
};

} // namespace Cabana

#endif
109 changes: 82 additions & 27 deletions core/unit_test/tstNeighborListArborX.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,13 +53,26 @@ void testArborXListHalf()
auto aosoa = createParticles( num_particle, box_min, box_max );
auto position = Cabana::slice<0>( aosoa );

// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist = Cabana::Experimental::makeNeighborList<device_type>(
Cabana::HalfNeighborTag{}, position, 0, aosoa.size(), test_radius );

// Check the neighbor list.
checkHalfNeighborList( nlist, position, test_radius );
{
// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist = Cabana::Experimental::makeNeighborList<device_type>(
Cabana::HalfNeighborTag{}, position, 0, aosoa.size(), test_radius );

// Check the neighbor list.
checkHalfNeighborList( nlist, position, test_radius );
}
{
// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist =
Cabana::Experimental::make2DNeighborList<device_type>(
Cabana::HalfNeighborTag{}, position, 0, aosoa.size(),
test_radius );

// Check the neighbor list.
checkHalfNeighborList( nlist, position, test_radius );
}
}

//---------------------------------------------------------------------------//
Expand All @@ -74,14 +87,28 @@ void testArborXListFullPartialRange()
auto aosoa = createParticles( num_particle, box_min, box_max );
auto position = Cabana::slice<0>( aosoa );

// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist = Cabana::Experimental::makeNeighborList<device_type>(
Cabana::FullNeighborTag{}, position, 0, num_ignore, test_radius );

// Check the neighbor list.
checkFullNeighborListPartialRange( nlist, position, test_radius,
num_ignore );
{
// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist = Cabana::Experimental::makeNeighborList<device_type>(
Cabana::FullNeighborTag{}, position, 0, num_ignore, test_radius );

// Check the neighbor list.
checkFullNeighborListPartialRange( nlist, position, test_radius,
num_ignore );
}
{
// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist =
Cabana::Experimental::make2DNeighborList<device_type>(
Cabana::FullNeighborTag{}, position, 0, num_ignore,
test_radius );

// Check the neighbor list.
checkFullNeighborListPartialRange( nlist, position, test_radius,
num_ignore );
}
}

//---------------------------------------------------------------------------//
Expand All @@ -95,14 +122,28 @@ void testNeighborArborXParallelFor()
auto aosoa = createParticles( num_particle, box_min, box_max );
auto position = Cabana::slice<0>( aosoa );

// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist = Cabana::Experimental::makeNeighborList<device_type>(
Cabana::FullNeighborTag{}, position, 0, aosoa.size(), test_radius );
{
// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist = Cabana::Experimental::makeNeighborList<device_type>(
Cabana::FullNeighborTag{}, position, 0, aosoa.size(), test_radius );

checkFirstNeighborParallelFor( nlist, position, test_radius );
checkFirstNeighborParallelFor( nlist, position, test_radius );

checkSecondNeighborParallelFor( nlist, position, test_radius );
checkSecondNeighborParallelFor( nlist, position, test_radius );
}
{
// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist =
Cabana::Experimental::make2DNeighborList<device_type>(
Cabana::FullNeighborTag{}, position, 0, aosoa.size(),
test_radius );

checkFirstNeighborParallelFor( nlist, position, test_radius );

checkSecondNeighborParallelFor( nlist, position, test_radius );
}
}

//---------------------------------------------------------------------------//
Expand All @@ -116,14 +157,28 @@ void testNeighborArborXParallelReduce()
auto aosoa = createParticles( num_particle, box_min, box_max );
auto position = Cabana::slice<0>( aosoa );

// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist = Cabana::Experimental::makeNeighborList<device_type>(
Cabana::FullNeighborTag{}, position, 0, aosoa.size(), test_radius );
{
// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist = Cabana::Experimental::makeNeighborList<device_type>(
Cabana::FullNeighborTag{}, position, 0, aosoa.size(), test_radius );

checkFirstNeighborParallelReduce( nlist, aosoa, test_radius );

checkSecondNeighborParallelReduce( nlist, aosoa, test_radius );
}
{
// Create the neighbor list.
using device_type = TEST_MEMSPACE; // sigh...
auto const nlist =
Cabana::Experimental::make2DNeighborList<device_type>(
Cabana::FullNeighborTag{}, position, 0, aosoa.size(),
test_radius );

checkFirstNeighborParallelReduce( nlist, aosoa, test_radius );
checkFirstNeighborParallelReduce( nlist, aosoa, test_radius );

checkSecondNeighborParallelReduce( nlist, aosoa, test_radius );
checkSecondNeighborParallelReduce( nlist, aosoa, test_radius );
}
}

//---------------------------------------------------------------------------//
Expand Down