WIP Add particle-grid periodic migrate and gather

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

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

@@ -199,60 +199,61 @@ 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.
+    // 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 ) );
+    };
+    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 Periodic_t>
+void sendBuffer( const Halo_t &halo, AoSoA_t &aosoa, View_t &send_buffer,
+                 const Periodic_t &periodic )
 {
-    // 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 ) );
+        periodic.modify_buffer( 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,68 @@ 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. 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, class Periodic_t>
+void gather( const Halo_t &halo, AoSoA_t &aosoa, const Periodic_t &periodic,
+             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, periodic );
+    Impl::recvBuffer( halo, aosoa, send_buffer );
+}
+
+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 +414,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;