Add particle-grid periodic migrate

cajita/unit_test/CMakeLists.txt

@@ -38,7 +38,7 @@ macro(Cajita_add_tests)
       list(APPEND CAJITA_UNIT_TEST_NUMTHREADS ${_nt})
     endif()
   endforeach()
-  set(CAJITA_UNIT_TEST_MAIN unit_test_main.cpp)
+  set(CAJITA_UNIT_TEST_MAIN ../../core/unit_test/mpi_unit_test_main.cpp)
   foreach(_device ${CAJITA_TEST_DEVICES})
     set(_dir ${CMAKE_CURRENT_BINARY_DIR}/${_device})
     file(MAKE_DIRECTORY ${_dir})

core/src/Cabana_PeriodicMigrate.hpp

@@ -0,0 +1,227 @@
+/****************************************************************************
+ * Copyright (c) 2018-2020 by the Cabana authors                            *
+ * All rights reserved.                                                     *
+ *                                                                          *
+ * This file is part of the Cabana library. Cabana is distributed under a   *
+ * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
+ * the top-level directory.                                                 *
+ *                                                                          *
+ * SPDX-License-Identifier: BSD-3-Clause                                    *
+ ****************************************************************************/
+
+#ifndef CABANA_PERIODICMIGRATE_HPP
+#define CABANA_PERIODICMIGRATE_HPP
+
+#include <Cabana_Distributor.hpp>
+
+#include <Cajita_GlobalGrid.hpp>
+#include <Cajita_GlobalMesh.hpp>
+#include <Cajita_LocalGrid.hpp>
+#include <Cajita_LocalMesh.hpp>
+
+#include <Kokkos_Core.hpp>
+
+#include <vector>
+
+namespace Cabana
+{
+
+//---------------------------------------------------------------------------//
+template <class LocalGridType, class PositionSliceType, class NeighborRankView,
+          class DestinationRankView>
+void periodic_shift( const LocalGridType &local_grid,
+                     const NeighborRankView &neighbor_ranks,
+                     DestinationRankView &destinations,
+                     PositionSliceType &positions )
+{
+    using execution_space = typename PositionSliceType::execution_space;
+
+    // Locate the particles in the global grid and get their destination
+    // rank. The particle halo should be constructed such that particles will
+    // only move to a location in the 26 neighbor halo or stay on this
+    // rank. If the particle crosses a periodic boundary update it's
+    // positions to represent the shift.
+    auto local_mesh = Cajita::createLocalMesh<Kokkos::HostSpace>( local_grid );
+    const auto &global_grid = local_grid.globalGrid();
+    const auto &global_mesh = global_grid.globalMesh();
+    const Kokkos::Array<double, 3> local_low = {
+        local_mesh.lowCorner( Cajita::Own(), Cajita::Dim::I ),
+        local_mesh.lowCorner( Cajita::Own(), Cajita::Dim::J ),
+        local_mesh.lowCorner( Cajita::Own(), Cajita::Dim::K )};
+    const Kokkos::Array<double, 3> local_high = {
+        local_mesh.highCorner( Cajita::Own(), Cajita::Dim::I ),
+        local_mesh.highCorner( Cajita::Own(), Cajita::Dim::J ),
+        local_mesh.highCorner( Cajita::Own(), Cajita::Dim::K )};
+    const Kokkos::Array<bool, 3> period = {
+        global_grid.isPeriodic( Cajita::Dim::I ),
+        global_grid.isPeriodic( Cajita::Dim::J ),
+        global_grid.isPeriodic( Cajita::Dim::K )};
+    const Kokkos::Array<double, 3> global_low = {
+        global_mesh.lowCorner( Cajita::Dim::I ),
+        global_mesh.lowCorner( Cajita::Dim::J ),
+        global_mesh.lowCorner( Cajita::Dim::K )};
+    const Kokkos::Array<double, 3> global_high = {
+        global_mesh.highCorner( Cajita::Dim::I ),
+        global_mesh.highCorner( Cajita::Dim::J ),
+        global_mesh.highCorner( Cajita::Dim::K )};
+    const Kokkos::Array<double, 3> global_span = {
+        global_mesh.extent( Cajita::Dim::I ),
+        global_mesh.extent( Cajita::Dim::J ),
+        global_mesh.extent( Cajita::Dim::K )};
+    Kokkos::parallel_for(
+        "periodic_shift",
+        Kokkos::RangePolicy<execution_space>( 0, positions.size() ),
+        KOKKOS_LAMBDA( const int p ) {
+            // Compute the logical index of the neighbor we are sending to.
+            int nid[3] = {1, 1, 1};
+            for ( int d = 0; d < 3; ++d )
+            {
+                if ( positions( p, d ) < local_low[d] )
+                    nid[d] = 0;
+                else if ( positions( p, d ) > local_high[d] )
+                    nid[d] = 2;
+            }
+
+            // Compute the destination MPI rank.
+            destinations( p ) = neighbor_ranks(
+                nid[Cajita::Dim::I] +
+                3 * ( nid[Cajita::Dim::J] + 3 * nid[Cajita::Dim::K] ) );
+
+            // Shift periodic coordinates if needed.
+            for ( int d = 0; d < 3; ++d )
+            {
+                if ( period[d] )
+                {
+                    if ( positions( p, d ) > global_high[d] )
+                        positions( p, d ) -= global_span[d];
+                    else if ( positions( p, d ) < global_low[d] )
+                        positions( p, d ) += global_span[d];
+                }
+            }
+        } );
+}
+
+namespace Impl
+{
+
+//---------------------------------------------------------------------------//
+template <class DeviceType, class LocalGridType, class PositionType>
+Cabana::Distributor<DeviceType>
+prepare_migrate( const LocalGridType &local_grid, PositionType &positions )
+{
+    using device_type = DeviceType;
+
+    // Of the 27 potential local grids figure out which are in our topology.
+    // Some of the ranks in this list may be invalid. We will update this list
+    // after we compute destination ranks so it is unique and valid.
+    std::vector<int> topology( 27, -1 );
+    int nr = 0;
+    for ( int k = -1; k < 2; ++k )
+        for ( int j = -1; j < 2; ++j )
+            for ( int i = -1; i < 2; ++i, ++nr )
+                topology[nr] = local_grid.neighborRank( i, j, k );
+
+    // Locate the particles in the global grid and get their destination
+    // rank and shift periodic coordinates if necessary.
+    Kokkos::View<int *, Kokkos::HostSpace, Kokkos::MemoryUnmanaged>
+        neighbor_ranks( topology.data(), topology.size() );
+    auto nr_mirror =
+        Kokkos::create_mirror_view_and_copy( device_type(), neighbor_ranks );
+    Kokkos::View<int *, device_type> destinations(
+        Kokkos::ViewAllocateWithoutInitializing( "destinations" ),
+        positions.size() );
+    periodic_shift( local_grid, nr_mirror, destinations, positions );
+
+    // Make the topology a list of unique and valid ranks.
+    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 ) );
+
+    // Create the Cabana distributor.
+    Cabana::Distributor<device_type> distributor(
+        local_grid.globalGrid().comm(), destinations, topology );
+
+    return distributor;
+}
+
+} // namespace Impl
+
+//---------------------------------------------------------------------------//
+/*!
+  \brief periodic_migrate uses a communication plan and migrates data from one
+  uniquely-owned decomposition to another uniquely-owned decomposition, taking
+  into account periodic boundary conditions with a Cajita grid. In-place
+  variant.
+*/
+template <class LocalGridType, class ParticleContainer,
+          std::size_t PositionIndex>
+void periodic_migrate( const LocalGridType &local_grid,
+                       ParticleContainer &particles,
+                       std::integral_constant<std::size_t, PositionIndex> )
+{
+    using device_type = typename ParticleContainer::device_type;
+
+    // Get the positions.
+    auto positions = Cabana::slice<PositionIndex>( particles );
+
+    // Periodic position shift.
+    auto distributor =
+        Impl::prepare_migrate<device_type>( local_grid, positions );
+
+    // Redistribute the particles.
+    Cabana::migrate( distributor, particles );
+}
+
+//---------------------------------------------------------------------------//
+/*!
+\brief periodic_migrate uses a communication plan and migrates data from one
+uniquely-owned decomposition to another uniquely-owned decomposition, taking
+into account periodic boundary conditions with a Cajita grid. Separate AoSoA
+variant.
+*/
+template <class LocalGridType, class ParticleContainer,
+          std::size_t PositionIndex>
+void periodic_migrate( const LocalGridType &local_grid,
+                       ParticleContainer &src_particles,
+                       std::integral_constant<std::size_t, PositionIndex>,
+                       ParticleContainer &dst_particles )
+{
+    using device_type = typename ParticleContainer::device_type;
+
+    // Get the positions.
+    auto positions = Cabana::slice<PositionIndex>( src_particles );
+
+    // Periodic position shift.
+    auto distributor =
+        Impl::prepare_migrate<device_type>( local_grid, positions );
+
+    // Redistribute the particles.
+    Cabana::migrate( distributor, src_particles, dst_particles );
+}
+
+//---------------------------------------------------------------------------//
+/*!
+\brief periodic_migrate uses a communication plan and migrates data from one
+uniquely-owned decomposition to another uniquely-owned decomposition, taking
+into account periodic boundary conditions with a Cajita grid. Separate
+slice/View variant.
+*/
+template <class LocalGridType, class ParticleContainer>
+void periodic_migrate( const LocalGridType &local_grid,
+                       ParticleContainer &src_positions,
+                       ParticleContainer &dst_positions )
+{
+    using device_type = typename ParticleContainer::device_type;
+
+    // Periodic position shift.
+    auto distributor =
+        Impl::prepare_migrate<device_type>( local_grid, src_positions );
+
+    // Redistribute the positions.
+    Cabana::migrate( distributor, src_positions, dst_positions );
+}
+
+} // end namespace Cabana
+
+#endif // end CABANA_PERIODICMIGRATE_HPP

core/unit_test/CMakeLists.txt

@@ -59,6 +59,9 @@ macro(Cabana_add_tests)
       target_include_directories(${_target} PRIVATE ${_dir}
         ${CMAKE_CURRENT_BINARY_DIR} ${CMAKE_CURRENT_SOURCE_DIR})
       target_link_libraries(${_target} cabanacore cabana_core_gtest)
+      if(_test STREQUAL PeriodicMigrate)
+        target_link_libraries(${_target} Cajita)
+      endif()
       if(CABANA_UNIT_TEST_MPI)
         foreach(_np ${CABANA_UNIT_TEST_MPIEXEC_NUMPROCS})
           # NOTE: When moving to CMake 3.10+ make sure to use MPIEXEC_EXECUTABLE instead
@@ -86,4 +89,7 @@ if(Cabana_ENABLE_ARBORX)
 endif()
 if(Cabana_ENABLE_MPI)
   Cabana_add_tests(MPI NAMES CommunicationPlan Distributor Halo)
+  if(Cabana_ENABLE_CAJITA)
+    Cabana_add_tests(MPI NAMES PeriodicMigrate)
+  endif()
 endif()