Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add support for migration with automatic copies #768

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add support for migration with automatic copies #768

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 26 additions & 3 deletions core/src/Cabana_Distributor.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 );
}

/*!
Expand Down Expand Up @@ -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 );
}

/*!
Expand Down
81 changes: 72 additions & 9 deletions grid/src/Cabana_Grid_ParticleGridDistributor.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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
Expand All @@ -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 )
{
Expand All @@ -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
Expand Down
25 changes: 24 additions & 1 deletion grid/src/Cabana_Grid_ParticleList.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand All @@ -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;
Expand Down
Loading