diff --git a/core/src/Cabana_Experimental_NeighborList.hpp b/core/src/Cabana_Experimental_NeighborList.hpp index 26ec7998e..a5d9aa3ba 100644 --- a/core/src/Cabana_Experimental_NeighborList.hpp +++ b/core/src/Cabana_Experimental_NeighborList.hpp @@ -245,7 +245,7 @@ struct CrsGraph Kokkos::View row_ptr; //! Neighbor offset shift. typename MemorySpace::size_type shift; - //! Total neighbors. + //! Total particles. typename MemorySpace::size_type total; }; @@ -368,7 +368,7 @@ struct Dense Kokkos::View val; //! Neighbor offset shift. typename MemorySpace::size_type shift; - //! Total neighbors. + //! Total particles. typename MemorySpace::size_type total; }; @@ -536,6 +536,21 @@ class NeighborList> 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 ) @@ -568,6 +583,21 @@ class NeighborList> 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 ) diff --git a/core/src/Cabana_NeighborList.hpp b/core/src/Cabana_NeighborList.hpp index dad7f9a6e..4055198b6 100644 --- a/core/src/Cabana_NeighborList.hpp +++ b/core/src/Cabana_NeighborList.hpp @@ -114,6 +114,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, @@ -134,6 +142,33 @@ class NeighborList //---------------------------------------------------------------------------// +namespace Impl +{ +//! Iterate to get the total number of neighbors. +template +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::numNeighbor( list, p ); + return total_n; +} + +//! Iterate to find the maximum number of neighbors. +template +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::numNeighbor( list, p ) > max_n ) + max_n = NeighborList::numNeighbor( list, p ); + return max_n; +} +} // namespace Impl + } // end namespace Cabana #endif // end CABANA_NEIGHBORLIST_HPP diff --git a/core/src/Cabana_VerletList.hpp b/core/src/Cabana_VerletList.hpp index 1e7f021fd..1bb06af15 100644 --- a/core/src/Cabana_VerletList.hpp +++ b/core/src/Cabana_VerletList.hpp @@ -90,6 +90,10 @@ struct VerletListData //! Neighbor list. Kokkos::View 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 @@ -146,8 +150,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, @@ -159,7 +163,7 @@ struct VerletListBuilder const std::size_t max_neigh ) : pid_begin( begin ) , pid_end( end ) - , max_n( max_neigh ) + , alloc_n( max_neigh ) { count = true; refill = false; @@ -341,13 +345,13 @@ struct VerletListBuilder void initCounts( VerletLayout2D ) { - if ( max_n > 0 ) + if ( alloc_n > 0 ) { count = false; _data.neighbors = Kokkos::View( Kokkos::ViewAllocateWithoutInitializing( "neighbors" ), - _data.counts.size(), max_n ); + _data.counts.size(), alloc_n ); } } @@ -384,8 +388,8 @@ struct VerletListBuilder { // Calculate the maximum number of neighbors. auto counts = _data.counts; - int max_num_neighbor = 0; - Kokkos::Max max_reduce( max_num_neighbor ); + int max; + Kokkos::Max max_reduce( max ); Kokkos::parallel_reduce( "Cabana::VerletListBuilder::reduce_max", Kokkos::RangePolicy( 0, _data.counts.size() ), @@ -395,16 +399,16 @@ struct VerletListBuilder }, max_reduce ); Kokkos::fence(); + _data.max_n = static_cast( 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( Kokkos::ViewAllocateWithoutInitializing( "neighbors" ), - _data.counts.size(), max_num_neighbor ); + _data.counts.size(), _data.max_n ); } } @@ -738,13 +742,22 @@ class NeighborList< using list_type = VerletList; - //! 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, @@ -778,11 +791,20 @@ class NeighborList< using list_type = VerletList; + //! 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. diff --git a/core/unit_test/neighbor_unit_test.hpp b/core/unit_test/neighbor_unit_test.hpp index f1f2f14ed..88cdf643e 100644 --- a/core/unit_test/neighbor_unit_test.hpp +++ b/core/unit_test/neighbor_unit_test.hpp @@ -30,6 +30,8 @@ struct TestNeighborList { Kokkos::View counts; Kokkos::View neighbors; + int max; + int total; }; template @@ -43,6 +45,8 @@ createTestListHostCopy( const TestNeighborList& 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; } @@ -56,15 +60,23 @@ copyListToHost( const ListType& list, const int num_particle, const int max_n ) Kokkos::View( "counts", num_particle ); list_copy.neighbors = Kokkos::View( "neighbors", num_particle, max_n ); - Kokkos::parallel_for( + Kokkos::Max max_reduce( list_copy.max ); + // Use max here because every rank should return the same value. + Kokkos::Max total_reduce( list_copy.total ); + Kokkos::parallel_reduce( "copy list", Kokkos::RangePolicy( 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::numNeighbor( list, p ); for ( int n = 0; n < list_copy.counts( p ); ++n ) list_copy.neighbors( p, n ) = Cabana::NeighborList::getNeighbor( list, p, n ); - } ); + + // Same for every particle, but we need to extract on device. + max_val = Cabana::NeighborList::maxNeighbor( list ); + total_val = Cabana::NeighborList::totalNeighbor( list ); + }, + max_reduce, total_reduce ); Kokkos::fence(); return createTestListHostCopy( list_copy ); } @@ -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( max_n ) ); + Kokkos::parallel_reduce( exec_policy, max_op, Kokkos::Max( list.max ), + Kokkos::Sum( list.total ) ); Kokkos::fence(); - list.neighbors = Kokkos::View( "test_neighbors", - num_particle, max_n ); + list.neighbors = Kokkos::View( + "test_neighbors", num_particle, list.max ); // Fill. auto fill_op = KOKKOS_LAMBDA( const int i ) @@ -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 ); } //---------------------------------------------------------------------------// @@ -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( N2_list_copy.total / 2.0 ), list_copy.total ); } //---------------------------------------------------------------------------// diff --git a/core/unit_test/tstNeighborListArborX.hpp b/core/unit_test/tstNeighborListArborX.hpp index b93c35d66..8861ed7bc 100644 --- a/core/unit_test/tstNeighborListArborX.hpp +++ b/core/unit_test/tstNeighborListArborX.hpp @@ -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(); }