Skip to content

Commit

Permalink
Add particle-grid periodic migrate and gather
Browse files Browse the repository at this point in the history
  • Loading branch information
streeve committed Jan 4, 2021
1 parent d46804f commit 7f5b0e4
Show file tree
Hide file tree
Showing 8 changed files with 1,431 additions and 50 deletions.
10 changes: 10 additions & 0 deletions core/src/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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}>
Expand Down
40 changes: 26 additions & 14 deletions core/src/Cabana_CommunicationPlan.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ namespace Cabana
{
namespace Impl
{

//---------------------------------------------------------------------------//
// Count sends and create steering algorithm tags.
struct CountSendsAndCreateSteeringDuplicated
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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 );
Expand Down
4 changes: 4 additions & 0 deletions core/src/Cabana_Core.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
164 changes: 128 additions & 36 deletions core/src/Cabana_Halo.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
#include <mpi.h>

#include <exception>
#include <type_traits>
#include <vector>

namespace Cabana
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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;
Expand Down
Loading

0 comments on commit 7f5b0e4

Please sign in to comment.