Skip to content

Commit

Permalink
Merge pull request #713 from streeve/total_neighbors
Browse files Browse the repository at this point in the history 
Add total neighbor interface and fix max neighbors
  • Loading branch information
@streeve
streeve authored Nov 10, 2023
2 parents a189f57 + 098aa73 commit f99c7db
Show file tree
Hide file tree
Showing 5 changed files with 136 additions and 26 deletions.
34 changes: 32 additions & 2 deletions core/src/Cabana_Experimental_NeighborList.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -245,7 +245,7 @@ struct CrsGraph
Kokkos::View<int*, MemorySpace> row_ptr;
//! Neighbor offset shift.
typename MemorySpace::size_type shift;
//! Total neighbors.
//! Total particles.
typename MemorySpace::size_type total;
};

Expand Down Expand Up @@ -368,7 +368,7 @@ struct Dense
Kokkos::View<int**, MemorySpace> val;
//! Neighbor offset shift.
typename MemorySpace::size_type shift;
//! Total neighbors.
//! Total particles.
typename MemorySpace::size_type total;
};

Expand Down Expand Up @@ -536,6 +536,21 @@ class NeighborList<Experimental::CrsGraph<MemorySpace, Tag>>
public:
//! Kokkos memory space.
using memory_space = MemorySpace;

//! Get the total number of neighbors across all particles.
KOKKOS_INLINE_FUNCTION
static size_type totalNeighbor( crs_graph_type const& crs_graph )
{
return Impl::totalNeighbor( crs_graph, crs_graph.total );
}

//! Get the maximum number of neighbors across all particles.
KOKKOS_INLINE_FUNCTION
static size_type maxNeighbor( crs_graph_type const& crs_graph )
{
return Impl::maxNeighbor( crs_graph, crs_graph.total );
}

//! Get the number of neighbors for a given particle index.
static KOKKOS_FUNCTION size_type
numNeighbor( crs_graph_type const& crs_graph, size_type p )
Expand Down Expand Up @@ -568,6 +583,21 @@ class NeighborList<Experimental::Dense<MemorySpace, Tag>>
public:
//! Kokkos memory space.
using memory_space = MemorySpace;

//! Get the total number of neighbors across all particles.
KOKKOS_INLINE_FUNCTION
static size_type totalNeighbor( specialization_type const& d )
{
return Impl::totalNeighbor( d, d.total );
}

//! Get the maximum number of neighbors across all particles.
KOKKOS_INLINE_FUNCTION
static size_type maxNeighbor( specialization_type const& d )
{
return Impl::maxNeighbor( d, d.total );
}

//! Get the number of neighbors for a given particle index.
static KOKKOS_FUNCTION size_type numNeighbor( specialization_type const& d,
size_type p )
Expand Down
35 changes: 35 additions & 0 deletions core/src/Cabana_NeighborList.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,14 @@ class NeighborList
//! Kokkos memory space.
using memory_space = typename NeighborListType::memory_space;

//! Get the total number of neighbors across all particles.
KOKKOS_INLINE_FUNCTION
static std::size_t totalNeighbor( const NeighborListType& list );

//! Get the maximum number of neighbors across all particles.
KOKKOS_INLINE_FUNCTION
static std::size_t maxNeighbor( const NeighborListType& list );

//! Get the number of neighbors for a given particle index.
KOKKOS_INLINE_FUNCTION
static std::size_t numNeighbor( const NeighborListType& list,
Expand All @@ -79,6 +87,33 @@ class NeighborList

//---------------------------------------------------------------------------//

namespace Impl
{
//! Iterate to get the total number of neighbors.
template <class ListType>
KOKKOS_INLINE_FUNCTION std::size_t
totalNeighbor( const ListType& list, const std::size_t num_particles )
{
std::size_t total_n = 0;
// Sum neighbors across all particles.
for ( std::size_t p = 0; p < num_particles; p++ )
total_n += NeighborList<ListType>::numNeighbor( list, p );
return total_n;
}

//! Iterate to find the maximum number of neighbors.
template <class ListType>
KOKKOS_INLINE_FUNCTION std::size_t
maxNeighbor( const ListType& list, const std::size_t num_particles )
{
std::size_t max_n = 0;
for ( std::size_t p = 0; p < num_particles; p++ )
if ( NeighborList<ListType>::numNeighbor( list, p ) > max_n )
max_n = NeighborList<ListType>::numNeighbor( list, p );
return max_n;
}
} // namespace Impl

} // end namespace Cabana

#endif // end CABANA_NEIGHBORLIST_HPP
48 changes: 35 additions & 13 deletions core/src/Cabana_VerletList.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,10 @@ struct VerletListData<MemorySpace, VerletLayout2D>
//! Neighbor list.
Kokkos::View<int**, memory_space> neighbors;

//! Actual maximum neighbors per particle (potentially less than allocated
//! space).
std::size_t max_n;

//! Add a neighbor to the list.
KOKKOS_INLINE_FUNCTION
void addNeighbor( const int pid, const int nid ) const
Expand Down Expand Up @@ -244,8 +248,8 @@ struct VerletListBuilder
bool refill;
bool count;

// Maximum neighbors per particle
std::size_t max_n;
// Maximum allocated neighbors per particle
std::size_t alloc_n;

// Constructor.
VerletListBuilder( PositionSlice slice, const std::size_t begin,
Expand All @@ -259,7 +263,7 @@ struct VerletListBuilder
, pid_end( end )
, cell_stencil( neighborhood_radius, cell_size_ratio, grid_min,
grid_max )
, max_n( max_neigh )
, alloc_n( max_neigh )
{
count = true;
refill = false;
Expand Down Expand Up @@ -438,13 +442,13 @@ struct VerletListBuilder

void initCounts( VerletLayout2D )
{
if ( max_n > 0 )
if ( alloc_n > 0 )
{
count = false;

_data.neighbors = Kokkos::View<int**, memory_space>(
Kokkos::ViewAllocateWithoutInitializing( "neighbors" ),
_data.counts.size(), max_n );
_data.counts.size(), alloc_n );
}
}

Expand Down Expand Up @@ -481,8 +485,8 @@ struct VerletListBuilder
{
// Calculate the maximum number of neighbors.
auto counts = _data.counts;
int max_num_neighbor = 0;
Kokkos::Max<int> max_reduce( max_num_neighbor );
int max;
Kokkos::Max<int> max_reduce( max );
Kokkos::parallel_reduce(
"Cabana::VerletListBuilder::reduce_max",
Kokkos::RangePolicy<execution_space>( 0, _data.counts.size() ),
Expand All @@ -492,16 +496,16 @@ struct VerletListBuilder
},
max_reduce );
Kokkos::fence();
_data.max_n = static_cast<std::size_t>( max );

// Reallocate the neighbor list if previous size is exceeded.
if ( count or ( std::size_t )
max_num_neighbor > _data.neighbors.extent( 1 ) )
if ( count or _data.max_n > _data.neighbors.extent( 1 ) )
{
refill = true;
Kokkos::deep_copy( _data.counts, 0 );
_data.neighbors = Kokkos::View<int**, memory_space>(
Kokkos::ViewAllocateWithoutInitializing( "neighbors" ),
_data.counts.size(), max_num_neighbor );
_data.counts.size(), _data.max_n );
}
}

Expand Down Expand Up @@ -833,13 +837,22 @@ class NeighborList<
using list_type =
VerletList<MemorySpace, AlgorithmTag, VerletLayoutCSR, BuildTag>;

//! Get the total number of neighbors (maximum size of CSR list).
//! Get the total number of neighbors across all particles.
KOKKOS_INLINE_FUNCTION
static std::size_t maxNeighbor( const list_type& list )
static std::size_t totalNeighbor( const list_type& list )
{
// Size of the allocated memory gives total neighbors.
return list._data.neighbors.extent( 0 );
}

//! Get the maximum number of neighbors across all particles.
KOKKOS_INLINE_FUNCTION
static std::size_t maxNeighbor( const list_type& list )
{
std::size_t num_p = list._data.counts.size();
return Impl::maxNeighbor( list, num_p );
}

//! Get the number of neighbors for a given particle index.
KOKKOS_INLINE_FUNCTION
static std::size_t numNeighbor( const list_type& list,
Expand Down Expand Up @@ -873,11 +886,20 @@ class NeighborList<
using list_type =
VerletList<MemorySpace, AlgorithmTag, VerletLayout2D, BuildTag>;

//! Get the total number of neighbors across all particles.
KOKKOS_INLINE_FUNCTION
static std::size_t totalNeighbor( const list_type& list )
{
std::size_t num_p = list._data.counts.size();
return Impl::totalNeighbor( list, num_p );
}

//! Get the maximum number of neighbors per particle.
KOKKOS_INLINE_FUNCTION
static std::size_t maxNeighbor( const list_type& list )
{
return list._data.neighbors.extent( 1 );
// Stored during neighbor search.
return list._data.max_n;
}

//! Get the number of neighbors for a given particle index.
Expand Down
39 changes: 31 additions & 8 deletions core/unit_test/neighbor_unit_test.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@ struct TestNeighborList
{
Kokkos::View<int*, Params...> counts;
Kokkos::View<int**, Params...> neighbors;
int max;
int total;
};

template <class KokkosMemorySpace>
Expand All @@ -43,6 +45,8 @@ createTestListHostCopy( const TestNeighborList<KokkosMemorySpace>& test_list )
Kokkos::resize( list_copy.neighbors, test_list.neighbors.extent( 0 ),
test_list.neighbors.extent( 1 ) );
Kokkos::deep_copy( list_copy.neighbors, test_list.neighbors );
list_copy.total = test_list.total;
list_copy.max = test_list.max;
return list_copy;
}

Expand All @@ -56,15 +60,23 @@ copyListToHost( const ListType& list, const int num_particle, const int max_n )
Kokkos::View<int*, TEST_MEMSPACE>( "counts", num_particle );
list_copy.neighbors =
Kokkos::View<int**, TEST_MEMSPACE>( "neighbors", num_particle, max_n );
Kokkos::parallel_for(
Kokkos::Max<int> max_reduce( list_copy.max );
// Use max here because every rank should return the same value.
Kokkos::Max<int> total_reduce( list_copy.total );
Kokkos::parallel_reduce(
"copy list", Kokkos::RangePolicy<TEST_EXECSPACE>( 0, num_particle ),
KOKKOS_LAMBDA( const int p ) {
KOKKOS_LAMBDA( const int p, int& max_val, int& total_val ) {
list_copy.counts( p ) =
Cabana::NeighborList<ListType>::numNeighbor( list, p );
for ( int n = 0; n < list_copy.counts( p ); ++n )
list_copy.neighbors( p, n ) =
Cabana::NeighborList<ListType>::getNeighbor( list, p, n );
} );

// Same for every particle, but we need to extract on device.
max_val = Cabana::NeighborList<ListType>::maxNeighbor( list );
total_val = Cabana::NeighborList<ListType>::totalNeighbor( list );
},
max_reduce, total_reduce );
Kokkos::fence();
return createTestListHostCopy( list_copy );
}
Expand Down Expand Up @@ -103,16 +115,19 @@ computeFullNeighborList( const PositionSlice& position,
Kokkos::fence();

// Allocate.
auto max_op = KOKKOS_LAMBDA( const int i, int& max_val )
auto max_op = KOKKOS_LAMBDA( const int i, int& max_val, int& total_val )
{
if ( max_val < list.counts( i ) )
{
max_val = list.counts( i );
}
total_val += list.counts( i );
};
int max_n;
Kokkos::parallel_reduce( exec_policy, max_op, Kokkos::Max<int>( max_n ) );
Kokkos::parallel_reduce( exec_policy, max_op, Kokkos::Max<int>( list.max ),
Kokkos::Sum<int>( list.total ) );
Kokkos::fence();
list.neighbors = Kokkos::View<int**, TEST_MEMSPACE>( "test_neighbors",
num_particle, max_n );
list.neighbors = Kokkos::View<int**, TEST_MEMSPACE>(
"test_neighbors", num_particle, list.max );

// Fill.
auto fill_op = KOKKOS_LAMBDA( const int i )
Expand Down Expand Up @@ -174,6 +189,10 @@ void checkFullNeighborList( const ListType& nlist,
for ( int n = 0; n < N2_list_copy.counts( p ); ++n )
EXPECT_EQ( computed_neighbors[n], actual_neighbors[n] );
}

// Check the total and max interfaces.
EXPECT_EQ( N2_list_copy.max, list_copy.max );
EXPECT_EQ( N2_list_copy.total, list_copy.total );
}

//---------------------------------------------------------------------------//
Expand Down Expand Up @@ -215,6 +234,10 @@ void checkHalfNeighborList( const ListType& nlist,
}
}
}

// Check the total and max interfaces (only approximate for max).
EXPECT_GE( N2_list_copy.max, list_copy.max );
EXPECT_EQ( static_cast<int>( N2_list_copy.total / 2.0 ), list_copy.total );
}

//---------------------------------------------------------------------------//
Expand Down
6 changes: 3 additions & 3 deletions core/unit_test/tstNeighborListArborX.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -290,13 +290,13 @@ void testNeighborArborXParallelReduce()
//---------------------------------------------------------------------------//
// TESTS
//---------------------------------------------------------------------------//
TEST( TEST_CATEGORY, verlet_list_full_test ) { testArborXListFull(); }
TEST( TEST_CATEGORY, neighbor_list_full_test ) { testArborXListFull(); }

//---------------------------------------------------------------------------//
TEST( TEST_CATEGORY, verlet_list_half_test ) { testArborXListHalf(); }
TEST( TEST_CATEGORY, neighbor_list_half_test ) { testArborXListHalf(); }

//---------------------------------------------------------------------------//
TEST( TEST_CATEGORY, verlet_list_full_range_test )
TEST( TEST_CATEGORY, neighbor_list_full_range_test )
{
testArborXListFullPartialRange();
}
Expand Down

0 comments on commit f99c7db

Please sign in to comment.