Cabana_Distributor.hpp

Go to the documentation of this file.

/****************************************************************************
 * Copyright (c) 2018-2023 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_DISTRIBUTOR_HPP
#define CABANA_DISTRIBUTOR_HPP
 
#include <Cabana_AoSoA.hpp>
#include <Cabana_CommunicationPlanBase.hpp>
#include <Cabana_Core_Config.hpp>
#include <Cabana_Slice.hpp>
 
#include <Kokkos_Core.hpp>
#include <Kokkos_Profiling_ScopedRegion.hpp>
 
#include <mpi.h>
 
#include <exception>
#include <vector>
 
namespace Cabana
{
//---------------------------------------------------------------------------//
template <class MemorySpace, class CommSpaceType = Mpi>
class Distributor : public CommunicationPlan<MemorySpace, CommSpaceType>
{
  public:
    using commspace_type = CommSpaceType;

    template <class ViewType>
    Distributor( MPI_Comm comm, const ViewType& element_export_ranks,
                 const std::vector<int>& neighbor_ranks )
        : CommunicationPlan<MemorySpace, CommSpaceType>( comm )
    {
        auto neighbor_ids = this->createWithTopology(
            Export(), element_export_ranks, neighbor_ranks );
        this->createExportSteering( neighbor_ids, element_export_ranks );
    }

    template <class ViewType>
    Distributor( MPI_Comm comm, const ViewType& element_export_ranks )
        : CommunicationPlan<MemorySpace, CommSpaceType>( comm )
    {
        auto neighbor_ids =
            this->createWithoutTopology( Export(), element_export_ranks );
        this->createExportSteering( neighbor_ids, element_export_ranks );
    }
};
 
//---------------------------------------------------------------------------//
template <typename>
struct is_distributor_impl : public std::false_type
{
};
 
template <typename MemorySpace, typename CommSpaceType>
struct is_distributor_impl<Distributor<MemorySpace, CommSpaceType>>
    : public std::true_type
{
};

template <class T>
struct is_distributor
    : public is_distributor_impl<typename std::remove_cv<T>::type>::type
{
};
 
} // end namespace Cabana
 
// Include communication backends from what is enabled in CMake.
#ifdef Cabana_ENABLE_MPI
#include <impl/Cabana_Migrate_Mpi.hpp>
#endif // Enable MPI
 
namespace Cabana
{
 
//---------------------------------------------------------------------------//
template <class ExecutionSpace, class Distributor_t, class AoSoA_t>
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 )
{
    // Check that src and dst are the right size.
    if ( src.size() != distributor.exportSize() )
        throw std::runtime_error( "Cabana::migrate: Source is "
                                  "the wrong size for migration! (Label: " +
                                  src.label() + ")" );
    if ( dst.size() != distributor.totalNumImport() )
        throw std::runtime_error( "Cabana::migrate: Destination "
                                  "is the wrong size for migration! (Label: " +
                                  dst.label() + ")" );
 
    // Move the data.
    Impl::migrateData( typename Distributor_t::commspace_type(), exec_space,
                       distributor, src, dst );
}

template <class Distributor_t, class AoSoA_t>
void migrate( 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 )
{
    migrate( typename Distributor_t::execution_space{}, distributor, src, dst );
}
 
//---------------------------------------------------------------------------//
template <class ExecutionSpace, class Distributor_t, class AoSoA_t>
void migrate( ExecutionSpace exec_space, const Distributor_t& distributor,
              AoSoA_t& aosoa,
              typename std::enable_if<( is_distributor<Distributor_t>::value &&
                                        is_aosoa<AoSoA_t>::value ),
                                      int>::type* = 0 )
{
    // Check that the AoSoA is the right size.
    if ( aosoa.size() != distributor.exportSize() )
        throw std::runtime_error(
            "Cabana::migrate (in-place): "
            "AoSoA is the wrong size for migration! (Label: " +
            aosoa.label() + ")" );
 
    // Determine if the source of destination decomposition has more data on
    // this rank.
    bool dst_is_bigger =
        ( distributor.totalNumImport() > distributor.exportSize() );
 
    // If the destination decomposition is bigger than the source
    // decomposition resize now so we have enough space to do the operation.
    if ( dst_is_bigger )
        aosoa.resize( distributor.totalNumImport() );
 
    // Move the data.
    Impl::migrateData( typename Distributor_t::commspace_type(), exec_space,
                       distributor, aosoa, aosoa );
 
    // 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() );
}

template <class Distributor_t, class AoSoA_t>
void migrate( const Distributor_t& distributor, AoSoA_t& aosoa,
              typename std::enable_if<( is_distributor<Distributor_t>::value &&
                                        is_aosoa<AoSoA_t>::value ),
                                      int>::type* = 0 )
{
    migrate( typename Distributor_t::execution_space{}, distributor, aosoa );
}

template <class ExecutionSpace, class Distributor_t, class Slice_t>
void migrate( ExecutionSpace exec_space, const Distributor_t& distributor,
              const Slice_t& src, Slice_t& dst,
              typename std::enable_if<( is_distributor<Distributor_t>::value &&
                                        is_slice<Slice_t>::value ),
                                      int>::type* = 0 )
{
    // Check that src and dst are the right size.
    if ( src.size() != distributor.exportSize() )
        throw std::runtime_error( "Cabana::migrate: Source Slice is the wrong "
                                  "size for migration! (Label: " +
                                  src.label() + ")" );
    if ( dst.size() != distributor.totalNumImport() )
        throw std::runtime_error( "Cabana::migrate: Destination Slice is the "
                                  "wrong size for migration! (Label: " +
                                  dst.label() + ")" );
 
    // Move the data.
    Impl::migrateSlice( typename Distributor_t::commspace_type(), exec_space,
                        distributor, src, dst );
}

template <class Distributor_t, class Slice_t>
void migrate( const Distributor_t& distributor, const Slice_t& src,
              Slice_t& dst,
              typename std::enable_if<( is_distributor<Distributor_t>::value &&
                                        is_slice<Slice_t>::value ),
                                      int>::type* = 0 )
{
    migrate( typename Distributor_t::execution_space{}, distributor, src, dst );
}
//---------------------------------------------------------------------------//
 
} // end namespace Cabana
 
#endif // end CABANA_DISTRIBUTOR_HPP