16#ifndef CABANA_DISTRIBUTOR_HPP
17#define CABANA_DISTRIBUTOR_HPP
23#include <Kokkos_Core.hpp>
24#include <Kokkos_Profiling_ScopedRegion.hpp>
61template <
class MemorySpace>
99 template <
class ViewType>
101 const std::vector<int>& neighbor_ranks )
105 Export(), element_export_ranks, neighbor_ranks );
136 template <
class ViewType>
149struct is_distributor_impl :
public std::false_type
153template <
typename MemorySpace>
154struct is_distributor_impl<
Distributor<MemorySpace>> :
public std::true_type
162 :
public is_distributor_impl<typename std::remove_cv<T>::type>::type
173template <
class ExecutionSpace,
class Distributor_t,
class AoSoA_t>
175 ExecutionSpace,
const Distributor_t& distributor,
const AoSoA_t& src,
181 Kokkos::Profiling::ScopedRegion region(
"Cabana::migrate" );
189 MPI_Comm_rank( distributor.comm(), &my_rank );
192 int num_n = distributor.numNeighbor();
198 std::size_t num_stay =
199 ( num_n > 0 && distributor.neighborRank( 0 ) == my_rank )
200 ? distributor.numExport( 0 )
204 std::size_t num_send = distributor.totalNumExport() - num_stay;
205 Kokkos::View<
typename AoSoA_t::tuple_type*,
206 typename Distributor_t::memory_space>
207 send_buffer( Kokkos::ViewAllocateWithoutInitializing(
208 "distributor_send_buffer" ),
212 Kokkos::View<
typename AoSoA_t::tuple_type*,
213 typename Distributor_t::memory_space>
214 recv_buffer( Kokkos::ViewAllocateWithoutInitializing(
215 "distributor_recv_buffer" ),
216 distributor.totalNumImport() );
219 auto steering = distributor.getExportSteering();
225 auto build_send_buffer_func = KOKKOS_LAMBDA(
const std::size_t i )
227 auto tpl = src.getTuple( steering( i ) );
229 recv_buffer( i ) = tpl;
231 send_buffer( i - num_stay ) = tpl;
233 Kokkos::RangePolicy<ExecutionSpace> build_send_buffer_policy(
234 0, distributor.totalNumExport() );
235 Kokkos::parallel_for(
"Cabana::Impl::distributeData::build_send_buffer",
236 build_send_buffer_policy, build_send_buffer_func );
240 const int mpi_tag = 1234;
243 std::vector<MPI_Request> requests;
244 requests.reserve( num_n );
245 std::pair<std::size_t, std::size_t> recv_range = { 0, 0 };
246 for (
int n = 0; n < num_n; ++n )
248 recv_range.second = recv_range.first + distributor.numImport( n );
250 if ( ( distributor.numImport( n ) > 0 ) &&
251 ( distributor.neighborRank( n ) != my_rank ) )
253 auto recv_subview = Kokkos::subview( recv_buffer, recv_range );
255 requests.push_back( MPI_Request() );
257 MPI_Irecv( recv_subview.data(),
258 recv_subview.size() *
259 sizeof(
typename AoSoA_t::tuple_type ),
260 MPI_BYTE, distributor.neighborRank( n ), mpi_tag,
261 distributor.comm(), &( requests.back() ) );
264 recv_range.first = recv_range.second;
268 std::pair<std::size_t, std::size_t> send_range = { 0, 0 };
269 for (
int n = 0; n < num_n; ++n )
271 if ( ( distributor.numExport( n ) > 0 ) &&
272 ( distributor.neighborRank( n ) != my_rank ) )
274 send_range.second = send_range.first + distributor.numExport( n );
276 auto send_subview = Kokkos::subview( send_buffer, send_range );
278 MPI_Send( send_subview.data(),
279 send_subview.size() *
280 sizeof(
typename AoSoA_t::tuple_type ),
281 MPI_BYTE, distributor.neighborRank( n ), mpi_tag,
282 distributor.comm() );
284 send_range.first = send_range.second;
289 std::vector<MPI_Status> status( requests.size() );
291 MPI_Waitall( requests.size(), requests.data(), status.data() );
292 if ( MPI_SUCCESS != ec )
293 throw std::logic_error(
294 "Cabana::Distributor: Failed MPI Communication" );
297 auto extract_recv_buffer_func = KOKKOS_LAMBDA(
const std::size_t i )
299 dst.setTuple( i, recv_buffer( i ) );
301 Kokkos::RangePolicy<ExecutionSpace> extract_recv_buffer_policy(
302 0, distributor.totalNumImport() );
303 Kokkos::parallel_for(
"Cabana::Impl::distributeData::extract_recv_buffer",
304 extract_recv_buffer_policy,
305 extract_recv_buffer_func );
309 MPI_Barrier( distributor.comm() );
336template <
class ExecutionSpace,
class Distributor_t,
class AoSoA_t>
337void migrate( ExecutionSpace exec_space,
const Distributor_t& distributor,
338 const AoSoA_t& src, AoSoA_t& dst,
344 if ( src.size() != distributor.exportSize() )
345 throw std::runtime_error(
"Cabana::migrate: Source is "
346 "the wrong size for migration! (Label: " +
348 if ( dst.size() != distributor.totalNumImport() )
349 throw std::runtime_error(
"Cabana::migrate: Destination "
350 "is the wrong size for migration! (Label: " +
354 Impl::distributeData( exec_space, distributor, src, dst );
374template <
class Distributor_t,
class AoSoA_t>
375void migrate(
const Distributor_t& distributor,
const AoSoA_t& src,
381 migrate(
typename Distributor_t::execution_space{}, distributor, src, dst );
408template <
class ExecutionSpace,
class Distributor_t,
class AoSoA_t>
409void migrate( ExecutionSpace exec_space,
const Distributor_t& distributor,
416 if ( aosoa.size() != distributor.exportSize() )
417 throw std::runtime_error(
418 "Cabana::migrate (in-place): "
419 "AoSoA is the wrong size for migration! (Label: " +
420 aosoa.label() +
")" );
425 ( distributor.totalNumImport() > distributor.exportSize() );
430 aosoa.resize( distributor.totalNumImport() );
433 Impl::distributeData( exec_space, distributor, aosoa, aosoa );
437 if ( !dst_is_bigger )
438 aosoa.resize( distributor.totalNumImport() );
462template <
class Distributor_t,
class AoSoA_t>
463void migrate(
const Distributor_t& distributor, AoSoA_t& aosoa,
468 migrate(
typename Distributor_t::execution_space{}, distributor, aosoa );
492template <
class ExecutionSpace,
class Distributor_t,
class Slice_t>
493void migrate( ExecutionSpace,
const Distributor_t& distributor,
494 const Slice_t& src, Slice_t& dst,
500 if ( src.size() != distributor.exportSize() )
501 throw std::runtime_error(
"Cabana::migrate: Source Slice is the wrong "
502 "size for migration! (Label: " +
504 if ( dst.size() != distributor.totalNumImport() )
505 throw std::runtime_error(
"Cabana::migrate: Destination Slice is the "
506 "wrong size for migration! (Label: " +
511 for (
size_t d = 2; d < src.viewRank(); ++d )
512 num_comp *= src.extent( d );
515 auto src_data = src.data();
516 auto dst_data = dst.data();
520 MPI_Comm_rank( distributor.comm(), &my_rank );
523 int num_n = distributor.numNeighbor();
529 std::size_t num_stay =
530 ( num_n > 0 && distributor.neighborRank( 0 ) == my_rank )
531 ? distributor.numExport( 0 )
536 std::size_t num_send = distributor.totalNumExport() - num_stay;
537 Kokkos::View<
typename Slice_t::value_type**, Kokkos::LayoutRight,
538 typename Distributor_t::memory_space>
539 send_buffer( Kokkos::ViewAllocateWithoutInitializing(
540 "distributor_send_buffer" ),
541 num_send, num_comp );
545 Kokkos::View<
typename Slice_t::value_type**, Kokkos::LayoutRight,
546 typename Distributor_t::memory_space>
547 recv_buffer( Kokkos::ViewAllocateWithoutInitializing(
548 "distributor_recv_buffer" ),
549 distributor.totalNumImport(), num_comp );
552 auto steering = distributor.getExportSteering();
557 auto build_send_buffer_func = KOKKOS_LAMBDA(
const std::size_t i )
559 auto s_src = Slice_t::index_type::s( steering( i ) );
560 auto a_src = Slice_t::index_type::a( steering( i ) );
561 std::size_t src_offset = s_src * src.stride( 0 ) + a_src;
563 for ( std::size_t n = 0; n < num_comp; ++n )
564 recv_buffer( i, n ) =
565 src_data[src_offset + n * Slice_t::vector_length];
567 for ( std::size_t n = 0; n < num_comp; ++n )
568 send_buffer( i - num_stay, n ) =
569 src_data[src_offset + n * Slice_t::vector_length];
571 Kokkos::RangePolicy<ExecutionSpace> build_send_buffer_policy(
572 0, distributor.totalNumExport() );
573 Kokkos::parallel_for(
"Cabana::migrate::build_send_buffer",
574 build_send_buffer_policy, build_send_buffer_func );
578 const int mpi_tag = 1234;
581 std::vector<MPI_Request> requests;
582 requests.reserve( num_n );
583 std::pair<std::size_t, std::size_t> recv_range = { 0, 0 };
584 for (
int n = 0; n < num_n; ++n )
586 recv_range.second = recv_range.first + distributor.numImport( n );
588 if ( ( distributor.numImport( n ) > 0 ) &&
589 ( distributor.neighborRank( n ) != my_rank ) )
592 Kokkos::subview( recv_buffer, recv_range, Kokkos::ALL );
594 requests.push_back( MPI_Request() );
596 MPI_Irecv( recv_subview.data(),
597 recv_subview.size() *
598 sizeof(
typename Slice_t::value_type ),
599 MPI_BYTE, distributor.neighborRank( n ), mpi_tag,
600 distributor.comm(), &( requests.back() ) );
603 recv_range.first = recv_range.second;
607 std::pair<std::size_t, std::size_t> send_range = { 0, 0 };
608 for (
int n = 0; n < num_n; ++n )
610 if ( ( distributor.numExport( n ) > 0 ) &&
611 ( distributor.neighborRank( n ) != my_rank ) )
613 send_range.second = send_range.first + distributor.numExport( n );
616 Kokkos::subview( send_buffer, send_range, Kokkos::ALL );
618 MPI_Send( send_subview.data(),
619 send_subview.size() *
620 sizeof(
typename Slice_t::value_type ),
621 MPI_BYTE, distributor.neighborRank( n ), mpi_tag,
622 distributor.comm() );
624 send_range.first = send_range.second;
629 std::vector<MPI_Status> status( requests.size() );
631 MPI_Waitall( requests.size(), requests.data(), status.data() );
632 if ( MPI_SUCCESS != ec )
633 throw std::logic_error(
"Cabana::migrate: Failed MPI Communication" );
636 auto extract_recv_buffer_func = KOKKOS_LAMBDA(
const std::size_t i )
638 auto s = Slice_t::index_type::s( i );
639 auto a = Slice_t::index_type::a( i );
640 std::size_t dst_offset = s * dst.stride( 0 ) + a;
641 for ( std::size_t n = 0; n < num_comp; ++n )
642 dst_data[dst_offset + n * Slice_t::vector_length] =
645 Kokkos::RangePolicy<ExecutionSpace> extract_recv_buffer_policy(
646 0, distributor.totalNumImport() );
647 Kokkos::parallel_for(
"Cabana::migrate::extract_recv_buffer",
648 extract_recv_buffer_policy,
649 extract_recv_buffer_func );
653 MPI_Barrier( distributor.comm() );
675template <
class Distributor_t,
class Slice_t>
676void migrate(
const Distributor_t& distributor,
const Slice_t& src,
682 migrate(
typename Distributor_t::execution_space{}, distributor, src, dst );
Array-of-Struct-of-Arrays particle data structure.
Multi-node communication patterns.
Slice a single particle property from an AoSoA.
Kokkos::View< size_type *, memory_space > createWithoutTopology(ExecutionSpace exec_space, Export, const RankViewType &element_export_ranks)
Export rank creator. Use this when you don't know who you will receiving from - only who you are send...
Definition Cabana_CommunicationPlan.hpp:770
Kokkos::View< size_type *, memory_space > createWithTopology(ExecutionSpace exec_space, Export, const RankViewType &element_export_ranks, const std::vector< int > &neighbor_ranks)
Neighbor and export rank creator. Use this when you already know which ranks neighbor each other (i....
Definition Cabana_CommunicationPlan.hpp:608
void createExportSteering(const PackViewType &neighbor_ids, const RankViewType &element_export_ranks)
Create the export steering vector.
Definition Cabana_CommunicationPlan.hpp:1504
MPI_Comm comm() const
Get the MPI communicator.
Definition Cabana_CommunicationPlan.hpp:480
CommunicationPlan(MPI_Comm comm)
Constructor.
Definition Cabana_CommunicationPlan.hpp:458
A communication plan for migrating data from one uniquely-owned decomposition to another uniquely own...
Definition Cabana_Distributor.hpp:63
Distributor(MPI_Comm comm, const ViewType &element_export_ranks)
Export rank constructor. Use this when you don't know who you will be receiving from - only who you a...
Definition Cabana_Distributor.hpp:137
Distributor(MPI_Comm comm, const ViewType &element_export_ranks, const std::vector< int > &neighbor_ranks)
Topology and export rank constructor. Use this when you already know which ranks neighbor each other ...
Definition Cabana_Distributor.hpp:100
Core: particle data structures and algorithms.
Definition Cabana_AoSoA.hpp:36
void migrate(ExecutionSpace exec_space, const Distributor_t &distributor, const AoSoA_t &src, AoSoA_t &dst, typename std::enable_if<(is_distributor< Distributor_t >::value &&is_aosoa< AoSoA_t >::value), int >::type *=0)
Synchronously migrate data between two different decompositions using the distributor forward communi...
Definition Cabana_Distributor.hpp:337
Export-based tag - default.
Definition Cabana_CommunicationPlan.hpp:42
Definition Cabana_Types.hpp:88
AoSoA static type checker.
Definition Cabana_AoSoA.hpp:61
Distributor static type checker.
Definition Cabana_Distributor.hpp:163
Slice static type checker.
Definition Cabana_Slice.hpp:868