From e54962a76ac3ed5d32ffa49055231e817d44df22 Mon Sep 17 00:00:00 2001 From: Sam Reeve <6740307+streeve@users.noreply.github.com> Date: Tue, 3 Sep 2024 11:58:18 -0400 Subject: [PATCH] Add support for host migration with device data; or vice versa --- core/src/Cabana_Distributor.hpp | 29 ++++++- .../Cabana_Grid_ParticleGridDistributor.hpp | 81 ++++++++++++++++--- grid/src/Cabana_Grid_ParticleList.hpp | 25 +++++- 3 files changed, 122 insertions(+), 13 deletions(-) diff --git a/core/src/Cabana_Distributor.hpp b/core/src/Cabana_Distributor.hpp index 0d55ab11e..f8dcdf296 100644 --- a/core/src/Cabana_Distributor.hpp +++ b/core/src/Cabana_Distributor.hpp @@ -182,6 +182,9 @@ void distributeData( static_assert( is_accessible_from{}, "" ); + static_assert( is_accessible_from{}, + "" ); // 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 ); } /*! diff --git a/grid/src/Cabana_Grid_ParticleGridDistributor.hpp b/grid/src/Cabana_Grid_ParticleGridDistributor.hpp index 3e3de1242..b925aa339 100644 --- a/grid/src/Cabana_Grid_ParticleGridDistributor.hpp +++ b/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 -Cabana::Distributor -createParticleGridDistributor( const LocalGridType& local_grid, - PositionSliceType& positions ) +template +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 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 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 +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 +template 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 +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 diff --git a/grid/src/Cabana_Grid_ParticleList.hpp b/grid/src/Cabana_Grid_ParticleList.hpp index db0d40152..a289aedf0 100644 --- a/grid/src/Cabana_Grid_ParticleList.hpp +++ b/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 @@ -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 + 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;