Particle-grid periodic communication

core/src/CMakeLists.txt

@@ -32,6 +32,12 @@ if(Cabana_ENABLE_MPI)
     )
 endif()
 
+if(Cabana_ENABLE_CAJITA)
+  list(APPEND HEADERS_PUBLIC
+    Cabana_ParticleGridCommunication.hpp
+    )
+endif()
+
 set(HEADERS_IMPL
   impl/Cabana_CartesianGrid.hpp
   impl/Cabana_Index.hpp
@@ -54,6 +60,10 @@ if(Cabana_ENABLE_MPI)
   target_link_libraries(cabanacore MPI::MPI_CXX)
 endif()
 
+if(Cabana_ENABLE_CAJITA)
+  target_link_libraries(cabanacore Cajita)
+endif()
+
 target_include_directories(cabanacore
   PUBLIC
   $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>

core/src/Cabana_CommunicationPlan.hpp

@@ -30,6 +30,7 @@ namespace Cabana
 {
 namespace Impl
 {
+
 //---------------------------------------------------------------------------//
 // Count sends and create steering algorithm tags.
 struct CountSendsAndCreateSteeringDuplicated
@@ -246,6 +247,29 @@ auto countSendsAndCreateSteering( const ExportRankView element_export_ranks,
     return std::make_pair( neighbor_counts, neighbor_ids );
 }
 
+//---------------------------------------------------------------------------//
+// Return unique neighbor ranks, with the current rank first.
+std::vector<int> getUniqueTopology( std::vector<int> topology )
+{
+    auto remove_end = std::remove( topology.begin(), topology.end(), -1 );
+    std::sort( topology.begin(), remove_end );
+    auto unique_end = std::unique( topology.begin(), remove_end );
+    topology.resize( std::distance( topology.begin(), unique_end ) );
+
+    // Put this rank first.
+    int my_rank = -1;
+    MPI_Comm_rank( MPI_COMM_WORLD, &my_rank );
+    for ( auto& n : topology )
+    {
+        if ( n == my_rank )
+        {
+            std::swap( n, topology[0] );
+            break;
+        }
+    }
+    return topology;
+}
+
 //---------------------------------------------------------------------------//
 
 } // end namespace Impl
@@ -464,11 +488,8 @@ class CommunicationPlan
         // Store the number of export elements.
         _num_export_element = element_export_ranks.size();
 
-        // Store the unique neighbors.
-        _neighbors = neighbor_ranks;
-        std::sort( _neighbors.begin(), _neighbors.end() );
-        auto unique_end = std::unique( _neighbors.begin(), _neighbors.end() );
-        _neighbors.resize( std::distance( _neighbors.begin(), unique_end ) );
+        // Store the unique neighbors (this rank first).
+        _neighbors = Impl::getUniqueTopology( neighbor_ranks );
         int num_n = _neighbors.size();
 
         // Get the size of this communicator.
@@ -483,15 +504,6 @@ class CommunicationPlan
         // communication space so any mpi tag will do.
         const int mpi_tag = 1221;
 
-        // If we are sending to ourself put that one first in the neighbor
-        // list.
-        for ( auto& n : _neighbors )
-            if ( n == my_rank )
-            {
-                std::swap( n, _neighbors[0] );
-                break;
-            }
-
         // Initialize import/export sizes.
         _num_export.assign( num_n, 0 );
         _num_import.assign( num_n, 0 );

core/src/Cabana_Core.hpp

@@ -33,6 +33,10 @@
 #include <Cabana_Halo.hpp>
 #endif
 
+#ifdef Cabana_ENABLE_CAJITA
+#include <Cabana_ParticleGridCommunication.hpp>
+#endif
+
 #ifdef Cabana_ENABLE_ARBORX
 #include <Cabana_Experimental_NeighborList.hpp>
 #endif

core/src/Cabana_Halo.hpp

@@ -21,6 +21,7 @@
 #include <mpi.h>
 
 #include <exception>
+#include <type_traits>
 #include <vector>
 
 namespace Cabana
@@ -199,60 +200,60 @@ struct is_halo : public is_halo_impl<typename std::remove_cv<T>::type>::type
 {
 };
 
-//---------------------------------------------------------------------------//
-/*!
-  \brief Synchronously gather data from the local decomposition to the ghosts
-  using the halo forward communication plan. AoSoA version. This is a
-  uniquely-owned to multiply-owned communication.
+namespace Impl
+{
 
-  A gather sends data from a locally owned elements to one or many ranks on
-  which they exist as ghosts. A locally owned element may be sent to as many
-  ranks as desired to be used as a ghost on those ranks. The value of the
-  element in the locally owned decomposition will be the value assigned to the
-  element in the ghosted decomposition.
+template <class Halo_t, class Container_t>
+void checkSize( const Halo_t& halo, Container_t& container )
+{
+    // Check that the AoSoA/slice is the right size.
+    if ( container.size() != halo.numLocal() + halo.numGhost() )
+        throw std::runtime_error( "AoSoA/slice is the wrong size for gather!" );
+}
 
-  \tparam Halo_t Halo type - must be a Halo.
+template <class Halo_t, class AoSoA_t, class View_t>
+void sendBuffer( const Halo_t& halo, AoSoA_t& aosoa, View_t& send_buffer )
+{
 
-  \tparam AoSoA_t AoSoA type - must be an AoSoA.
+    // Get the steering vector for the sends.
+    auto steering = halo.getExportSteering();
 
-  \param halo The halo to use for the gather.
+    // Gather from the local data into a tuple-contiguous send buffer.
+    auto gather_send_buffer_func = KOKKOS_LAMBDA( const std::size_t i )
+    {
+        send_buffer( i ) = aosoa.getTuple( steering( i ) );
+    };
+    Kokkos::RangePolicy<typename Halo_t::execution_space>
+        gather_send_buffer_policy( 0, halo.totalNumExport() );
+    Kokkos::parallel_for( "Cabana::gather::gather_send_buffer",
+                          gather_send_buffer_policy, gather_send_buffer_func );
+    Kokkos::fence();
+}
 
-  \param aosoa The AoSoA on which to perform the gather. The AoSoA should have
-  a size equivalent to halo.numGhost() + halo.numLocal(). The locally owned
-  elements are expected to appear first (i.e. in the first halo.numLocal()
-  elements) and the ghosted elements are expected to appear second (i.e. in
-  the next halo.numGhost() elements()).
-*/
-template <class Halo_t, class AoSoA_t>
-void gather( const Halo_t& halo, AoSoA_t& aosoa,
-             typename std::enable_if<( is_halo<Halo_t>::value &&
-                                       is_aosoa<AoSoA_t>::value ),
-                                     int>::type* = 0 )
+template <class Halo_t, class AoSoA_t, class View_t, class Modify_t>
+void sendBuffer( const Halo_t& halo, AoSoA_t& aosoa, View_t& send_buffer,
+                 const Modify_t& modify_functor )
 {
-    // Check that the AoSoA is the right size.
-    if ( aosoa.size() != halo.numLocal() + halo.numGhost() )
-        throw std::runtime_error( "AoSoA is the wrong size for gather!" );
-
-    // Allocate a send buffer.
-    Kokkos::View<typename AoSoA_t::tuple_type*, typename Halo_t::memory_space>
-        send_buffer(
-            Kokkos::ViewAllocateWithoutInitializing( "halo_send_buffer" ),
-            halo.totalNumExport() );
-
     // Get the steering vector for the sends.
     auto steering = halo.getExportSteering();
 
     // Gather from the local data into a tuple-contiguous send buffer.
+    // Pass send buffer to user modification functor class to add shifts.
     auto gather_send_buffer_func = KOKKOS_LAMBDA( const std::size_t i )
     {
         send_buffer( i ) = aosoa.getTuple( steering( i ) );
+        modify_functor( send_buffer, i );
     };
     Kokkos::RangePolicy<typename Halo_t::execution_space>
         gather_send_buffer_policy( 0, halo.totalNumExport() );
     Kokkos::parallel_for( "Cabana::gather::gather_send_buffer",
                           gather_send_buffer_policy, gather_send_buffer_func );
     Kokkos::fence();
+}
 
+template <class Halo_t, class AoSoA_t, class View_t>
+void recvBuffer( const Halo_t& halo, AoSoA_t& aosoa, const View_t& send_buffer )
+{
     // Allocate a receive buffer.
     Kokkos::View<typename AoSoA_t::tuple_type*, typename Halo_t::memory_space>
         recv_buffer(
@@ -320,6 +321,98 @@ void gather( const Halo_t& halo, AoSoA_t& aosoa,
     MPI_Barrier( halo.comm() );
 }
 
+} // namespace Impl
+
+//---------------------------------------------------------------------------//
+/*!
+  \brief Synchronously gather data from the local decomposition to the ghosts
+  using the halo forward communication plan. AoSoA version, where the buffer is
+  modified before being sent. This is a uniquely-owned to multiply-owned
+  communication.
+
+  A gather sends data from a locally owned elements to one or many ranks on
+  which they exist as ghosts. A locally owned element may be sent to as many
+  ranks as desired to be used as a ghost on those ranks. The value of the
+  element in the locally owned decomposition will be the value assigned to the
+  element in the ghosted decomposition.
+
+  \tparam Halo_t Halo type - must be a Halo.
+
+  \tparam AoSoA_t AoSoA type - must be an AoSoA.
+
+  \tparam Modify_t Buffer modification type.
+
+  \param halo The halo to use for the gather.
+
+  \param aosoa The AoSoA on which to perform the gather. The AoSoA should have
+  a size equivalent to halo.numGhost() + halo.numLocal(). The locally owned
+  elements are expected to appear first (i.e. in the first halo.numLocal()
+  elements) and the ghosted elements are expected to appear second (i.e. in
+  the next halo.numGhost() elements()).
+
+  \param modify_functor Class containing functor to modify the send buffer
+  before being sent (e.g. for periodic coordinate update).
+*/
+template <class Halo_t, class AoSoA_t, class Modify_t>
+void gather( const Halo_t& halo, AoSoA_t& aosoa, const Modify_t& modify_functor,
+             typename std::enable_if<( is_halo<Halo_t>::value &&
+                                       is_aosoa<AoSoA_t>::value ),
+                                     int>::type* = 0 )
+{
+    Impl::checkSize( halo, aosoa );
+
+    // Allocate a send buffer.
+    Kokkos::View<typename AoSoA_t::tuple_type*, typename Halo_t::memory_space>
+        send_buffer(
+            Kokkos::ViewAllocateWithoutInitializing( "halo_send_buffer" ),
+            halo.totalNumExport() );
+
+    Impl::sendBuffer( halo, aosoa, send_buffer, modify_functor );
+    Impl::recvBuffer( halo, aosoa, send_buffer );
+}
+
+//---------------------------------------------------------------------------//
+/*!
+  \brief Synchronously gather data from the local decomposition to the ghosts
+  using the halo forward communication plan. AoSoA version. This is a
+  uniquely-owned to multiply-owned communication.
+
+  A gather sends data from a locally owned elements to one or many ranks on
+  which they exist as ghosts. A locally owned element may be sent to as many
+  ranks as desired to be used as a ghost on those ranks. The value of the
+  element in the locally owned decomposition will be the value assigned to the
+  element in the ghosted decomposition.
+
+  \tparam Halo_t Halo type - must be a Halo.
+
+  \tparam AoSoA_t AoSoA type - must be an AoSoA.
+
+  \param halo The halo to use for the gather.
+
+  \param aosoa The AoSoA on which to perform the gather. The AoSoA should have
+  a size equivalent to halo.numGhost() + halo.numLocal(). The locally owned
+  elements are expected to appear first (i.e. in the first halo.numLocal()
+  elements) and the ghosted elements are expected to appear second (i.e. in
+  the next halo.numGhost() elements()).
+*/
+template <class Halo_t, class AoSoA_t>
+void gather( const Halo_t& halo, AoSoA_t& aosoa,
+             typename std::enable_if<( is_halo<Halo_t>::value &&
+                                       is_aosoa<AoSoA_t>::value ),
+                                     int>::type* = 0 )
+{
+    Impl::checkSize( halo, aosoa );
+
+    // Allocate a send buffer.
+    Kokkos::View<typename AoSoA_t::tuple_type*, typename Halo_t::memory_space>
+        send_buffer(
+            Kokkos::ViewAllocateWithoutInitializing( "halo_send_buffer" ),
+            halo.totalNumExport() );
+
+    Impl::sendBuffer( halo, aosoa, send_buffer );
+    Impl::recvBuffer( halo, aosoa, send_buffer );
+}
+
 //---------------------------------------------------------------------------//
 /*!
   \brief Synchronously gather data from the local decomposition to the ghosts
@@ -351,8 +444,7 @@ void gather( const Halo_t& halo, Slice_t& slice,
                                      int>::type* = 0 )
 {
     // Check that the Slice is the right size.
-    if ( slice.size() != halo.numLocal() + halo.numGhost() )
-        throw std::runtime_error( "Slice is the wrong size for gather!" );
+    Impl::checkSize( halo, slice );
 
     // Get the number of components in the slice.
     std::size_t num_comp = 1;