Merge pull request #793 from streeve/fixup_lcl_neigh_interface

Fixup LinkedCellList neighbor interface
ECP-copa · Oct 21, 2024 · 86b41a2 · 86b41a2
2 parents facdb90 + e00db62
commit 86b41a2
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 28 deletions.
diff --git a/core/src/Cabana_LinkedCellList.hpp b/core/src/Cabana_LinkedCellList.hpp
@@ -761,18 +761,29 @@ class NeighborList<LinkedCellList<MemorySpace, Scalar>>
     //! Neighbor list type.
     using list_type = LinkedCellList<MemorySpace, Scalar>;
 
-    //! Get the maximum number of neighbors per particle.
+    //! Get the total number of neighbors across all particles.
     KOKKOS_INLINE_FUNCTION static std::size_t
     totalNeighbor( const list_type& list )
     {
-        return Impl::totalNeighbor( list, list.numParticles() );
+        std::size_t total_n = 0;
+        // Sum neighbors across all particles in range.
+        for ( std::size_t p = list.getParticleBegin();
+              p < list.getParticleEnd(); p++ )
+            total_n += numNeighbor( list, p );
+        return total_n;
     }
 
-    //! Get the maximum number of neighbors across all particles.
+    //! Get the maximum number of neighbors per particles.
     KOKKOS_INLINE_FUNCTION
     static std::size_t maxNeighbor( const list_type& list )
     {
-        return Impl::maxNeighbor( list, list.numParticles() );
+        std::size_t max_n = 0;
+        // Max neighbors across all particles in range.
+        for ( std::size_t p = list.getParticleBegin();
+              p < list.getParticleEnd(); p++ )
+            if ( numNeighbor( list, p ) > max_n )
+                max_n = numNeighbor( list, p );
+        return max_n;
     }
 
     //! Get the number of neighbors for a given particle index.
@@ -817,14 +828,7 @@ class NeighborList<LinkedCellList<MemorySpace, Scalar>>
                     {
                         int particle_id = list.binOffset( i, j, k ) +
                                           ( neighbor_index - previous_count );
-                        if ( list.sorted() )
-                        {
-                            return particle_id + list.getParticleBegin();
-                        }
-                        else
-                        {
-                            return list.permutation( particle_id );
-                        }
+                        return list.getParticle( particle_id );
                     }
                     previous_count = total_count;
                 }

diff --git a/core/src/Cabana_Parallel.hpp b/core/src/Cabana_Parallel.hpp
@@ -1283,18 +1283,15 @@ struct LinkedCellParallelReduce
                     // neighbors.
                     auto offset = _list.binOffset( gi, gj, gk );
                     auto size = _list.binSize( gi, gj, gk );
-                    for ( std::size_t j = offset; j < offset + size; ++j )
+                    for ( std::size_t n = offset; n < offset + size; ++n )
                     {
                         // Get the true id of the candidate neighbor.
-                        std::size_t jj;
-                        if ( !_list.sorted() )
-                            jj = _list.permutation( j );
-                        else
-                            jj = j + _begin;
+                        auto j = _list.getParticle( n );
+
                         // Avoid self interactions (dummy position args).
-                        if ( _discriminator.isValid( i, 0, 0, 0, jj, 0, 0, 0 ) )
+                        if ( _discriminator.isValid( i, 0, 0, 0, j, 0, 0, 0 ) )
                         {
-                            Impl::functorTagDispatch<WorkTag>( _functor, i, jj,
+                            Impl::functorTagDispatch<WorkTag>( _functor, i, j,
                                                                ival );
                         }
                     }
@@ -1322,20 +1319,17 @@ struct LinkedCellParallelReduce
                     auto size = _list.binSize( gi, gj, gk );
                     Kokkos::parallel_for(
                         Kokkos::TeamThreadRange( team, offset, offset + size ),
-                        [&]( const index_type j )
+                        [&]( const index_type n )
                         {
                             // Get the true id of the candidate neighbor.
-                            std::size_t jj;
-                            if ( !_list.sorted() )
-                                jj = _list.permutation( j );
-                            else
-                                jj = j + _begin;
+                            auto j = _list.getParticle( n );
+
                             // Avoid self interactions (dummy position args).
-                            if ( _discriminator.isValid( i, 0, 0, 0, jj, 0, 0,
+                            if ( _discriminator.isValid( i, 0, 0, 0, j, 0, 0,
                                                          0 ) )
                             {
                                 Impl::functorTagDispatch<WorkTag>( _functor, i,
-                                                                   jj, ival );
+                                                                   j, ival );
                             }
                         } );
                 }