Add support for migration with automatic copies

core/src/Cabana_Distributor.hpp

@@ -182,6 +182,9 @@ void distributeData(
     static_assert( is_accessible_from<typename Distributor_t::memory_space,
                                       ExecutionSpace>{},
                    "" );
+    static_assert( is_accessible_from<typename Distributor_t::memory_space,
+                                      typename AoSoA_t::memory_space>{},
+                   "" );
 
     // Get the MPI rank we are currently on.
     int my_rank = -1;
@@ -345,8 +348,18 @@ void migrate( ExecutionSpace exec_space, const Distributor_t& distributor,
         throw std::runtime_error(
             "Destination is the wrong size for migration!" );
 
+    // Enable cases where the communication happens in a different space than
+    // the particle data, e.g. device data communicated on the host.
+    using comm_space = typename Distributor_t::memory_space;
+    auto src_copy = Cabana::create_mirror_view_and_copy( comm_space{}, src );
+    auto dst_copy = Cabana::create_mirror_view_and_copy( comm_space{}, dst );
+
     // Move the data.
-    Impl::distributeData( exec_space, distributor, src, dst );
+    Impl::distributeData( exec_space, distributor, src_copy, dst_copy );
+
+    // Copy back, if needed.
+    using dst_space = typename AoSoA_t::memory_space;
+    dst = Cabana::create_mirror_view_and_copy( dst_space{}, dst_copy );
 }
 
 /*!
@@ -421,13 +434,23 @@ void migrate( ExecutionSpace exec_space, const Distributor_t& distributor,
     if ( dst_is_bigger )
         aosoa.resize( distributor.totalNumImport() );
 
+    // Enable cases where the communication happens in a different space than
+    // the particle data, e.g. device data communicated on the host.
+    using comm_space = typename Distributor_t::memory_space;
+    auto aosoa_copy =
+        Cabana::create_mirror_view_and_copy( comm_space{}, aosoa );
+
     // Move the data.
-    Impl::distributeData( exec_space, distributor, aosoa, aosoa );
+    Impl::distributeData( exec_space, distributor, aosoa_copy, aosoa_copy );
 
     // If the destination decomposition is smaller than the source
     // decomposition resize after we have moved the data.
     if ( !dst_is_bigger )
-        aosoa.resize( distributor.totalNumImport() );
+        aosoa_copy.resize( distributor.totalNumImport() );
+
+    // Copy back, if needed.
+    using aosoa_space = typename AoSoA_t::memory_space;
+    aosoa = Cabana::create_mirror_view_and_copy( aosoa_space{}, aosoa_copy );
 }
 
 /*!

grid/src/Cabana_Grid_ParticleGridDistributor.hpp

@@ -234,13 +234,15 @@ int migrateCount( const LocalGridType& local_grid,
   \param local_grid The local grid containing periodicity and system bound
   information.
   \param positions The particle positions.
+  \note CommMemorySpace Memory space which may not be accessible from the
+  particle memory space with automatic copies within the Distributor.
 
   \return Distributor for later migration.
 */
-template <class LocalGridType, class PositionSliceType>
-Cabana::Distributor<typename PositionSliceType::memory_space>
-createParticleGridDistributor( const LocalGridType& local_grid,
-                               PositionSliceType& positions )
+template <class LocalGridType, class PositionSliceType, class CommMemorySpace>
+auto createParticleGridDistributor( const LocalGridType& local_grid,
+                                    PositionSliceType& positions,
+                                    CommMemorySpace )
 {
     using memory_space = typename PositionSliceType::memory_space;
 
@@ -261,11 +263,37 @@ createParticleGridDistributor( const LocalGridType& local_grid,
                                   positions );
 
     // Create the Cabana distributor.
-    Cabana::Distributor<memory_space> distributor(
-        local_grid.globalGrid().comm(), destinations, topology );
+    // Copy the destinations in case the space is not accessible.
+    auto destinations_copy =
+        Kokkos::create_mirror_view_and_copy( CommMemorySpace{}, destinations );
+    Cabana::Distributor<CommMemorySpace> distributor(
+        local_grid.globalGrid().comm(), destinations_copy, topology );
     return distributor;
 }
 
+/*!
+  \brief Determine which data should be migrated from one uniquely-owned
+  decomposition to another uniquely-owned decomposition, using bounds of the
+  grid and taking periodic boundaries into account.
+
+  \tparam LocalGridType LocalGrid type.
+  \tparam PositionSliceType Position type.
+
+  \param local_grid The local grid containing periodicity and system bound
+  information.
+  \param positions The particle positions.
+
+  \return Distributor for later migration.
+*/
+template <class LocalGridType, class PositionSliceType>
+auto createParticleGridDistributor( const LocalGridType& local_grid,
+                                    PositionSliceType& positions )
+{
+    using particle_space = typename PositionSliceType::memory_space;
+    return createParticleGridDistributor( local_grid, positions,
+                                          particle_space{} );
+}
+
 //---------------------------------------------------------------------------//
 /*!
   \brief Migrate data from one uniquely-owned decomposition to another
@@ -279,16 +307,19 @@ createParticleGridDistributor( const LocalGridType& local_grid,
   \param local_grid The local grid containing periodicity and system bounds.
   \param positions Particle positions.
   \param particles The particle AoSoA.
+  \param comm_space Memory space in which to create the Distributor and
+  communicate data.
   \param min_halo_width Number of halo mesh widths to allow particles before
   migrating.
   \param force_migrate Migrate particles outside the local domain regardless of
   ghosted halo.
   \return Whether any particle migration occurred.
 */
-template <class LocalGridType, class ParticlePositions, class ParticleContainer>
+template <class LocalGridType, class ParticlePositions, class ParticleContainer,
+          class CommSpace>
 bool particleGridMigrate( const LocalGridType& local_grid,
                           const ParticlePositions& positions,
-                          ParticleContainer& particles,
+                          ParticleContainer& particles, CommSpace comm_space,
                           const int min_halo_width,
                           const bool force_migrate = false )
 {
@@ -306,13 +337,45 @@ bool particleGridMigrate( const LocalGridType& local_grid,
             return false;
     }
 
-    auto distributor = createParticleGridDistributor( local_grid, positions );
+    auto distributor =
+        createParticleGridDistributor( local_grid, positions, comm_space );
 
     // Redistribute the particles.
     migrate( distributor, particles );
     return true;
 }
 
+//---------------------------------------------------------------------------//
+/*!
+  \brief Migrate data from one uniquely-owned decomposition to another
+  uniquely-owned decomposition, using the bounds and periodic boundaries of the
+  grid to determine which particles should be moved. In-place variant.
+
+  \tparam LocalGridType LocalGrid type.
+  \tparam ParticlePositions Particle position type.
+  \tparam PositionContainer AoSoA type.
+
+  \param local_grid The local grid containing periodicity and system bounds.
+  \param positions Particle positions.
+  \param particles The particle AoSoA.
+  \param min_halo_width Number of halo mesh widths to allow particles before
+  migrating.
+  \param force_migrate Migrate particles outside the local domain regardless of
+  ghosted halo.
+  \return Whether any particle migration occurred.
+*/
+template <class LocalGridType, class ParticlePositions, class ParticleContainer>
+bool particleGridMigrate( const LocalGridType& local_grid,
+                          const ParticlePositions& positions,
+                          ParticleContainer& particles,
+                          const int min_halo_width,
+                          const bool force_migrate = false )
+{
+    using memory_space = typename ParticlePositions::memory_space;
+    return particleGridMigrate( local_grid, positions, particles,
+                                memory_space{}, min_halo_width, force_migrate );
+}
+
 //---------------------------------------------------------------------------//
 /*!
   \brief Migrate data from one uniquely-owned decomposition to another

grid/src/Cabana_Grid_ParticleList.hpp

@@ -84,7 +84,10 @@ class ParticleList
 
     /*!
       \brief Redistribute particles to new owning grids with explicit field.
-      \tparam PositionFieldTag Field tag for position data.
+      \param local_grid The local grid containing system bounds.
+      \param force_redistribute Whether to migrate particles within the halo
+      region immediately.
+       \note PositionFieldTag Field tag for position data.
       \return Whether the particles were actually redistributed.
     */
     template <class PositionFieldTag, class LocalGridType>
@@ -96,6 +99,26 @@ class ParticleList
             local_grid.haloCellWidth(), force_redistribute );
     }
 
+    /*!
+       \brief Redistribute particles to new owning grids with explicit field.
+       \param local_grid The local grid containing system bounds.
+       \param comm_space Memory space in which to create the Distributor and
+       communicate data.
+       \param force_redistribute Whether to migrate particles within the halo
+       region immediately.
+       \note PositionFieldTag Field tag for position data.
+       \return Whether the particles were actually redistributed.
+     */
+    template <class PositionFieldTag, class LocalGridType, class CommSpace>
+    bool redistribute( const LocalGridType& local_grid, PositionFieldTag,
+                       CommSpace comm_space,
+                       const bool force_redistribute = false )
+    {
+        return particleGridMigrate(
+            local_grid, this->slice( PositionFieldTag() ), _aosoa, comm_space,
+            local_grid.haloCellWidth(), force_redistribute );
+    }
+
   protected:
     //! Particle AoSoA.
     using base::_aosoa;